From 1550ab806519b92d8782eb9cf604011e4d3c4ab6 Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Thu, 13 Feb 2025 09:54:44 +0100
Subject: [PATCH 1/2] Use invariant.load for ldg

---
 src/device/pointer.jl | 52 ++++++++++++++++++++++++++-----------------
 1 file changed, 32 insertions(+), 20 deletions(-)

diff --git a/src/device/pointer.jl b/src/device/pointer.jl
index 0be5b0d4f4..c9d514c81e 100644
--- a/src/device/pointer.jl
+++ b/src/device/pointer.jl
@@ -15,29 +15,41 @@ const Local    = 5
 end
 
 
-## ldg
+@inline @generated function pointerref_ldg(ptr::LLVMPtr{T,A}, i::I, ::Val{align}) where {T,A,I,align}
+    sizeof(T) == 0 && return T.instance
+    LLVM.@dispose ctx=LLVM.Context() begin
+        eltyp = convert(LLVM.LLVMType, T)
 
-const LDGTypes = (UInt8, UInt16, UInt32, UInt64, Int8, Int16, Int32, Int64,
-                  Float32, Float64)
+        T_idx = convert(LLVM.LLVMType, I)
+        T_ptr = convert(LLVM.LLVMType, ptr)
 
-# TODO: this functionality should throw <sm_32
-# NOTE: CUDA 8.0 supports more caching modifiers, but those aren't supported by LLVM yet
-for T in LDGTypes
-    class = if T <: Integer
-        :i
-    elseif T <: AbstractFloat
-        :f
-    end
-    # TODO: p class
-    width = sizeof(T)*8 # in bits
-    typ = Symbol(class, width)
+        T_typed_ptr = LLVM.PointerType(eltyp, A)
 
-    intr = "llvm.nvvm.ldg.global.$class.$typ.p1$typ"
-    @eval @inline function pointerref_ldg(base_ptr::LLVMPtr{$T,AS.Global}, i::Integer,
-                                          ::Val{align}) where align
-        offset = i-one(i) # in elements
-        ptr = base_ptr + offset*sizeof($T)
-        @typed_ccall($intr, llvmcall, $T, (LLVMPtr{$T,AS.Global}, Int32), ptr, Val(align))
+        # create a function
+        param_types = [T_ptr, T_idx]
+        llvm_f, _ = LLVM.Interop.create_function(eltyp, param_types)
+
+        # generate IR
+        LLVM.@dispose builder=LLVM.IRBuilder() begin
+            entry = LLVM.BasicBlock(llvm_f, "entry")
+            LLVM.position!(builder, entry)
+            ptr = if LLVM.supports_typed_pointers(ctx)
+                typed_ptr = LLVM.bitcast!(builder, LLVM.parameters(llvm_f)[1], T_typed_ptr)
+                LLVM.inbounds_gep!(builder, eltyp, typed_ptr, [LLVM.parameters(llvm_f)[2]])
+            else
+                LLVM.inbounds_gep!(builder, eltyp, LLVM.parameters(llvm_f)[1], [LLVM.parameters(llvm_f)[2]])
+            end
+            ld = LLVM.load!(builder, eltyp, ptr)
+            if A != 0
+                LLVM.metadata(ld)[LLVM.MD_tbaa] = LLVM.Interop.tbaa_addrspace(A)
+            end
+            LLVM.alignment!(ld, align)
+            LLVM.metadata(ld)[LLVM.MD_invariant_load] = LLVM.MDNode(LLVM.Metadata[nothing])
+
+            LLVM.ret!(builder, ld)
+        end
+
+        LLVM.Interop.call_function(llvm_f, T, Tuple{LLVMPtr{T,A}, I}, :ptr, :(i-one(I)))
     end
 end
 

From 49fefeaa846c05d03d7313c5ae8a4ac3640454fe Mon Sep 17 00:00:00 2001
From: Valentin Churavy <v.churavy@gmail.com>
Date: Thu, 13 Feb 2025 09:58:14 +0100
Subject: [PATCH 2/2] fixup! Use invariant.load for ldg

---
 src/device/pointer.jl | 6 +-----
 1 file changed, 1 insertion(+), 5 deletions(-)

diff --git a/src/device/pointer.jl b/src/device/pointer.jl
index c9d514c81e..baa7709f04 100644
--- a/src/device/pointer.jl
+++ b/src/device/pointer.jl
@@ -78,9 +78,5 @@ end
 
 export unsafe_cached_load
 
-unsafe_cached_load(p::LLVMPtr{<:Union{LDGTypes...},AS.Global}, i::Integer=1, align::Val=Val(1)) =
-    pointerref_ldg(p, i, align)
-# NOTE: fall back to normal unsafe_load for unsupported types. we could be smarter here,
-#       e.g. destruct/load/reconstruct, but that's too complicated for what it's worth.
 unsafe_cached_load(p::LLVMPtr, i::Integer=1, align::Val=Val(1)) =
-    unsafe_load(p, i, align)
+    pointerref_ldg(p, i, align)