Implement KernelInterface, and support KernelAbstractions 0.10

.buildkite/pipeline.yml

@@ -24,6 +24,16 @@ steps:
         agents:
           queue: "oneapi"
         commands: |
+          # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet, and
+          #      neither is an AcceleratedKernels that supports them. Julia 1.10 ignores
+          #      [sources], and Julia 1.11 doesn't pick them up for the test environment,
+          #      so check them out and develop them together, as the registered
+          #      AcceleratedKernels excludes KernelAbstractions 0.10.
+          if [[ "{{matrix.julia}}" == "1.10" || "{{matrix.julia}}" == "1.11" ]]; then
+            git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka
+            git clone --depth 1 https://github.com/JuliaGPU/AcceleratedKernels.jl ak
+            julia --project -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in ("ka", "ka/lib/KernelInterface", "ak")])'
+          fi
           julia --project=deps deps/build_ci.jl
         if: |
           build.message !~ /\[skip [^\]]*(tests|julia)/ &&

.github/workflows/docs.yml

@@ -29,7 +29,9 @@ jobs:
           ssh-strict: false
       - uses: julia-actions/setup-julia@latest
         with:
-          version: 'lts'
+          # XXX: needs Julia 1.11+ to pick up the unregistered KernelAbstractions 0.10 and
+          #      KernelInterface 0.4 from [sources]
+          version: '1.12'
       - uses: julia-actions/cache@v3
       - uses: julia-actions/julia-buildpkg@latest
       - run: julia --project=docs/ docs/make.jl

.gitlab-ci.yml

@@ -90,6 +90,18 @@ stages: [build, test, coverage]
   script:
     - !reference [.aurora-env, script]
     - julia --color=yes --project=deps deps/build_local.jl
+    # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet, nor is an
+    #      AcceleratedKernels that supports them, and Julia 1.10 ignores [sources]: check
+    #      them out and develop them (the checkouts travel to the test job as artifacts,
+    #      like test/Manifest.toml)
+    - |
+      if [ "$JULIA_VERSION" = "1.10" ]; then
+        rm -rf ka ak
+        git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka
+        git clone --depth 1 https://github.com/JuliaGPU/AcceleratedKernels.jl ak
+        julia --color=yes --project=. -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in ("ka", "ka/lib/KernelInterface", "ak")])'
+        julia --color=yes --project=test -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "ka", "ka/lib/KernelInterface", "ak")])'
+      fi
     # Instantiate (and thereby precompile) both environments here so the 1 h batch job
     # spends its walltime on tests, not on Pkg. Manifests are gitignored, so the test
     # env must be resolved here with oneAPI dev'ed at the checkout (path "..") — a plain
@@ -113,6 +125,8 @@ stages: [build, test, coverage]
       - LocalPreferences.toml
       - test/LocalPreferences.toml
       - test/Manifest.toml
+      - ka
+      - ak
     expire_in: 1 week
 
 .test:lts:

Project.toml

@@ -13,6 +13,7 @@ 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"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
@@ -32,6 +33,14 @@ oneAPI_Level_Zero_Headers_jll = "f4bc562b-d309-54f8-9efb-476e56f0410d"
 oneAPI_Level_Zero_Loader_jll = "13eca655-d68d-5b81-8367-6d99d727ab01"
 oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36"
 
+[sources]
+AcceleratedKernels = {url = "https://github.com/JuliaGPU/AcceleratedKernels.jl", rev = "main"}
+KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"}
+KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"}
+
+[extensions]
+KernelAbstractionsExt = "KernelAbstractions"
+
 [compat]
 AbstractFFTs = "1.5.0"
 AcceleratedKernels = "0.3.1, 0.4"
@@ -41,12 +50,13 @@ ExprTools = "0.1"
 GPUArrays = "11.5.14"
 GPUCompiler = "2.9"
 GPUToolbox = "3.1"
-KernelAbstractions = "0.9.39"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "6, 7, 8, 9"
 NEO_jll = "=26.18.38308"
 PrecompileTools = "1"
 Preferences = "1"
