adapt to JuliaGPU/GPUCompiler.jl#718

Author: simeonschaub

.github/workflows/Test.yml

@@ -136,7 +136,7 @@ jobs:
           julia --project -e '
             using Pkg
             Pkg.develop(path="lib/intrinsics")
-            Pkg.add(name="GPUCompiler", rev="sds/additional_args")'
+            Pkg.add(name="GPUCompiler", rev="sds/add_input_args")'
 
       - name: Test OpenCL.jl
         uses: julia-actions/julia-runtest@v1

Project.toml

@@ -22,7 +22,7 @@ SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [sources]
-GPUCompiler = {rev = "sds/additional_args", url = "https://github.com/JuliaGPU/GPUCompiler.jl.git"}
+GPUCompiler = {rev = "sds/add_input_args", url = "https://github.com/JuliaGPU/GPUCompiler.jl.git"}
 SPIRVIntrinsics = {path = "lib/intrinsics"}
 
 [compat]

src/compiler/compilation.jl

@@ -19,21 +19,27 @@ GPUCompiler.isintrinsic(job::OpenCLCompilerJob, fn::String) =
     contains(fn, "__spirv_")
 
 GPUCompiler.kernel_state_type(job::OpenCLCompilerJob) = KernelState
-GPUCompiler.additional_arg_types(job::OpenCLCompilerJob) = (; random_keys = LLVMPtr{UInt32, AS.Workgroup}, random_counters = LLVMPtr{UInt32, AS.Workgroup})
 
 function GPUCompiler.finish_module!(@nospecialize(job::OpenCLCompilerJob),
                                     mod::LLVM.Module, entry::LLVM.Function)
     entry = invoke(GPUCompiler.finish_module!,
                    Tuple{CompilerJob{SPIRVCompilerTarget}, LLVM.Module, LLVM.Function},
                    job, mod, entry)
 
+    kernel_intrinsics = Dict(
+        "julia.spirv.random_keys" => (; name = "random_keys", typ = LLVMPtr{UInt32, AS.Workgroup}),
+        "julia.spirv.random_counters" => (; name = "random_counters", typ = LLVMPtr{UInt32, AS.Workgroup}),
+    )
+    entry′ = GPUCompiler.add_input_arguments!(job, mod, entry, kernel_intrinsics)
+
     # if this kernel uses our RNG, we should prime the shared state.
     # XXX: these transformations should really happen at the Julia IR level...
-    if callconv(entry) == LLVM.API.LLVMSPIRKERNELCallConv && haskey(functions(mod), "julia.gpu.additional_arg_getter")
+    if job.config.kernel && entry !== entry′
+        entry = entry′
         # insert call to `initialize_rng_state`
         f = initialize_rng_state
         ft = typeof(f)
-        tt = Tuple{}
+        tt = NTuple{2, LLVMPtr{UInt32, AS.Workgroup}}
 
         # create a deferred compilation job for `initialize_rng_state`
         src = methodinstance(ft, tt, GPUCompiler.tls_world_age())
@@ -74,16 +80,18 @@ function GPUCompiler.finish_module!(@nospecialize(job::OpenCLCompilerJob),
             # call the `initialize_rng_state` function
             rt = Core.Compiler.return_type(f, tt)
             llvm_rt = convert(LLVMType, rt)
-            llvm_ft = LLVM.FunctionType(llvm_rt)
+            llvm_ft = LLVM.FunctionType(llvm_rt, [convert(LLVMType, LLVMPtr{UInt32, AS.Workgroup}) for _ in 1:2])
             fptr = inttoptr!(builder, fptr, LLVM.PointerType(llvm_ft))
-            call!(builder, llvm_ft, fptr)
+            random_keys = findfirst(arg -> name(arg) == "random_keys", parameters(entry))
+            random_counters = findfirst(arg -> name(arg) == "random_counters", parameters(entry))
+            call!(builder, llvm_ft, fptr, parameters(entry)[[random_keys, random_counters]])
             br!(builder, top_bb)
         end
 
         # XXX: put some of the above behind GPUCompiler abstractions
         #      (e.g., a compile-time version of `deferred_codegen`)
     end
-    return entry
+    return entry′
 end
 
 ## compiler implementation (cache, configure, compile, and link)

src/device/random.jl

@@ -10,21 +10,22 @@ import RandomNumbers
 # local memory with the actual seed, per subgroup, set by `initialize_rng_state`` or overridden by calling `seed!`
 @inline function global_random_keys()
     n = get_num_sub_groups()
-    ptr = additional_args(Val{1}())::LLVMPtr{UInt32, AS.Workgroup}
+    ptr = random_keys()::LLVMPtr{UInt32, AS.Workgroup}
     return CLDeviceArray{UInt32, 1, AS.Workgroup}((n,), ptr)
 end
 
 # local memory with per-subgroup counters, incremented when generating numbers
 @inline function global_random_counters()
     n = get_num_sub_groups()
-    ptr = additional_args(Val{2}())::LLVMPtr{UInt32, AS.Workgroup}
+    ptr = random_counters()::LLVMPtr{UInt32, AS.Workgroup}
     return CLDeviceArray{UInt32, 1, AS.Workgroup}((n,), ptr)
 end
 
 # initialization function, called automatically at the start of each kernel
-function initialize_rng_state()
-    random_keys = global_random_keys()
-    random_counters = global_random_counters()
+function initialize_rng_state(random_keys_ptr::LLVMPtr{UInt32, AS.Workgroup}, random_counters_ptr::LLVMPtr{UInt32, AS.Workgroup})
+    n = get_num_sub_groups()
+    random_keys = CLDeviceArray{UInt32, 1, AS.Workgroup}((n,), random_keys_ptr)
+    random_counters = CLDeviceArray{UInt32, 1, AS.Workgroup}((n,), random_counters_ptr)
 
     subgroup_id = get_sub_group_id()
     @inbounds random_keys[subgroup_id] = kernel_state().random_seed

src/device/runtime.jl

@@ -20,4 +20,5 @@ struct KernelState
 end
 
 @inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)
-@inline @generated additional_args(::Val{i}) where {i} = GPUCompiler.additional_arg_value(LLVMPtr{UInt32, AS.Workgroup}, i)
+@inline @generated random_keys() = GPUCompiler.call_custom_intrinsic(LLVMPtr{UInt32, AS.Workgroup}, "julia.spirv.random_keys", "random_keys")
+@inline @generated random_counters() = GPUCompiler.call_custom_intrinsic(LLVMPtr{UInt32, AS.Workgroup}, "julia.spirv.random_counters", "random_counters")

test/Project.toml

@@ -24,7 +24,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 pocl_jll = "627d6b7a-bbe6-5189-83e7-98cc0a5aeadd"
 
 [sources]
-GPUCompiler = {rev = "sds/additional_args", url = "https://github.com/JuliaGPU/GPUCompiler.jl.git"}
+GPUCompiler = {rev = "sds/add_input_args", url = "https://github.com/JuliaGPU/GPUCompiler.jl.git"}
 
 [compat]
 pocl_jll = "7.0"