Implement KernelInterface, and support KernelAbstractions 0.10

Project.toml

@@ -11,7 +11,7 @@ ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 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"
@@ -26,11 +26,16 @@ SPIRV_LLVM_Translator_jll = "4a5d46fc-d8cf-5151-a261-86b458210efb"
 SPIRV_Tools_jll = "6ac6d60f-d740-5983-97d7-a4482c0689f4"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 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"
 
+[weakdeps]
+KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+
+[extensions]
+KernelAbstractionsExt = "KernelAbstractions"
+
 [compat]
 AbstractFFTs = "1.5.0"
 Adapt = "4"
@@ -39,7 +44,8 @@ ExprTools = "0.1"
 GPUArrays = "12.1"
 GPUCompiler = "2.13"
 GPUToolbox = "3.3.1"
-KernelAbstractions = "0.9.39"
+KernelAbstractions = "0.10"
+KernelInterface = "0.4"
 LLVM = "10"
 NEO_jll = "=26.18.38308"
 PrecompileTools = "1"
@@ -49,7 +55,6 @@ SPIRV_LLVM_Backend_jll = "23"
 SPIRV_LLVM_Translator_jll = "23"
 SPIRV_Tools_jll = "2025.4.0"
 SpecialFunctions = "1.3, 2"
-StaticArrays = "1"
 julia = "1.10"
 oneAPI_Level_Zero_Loader_jll = "1.29"
 oneAPI_Support_jll = "0.10.0"

ext/KernelAbstractionsExt.jl

@@ -0,0 +1,22 @@
+module KernelAbstractionsExt
+
+using oneAPI
+
+import KernelAbstractions as KA
+
+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
+
+
+## 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

@@ -124,9 +124,9 @@ end
 include("context.jl")
 include("cmdqueue.jl")
 include("cmdlist.jl")
-include("synchronization.jl")
 include("fence.jl")
 include("event.jl")
+include("synchronization.jl")
 include("barrier.jl")
 include("module.jl")
 include("memory.jl")

lib/level-zero/synchronization.jl

@@ -1,16 +1,16 @@
 # cooperative synchronization
 #
-# `zeCommandListHostSynchronize` and `zeCommandQueueSynchronize` block the calling thread
-# until the work has completed, so no other task can run on it in the meantime. Instead, wait
-# using GPUToolbox's `cooperative_wait`: first poll, which keeps the latency of short
-# operations low, and then block in the driver on a separate thread, while the calling task
-# yields.
+# `zeCommandListHostSynchronize`, `zeCommandQueueSynchronize` and `zeEventHostSynchronize`
+# block the calling thread until the work has completed, so no other task can run on it in
+# the meantime. Instead, wait using GPUToolbox's `cooperative_wait`: first poll, which keeps
+# the latency of short operations low, and then block in the driver on a separate thread,
+# while the calling task yields.
 
 using GPUToolbox: cooperative_wait
 
 export nonblocking_synchronize
 
-const SyncObject = Union{ZeImmediateCommandList, ZeCommandQueue}
+const SyncObject = Union{ZeImmediateCommandList, ZeCommandQueue, ZeEvent}
 
 # with a zero timeout, a synchronization is a query
 function check_done(res::ze_result_t)
@@ -34,13 +34,17 @@ gcsafe_synchronize(list::ZeImmediateCommandList) =
 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
 
 """
     nonblocking_synchronize(list_or_queue)
+    nonblocking_synchronize(event)
 
-Wait for the work on an immediate command list or command queue to complete, like
-[`synchronize`](@ref), but without blocking the calling thread: other tasks keep running
-while this one waits.
+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 calling
+thread: other tasks keep running while this one waits.
 """
 function nonblocking_synchronize(obj::SyncObject)
     # when polling found the work to be done, synchronize again to check for errors

src/compiler/compilation.jl

@@ -1,6 +1,9 @@
 ## gpucompiler interface implementation
 
-struct oneAPICompilerParams <: AbstractCompilerParams end
+Base.@kwdef struct oneAPICompilerParams <: AbstractCompilerParams
+    sub_group_size::Union{Nothing,Int} = nothing
+end
+
 const oneAPICompilerConfig = CompilerConfig{SPIRVCompilerTarget, oneAPICompilerParams}
 const oneAPICompilerJob = CompilerJob{SPIRVCompilerTarget,oneAPICompilerParams}
 
@@ -56,6 +59,11 @@ function GPUCompiler.finish_module!(job::oneAPICompilerJob, mod::LLVM.Module,
                    Tuple{CompilerJob{SPIRVCompilerTarget}, typeof(mod), typeof(entry)},
                    job, mod, entry)
 
+    # require a sub-group size, instead of letting the compiler choose one
+    if job.config.params.sub_group_size !== nothing
+        entry.metadata["intel_reqd_sub_group_size"] = MDNode([ConstantInt(Int32(job.config.params.sub_group_size))])
+    end
+
     # OpenCL 2.0
     push!(
         get!(mod.metadata, "opencl.ocl.version").operands,
@@ -290,11 +298,18 @@ function _atomics_capabilities(dev, properties; supports_fp16, supports_fp64)
     )
 end
 
