Cut allocations when launching CPU kernels

src/pocl/backend.jl

@@ -185,7 +185,7 @@ function threads_to_workgroupsize(threads, ndrange)
     end
 end
 
-function (obj::KA.Kernel{POCLBackend})(args...; ndrange = nothing, workgroupsize = nothing)
+function (obj::KA.Kernel{POCLBackend})(args::Vararg{Any, N}; ndrange = nothing, workgroupsize = nothing) where {N}
     ndrange, workgroupsize, iterspace, dynamic =
         KA.launch_config(obj, ndrange, workgroupsize)
 

src/pocl/compiler/execution.jl

@@ -86,7 +86,7 @@ end
 ## argument conversion
 
 struct KernelAdaptor
-    svm_pointers::Vector{Ptr{Cvoid}}
+    svm_pointers::Union{Nothing, Vector{Ptr{Cvoid}}}
 end
 
 # # assume directly-passed pointers are SVM pointers
@@ -137,7 +137,7 @@ register methods for the the `OpenCL.KernelAdaptor` type.
 The `pointers` argument is used to collect pointers to indirect SVM buffers, which need to
 be registered with OpenCL before invoking the kernel.
 """
-function clconvert(arg, pointers::Vector{Ptr{Cvoid}} = Ptr{Cvoid}[])
+function clconvert(arg, pointers::Union{Nothing, Vector{Ptr{Cvoid}}} = nothing)
     return adapt(KernelAdaptor(pointers), arg)
 end
 
@@ -149,11 +149,11 @@ abstract type AbstractKernel{F, TT} end
 pass_arg(@nospecialize dt) = !(GPUCompiler.isghosttype(dt) || Core.Compiler.isconstType(dt))
 
 @inline @generated function (kernel::AbstractKernel{F, TT})(
-        args...;
-        call_kwargs...
-    ) where {F, TT}
+        args::Vararg{Any, N};
+        global_size = (1,), local_size = nothing
+    ) where {F, TT, N}
     sig = Tuple{F, TT.parameters...}    # Base.signature_type with a function type
-    args = (:(kernel.f), (:(clconvert(args[$i], svm_pointers)) for i in 1:length(args))...)
+    args = (:(kernel.f), (:(clconvert(args[$i])) for i in 1:length(args))...)
 
     # filter out ghost arguments that shouldn't be passed
     to_pass = map(pass_arg, sig.parameters)
@@ -175,8 +175,7 @@ pass_arg(@nospecialize dt) = !(GPUCompiler.isghosttype(dt) || Core.Compiler.isco
     call_tt = Base.to_tuple_type(call_t)
 
     return quote
-        svm_pointers = Ptr{Cvoid}[]
-        $cl.clcall(kernel.fun, $call_tt, $(call_args...); svm_pointers, kernel.rng_state, call_kwargs...)
+        $cl.clcall(kernel.fun, $call_tt, $(call_args...); global_size, local_size, kernel.rng_state)
     end
 end
 
@@ -194,8 +193,37 @@ end
 
 const clfunction_lock = ReentrantLock()
 
+# `HostKernel` with the world age and context it was resolved in; valid as long as no method
+# has been defined since and the context is unchanged.
+struct ResolvedKernel
+    world::UInt
+    context::nanoOpenCL.Context
+    kernel::Any
+end
+
+# `HostKernel{F, tt}` instances keyed by their type
+const _kernel_fastpath = Dict{DataType, ResolvedKernel}()
+
+# On Julia 1.11 to 1.13 reading a `ScopedValue` allocates; outside of any dynamic scope it
+# holds its default, so the read is skipped there.
+@static if v"1.11" <= VERSION < v"1.14-"
+    @inline compile_hook_set() = Core.current_scope() !== nothing && GPUCompiler.compile_hook[] !== nothing
+else
+    @inline compile_hook_set() = GPUCompiler.compile_hook[] !== nothing
+end
+
 function clfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F, TT}
     Base.@lock clfunction_lock begin
+        ctx = context()
+        world = Base.get_world_counter()
+        cacheable = isempty(kwargs) && !compile_hook_set()
+        if cacheable
+            entry = get(_kernel_fastpath, HostKernel{F, tt}, nothing)
+            if entry !== nothing && entry.world == world && entry.context === ctx
+                return entry.kernel::HostKernel{F, tt}
+            end
+        end
+
         config = compiler_config(device(); kwargs...)::OpenCLCompilerConfig
         source = methodinstance(F, tt)
         job = CompilerJob(source, config)
@@ -204,28 +232,34 @@ function clfunction(f::F, tt::TT = Tuple{}; kwargs...) where {F, TT}
 
         # Resolve the cl.Kernel for the active context. Linear scan over the
         # session-local cache; almost always n=1, so this is one `===` compare.
-        ctx = context()
-        kernel = Ref{nanoOpenCL.Kernel}()
+        cached = nothing
         @inbounds for (cached_ctx, cached_kernel) in res.kernels
             if cached_ctx === ctx
-                kernel[] = cached_kernel
+                cached = cached_kernel
                 break
             end
         end
-        if !isassigned(kernel)
-            kernel[] = link_kernel(job, res.obj::Vector{UInt8}, res.entry::String)
+        kernel = if cached === nothing
+            linked = link_kernel(job, res.obj::Vector{UInt8}, res.entry::String)
             # Don't cache session-local kernel handles while precompiling: the
             # results struct is serialized into the package image along with its
             # CodeInstance, and the handles would come back dangling.
             if ccall(:jl_generating_output, Cint, ()) != 1
-                push!(res.kernels, (ctx, kernel[]))
+                push!(res.kernels, (ctx, linked))
             end
+            linked
+        else
+            cached
         end
 
-        h = hash(kernel[], hash(f, hash(tt)))
-        return get!(_kernel_instances, h) do
-            HostKernel{F, tt}(f, kernel[], res.device_rng)
+        h = hash(kernel, hash(f, hash(tt)))
+        hostkernel = get!(_kernel_instances, h) do
+            HostKernel{F, tt}(f, kernel, res.device_rng)
         end::HostKernel{F, tt}
+        if cacheable
+            _kernel_fastpath[HostKernel{F, tt}] = ResolvedKernel(world, ctx, hostkernel)
+        end
+        return hostkernel
     end
 end
 
@@ -242,7 +276,9 @@ function compile_or_lookup(@nospecialize(job::CompilerJob))::OpenCLResults
     res = GPUCompiler.cached_results(OpenCLResults, job)
     if res === nothing || res.obj === nothing || GPUCompiler.compile_hook[] !== nothing
         compiled = compile_to_obj(job)
-        res = @something res GPUCompiler.cached_results(OpenCLResults, job)
+        if res === nothing
+            res = GPUCompiler.cached_results(OpenCLResults, job)::OpenCLResults
+        end
         res.obj = compiled.obj
         res.entry = compiled.entry
         res.device_rng = compiled.device_rng

src/pocl/nanoOpenCL.jl

@@ -665,6 +665,23 @@ end
     )::cl_int
 end
 
+# sizes passed as tuples, which `ccall` copies to the stack
+@checked function clEnqueueNDRangeKernel(
+        command_queue, kernel, work_dim,
+        global_work_size::NTuple{3, Csize_t}, local_work_size::NTuple{3, Csize_t}, event::Ref{cl_event}
+    )
+    @ccall libopencl.POclEnqueueNDRangeKernel(
+        command_queue::cl_command_queue,
+        kernel::cl_kernel, work_dim::cl_uint,
+        C_NULL::Ptr{Csize_t},
+        global_work_size::Ref{NTuple{3, Csize_t}},
+        local_work_size::Ref{NTuple{3, Csize_t}},
+        0::cl_uint,
+        C_NULL::Ptr{cl_event},
+        event::Ref{cl_event}
+    )::cl_int
+end
+
 @checked function clEnqueueNDRangeKernel(
         command_queue, kernel, work_dim,
         global_work_offset, global_work_size,
@@ -1254,9 +1271,10 @@ function set_arg!(k::Kernel, idx::Integer, arg::LocalMem)
 end
 
 function set_arg!(k::Kernel, idx::Integer, arg::T) where {T}
-    ref = Ref(arg)
-    tsize = sizeof(ref)
-    err = unchecked_clSetKernelArg(k, cl_uint(idx - 1), tsize, ref)
+    # `Ref{T}` makes `ccall` pass a pointer to a stack copy of `arg`
+    err = @ccall libopencl.POclSetKernelArg(
+        k::cl_kernel, cl_uint(idx - 1)::cl_uint, sizeof(T)::Csize_t, arg::Ref{T}
+    )::cl_int
     if err == CL_INVALID_ARG_SIZE
         error(
             """Mismatch between Julia and OpenCL type for kernel argument $idx.
