address review comments

Co-authored-by: Tim Besard <[email protected]>

Author: simeonschaub

lib/cl/kernel.jl

@@ -158,7 +158,7 @@ end
 
 function enqueue_kernel(k::Kernel, global_work_size, local_work_size=nothing;
                         global_work_offset=nothing, wait_on::Vector{Event}=Event[],
-                        device_rng=false)
+                        rng_state=false)
     max_work_dim = device().max_work_item_dims
     work_dim     = length(global_work_size)
     if work_dim > max_work_dim
@@ -201,7 +201,7 @@ function enqueue_kernel(k::Kernel, global_work_size, local_work_size=nothing;
         # null local size means OpenCL decides
     end
 
-    if device_rng
+    if rng_state
         if local_work_size !== nothing
             num_sub_groups = KernelSubGroupInfo(k, device(), lsize).sub_group_count
         else
@@ -249,7 +249,7 @@ function call(
         k::Kernel, args...; global_size = (1,), local_size = nothing,
         global_work_offset = nothing, wait_on::Vector{Event} = Event[],
         indirect_memory::Vector{AbstractMemory} = AbstractMemory[],
-        device_rng=false,
+        rng_state=false,
     )
     set_args!(k, args...)
     if !isempty(indirect_memory)
@@ -303,7 +303,7 @@ function call(
             clSetKernelExecInfo(k, CL_KERNEL_EXEC_INFO_USM_PTRS_INTEL, sizeof(usm_pointers), usm_pointers)
         end
     end
-    enqueue_kernel(k, global_size, local_size; global_work_offset, wait_on, device_rng)
+    enqueue_kernel(k, global_size, local_size; global_work_offset, wait_on, rng_state)
 end
 
 # From `julia/base/reflection.jl`, adjusted to add specialization on `t`.

src/compiler/compilation.jl

@@ -28,7 +28,7 @@ function GPUCompiler.finish_module!(@nospecialize(job::OpenCLCompilerJob),
 
     # if this kernel uses our RNG, we should prime the shared state.
     # XXX: these transformations should really happen at the Julia IR level...
-    if haskey(functions(mod), "julia.spirv.random_keys") && job.config.kernel
+    if haskey(functions(mod), "julia.opencl.random_keys") && job.config.kernel
         # insert call to `initialize_rng_state`
         f = initialize_rng_state
         ft = typeof(f)
@@ -88,8 +88,8 @@ end
 function GPUCompiler.finish_linked_module!(@nospecialize(job::OpenCLCompilerJob), mod::LLVM.Module)
     for f in GPUCompiler.kernels(mod)
         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}),
+            "julia.opencl.random_keys" => (; name = "random_keys", typ = LLVMPtr{UInt32, AS.Workgroup}),
+            "julia.opencl.random_counters" => (; name = "random_counters", typ = LLVMPtr{UInt32, AS.Workgroup}),
         )
         GPUCompiler.add_input_arguments!(job, mod, f, kernel_intrinsics)
     end

src/compiler/execution.jl

@@ -159,8 +159,7 @@ abstract type AbstractKernel{F, TT} end
 
     quote
         indirect_memory = cl.AbstractMemory[]
-        device_rng = kernel.fun.num_args == $(length(call_args) + 2)
-        clcall(kernel.fun, $call_tt, $(call_args...); indirect_memory, device_rng, call_kwargs...)
+        clcall(kernel.fun, $call_tt, $(call_args...); indirect_memory, kernel.rng_state, call_kwargs...)
     end
 end
 
@@ -171,6 +170,7 @@ end
 struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     f::F
     fun::cl.Kernel
+    rng_state::Bool
 end
 
 
@@ -194,8 +194,10 @@ function clfunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
         h = hash(fun, hash(f, hash(tt)))
         kernel = get(_kernel_instances, h, nothing)
         if kernel === nothing
+            # TODO: move the `rng_state` check into `OpenCL.compile` so we avoid the API call?
+            rng_state = fun.num_args == length(tt.parameters) + 2
             # create the kernel state object
-            kernel = HostKernel{F,tt}(f, fun)
+            kernel = HostKernel{F,tt}(f, fun, rng_state)
             _kernel_instances[h] = kernel
         end
         return kernel::HostKernel{F,tt}

src/device/runtime.jl

@@ -21,11 +21,19 @@ end
 
 @inline @generated kernel_state() = GPUCompiler.kernel_state_value(KernelState)
 
+## intrinsics for adding and accessing additional kernel arguments
+
+# The amount of local shared memory we need for storing RNG state is determined
+# dynamically at kernel launch time, so needs to be passed as additional arguments
+# to the kernel.
+# We define intrinsics that get transformed into additional kernel arguments which
+# then get propagated across function calls to the caller.
+
 function additional_arg_intr(mod::LLVM.Module, T_state, name)
-    state_intr = if haskey(functions(mod), "julia.spirv.$name")
-        functions(mod)["julia.spirv.$name"]
+    state_intr = if haskey(functions(mod), "julia.opencl.$name")
+        functions(mod)["julia.opencl.$name"]
     else
-        LLVM.Function(mod, "julia.spirv.$name", LLVM.FunctionType(T_state))
+        LLVM.Function(mod, "julia.opencl.$name", LLVM.FunctionType(T_state))
     end
     push!(function_attributes(state_intr), EnumAttribute("readnone", 0))
 

test/Project.toml

@@ -3,7 +3,6 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
-GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 IOCapture = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"