POCL: reduce launch memory allocations

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelAbstractions"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
-version = "0.10.0"
+version = "0.10.1"
 authors = ["Valentin Churavy <v.churavy@gmail.com> and contributors"]
 
 [workspace]
@@ -43,7 +43,7 @@ StaticArraysExt = "StaticArrays"
 Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "1.5"
 GPUCompiler = "2.13"
-GPUToolbox = "3.3.3"
+GPUToolbox = "3.4.1"
 KernelInterface = "0.4"
 LLVM = "10"
 LinearAlgebra = "1.6"

src/pocl/backend.jl

@@ -165,11 +165,10 @@ struct POCLKernel{K, F}
 end
 
 function KI.kernel_function(backend::POCLBackend, f::F, tt::TT = Tuple{}; name = nothing, kwargs...) where {F, TT}
-    # fix the sub-group width, as `KI.sub_group_size` promises. pass it even if the device
-    # has no sub-groups, so that `clfunction` is only compiled for one set of keywords.
-    sub_group_size = device_limits().sub_group_size
-    sub_group_size = sub_group_size > 0 ? sub_group_size : nothing
-    kernel = clfunction(clconvert(f), tt; name, sub_group_size, kwargs...)
+    # the sub-group width defaults to the one `KI.sub_group_size` reports. it is determined
+    # when creating the compiler configuration, since passing it here would make the keyword
+    # arguments' type depend on the device, and the call to `clfunction` dynamic.
+    kernel = clfunction(clconvert(f), tt; name, kwargs...)
     kern = POCLKernel(kernel, f)
     return KI.Kernel{POCLBackend, typeof(kern)}(backend, kern)
 end

src/pocl/compiler/compilation.jl

@@ -165,7 +165,8 @@ end
 
 The GPUCompiler configuration for compiling kernels for `dev`, cached per device and
 keyword arguments. Besides those of `CompilerConfig` (`kernel`, `name`, `always_inline`,
-`debug_level`) and `sub_group_size`, it takes:
+`debug_level`) and `sub_group_size` (the sub-group width to require, or `nothing` for none,
+defaulting to the width `KI.sub_group_size` reports), it takes:
 
 - `atomics`: override the atomic capabilities GPUCompiler may select directly. This
   replaces the whole device-derived `SPIRVAtomics` (see `spirv_atomics`), and defaults to
@@ -190,7 +191,8 @@ const default_extensions = "+SPV_KHR_no_integer_wrap_decoration"
 const _toolchain = Ref{Any}()
 const _compiler_configs = Dict{UInt, OpenCLCompilerConfig}()
 function compiler_config(dev::cl.Device; kwargs...)
-    h = hash(dev, hash(kwargs))
+    # (as a `NamedTuple`: hashing the `Pairs` boxes each value, on every launch)
+    h = hash(dev, hash(values(kwargs)))
     # launches already hold this (reentrant) lock, but reflection doesn't
     return @lock clfunction_lock begin
         config = get(_compiler_configs, h, nothing)
@@ -204,7 +206,7 @@ end
 @noinline function _compiler_config(
         dev; kernel = true, name = nothing, always_inline = false,
         debug_level = Base.JLOptions().debug_level,
-        sub_group_size::Union{Nothing, Int} = 32,
+        sub_group_size::Union{Nothing, Int} = default_sub_group_size(dev),
         atomics::SPIRVAtomics = spirv_atomics(dev),
         extensions::Union{Nothing, String} = nothing, kwargs...
     )

src/pocl/nanoOpenCL.jl

@@ -426,7 +426,8 @@ end
     throw(CLError(res))
 end
 
-function check(f)
+# inlined, so that `Ref`s passed to API calls in `f` can be allocated on the stack
+@inline function check(f)
     res = f()
 
     if res != CL_SUCCESS