-SPIRVIntrinsics = "1"
+SPIRVIntrinsics = "1.1.3"
 SPIRV_LLVM_Backend_jll = "23"
 SPIRV_LLVM_Translator_jll = "23"
 SPIRV_Tools_jll = "2025.4.0"

ext/KernelAbstractionsExt.jl

@@ -0,0 +1,32 @@
+module KernelAbstractionsExt
+
+using oneAPI
+using oneAPI: @device_override, method_table
+
+import KernelAbstractions as KA
+
+import StaticArrays
+
+import Adapt
+
+Adapt.adapt_storage(::KA.CPU, a::oneArray) = convert(Array, a)
+
+# sparse arrays (oneMKL is only available on Linux)
+@static if Sys.islinux()
+    import SparseArrays
+    Adapt.adapt_storage(::KA.CPU, a::oneAPI.oneMKL.oneAbstractSparseMatrix) = SparseArrays.SparseMatrixCSC(a)
+end
+
+
+## scratch memory
+
+@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims}
+    StaticArrays.MArray{Tuple{Dims...}, T}(undef)
+end
+
+
+## other
+
+Adapt.adapt_storage(to::KA.ConstAdaptor, a::oneDeviceArray) = Base.Experimental.Const(a)
+
+end

lib/level-zero/event.jl

@@ -36,8 +36,8 @@ mutable struct ZeEvent
     handle::ze_event_handle_t
     pool::ZeEventPool
 
-    function ZeEvent(pool, index::Integer)
-        desc_ref = Ref(ze_event_desc_t(; index=index-1))
+    function ZeEvent(pool, index::Integer; signal=0, wait=0)
+        desc_ref = Ref(ze_event_desc_t(; index=index-1, signal, wait))
         handle_ref = Ref{ze_event_handle_t}()
         zeEventCreate(pool, desc_ref, handle_ref)
         obj = new(handle_ref[], pool)

lib/level-zero/module.jl

@@ -85,13 +85,17 @@ mutable struct ZeKernel
     # Read on every launch by the scratch hedge, so it must not cost an API call.
     spill::Int
 
+    # cached maxGroupSize, seeded by `properties`; -1 while unqueried, and 0 without the
+    # MAX_GROUP_SIZE extension. Read on every KernelInterface launch.
+    max_group_size::Int
+
     function ZeKernel(mod, name)
         GC.@preserve name begin
             desc_ref = Ref(ze_kernel_desc_t(; pKernelName=pointer(name)))
             handle_ref = Ref{ze_kernel_handle_t}()
             zeKernelCreate(mod, desc_ref, handle_ref)
         end
-        obj = new(mod, handle_ref[], ReentrantLock(), -1)
+        obj = new(mod, handle_ref[], ReentrantLock(), -1, -1)
 
         finalizer(obj) do obj
             zeKernelDestroy(obj)
@@ -268,6 +272,8 @@ function properties(kernel::ZeKernel)
 
     props = props_ref[]
     kernel.spill = Int(props.spillMemSize)
