Implement KernelInterface, and support KernelAbstractions 0.10

.buildkite/pipeline.yml

@@ -47,7 +47,17 @@ steps:
             # SPIRVIntrinsics first; otherwise the Pkg.add below resolves
             # against the registry, which has no SPIRVIntrinsics 1. Pkg.test
             # then carries the in-tree copy into the test sandbox.
-            Pkg.develop(path="lib/intrinsics")
+            if VERSION < v"1.12"
+              # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 are not registered yet,
+              #      so develop them from their branch, in the same Pkg operation
+              run(`git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka`)
+              Pkg.develop([PackageSpec(; path) for path in ("lib/intrinsics", "ka", "ka/lib/KernelInterface")])
+              # Pkg < 1.12 turns the developed weak dependency into a strong one; restore
+              # the project, keeping the developed package in the manifest
+              run(`git checkout Project.toml`)
+            else
+              Pkg.develop(path="lib/intrinsics")
+            end
             Pkg.add("{{matrix.pocl}}_jll")
             Pkg.add("InteractiveUtils")
 

.github/workflows/Test.yml

@@ -142,7 +142,19 @@ jobs:
             using Pkg
             # Julia 1.10 does not support [sources], so dev the in-tree
             # SPIRVIntrinsics; Pkg.test then carries it into the test sandbox.
-            Pkg.develop(path="lib/intrinsics")'
+            if VERSION < v"1.12"
+              # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 are not registered yet,
+              #      and Julia 1.11 does not pick them up from the [sources] of the test project
+              #      without workspaces, so develop them from their branch, in the same
+              #      Pkg operation
+              run(`git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka`)
+              Pkg.develop([PackageSpec(; path) for path in ("lib/intrinsics", "ka", "ka/lib/KernelInterface")])
+              # Pkg < 1.12 turns the developed weak dependency into a strong one; restore
+              # the project, keeping the developed package in the manifest
+              run(`git checkout Project.toml`)
+            else
+              Pkg.develop(path="lib/intrinsics")
+            end'
 
       - name: Test OpenCL.jl
         uses: julia-actions/julia-runtest@v1

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,18 +23,25 @@ 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"
 spirv2clc_jll = "f0274c0c-8c8a-59f1-85b7-f7d60330c5fb"
 
+[weakdeps]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+
+[extensions]
+KernelAbstractionsExt = "KernelAbstractions"
+
 [sources]
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"}
 SPIRVIntrinsics = {path = "lib/intrinsics"}
 
 [compat]
 Adapt = "4"
 GPUArrays = "11.2.1"
 GPUCompiler = "2.12"
 GPUToolbox = "3.3.2"
-KernelAbstractions = "0.9.38"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "10"
 LinearAlgebra = "1"
 OpenCL_jll = "=2024.10.24"
@@ -47,6 +54,9 @@ Reexport = "1"
 SPIRVIntrinsics = "1.2"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_Tools_jll = "2025.1"
-StaticArrays = "1"
 julia = "1.10"
 spirv2clc_jll = "0.2.0"
+
+# XXX: lets Pkg < 1.12 develop the weak dependency (see CI)
+[extras]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"

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/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

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