@@ -1355,8 +1356,9 @@ function enqueue_kernel(
         if length(local_work_size) != work_dim
             throw(ArgumentError("global_work_size and local_work_size have differing dims"))
         end
+        # inlined, so that `ret_event` can be allocated on the stack
         ret_event = Ref{cl_event}()
-        clEnqueueNDRangeKernel(
+        @inline clEnqueueNDRangeKernel(
             queue(), k, cl_uint(work_dim),
             work_size(global_work_size), work_size(local_work_size), ret_event
         )

src/pocl/pocl.jl

@@ -23,6 +23,15 @@ struct Session
     }
 end
 
+# the sub-group width kernels are compiled for, or `nothing` if the device has no sub-groups.
+# POCL can technically support any sub-group size; prefer the common GPU ones, and only
+# widths of up to 64, so that `KI.sub_group_ballot` (a `UInt64` mask) covers every lane.
+function default_sub_group_size(device::cl.Device)
+    sg_sizes = filter(<=(64), device.sub_group_sizes)
+    common = filter(in(sg_sizes), [32, 64, 16, sg_sizes...])
+    return isempty(common) ? nothing : first(common)
+end
+
 function Session()
     cl.pocl_standalone_jll.is_available() ||
         error("The CPU back-end needs PoCL, which is not available for $(Sys.MACHINE)")
@@ -39,15 +48,12 @@ function Session()
     context = cl.Context(device)
 
     sizes = device.max_work_item_size
-    # POCL can technically support any sub-group size; prefer the common GPU ones, and only
-    # widths of up to 64, so that `KI.sub_group_ballot` (a `UInt64` mask) covers every lane
-    sg_sizes = filter(<=(64), device.sub_group_sizes)
-    common = filter(in(sg_sizes), [32, 64, 16, sg_sizes...])
+    sub_group_size = default_sub_group_size(device)
     limits = (;
         max_work_group_size = Int(device.max_work_group_size),
         max_work_group_dims = ntuple(d -> d <= length(sizes) ? sizes[d] : 1, 3),
         # 0 if the device has no sub-groups
-        sub_group_size = isempty(common) ? 0 : first(common),
+        sub_group_size = something(sub_group_size, 0),
         # the threads PoCL runs work-groups on
         compute_units = Int(device.max_compute_units),
     )

test/runtests.jl

@@ -423,6 +423,64 @@ end
     end
 end
 
+# launching a compiled kernel doesn't allocate
+@testset "POCL launch allocations" begin
+    mod = @eval module $(gensym())
+    import KernelAbstractions, KernelAbstractions.POCL
+    using KernelAbstractions: @kernel, @index
+    kernel(A) = (A[POCL.get_global_id()] = 1; return)
+    compiler_config(dev) = POCL.compiler_config(dev; name = nothing)
+    # `tt` is only specialized on as a type parameter
+    clfunction(::Type{TT}) where {TT} = POCL.clfunction(kernel, TT; name = nothing)
+    kernel_function(::Type{TT}) where {TT} =
+        KernelAbstractions.KernelInterface.kernel_function(POCL.POCLBackend(), kernel, TT)
+    function launch_tuple(kernel, A, info)
+        event = POCL.launch_tuple(kernel, (A,), info; global_size = (4,), local_size = (4,))
+        wait(event)
+        POCL.cl.clReleaseEvent(event)
+    end
+    @kernel function ka_kernel!(A, n)
+        I = @index(Global, Linear)
+        x = A[I]
+        for _ in 1:n
+            x = sin(x)
+        end
+        @inbounds A[I] = x
+    end
+    launch(kernel, A, n) = kernel(A, n; ndrange = length(A))
+    end
+    # with `args...`, Julia 1.11 and older box the result
+    allocated(f::F, args::Vararg{Any, N}) where {F, N} = (f(args...); @allocated f(args...))
+
+    # these allocate on Julia 1.10, and those that look up compiled code in GPUCompiler on 1.11
+    skip = VERSION < v"1.11"
+    skip_lookup = VERSION < v"1.12"
+    tt = Tuple{typeof(POCL.clconvert(zeros(Int, 4)))}
+    @test allocated(mod.compiler_config, POCL.device()) == 0 skip = skip
+    @test allocated(mod.clfunction, tt) == 0 skip = skip_lookup
+    @test allocated(mod.kernel_function, tt) == 0 skip = skip_lookup
+
+    # neither does waiting for the kernel, also when it takes long enough for the wait to
+    # involve a completion callback (except on 1.11, and to shield it from task cancellation
+    # on 1.14+)
+    skip_wait = VERSION < v"1.12" || isdefined(Base, :CANCEL_TOKEN)
+    A = zeros(Int, 4)
+    info = POCL.exception_info()
+    GC.@preserve info begin
+        info_ptr = Base.unsafe_convert(Ptr{POCL.ExceptionInfo_st}, info)
+        kernel = mod.clfunction(tt)
+        mod.launch_tuple(kernel, A, info_ptr)
+        @test all(==(1), A)
+        @test allocated(mod.launch_tuple, kernel, A, info_ptr) == 0 skip = skip_wait
+    end
+    B = ones(4)
+    kernel = mod.ka_kernel!(CPU(), 4)
+    for n in (1, 100_000)
+        mod.launch(kernel, B, n)
+        @test allocated(mod.launch, kernel, B, n) == 0 skip = skip_wait
+    end
+end
+
 @testset "POCL exceptions" begin
     @kernel function exception_kernel!(A)
         I = @index(Global, Linear)