@@ -1275,19 +1293,36 @@ function set_arg!(k::Kernel, idx::Integer, arg::T) where {T}
     return k
 end
 
-function set_args!(k::Kernel, args...)
-    for (i, a) in enumerate(args)
-        set_arg!(k, i, a)
-    end
-    return
+set_args!(k::Kernel, args::Vararg{Any, N}) where {N} = set_args!(k, 1, args...)
+@inline set_args!(k::Kernel, i::Int) = nothing
+@inline function set_args!(k::Kernel, i::Int, arg, args::Vararg{Any, N}) where {N}
+    set_arg!(k, i, arg)
+    return set_args!(k, i + 1, args...)
 end
 
+# work sizes padded to the three dimensions OpenCL devices support
+const WorkSize = NTuple{3, Csize_t}
+@inline work_size(sizes) = ntuple(i -> i <= length(sizes) ? Csize_t(sizes[i]) : Csize_t(0), Val(3))
+
 function enqueue_kernel(
         k::Kernel, global_work_size, local_work_size = nothing;
         global_work_offset = nothing, rng_state = false, nargs = nothing
     )
-    max_work_dim = device().max_work_item_dims
     work_dim = length(global_work_size)
+
+    if global_work_offset === nothing && local_work_size !== nothing && !rng_state && work_dim <= 3
+        if length(local_work_size) != work_dim
+            throw(ArgumentError("global_work_size and local_work_size have differing dims"))
+        end
+        ret_event = Ref{cl_event}()
+        clEnqueueNDRangeKernel(
+            queue(), k, cl_uint(work_dim),
+            work_size(global_work_size), work_size(local_work_size), ret_event
+        )
+        return Event(ret_event[])
+    end
+
+    max_work_dim = device().max_work_item_dims
     if work_dim > max_work_dim
         throw(ArgumentError("global_work_size has max dim of $max_work_dim"))
     end
@@ -1354,19 +1389,19 @@ function enqueue_kernel(
 end
 
 function call(
-        k::Kernel, args...; global_size = (1,), local_size = nothing,
+        k::Kernel, args::Vararg{Any, N}; global_size = (1,), local_size = nothing,
         global_work_offset = nothing,
-        svm_pointers::Vector{Ptr{Cvoid}} = Ptr{Cvoid}[],
+        svm_pointers::Union{Nothing, Vector{Ptr{Cvoid}}} = nothing,
         rng_state = false
-    )
+    ) where {N}
     set_args!(k, args...)
-    if !isempty(svm_pointers)
+    if svm_pointers !== nothing && !isempty(svm_pointers)
         clSetKernelExecInfo(
             k, CL_KERNEL_EXEC_INFO_SVM_PTRS,
             sizeof(svm_pointers), svm_pointers
         )
     end
-    return enqueue_kernel(k, global_size, local_size; global_work_offset, rng_state, nargs = length(args))
+    return enqueue_kernel(k, global_size, local_size; global_work_offset, rng_state, nargs = N)
 end
 
 # convert the argument values to match the kernel's signature (specified by the user)