+    kernel.max_group_size = max_group_size_props_ref === nothing ? 0 :
+        Int(max_group_size_props_ref[].maxGroupSize)
     return (
         numKernelArgs=Int(props.numKernelArgs),
         requiredGroupSize=ZeDim3(props.requiredGroupSizeX,
@@ -295,6 +301,13 @@ function spill_mem_size(kernel::ZeKernel)
     return s >= 0 ? s : Int(properties(kernel).spillMemSize)
 end
 
+# Cached access to a kernel's maxGroupSize, `missing` without the MAX_GROUP_SIZE extension.
+function max_group_size(kernel::ZeKernel)
+    s = kernel.max_group_size
+    s < 0 && (s = coalesce(properties(kernel).maxGroupSize, 0))
+    return s > 0 ? s : missing
+end
+
 
 ## execution
 

lib/level-zero/oneL0.jl

@@ -126,6 +126,7 @@ include("cmdqueue.jl")
 include("cmdlist.jl")
 include("fence.jl")
 include("event.jl")
+include("synchronization.jl")
 include("barrier.jl")
 include("module.jl")
 include("memory.jl")

lib/level-zero/synchronization.jl

@@ -0,0 +1,204 @@
+# cooperative synchronization
+#
+# `zeCommandListHostSynchronize`, `zeCommandQueueSynchronize` and `zeEventHostSynchronize`
+# block the calling thread until the work has completed, so no other task can run on it in the meantime. As CUDA.jl
+# does, first busy-wait on a non-blocking query, which keeps the latency of short
+# operations low, and then block in the driver on a separate thread, while the calling
+# task waits for that thread without blocking the scheduler.
+
+export nonblocking_synchronize
+
+const SyncObject = Union{ZeImmediateCommandList, ZeCommandQueue, ZeEvent}
+
+# with a zero timeout, a synchronization is a query
+function check_done(res::ze_result_t)
+    if res == RESULT_NOT_READY
+        return false
+    elseif res == RESULT_SUCCESS
+        return true
+    else
+        throw_api_error(res)
+    end
+end
+Base.isdone(list::ZeImmediateCommandList) =
+    check_done(unchecked_zeCommandListHostSynchronize(list, 0))
+Base.isdone(queue::ZeCommandQueue) =
+    check_done(unchecked_zeCommandQueueSynchronize(queue, 0))
+
+# the blocking synchronization, marked GC-safe so that it doesn't keep the GC from running
+gcsafe_synchronize(list::ZeImmediateCommandList) =
+    @gcsafe_ccall libze_loader.zeCommandListHostSynchronize(
+        list::ze_command_list_handle_t, typemax(UInt64)::UInt64)::ze_result_t
+gcsafe_synchronize(queue::ZeCommandQueue) =
+    @gcsafe_ccall libze_loader.zeCommandQueueSynchronize(
+        queue::ze_command_queue_handle_t, typemax(UInt64)::UInt64)::ze_result_t
+gcsafe_synchronize(event::ZeEvent) =
+    @gcsafe_ccall libze_loader.zeEventHostSynchronize(
+        event::ze_event_handle_t, typemax(UInt64)::UInt64)::ze_result_t
+
+# once the work has completed, synchronizing a list or queue doesn't block anymore; a
+# signaled event needs nothing more
+finish_synchronization(obj::Union{ZeImmediateCommandList, ZeCommandQueue}) = synchronize(obj)
+finish_synchronization(::ZeEvent) = nothing
+
+
+## bidirectional channel
+
+# custom, unbuffered channel that supports returning a value to the sender
+# without the need for a second channel
+struct BidirectionalChannel{I,O} <: AbstractChannel{I}
+    cond_take::Threads.Condition                 # waiting for data to become available
+    cond_put::Threads.Condition                  # waiting for a writeable slot
+    cond_ret::Threads.Condition                  # waiting for a data to be returned
+
+    function BidirectionalChannel{I,O}() where {I,O}
+        lock = ReentrantLock()
+        cond_put = Threads.Condition(lock)
+        cond_take = Threads.Condition(lock)
+        cond_ret = Threads.Condition(lock)
+        return new(cond_take, cond_put, cond_ret)
+    end
+end
+
+Base.put!(c::BidirectionalChannel{I}, v) where {I} = put!(c, convert(I, v))
+function Base.put!(c::BidirectionalChannel{I,O}, v::I) where {I,O}
+    lock(c)
+    try
+        # wait for a slot to be available
+        while isempty(c.cond_take)
+            Base.wait(c.cond_put)
+        end
+
+        # pass a value to the consumer
+        notify(c.cond_take, v, false, false)
+
+        # wait for a return value to be produced
+        Base.wait(c.cond_ret)::O
+    finally
+        unlock(c)
+    end
+end
+
+function Base.take!(f::Base.Callable, c::BidirectionalChannel{I,O}) where {I,O}
+    lock(c)
+    try
+        # notify the producer that we're ready to accept a value
+        notify(c.cond_put, nothing, false, false)
+
+        # receive a value from the producer
+        v = Base.wait(c.cond_take)::I
+
+        # return a value to the producer
+        ret = f(v)::O
+        notify(c.cond_ret, ret, false, false)
+    finally
+        unlock(c)
+    end
+end
+
+Base.lock(c::BidirectionalChannel) = lock(c.cond_take)
+Base.unlock(c::BidirectionalChannel) = unlock(c.cond_take)
+
+
+## fast path
+
+# before blocking on a separate thread, which has some overhead, busy-wait on a query of
+# the object to synchronize. when this returns true, the object still has to be
+# synchronized, but that won't block anymore.
+function spinning_synchronization(f, obj)
+    # fast path
+    f(obj) && return true
+
+    # minimize latency of short operations by busy-waiting,
+    # initially without even yielding to other tasks
+    spins = 0
+    while spins < 256
+        if spins < 32
+            ccall(:jl_cpu_pause, Cvoid, ())
+            # temporary solution before we have gc transition support in codegen.
+            ccall(:jl_gc_safepoint, Cvoid, ())
+        else
+            yield()
+        end
+        f(obj) && return true
+        spins += 1
+    end
+
+    return false
+end
+
+
+## slow path: synchronize on a separate thread
+
+const MAX_SYNC_THREADS = 4
+const sync_channels = Array{BidirectionalChannel{SyncObject,ze_result_t}}(undef, MAX_SYNC_THREADS)
+const sync_channel_cursor = Threads.Atomic{UInt32}(1)
+const sync_channel_lock = Base.ReentrantLock()
+
+function synchronization_worker(data)
+    i = Int(data)
+    chan = sync_channels[i]
+
+    while true
+        # wait for work
+        take!(gcsafe_synchronize, chan)
+    end
+end
+
+@noinline function create_synchronization_worker(i)
+    lock(sync_channel_lock) do
+        # test and test-and-set
+        if isassigned(sync_channels, i)
+            return
+        end
+
+        # should be safe to assign before threads are running;
+        # any user will just submit work that makes it block
+        sync_channels[i] = BidirectionalChannel{SyncObject,ze_result_t}()
+
+        # we don't know what the size of uv_thread_t is, so reserve enough space
+        tid = Ref{NTuple{32, UInt8}}(ntuple(i -> 0, 32))
+
+        cb = @cfunction(synchronization_worker, Cvoid, (Ptr{Cvoid},))
+        err = @ccall uv_thread_create(tid::Ptr{Cvoid}, cb::Ptr{Cvoid}, Ptr{Cvoid}(i)::Ptr{Cvoid})::Cint
+        err == 0 || Base.uv_error("uv_thread_create", err)
+        err = @ccall uv_thread_detach(tid::Ptr{Cvoid})::Cint
+        err == 0 || Base.uv_error("uv_thread_detach", err)
+    end
+
+    return
+end
+
+"""
+    nonblocking_synchronize(list_or_queue)
+    nonblocking_synchronize(event)
+
+Wait for the work on an immediate command list or command queue to complete, or for an
+event to be signaled, like [`synchronize`](@ref) or `wait`, but without blocking the Julia
+scheduler: other tasks keep running while this one waits.
+"""
+function nonblocking_synchronize(obj::SyncObject)
+    if spinning_synchronization(Base.isdone, obj)
+        # done, so this doesn't block
+        finish_synchronization(obj)
+        return
+    end
+
+    # pick a worker channel: sticky per task, so repeated synchronizations from the
+    # same task always hit the same, already running worker thread.
+    tls = task_local_storage()
+    i = get!(tls, :ZeSyncChannel) do
+        mod1(Threads.atomic_add!(sync_channel_cursor, UInt32(1)), MAX_SYNC_THREADS)
+    end::Int
+    if !isassigned(sync_channels, i)
+        create_synchronization_worker(i)
+    end
+    chan = @inbounds sync_channels[i]
+
+    # submit the object to synchronize; unlike with regular channels, this `put!` blocks
+    # until the worker has synchronized it and returned the result
+    res = put!(chan, obj)
+    res == RESULT_SUCCESS || throw_api_error(res)
+
+    return
+end