-@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false, kwargs...)
+@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false,
+                                    sub_group_size=nothing, kwargs...)
     properties = oneL0.module_properties(dev)
     supports_fp16 = properties.fp16flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP16 == oneL0.ZE_DEVICE_MODULE_FLAG_FP16
     supports_fp64 = properties.fp64flags & oneL0.ZE_DEVICE_MODULE_FLAG_FP64 == oneL0.ZE_DEVICE_MODULE_FLAG_FP64
 
+    if sub_group_size !== nothing
+        sizes = oneL0.compute_properties(dev).subGroupSizes
+        sub_group_size in sizes ||
+            throw(ArgumentError("Sub-group size $sub_group_size is not supported by this device, which supports $(join(sizes, ", "))"))
+    end
+
     # SPIR-V codegen path. The Aurora LTS NEO/IGC runtime only accepts SPIR-V from the
     # Khronos translator; the rolling stack uses the LLVM SPIR-V back-end. GPUCompiler picks
     # the tool from the target's `backend` field and loads the JLL lazily, so both can be
@@ -329,7 +344,7 @@ end
     target = SPIRVCompilerTarget(; backend, extensions = extensions_str, atomics,
                                    supports_fp16, supports_fp64, supports_bfloat16,
                                    driver = :intel, kwargs...)
-    params = oneAPICompilerParams()
+    params = oneAPICompilerParams(; sub_group_size)
     CompilerConfig(target, params; kernel, name, always_inline)
 end
 

src/compiler/execution.jl

@@ -4,7 +4,7 @@ export @oneapi, zefunction, kernel_convert
 ## high-level @oneapi interface
 
 const MACRO_KWARGS = [:launch]
-const COMPILER_KWARGS = [:kernel, :name, :always_inline, :atomics]
+const COMPILER_KWARGS = [:kernel, :name, :always_inline, :atomics, :sub_group_size]
 const LAUNCH_KWARGS = [:groups, :items, :queue]
 
 """
@@ -33,6 +33,9 @@ launches the kernel on the GPU.
   additions on Xe-LP). Disabled floating-point operations are implemented with integer
   compare-and-swap loops instead, and disabling `int64` makes 64-bit atomic operations an
   error.
+- `sub_group_size::Union{Int,Nothing}=nothing`: The sub-group size the kernel has to be
+  compiled for, one of the device's `oneL0.compute_properties(dev).subGroupSizes`. By
+  default, the compiler chooses one.
 
 ## Launch Keywords (runtime)
 - `groups`: Number of workgroups (required). Can be an integer or tuple.
@@ -249,9 +252,9 @@ function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
     # configurations, so roll our own version that behaves like CUDA's
     # occupancy API and assumes the kernel still does bounds checking.
 
-    kernel_props = oneL0.properties(kernel.fun)
-    group_size = if kernel_props.maxGroupSize !== missing
-        kernel_props.maxGroupSize
+    max_group_size = oneL0.max_group_size(kernel.fun)
+    group_size = if max_group_size !== missing
+        max_group_size
     else
         # without the MAX_GROUP_SIZE extension, we need to be conservative
         dev = kernel.fun.mod.device
@@ -269,7 +272,7 @@ function launch_configuration(kernel::HostKernel{F,TT}) where {F,TT}
     # size but does not fold it into `maxGroupSize`, so account for it here. Rounded down to
     # a power of two, both because group sizes want to be anyway and to stay clear of the
     # limit rather than right at it.
-    spill = kernel_props.spillMemSize
+    spill = oneL0.spill_mem_size(kernel.fun)
     if spill > 0 && group_size * spill > MAX_GROUP_SCRATCH
         group_size = max(1, prevpow(2, max(1, MAX_GROUP_SCRATCH ÷ spill)))
     end

src/context.jl

@@ -284,7 +284,7 @@ global_queue(ctx::ZeContext, dev::ZeDevice) = stream_queue(global_stream(ctx, de
 # Register `stream` as a stream targeting (ctx, dev) so `synchronize_all_streams`/
 # `release` can find and drain it before freeing buffers whose in-flight work it may
 # still reference. EVERY stream that becomes a task's active stream must go through here
-# — not just the one `global_stream` creates but also the replacement `KA.priority!`
+# — not just the one `global_stream` creates but also the replacement `KI.priority!`
 # installs — or the unregistered stream's in-flight work can outlive a freed buffer (a
 # use-after-free that faults and bans the context on the LTS NEO stack). Only the LTS
 # stack maintains the registry; on the rolling stack this is a no-op. Returns `stream`.

src/oneAPI.jl

@@ -12,7 +12,7 @@ using SpecialFunctions
 
 import Preferences
 
-import KernelAbstractions: KernelAbstractions
+import KernelInterface
 
 using LLVM, LLVM.IR, LLVM.Build, LLVM.Interop
 using Core: LLVMPtr