Implement KernelInterface, and support KernelAbstractions 0.10

Project.toml

@@ -10,7 +10,7 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 GPUToolbox = "096a3bc2-3ced-46d0-87f4-dd12716f4bfc"
-KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf"
 LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 OpenCL_jll = "6cb37087-e8b6-5417-8430-1f242f1e46e4"
@@ -23,10 +23,15 @@ Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SPIRVIntrinsics = "71d1d633-e7e8-4a92-83a1-de8814b09ba8"
 SPIRV_LLVM_Backend_jll = "4376b9bf-cff8-51b6-bb48-39421dff0d0c"
 SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UnsafeAtomics = "013be700-e6cd-48c3-b4a1-df204f14c38f"
 spirv2clc_jll = "f0274c0c-8c8a-59f1-85b7-f7d60330c5fb"
 
+[weakdeps]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+
+[extensions]
+KernelAbstractionsExt = "KernelAbstractions"
+
 [sources]
 SPIRVIntrinsics = {path = "lib/intrinsics"}
 
@@ -35,7 +40,8 @@ Adapt = "4"
 GPUArrays = "12.1"
 GPUCompiler = "2.13"
 GPUToolbox = "3.3.2"
-KernelAbstractions = "0.9.38"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "10"
 LinearAlgebra = "1"
 OpenCL_jll = "=2024.10.24"
@@ -48,7 +54,6 @@ Reexport = "1"
 SPIRVIntrinsics = "1.4"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_Tools_jll = "2025.1"
-StaticArrays = "1"
 UnsafeAtomics = "0.4"
 julia = "1.10"
-spirv2clc_jll = "0.2.2"
+spirv2clc_jll = "0.2.5"

docs/src/95-reference.md

@@ -13,5 +13,5 @@ Pages = ["95-reference.md"]
 ```
 
 ```@autodocs
-Modules = [OpenCL, OpenCL.cl]
+Modules = [OpenCL, OpenCL.cl, OpenCL.OpenCLKernels]
 ```

ext/KernelAbstractionsExt.jl

@@ -0,0 +1,16 @@
+module KernelAbstractionsExt
+
+using OpenCL
+
+import KernelAbstractions as KA
+
+import Adapt
+
+Adapt.adapt_storage(::KA.CPU, a::CLArray) = convert(Array, a)
+
+# `@Const` applies `constify` inside the kernel, where arguments have already been
+# converted to device arrays, so the rule has to be registered for `CLDeviceArray`
+# rather than for `CLArray`.
+Adapt.adapt_storage(::KA.ConstAdaptor, a::CLDeviceArray) = Base.Experimental.Const(a)
+
+end

lib/cl/device.jl

@@ -369,7 +369,15 @@ function cl_device_type(dtype::Symbol)
     return cl_dtype
 end
 
-sub_groups_supported(d::Device) = "cl_khr_subgroups" in d.extensions || "cl_intel_subgroups" in d.extensions
+# sub-groups are an extension before OpenCL 3.0, and an optional core feature since, which
+# devices without it report with a maximum of 0 sub-groups
+function sub_groups_supported(d::Device)
+    ("cl_khr_subgroups" in d.extensions || "cl_intel_subgroups" in d.extensions) && return true
+    d.opencl_version >= v"3.0" || return false
+    num = Ref{cl_uint}(0)
+    err = unchecked_clGetDeviceInfo(d, CL_DEVICE_MAX_NUM_SUB_GROUPS, sizeof(cl_uint), num, C_NULL)
+    return err == CL_SUCCESS && num[] > 0
+end
 function sub_group_size(d::Device)
     sub_groups_supported(d) || return 0
     if "cl_amd_device_attribute_query" in d.extensions

lib/cl/kernel.jl

@@ -260,14 +260,15 @@ function enqueue_task(k::Kernel; wait_for=nothing)
     return ret_event[]
 end
 
+# the arguments are passed as a tuple, see `set_args!`
 function call(
-        k::Kernel, args...; global_size = (1,), local_size = nothing,
+        k::Kernel, args::Tuple; global_size = (1,), local_size = nothing,
         global_work_offset = nothing, wait_on::Vector{Event} = Event[],
         indirect_memory::Vector{AbstractMemory} = AbstractMemory[],
         rng_state=false,
     )
     return Base.@lock k begin
-    set_args!(k, args...)
+    _set_args!(k, args)
     if !isempty(indirect_memory)
         svm_pointers = CLPtr{Cvoid}[]
         usm_pointers = CLPtr{Cvoid}[]
@@ -360,7 +361,7 @@ clcall(f::F, types::Tuple, args::Vararg{Any,N}; kwargs...) where {N,F} =
 
 function clcall(k::Kernel, types::Type{T}, args::Vararg{Any,N}; kwargs...) where {T,N}
     call_closure = function (converted_args::Vararg{Any,N})
-        call(k, converted_args...; kwargs...)
+        call(k, converted_args; kwargs...)
     end
     convert_arguments(call_closure, types, args...)
 end
@@ -415,10 +416,12 @@ unlock_arguments(locks) = foreach(unlock, Iterators.reverse(locks))
     return ex
 end
 
-function set_args!(k::Kernel, args...)
-    for (i, a) in enumerate(args)
-        set_arg!(k, i, a)
-    end
+set_args!(k::Kernel, args::Vararg{Any,N}) where {N} = _set_args!(k, args)
+
+# one call per argument: a splat of more than 32 arguments isn't a direct call
+@inline @generated function _set_args!(k::Kernel, args::Tuple)
+    calls = (:(set_arg!(k, $i, args[$i])) for i in 1:fieldcount(args))
+    return :($(calls...); nothing)
 end
 
 function set_arg!(k::Kernel, idx::Integer, arg::T) where {T}

src/OpenCL.jl

@@ -11,7 +11,7 @@ using GPUArrays
 using Random
 using Preferences
 
-import KernelAbstractions: KernelAbstractions
+import KernelInterface
 
 using Core: LLVMPtr