From 45eff3fe624b94199a57bd0908b13e92be4d384f Mon Sep 17 00:00:00 2001 From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> Date: Sat, 22 Aug 2026 15:32:34 -0300 Subject: [PATCH 01/12] Only list KernelAbstractions in versioninfo when it is loaded In preparation for KernelAbstractions becoming a weak dependency. --- src/utils.jl | 14 +++++++++++--- 1 file changed, 11 insertions(+), 3 deletions(-) diff --git a/src/utils.jl b/src/utils.jl index ddf4b1f91..03d4de12c 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -27,11 +27,19 @@ function versioninfo(io::IO=stdout) println(io, "- LLVM: $(LLVM.version())") println(io) + + get_module(name::Symbol) = (name, getfield(oneAPI, name)) + function get_module(pkg::Tuple{String, String}) + id = Base.PkgId(Base.UUID(pkg[1]), pkg[2]) + (pkg[2], get(Base.loaded_modules, id, nothing)) + end + println(io, "Julia packages:") println(io, "- oneAPI.jl: $(Base.pkgversion(oneAPI))") - for name in [:GPUArrays, :GPUCompiler, :KernelAbstractions, :LLVM, :SPIRVIntrinsics] - mod = getfield(oneAPI, name) - println(io, "- $(name): $(Base.pkgversion(mod))") + for pkg in [:GPUArrays, :GPUCompiler, ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"), + :LLVM, :SPIRVIntrinsics] + name, mod = get_module(pkg) + isnothing(mod) || println(io, "- $(name): $(Base.pkgversion(mod))") end println(io) From a1f8d7bf54122509ea76aac7a049c71698197ff6 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Mon, 28 Sep 2026 10:12:28 +0200 Subject: [PATCH 02/12] Cache a kernel's maximum group size Like the spill size, cache the kernel's `maxGroupSize` in the `ZeKernel`, so that `launch_configuration` (and the KernelInterface launch validation, which checks every explicit work-group size against it) don't query the kernel properties, which allocates, on every launch. --- lib/level-zero/module.jl | 15 ++++++++++++++- src/compiler/execution.jl | 8 ++++---- 2 files changed, 18 insertions(+), 5 deletions(-) diff --git a/lib/level-zero/module.jl b/lib/level-zero/module.jl index 91cc0dfcf..6e0635df6 100644 --- a/lib/level-zero/module.jl +++ b/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 diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index c5fb4b927..04c2501ff 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -249,9 +249,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 +269,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 From 906f0dbaf15e9a37a3d1f1193ac4e1f28add54d0 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:03:56 +0200 Subject: [PATCH 03/12] Implement KernelInterface, and support KernelAbstractions 0.10 KernelAbstractions 0.10 builds on KernelInterface, which defines what a back end provides: memory and device management, compiling and launching kernels, and the device-side intrinsics. KernelAbstractions then launches `@kernel` kernels itself on any KernelInterface back end, which replaces oneAPI's copy of that launch path: partitioning the ndrange, building the kernel's context, and tuning the work-group size. `oneAPIBackend` now implements KernelInterface, and oneAPI depends on it instead of on KernelAbstractions. What KernelAbstractions still needs from a back end moves to an extension: the `MArray` behind `@private`, the Adapt rule for `@Const`, and adapting to KernelAbstractions' CPU back end. `prefer_blocks` now applies in `KI.launch_configuration`, which receives the number of work-items. The index queries are computed in the requested type with `% T`, which needs SPIRVIntrinsics 1.1.3 to truncate the 3-D built-ins without producing illegal vector types. `KI.launch` rejects `items` and `groups`, which would override the launch geometry that KernelInterface validated, and refuses to launch a kernel on another context or device than the one it was compiled for. `KI.kernel_function` converts the callable to compile it, and the kernel keeps the original alive, as the converted form of a closure only holds pointers to the arrays it captures. Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> --- Project.toml | 11 +- ext/KernelAbstractionsExt.jl | 32 +++++ src/context.jl | 2 +- src/oneAPI.jl | 2 +- src/oneAPIKernels.jl | 266 ++++++++++++++++------------------- src/utils.jl | 2 +- test/Project.toml | 1 + test/kernelabstractions.jl | 76 +++++++++- test/kernelinterface.jl | 27 ++++ 9 files changed, 271 insertions(+), 148 deletions(-) create mode 100644 ext/KernelAbstractionsExt.jl create mode 100644 test/kernelinterface.jl diff --git a/Project.toml b/Project.toml index 569f884ff..24b549704 100644 --- a/Project.toml +++ b/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" @@ -31,6 +31,12 @@ 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 +45,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" diff --git a/ext/KernelAbstractionsExt.jl b/ext/KernelAbstractionsExt.jl new file mode 100644 index 000000000..73f4a14da --- /dev/null +++ b/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 diff --git a/src/context.jl b/src/context.jl index 6c9e188fa..521adb3ec 100644 --- a/src/context.jl +++ b/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`. diff --git a/src/oneAPI.jl b/src/oneAPI.jl index 1f098f680..f6489f6c2 100644 --- a/src/oneAPI.jl +++ b/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 diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index ae693caac..a4089e6d9 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -1,11 +1,9 @@ module oneAPIKernels using ..oneAPI -using ..oneAPI: @device_override, SPIRVIntrinsics, method_table +using ..oneAPI: @device_override, SPIRVIntrinsics, method_table, kernel_convert, zefunction -import KernelAbstractions as KA - -import StaticArrays +import KernelInterface as KI import Adapt @@ -14,215 +12,201 @@ import Adapt export oneAPIBackend -struct oneAPIBackend <: KA.GPU +""" + oneAPIBackend(; prefer_blocks=false, always_inline=false) + +KernelInterface back end for oneAPI, which KernelAbstractions uses to launch kernels on +Intel GPUs. `prefer_blocks=true` spreads small launches over more, smaller work-groups. +""" +struct oneAPIBackend <: KI.Backend prefer_blocks::Bool always_inline::Bool end oneAPIBackend(; prefer_blocks = false, always_inline = false) = oneAPIBackend(prefer_blocks, always_inline) -@inline KA.allocate(::oneAPIBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where {T} = oneArray{T, length(dims), unified ? oneAPI.oneL0.SharedBuffer : oneAPI.oneL0.DeviceBuffer}(undef, dims) -@inline KA.zeros(::oneAPIBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where {T} = fill!(oneArray{T, length(dims), unified ? oneAPI.oneL0.SharedBuffer : oneAPI.oneL0.DeviceBuffer}(undef, dims), zero(T)) -@inline KA.ones(::oneAPIBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where {T} = fill!(oneArray{T, length(dims), unified ? oneAPI.oneL0.SharedBuffer : oneAPI.oneL0.DeviceBuffer}(undef, dims), one(T)) +@inline KI.allocate(::oneAPIBackend, ::Type{T}, dims::Tuple; unified::Bool = false) where {T} = oneArray{T, length(dims), unified ? oneAPI.oneL0.SharedBuffer : oneAPI.oneL0.DeviceBuffer}(undef, dims) -KA.get_backend(::oneArray) = oneAPIBackend() -KA.synchronize(::oneAPIBackend) = oneAPI.oneL0.synchronize() -KA.supports_float64(::oneAPIBackend) = false # TODO: Check if this is device dependent -KA.supports_unified(::oneAPIBackend) = true +KI.get_backend(::oneArray) = oneAPIBackend() +KI.synchronize(::oneAPIBackend) = oneAPI.oneL0.synchronize() +KI.supports_float64(::oneAPIBackend) = device_limits().supports_float64 +KI.supports_unified(::oneAPIBackend) = true +KI.supports_atomics(::oneAPIBackend) = true -KA.functional(::oneAPIBackend) = oneAPI.functional() +KI.functional(::oneAPIBackend) = oneAPI.functional() Adapt.adapt_storage(::oneAPIBackend, a::AbstractArray) = Adapt.adapt(oneArray, a) Adapt.adapt_storage(::oneAPIBackend, a::oneArray) = a -Adapt.adapt_storage(::KA.CPU, a::oneArray) = convert(Array, a) # sparse arrays (oneMKL is only available on Linux) @static if Sys.islinux() - import GPUArrays, SparseArrays - KA.get_backend(::oneAPI.oneMKL.oneAbstractSparseMatrix) = oneAPIBackend() + import GPUArrays + KI.get_backend(::oneAPI.oneMKL.oneAbstractSparseMatrix) = oneAPIBackend() # without this, `adapt_storage(::oneAPIBackend, ::AbstractArray)` would densify sparse arrays Adapt.adapt_storage(::oneAPIBackend, a::GPUArrays.AbstractGPUSparseArray) = a - Adapt.adapt_storage(::KA.CPU, a::oneAPI.oneMKL.oneAbstractSparseMatrix) = SparseArrays.SparseMatrixCSC(a) end ## Memory Operations -function KA.copyto!(::oneAPIBackend, A, B) +function KI.copyto!(::oneAPIBackend, A, B) + length(A) == length(B) || + throw(ArgumentError("Arrays must have the same length, got $(length(A)) and $(length(B))")) copyto!(A, B) # TODO: Address device to host copies in jl being synchronizing + return A end +KI.unsafe_free!(A::oneArray) = oneAPI.unsafe_free!(A) + ## Device Operations -function KA.ndevices(::oneAPIBackend) +function KI.ndevices(::oneAPIBackend) return length(oneAPI.devices()) end -function KA.device(::oneAPIBackend)::Int - dev = oneAPI.device() +function device_index(dev)::Int devs = oneAPI.devices() idx = findfirst(==(dev), devs) return idx === nothing ? 1 : idx end +KI.device(::oneAPIBackend)::Int = device_index(oneAPI.device()) +KI.device(::oneAPIBackend, A::oneArray)::Int = device_index(oneAPI.device(A)) -function KA.device!(backend::oneAPIBackend, id::Int) - return oneAPI.device!(id) +function KI.device!(backend::oneAPIBackend, id::Int) + oneAPI.device!(id) + return end ## Kernel Launch -function KA.mkcontext(kernel::KA.Kernel{oneAPIBackend}, _ndrange, iterspace) - KA.CompilerMetadata{KA.ndrange(kernel), KA.DynamicCheck}(_ndrange, iterspace) -end -function KA.mkcontext(kernel::KA.Kernel{oneAPIBackend}, I, _ndrange, iterspace, - ::Dynamic) where Dynamic - KA.CompilerMetadata{KA.ndrange(kernel), Dynamic}(I, _ndrange, iterspace) -end - -function KA.launch_config(kernel::KA.Kernel{oneAPIBackend}, ndrange, workgroupsize) - if ndrange isa Integer - ndrange = (ndrange,) - end - if workgroupsize isa Integer - workgroupsize = (workgroupsize, ) - end +KI.argconvert(::oneAPIBackend, arg) = kernel_convert(arg) - # partition checked that the ndrange's agreed - if KA.ndrange(kernel) <: KA.StaticSize - ndrange = nothing - end - - iterspace, dynamic = if KA.workgroupsize(kernel) <: KA.DynamicSize && - workgroupsize === nothing - # use ndrange as preliminary workgroupsize for autotuning - # (clamped to 1, since an empty ndrange cannot serve as a workgroup size) - KA.partition(kernel, ndrange, max.(ndrange, 1)) - else - KA.partition(kernel, ndrange, workgroupsize) - end +# a compiled kernel, and the callable it was compiled from: the converted callable only holds +# pointers to the arrays a closure captures, so the kernel has to keep the original alive +struct oneAPIKernel{F, H <: oneAPI.HostKernel} + f::F + host::H +end - return ndrange, workgroupsize, iterspace, dynamic +function KI.kernel_function(backend::oneAPIBackend, f::F, tt::TT=Tuple{}; name = nothing, kwargs...) where {F,TT} + host = zefunction(kernel_convert(f), tt; name, backend.always_inline, kwargs...) + kern = oneAPIKernel(f, host) + KI.Kernel{oneAPIBackend, typeof(kern)}(backend, kern) end -function threads_to_workgroupsize(threads, ndrange) - total = 1 - return map(ndrange) do n - x = max(1, min(div(threads, total), n)) - total *= x - return x +# XXX: calling a `HostKernel` takes the arguments as varargs, so this splats them +function KI.launch(obj::KI.Kernel{oneAPIBackend}, groups::Dims{3}, items::Dims{3}, + args::Tuple; kwargs...) + # KernelInterface has validated the launch geometry + if haskey(kwargs, :items) || haskey(kwargs, :groups) + throw(ArgumentError("KernelInterface kernels take `numgroups`, `workgroupsize` or `ndrange`, not `items` or `groups`")) end + # kernels are compiled for a context and device, and by default launched on the task's + # stream for the active ones + (; f, host) = obj.kern + check_target(host.fun.mod, get(kwargs, :queue, nothing)) + GC.@preserve f host(args...; items, groups, kwargs...) + return end -function (obj::KA.Kernel{oneAPIBackend})(args...; ndrange=nothing, workgroupsize=nothing) - backend = KA.backend(obj) - - ndrange, workgroupsize, iterspace, dynamic = KA.launch_config(obj, ndrange, workgroupsize) - # this might not be the final context, since we may tune the workgroupsize - ctx = KA.mkcontext(obj, ndrange, iterspace) - - # If the kernel is statically sized we can tell the compiler about that - if KA.workgroupsize(obj) <: KA.StaticSize - # TODO: maxthreads - # maxthreads = prod(KA.get(KA.workgroupsize(obj))) +function check_target(mod::oneAPI.oneL0.ZeModule, queue) + ctx, dev = if queue === nothing + context(), device() + elseif queue isa oneAPI.oneStream + queue.ctx, queue.dev else - # maxthreads = nothing + queue.context, queue.device end + mod.context == ctx && mod.device == dev || + throw(ArgumentError("Cannot launch a kernel on another context or device than the one it was compiled for")) + return +end - kernel = @oneapi launch = false always_inline = backend.always_inline obj.f(ctx, args...) - - # figure out the optimal workgroupsize automatically - if KA.workgroupsize(obj) <: KA.DynamicSize && workgroupsize === nothing - items = oneAPI.launch_configuration(kernel) - - if backend.prefer_blocks - # Prefer blocks over threads: - # Reducing the workgroup size (items) increases the number of workgroups (blocks). - # We use a simple heuristic here since we lack full occupancy info (max_blocks) from launch_configuration. - - # If the total range is large enough, full workgroups are fine. - # If the range is small, we might want to reduce 'items' to create more blocks to fill the GPU. - # (Simplified logic compared to CUDA.jl which uses explicit occupancy calculators) - total_items = prod(ndrange) - if total_items < items * 16 # Heuristic factor - # Force at least a few blocks if possible by reducing items per block - target_blocks = 16 # Target at least 16 blocks - items = max(1, min(items, cld(total_items, target_blocks))) - end +function KI.max_work_group_size(kernel::KI.Kernel{oneAPIBackend})::Int + fun = kernel.kern.host.fun + max_group_size = oneAPI.oneL0.max_group_size(fun) + # without the MAX_GROUP_SIZE extension, the device limit is all we know + return coalesce(max_group_size, device_limits(fun.mod.device).max_work_group_size) +end +function KI.launch_configuration( + kernel::KI.Kernel{oneAPIBackend}; nitems::Union{Integer, Nothing} = nothing, + max_work_group_size::Integer = typemax(Int) + ) + items = min(oneAPI.launch_configuration(kernel.kern.host), max_work_group_size) + if kernel.backend.prefer_blocks && nitems !== nothing + # prefer more, smaller work-groups: without an occupancy API that tells how many + # work-groups fill the device, make sure that a small launch uses at least a few + target_groups = 16 + if nitems < items * target_groups + items = max(1, min(items, cld(nitems, target_groups))) end - - workgroupsize = threads_to_workgroupsize(items, ndrange) - iterspace, dynamic = KA.partition(obj, ndrange, workgroupsize) - ctx = KA.mkcontext(obj, ndrange, iterspace) end - - groups = length(KA.blocks(iterspace)) - items = length(KA.workitems(iterspace)) - - if groups == 0 - return nothing + return (; workgroupsize = Int(items)) +end +# querying the device allocates, so cache what every launch needs +const DeviceLimits = @NamedTuple{ + max_work_group_size::Int, max_work_group_dims::NTuple{3, Int}, max_num_groups::NTuple{3, Int}, + supports_float64::Bool, +} +function device_limits(dev::oneAPI.oneL0.ZeDevice = device()) + limits = get!(task_local_storage(), :oneAPIDeviceLimits) do + Dict{oneAPI.oneL0.ZeDevice, DeviceLimits}() + end::Dict{oneAPI.oneL0.ZeDevice, DeviceLimits} + get!(limits, dev) do + props = oneAPI.oneL0.compute_properties(dev) + module_props = oneAPI.oneL0.module_properties(dev) + (; max_work_group_size = props.maxTotalGroupSize, + max_work_group_dims = (props.maxGroupSizeX, props.maxGroupSizeY, props.maxGroupSizeZ), + max_num_groups = (props.maxGroupCountX, props.maxGroupCountY, props.maxGroupCountZ), + supports_float64 = module_props.flags & oneAPI.oneL0.ZE_DEVICE_MODULE_FLAG_FP64 != 0) end - - # Launch kernel - kernel(ctx, args...; items, groups) - - return nothing +end +KI.max_work_group_size(::oneAPIBackend)::Int = device_limits().max_work_group_size +KI.max_work_group_dims(::oneAPIBackend)::NTuple{3, Int} = device_limits().max_work_group_dims +KI.max_num_groups(::oneAPIBackend)::NTuple{3, Int} = device_limits().max_num_groups +# Level Zero calls Xe cores sub-slices +function KI.multiprocessor_count(::oneAPIBackend)::Int + props = oneAPI.oneL0.properties(device()) + return props.numSlices * props.numSubslicesPerSlice end ## Indexing Functions +## COV_EXCL_START -@device_override @inline function KA.__index_Local_Linear(ctx) - return get_local_id() -end +# computed with `% T`, which unlike `T(x)` has no error path -@device_override @inline function KA.__index_Group_Linear(ctx) - return get_group_id() +@device_override @inline function KI.get_local_id(::Type{T}) where {T} + return (; x = get_local_id(1) % T, y = get_local_id(2) % T, z = get_local_id(3) % T) end -@device_override @inline function KA.__index_Global_Linear(ctx) - return get_global_id() +@device_override @inline function KI.get_group_id(::Type{T}) where {T} + return (; x = get_group_id(1) % T, y = get_group_id(2) % T, z = get_group_id(3) % T) end -@device_override @inline function KA.__index_Local_Cartesian(ctx) - @inbounds KA.workitems(KA.__iterspace(ctx))[get_local_id()] +@device_override @inline function KI.get_local_size(::Type{T}) where {T} + return (; x = get_local_size(1) % T, y = get_local_size(2) % T, z = get_local_size(3) % T) end -@device_override @inline function KA.__index_Group_Cartesian(ctx) - @inbounds KA.blocks(KA.__iterspace(ctx))[get_group_id()] +@device_override @inline function KI.get_num_groups(::Type{T}) where {T} + return (; x = get_num_groups(1) % T, y = get_num_groups(2) % T, z = get_num_groups(3) % T) end -@device_override @inline function KA.__index_Global_Cartesian(ctx) - return @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(), get_local_id()) -end - -@device_override @inline function KA.__validindex(ctx) - if KA.__dynamic_checkbounds(ctx) - I = @inbounds KA.expand(KA.__iterspace(ctx), get_group_id(), get_local_id()) - return I in KA.__ndrange(ctx) - else - return true - end -end +## Shared Memory -## Shared and Scratch Memory - -@device_override @inline function KA.SharedMemory(::Type{T}, ::Val{Dims}, ::Val{Id}) where {T, Dims, Id} +@device_override @inline function KI.localmemory(::Type{T}, ::Val{Dims}) where {T, Dims} ptr = oneAPI.emit_localmemory(T, Val(prod(Dims))) oneDeviceArray(Dims, ptr) end -@device_override @inline function KA.Scratchpad(ctx, ::Type{T}, ::Val{Dims}) where {T, Dims} - StaticArrays.MArray{KA.__size(Dims), T}(undef) -end - ## Synchronization and Printing -@device_override @inline function KA.__synchronize() +@device_override @inline function KI.barrier() # Fence both local and global memory across the workgroup barrier, matching CUDA # `__syncthreads` semantics. `barrier(0)` lowers to `OpControlBarrier` with # `SequentiallyConsistent` but WITHOUT any storage-class bit, which the SPIR-V spec @@ -232,18 +216,16 @@ end barrier(SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE) end -@device_override @inline function KA.__print(args...) +@device_override @inline function KI._print(args...) oneAPI._print(args...) end +## COV_EXCL_STOP -## Other - -Adapt.adapt_storage(to::KA.ConstAdaptor, a::oneDeviceArray) = Base.Experimental.Const(a) -KA.argconvert(::KA.Kernel{oneAPIBackend}, arg) = kernel_convert(arg) +## Other -function KA.priority!(::oneAPIBackend, prio::Symbol) +function KI.priority!(::oneAPIBackend, prio::Symbol) if !(prio in (:high, :normal, :low)) error("priority must be one of :high, :normal, :low") end diff --git a/src/utils.jl b/src/utils.jl index 03d4de12c..05bceb8af 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -37,7 +37,7 @@ function versioninfo(io::IO=stdout) println(io, "Julia packages:") println(io, "- oneAPI.jl: $(Base.pkgversion(oneAPI))") for pkg in [:GPUArrays, :GPUCompiler, ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"), - :LLVM, :SPIRVIntrinsics] + :KernelInterface, :LLVM, :SPIRVIntrinsics] name, mod = get_module(pkg) isnothing(mod) || println(io, "- $(name): $(Base.pkgversion(mod))") end diff --git a/test/Project.toml b/test/Project.toml index beab523b5..af23a69ba 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,6 +9,7 @@ GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7" InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +KernelInterface = "4ee993da-d684-4d17-a7dd-4e58e78d92bf" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" NEO_jll = "700fe977-ac61-5f37-bbc8-c6c4b2b6a9fd" ParallelTestRunner = "d3525ed8-44d0-4b2c-a655-542cee43accc" diff --git a/test/kernelabstractions.jl b/test/kernelabstractions.jl index 48bbe1cb7..e3f3b56cb 100644 --- a/test/kernelabstractions.jl +++ b/test/kernelabstractions.jl @@ -1,8 +1,82 @@ import KernelAbstractions +import KernelAbstractions as KA +import KernelInterface as KI include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.jl")) skip_tests=Set([ "sparse", "Convert", # Need to opt out of i128 + "Random", # oneAPI doesn't support Random's default RNG in kernels yet + # these run kernels on KernelAbstractions' POCL-based CPU back-end + "CPU synchronization", + "fallback test: callable types", ]) -Testsuite.testsuite(oneAPIBackend, "oneAPI", oneAPI, oneArray, oneDeviceArray; skip_tests) +Testsuite.testsuite(()->oneAPIBackend(), "oneAPI", oneAPI, oneArray, oneDeviceArray; skip_tests) + +KA.@kernel function store_global_linear!(A) + I = KA.@index(Global, Linear) + @inbounds A[I] = I +end + +KA.@kernel function store_last_index!(A) + I = KA.@index(Global, Linear) + if I == prod(KA.@ndrange()) + @inbounds A[1] = I + @inbounds A[2] = KA.@index(Global, Cartesian)[2] + end +end + +@testset "launch configuration" begin + backend = oneAPIBackend() + function select(kernel, ndrange, workgroupsize=nothing) + ndrange, workgroupsize, iterspace, _ = KA.launch_config(kernel, ndrange, workgroupsize) + KA.select_launch(kernel, workgroupsize, iterspace) + end + + # kernels are launched on an N-d grid, computing indices in 32 bits + kernel = store_global_linear!(backend) + @test select(kernel, (64, 32, 16)) === KA.NDLaunch{Int32}() + @test select(kernel, (4, 4, 4, 4)) === KA.LinearLaunch{Int32}() + + # which doesn't need divisions to compute the index of a dynamic N-d range + A = oneAPI.zeros(Int, 64, 32, 16) + ir = sprint(io -> oneAPI.@device_code_llvm io=io kernel(A; ndrange=size(A))) + @test !occursin(r"\b[us](div|rem) ", ir) + @test Array(A) == LinearIndices(A) + + # tuning for more work-groups (the testsuite above uses the default) + Testsuite.launch_testsuite(()->oneAPIBackend(; prefer_blocks=true), oneArray) + + # iteration spaces that don't fit 32 bits use 64-bit indices + kernel = store_last_index!(backend) + A = oneAPI.zeros(Int, 2) + for (dims, launch) in (((2^16 + 1, 2^15), KA.NDLaunch{Int}()), + ((2^11 + 1, 2^10, 2^10, 1), KA.LinearLaunch{Int}())) + @test select(kernel, dims) === launch + kernel(A; ndrange=dims) + @test Array(A) == [prod(dims), dims[2]] + end +end + +function ki_store_index!(A) + i = KI.get_global_id().x + if i <= length(A) + @inbounds A[i] = i + end + return +end + +@testset "tuning" begin + # tuning receives the number of work-items; prefer_blocks launches more, smaller groups + A = oneAPI.zeros(Int, 1024) + kernel = KI.@launch oneAPIBackend() launch=false ki_store_index!(A) + items = KI.launch_configuration(kernel; nitems=1024).workgroupsize + @test items <= KI.max_work_group_size(kernel) + kernel = KI.@launch oneAPIBackend(; prefer_blocks=true) launch=false ki_store_index!(A) + fewer = KI.launch_configuration(kernel; nitems=1024).workgroupsize + @test fewer < items + # ... but not for a bound on the work-group size alone + @test KI.launch_configuration(kernel; max_work_group_size=1024).workgroupsize == items + KI.@launch oneAPIBackend(; prefer_blocks=true) ndrange=length(A) ki_store_index!(A) + @test Array(A) == 1:1024 +end diff --git a/test/kernelinterface.jl b/test/kernelinterface.jl new file mode 100644 index 000000000..30f00df2b --- /dev/null +++ b/test/kernelinterface.jl @@ -0,0 +1,27 @@ +import KernelInterface +import KernelInterface as KI + +include(joinpath(dirname(pathof(KernelInterface)), "..", "test", "testsuite.jl")) + +Testsuite.testsuite(oneAPIBackend(), oneArray) + +function ki_fill!(A) + i = KI.get_global_id().x + if i <= length(A) + @inbounds A[i] = i + end + return +end + +@testset "launch keywords" begin + A = oneAPI.zeros(Int, 4) + kernel = KI.@launch oneAPIBackend() launch=false ki_fill!(A) + + # oneAPI's launch options are passed on + kernel(A; ndrange=4, queue=oneAPI.global_stream(oneAPI.context(), oneAPI.device())) + @test Array(A) == 1:4 + + # but not ones that would override the launch geometry + @test_throws ArgumentError kernel(A; ndrange=4, items=8) + @test_throws ArgumentError kernel(A; ndrange=4, groups=2) +end From 93c2acb048d1dcd27d03c5cd0efed061787cc991 Mon Sep 17 00:00:00 2001 From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> Date: Tue, 6 Jan 2026 15:40:48 -0400 Subject: [PATCH 04/12] Allow setting the sub-group size of a kernel The `sub_group_size` compiler keyword sets the `intel_reqd_sub_group_size` metadata of the kernel. By default, the compiler still chooses the sub-group size, and a size the device doesn't support is an error. --- src/compiler/compilation.jl | 21 ++++++++++++++++++--- src/compiler/execution.jl | 5 ++++- test/execution.jl | 16 ++++++++++++++++ 3 files changed, 38 insertions(+), 4 deletions(-) diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl index 33b5537b8..f8aa0cb0c 100644 --- a/src/compiler/compilation.jl +++ b/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 diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 04c2501ff..6dd470934 100644 --- a/src/compiler/execution.jl +++ b/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. diff --git a/test/execution.jl b/test/execution.jl index 00af8e3a5..19d71ae65 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -41,6 +41,22 @@ end end +function store_max_sub_group_size(a) + @inbounds a[1] = get_max_sub_group_size() + return +end + +@testset "sub-group size" begin + a = oneArray{UInt32}(undef, 1) + sizes = oneL0.compute_properties(device()).subGroupSizes + for sub_group_size in sizes + @oneapi items=64 sub_group_size store_max_sub_group_size(a) + @test Array(a)[1] == sub_group_size + end + @test_throws ArgumentError @oneapi sub_group_size=maximum(sizes)+1 dummy() +end + + @testset "inference" begin foo() = @oneapi dummy() @inferred foo() From e1a619ce275b2f3f8ffaf1b923aeb32d83b25dc9 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:15:39 +0200 Subject: [PATCH 05/12] Support KernelInterface's sub-groups Implement the sub-group queries, `sub_group_barrier` and `shfl_down`, and report their support to KernelInterface, which requires kernels to execute with a fixed sub-group width: `kernel_function` compiles for 32 work-items, like a CUDA warp, if the device supports that width, and for its widest one otherwise. Shuffles of `Float16` and `Float64` values depend on the device supporting those types. KernelInterface leaves unspecified how work-items are grouped into sub-groups; Intel GPUs form them from consecutive linear work-item indices, so the last sub-group of a work-group can be partial, which a test checks. Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> --- src/oneAPIKernels.jl | 45 +++++++++++++++++++++++++++++++++++++++-- test/kernelinterface.jl | 36 +++++++++++++++++++++++++++++++++ 2 files changed, 79 insertions(+), 2 deletions(-) diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index a4089e6d9..b1507d107 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -32,6 +32,13 @@ KI.synchronize(::oneAPIBackend) = oneAPI.oneL0.synchronize() KI.supports_float64(::oneAPIBackend) = device_limits().supports_float64 KI.supports_unified(::oneAPIBackend) = true KI.supports_atomics(::oneAPIBackend) = true +KI.supports_subgroups(::oneAPIBackend) = device_limits().sub_group_size > 0 +function KI.supports_shuffle(::oneAPIBackend, ::Type{T}) where {T} + T in SPIRVIntrinsics.gentypes || return false + T === Float64 && return device_limits().supports_float64 + T === Float16 && return device_limits().supports_float16 + return true +end KI.functional(::oneAPIBackend) = oneAPI.functional() @@ -92,7 +99,16 @@ struct oneAPIKernel{F, H <: oneAPI.HostKernel} end function KI.kernel_function(backend::oneAPIBackend, f::F, tt::TT=Tuple{}; name = nothing, kwargs...) where {F,TT} - host = zefunction(kernel_convert(f), tt; name, backend.always_inline, kwargs...) + # compile for the sub-group width that `KI.sub_group_size` promises + sub_group_size = KI.sub_group_size(backend) + if get(kwargs, :sub_group_size, sub_group_size) != sub_group_size + throw(ArgumentError("KernelInterface kernels are compiled for a sub-group size of $sub_group_size, got $(kwargs[:sub_group_size])")) + end + host = if sub_group_size > 0 + zefunction(kernel_convert(f), tt; name, backend.always_inline, kwargs..., sub_group_size) + else + zefunction(kernel_convert(f), tt; name, backend.always_inline, kwargs...) + end kern = oneAPIKernel(f, host) KI.Kernel{oneAPIBackend, typeof(kern)}(backend, kern) end @@ -149,7 +165,7 @@ end # querying the device allocates, so cache what every launch needs const DeviceLimits = @NamedTuple{ max_work_group_size::Int, max_work_group_dims::NTuple{3, Int}, max_num_groups::NTuple{3, Int}, - supports_float64::Bool, + sub_group_size::Int, supports_float16::Bool, supports_float64::Bool, } function device_limits(dev::oneAPI.oneL0.ZeDevice = device()) limits = get!(task_local_storage(), :oneAPIDeviceLimits) do @@ -158,15 +174,22 @@ function device_limits(dev::oneAPI.oneL0.ZeDevice = device()) get!(limits, dev) do props = oneAPI.oneL0.compute_properties(dev) module_props = oneAPI.oneL0.module_properties(dev) + # the sub-group width that `kernel_function` compiles for: 32, like a CUDA warp, if + # the device supports it, and 0 if the device has no sub-groups + sg_sizes = props.subGroupSizes + sub_group_size = 32 in sg_sizes ? 32 : maximum(sg_sizes; init = 0) (; max_work_group_size = props.maxTotalGroupSize, max_work_group_dims = (props.maxGroupSizeX, props.maxGroupSizeY, props.maxGroupSizeZ), max_num_groups = (props.maxGroupCountX, props.maxGroupCountY, props.maxGroupCountZ), + sub_group_size, + supports_float16 = module_props.flags & oneAPI.oneL0.ZE_DEVICE_MODULE_FLAG_FP16 != 0, supports_float64 = module_props.flags & oneAPI.oneL0.ZE_DEVICE_MODULE_FLAG_FP64 != 0) end end KI.max_work_group_size(::oneAPIBackend)::Int = device_limits().max_work_group_size KI.max_work_group_dims(::oneAPIBackend)::NTuple{3, Int} = device_limits().max_work_group_dims KI.max_num_groups(::oneAPIBackend)::NTuple{3, Int} = device_limits().max_num_groups +KI.sub_group_size(::oneAPIBackend)::Int = device_limits().sub_group_size # Level Zero calls Xe cores sub-slices function KI.multiprocessor_count(::oneAPIBackend)::Int props = oneAPI.oneL0.properties(device()) @@ -195,6 +218,16 @@ end return (; x = get_num_groups(1) % T, y = get_num_groups(2) % T, z = get_num_groups(3) % T) end +@device_override KI.get_sub_group_size(::Type{T}) where {T} = get_sub_group_size() % T + +@device_override KI.get_max_sub_group_size(::Type{T}) where {T} = get_max_sub_group_size() % T + +@device_override KI.get_num_sub_groups(::Type{T}) where {T} = get_num_sub_groups() % T + +@device_override KI.get_sub_group_id(::Type{T}) where {T} = get_sub_group_id() % T + +@device_override KI.get_sub_group_local_id(::Type{T}) where {T} = get_sub_group_local_id() % T + ## Shared Memory @@ -216,6 +249,14 @@ end barrier(SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE) end +@device_override @inline function KI.sub_group_barrier() + sub_group_barrier(SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE) +end + +@device_override function KI.shfl_down(val::T, offset::Integer) where T + sub_group_shuffle(val, get_sub_group_local_id() + offset) +end + @device_override @inline function KI._print(args...) oneAPI._print(args...) end diff --git a/test/kernelinterface.jl b/test/kernelinterface.jl index 30f00df2b..df0287b55 100644 --- a/test/kernelinterface.jl +++ b/test/kernelinterface.jl @@ -5,6 +5,42 @@ include(joinpath(dirname(pathof(KernelInterface)), "..", "test", "testsuite.jl") Testsuite.testsuite(oneAPIBackend(), oneArray) +function ki_subgroup_kernel(num, sizes, max, id, lane) + l = KI.get_local_id() + s = KI.get_local_size() + i = l.x + (l.y - 1) * s.x + @inbounds begin + num[i] = KI.get_num_sub_groups() + sizes[i] = KI.get_sub_group_size() + max[i] = KI.get_max_sub_group_size() + id[i] = KI.get_sub_group_id() + lane[i] = KI.get_sub_group_local_id() + end + return +end + +# KernelInterface leaves the formation of sub-groups unspecified; Intel GPUs form them from +# consecutive linear work-item indices +@testset "partial sub-groups" begin + backend = oneAPIBackend() + @test KI.supports_subgroups(backend) + width = KI.sub_group_size(backend) + + # a (width+1)x2 work-group is made up of 3 sub-groups, the last one only partially filled + workgroupsize = (width + 1, 2) + n = prod(workgroupsize) + num, sizes, max, id, lane = (oneArray{UInt32}(undef, n) for _ in 1:5) + KI.@launch backend workgroupsize=workgroupsize ki_subgroup_kernel(num, sizes, max, id, lane) + @test all(==(3), Array(num)) + @test all(==(width), Array(max)) + @test Array(sizes) == [i < 2width ? width : 2 for i in 0:n-1] + @test Array(id) == [div(i, width) + 1 for i in 0:n-1] + @test Array(lane) == [rem(i, width) + 1 for i in 0:n-1] + + # the width is fixed + @test_throws ArgumentError KI.@launch backend launch=false sub_group_size=(width == 16 ? 8 : 16) ki_subgroup_kernel(num, sizes, max, id, lane) +end + function ki_fill!(A) i = KI.get_global_id().x if i <= length(A) From cc65d1e18c9a5c4bdd7dba3c8475020e57acfe69 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:18:05 +0200 Subject: [PATCH 06/12] Order work across tasks with Level Zero events Implement `KI.record_event` with an event signaled on the task's stream, and `KI.wait_event` by waiting for it on the host, cooperatively, like `synchronize` does. `KernelAbstractions.@spawn` uses them to order a new task's work after the work its parent had queued, without synchronizing the parent; the default is a full synchronization. The wait happens on the host, not by making the waiting task's stream wait on the device, because a Level Zero event must not be destroyed while a command list still references it, and the stream would reference it until it executes the wait. For the same reason, recorded events are kept alive until they have been signaled, independently of the recording task, which may end before that. KernelInterface's testsuite checks them once `record_event` returns an event. Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> --- lib/level-zero/event.jl | 4 +-- lib/level-zero/oneL0.jl | 2 +- lib/level-zero/synchronization.jl | 22 +++++++------ src/oneAPIKernels.jl | 52 +++++++++++++++++++++++++++++++ 4 files changed, 68 insertions(+), 12 deletions(-) diff --git a/lib/level-zero/event.jl b/lib/level-zero/event.jl index 92460f2b2..d805cf129 100644 --- a/lib/level-zero/event.jl +++ b/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) diff --git a/lib/level-zero/oneL0.jl b/lib/level-zero/oneL0.jl index c36b9eee3..f8e32e8ad 100644 --- a/lib/level-zero/oneL0.jl +++ b/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") diff --git a/lib/level-zero/synchronization.jl b/lib/level-zero/synchronization.jl index e3a3a4e20..5d5235180 100644 --- a/lib/level-zero/synchronization.jl +++ b/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 diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index b1507d107..6cc2c6d3e 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -264,6 +264,58 @@ end ## COV_EXCL_STOP +## Events + +# Level Zero has no reusable timeline events, so every `record_event` creates a fresh +# one-shot event in its own host-visible pool. Its signal has `HOST` scope, the broadest: +# it flushes caches far enough for the host *and* other devices. +# +# `wait_event` waits for the event on the host, cooperatively: making the waiting task's +# stream wait for it on the device would require keeping the event alive until that stream +# has executed the wait. Recorded events are kept alive until they have been signaled, as +# destroying an event that the device still has to signal is undefined; not by the +# recording task, which may end before that. +const pending_events = oneAPI.oneL0.ZeEvent[] +const pending_events_lock = ReentrantLock() + +function KI.record_event(::oneAPIBackend) + ctx = oneAPI.context() + dev = oneAPI.device() + + pool = oneAPI.oneL0.ZeEventPool(ctx, 1; + flags = oneAPI.oneL0.ZE_EVENT_POOL_FLAG_HOST_VISIBLE) + ev = oneAPI.oneL0.ZeEvent(pool, 1; + signal = oneAPI.oneL0.ZE_EVENT_SCOPE_FLAG_HOST, + wait = oneAPI.oneL0.ZE_EVENT_SCOPE_FLAG_HOST) + + @lock pending_events_lock begin + filter!(!Base.isdone, pending_events) + push!(pending_events, ev) + end + + # the task's stream is an in-order immediate command list, so the signal fires once + # everything appended before it has completed. `execute!` first drains any pending + # oneMKL work on the companion queue (`mkl_wait!`), so that is captured as well. + try + oneAPI.oneL0.execute!(oneAPI.global_stream(ctx, dev)) do list + oneAPI.oneL0.append_signal!(list, ev) + end + catch + # the signal wasn't submitted, so nothing references the event + @lock pending_events_lock filter!(!=(ev), pending_events) + rethrow() + end + + return ev +end + +# the event can come from any device, also one of another driver +function KI.wait_event(::oneAPIBackend, ev::oneAPI.oneL0.ZeEvent) + oneAPI.oneL0.nonblocking_synchronize(ev) + return +end + + ## Other function KI.priority!(::oneAPIBackend, prio::Symbol) From 23f2743243c2d278d61b3b7d70addca952f4a8bc Mon Sep 17 00:00:00 2001 From: Christian Guinard <28689358+christiangnrd@users.noreply.github.com> Date: Wed, 30 Sep 2026 18:18:12 +0200 Subject: [PATCH 07/12] Report oneAPI's version info through KernelInterface `KI.versioninfo(oneAPIBackend())` prints `oneAPI.versioninfo()`. --- src/oneAPIKernels.jl | 2 ++ test/kernelinterface.jl | 4 ++++ 2 files changed, 6 insertions(+) diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index 6cc2c6d3e..8bc28dbdc 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -318,6 +318,8 @@ end ## Other +KI.versioninfo(io::IO, ::oneAPIBackend) = oneAPI.versioninfo(io) + function KI.priority!(::oneAPIBackend, prio::Symbol) if !(prio in (:high, :normal, :low)) error("priority must be one of :high, :normal, :low") diff --git a/test/kernelinterface.jl b/test/kernelinterface.jl index df0287b55..27c019482 100644 --- a/test/kernelinterface.jl +++ b/test/kernelinterface.jl @@ -61,3 +61,7 @@ end @test_throws ArgumentError kernel(A; ndrange=4, items=8) @test_throws ArgumentError kernel(A; ndrange=4, groups=2) end + +@testset "versioninfo" begin + @test occursin("oneAPI.jl", sprint(KI.versioninfo, oneAPIBackend())) +end From 67bd1314e7817ac9b199d1c70ed475e9ba5c5e0d Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Wed, 30 Sep 2026 18:10:14 +0200 Subject: [PATCH 08/12] [TEMP] Get KernelAbstractions 0.10 and KernelInterface 0.4 from their main branch Neither is registered yet. KernelAbstractions is an indirect dependency too (through GPUArrays and AcceleratedKernels), so resolving oneAPI needs it from [sources], which Pkg only accepts for a regular dependency: it is one until this commit is dropped. The test project takes oneAPI from the checkout, so that it doesn't resolve the registered release. Julia 1.10 ignores [sources], and Julia 1.11 doesn't pick them up for the test environment, so CI develops both explicitly there (and runs the 1.10 tests without the julia-test plugin, which instantiates the project first), and the documentation is built with Julia 1.12. Drop this commit once they are registered. --- .buildkite/pipeline.yml | 20 ++++++++++++++++++++ .github/workflows/docs.yml | 4 +++- .gitlab-ci.yml | 12 ++++++++++++ Project.toml | 6 ++++-- test/Project.toml | 5 +++++ 5 files changed, 44 insertions(+), 3 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 903fdc7aa..2af980966 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -26,7 +26,19 @@ steps: agents: queue: "oneapi" commands: | + # XXX: KernelAbstractions 0.10 and KernelInterface 0.4 aren't registered yet. + # 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. + if [[ "{{matrix.julia}}" == "1.10" || "{{matrix.julia}}" == "1.11" ]]; then + git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka + julia --project -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in ("ka", "ka/lib/KernelInterface")])' + fi julia --project=deps deps/build_ci.jl + # The julia-test plugin instantiates the project before these commands, which fails + # on Julia 1.10, so that step runs the tests here instead. + if [[ "{{matrix.julia}}" == "1.10" ]]; then + julia --project -e 'using Pkg; Pkg.test(; coverage=true, test_args=`--quickfail`)' + fi if: | build.message !~ /\[skip [^\]]*(tests|julia)/ && (build.message =~ /\[only [^\]]*(tests|julia)/ || build.message !~ /\[only /) @@ -38,6 +50,14 @@ steps: # the oldest supported release - <<: *julia + plugins: + - JuliaCI/julia#v1: + version: "{{matrix.julia}}" + - JuliaCI/julia-coverage#v1: + dirs: + - src + - lib + - examples if: | build.message !~ /\[skip [^\]]*(tests|julia)/ && (build.message =~ /\[only [^\]]*(tests|julia)/ || diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index c853602fa..338af8596 100644 --- a/.github/workflows/docs.yml +++ b/.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 diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index cf54dd92e..e64e9ab3f 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -90,6 +90,17 @@ 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, 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 + git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka + julia --color=yes --project=. -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in ("ka", "ka/lib/KernelInterface")])' + julia --color=yes --project=test -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "ka", "ka/lib/KernelInterface")])' + 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 +124,7 @@ stages: [build, test, coverage] - LocalPreferences.toml - test/LocalPreferences.toml - test/Manifest.toml + - ka expire_in: 1 week .test:lts: diff --git a/Project.toml b/Project.toml index 24b549704..bf83460d3 100644 --- a/Project.toml +++ b/Project.toml @@ -11,6 +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" @@ -31,8 +32,9 @@ 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" +[sources] +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" diff --git a/test/Project.toml b/test/Project.toml index af23a69ba..ef9f63055 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -27,5 +27,10 @@ libigc_jll = "94295238-5935-5bd7-bb0f-b00942e9bdd5" oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36" +[sources] +KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"} +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"} +oneAPI = {path = ".."} + [compat] ParallelTestRunner = "2.8" From 90c6de91aaafbdedb27afc8366fcc21030c6c29d Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 6 Oct 2026 11:01:01 +0200 Subject: [PATCH 09/12] Use KernelAbstractions' @private storage KernelAbstractions now backs @private with stack memory from GPUCompiler.alloca, defined as an overlay in GPUCompiler's shared method table, so drop our StaticArrays-based Scratchpad override and with it the StaticArrays dependency. This relies on our compiler consulting the shared method table, which oneAPI.jl#660 implements. --- Project.toml | 2 -- ext/KernelAbstractionsExt.jl | 10 ---------- 2 files changed, 12 deletions(-) diff --git a/Project.toml b/Project.toml index bf83460d3..495dc2ca2 100644 --- a/Project.toml +++ b/Project.toml @@ -27,7 +27,6 @@ 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" @@ -58,7 +57,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" diff --git a/ext/KernelAbstractionsExt.jl b/ext/KernelAbstractionsExt.jl index 73f4a14da..f9a0d9fcd 100644 --- a/ext/KernelAbstractionsExt.jl +++ b/ext/KernelAbstractionsExt.jl @@ -1,12 +1,9 @@ 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) @@ -18,13 +15,6 @@ Adapt.adapt_storage(::KA.CPU, a::oneArray) = convert(Array, 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) From 021f9610b28cb2295c074116a69b3e68e7a0c0a9 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 10 Oct 2026 18:50:06 +0200 Subject: [PATCH 10/12] [TEMP] Test against KernelAbstractions' vc/ki-subgroup-ops branch Get KernelAbstractions and KernelInterface from JuliaGPU/KernelAbstractions.jl#831, which adds the sub-group operations implemented here. Drop this commit when it is merged. --- .buildkite/pipeline.yml | 2 +- .gitlab-ci.yml | 2 +- Project.toml | 4 ++-- test/Project.toml | 4 ++-- 4 files changed, 6 insertions(+), 6 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index 2af980966..dcea82f18 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -30,7 +30,7 @@ steps: # 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. 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 --branch vc/ki-subgroup-ops https://github.com/JuliaGPU/KernelAbstractions.jl ka julia --project -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in ("ka", "ka/lib/KernelInterface")])' fi julia --project=deps deps/build_ci.jl diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml index e64e9ab3f..6f415995c 100644 --- a/.gitlab-ci.yml +++ b/.gitlab-ci.yml @@ -97,7 +97,7 @@ stages: [build, test, coverage] - | if [ "$JULIA_VERSION" = "1.10" ]; then rm -rf ka - git clone --depth 1 --branch main https://github.com/JuliaGPU/KernelAbstractions.jl ka + git clone --depth 1 --branch vc/ki-subgroup-ops https://github.com/JuliaGPU/KernelAbstractions.jl ka julia --color=yes --project=. -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in ("ka", "ka/lib/KernelInterface")])' julia --color=yes --project=test -e 'using Pkg; Pkg.develop([PackageSpec(; path) for path in (".", "ka", "ka/lib/KernelInterface")])' fi diff --git a/Project.toml b/Project.toml index 495dc2ca2..67e5f00e4 100644 --- a/Project.toml +++ b/Project.toml @@ -32,8 +32,8 @@ oneAPI_Level_Zero_Loader_jll = "13eca655-d68d-5b81-8367-6d99d727ab01" oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36" [sources] -KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"} -KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"} +KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"} +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"} [extensions] KernelAbstractionsExt = "KernelAbstractions" diff --git a/test/Project.toml b/test/Project.toml index ef9f63055..6997af150 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -28,8 +28,8 @@ oneAPI = "8f75cd03-7ff8-4ecb-9b8f-daf728133b1b" oneAPI_Support_jll = "b049733a-a71d-5ed3-8eba-7d323ac00b36" [sources] -KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main"} -KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "main", subdir = "lib/KernelInterface"} +KernelAbstractions = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops"} +KernelInterface = {url = "https://github.com/JuliaGPU/KernelAbstractions.jl", rev = "vc/ki-subgroup-ops", subdir = "lib/KernelInterface"} oneAPI = {path = ".."} [compat] From bbd258956fb52d5355d0ea8ca8e1c44d7343b205 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sat, 10 Oct 2026 18:50:10 +0200 Subject: [PATCH 11/12] Implement KernelInterface's sub-group shuffles, votes and collectives Implement `shfl`, `shfl_up`, `shfl_down` and `shfl_xor` for the types SPIR-V shuffles natively (leaving others to KernelInterface's word and field fallbacks), the votes with the native ballot, and the native reductions and scans where they have Julia's semantics: `+` on integers and floats, and `min`/`max` on integers. `shfl_down` and `shfl_up` return the work-item's own value past the sub-group width, which SPIR-V leaves undefined. Intel's GPU compiler forms sub-groups from consecutive work-items, and they execute independently, so report both `supports_linear_subgroups` and `supports_independent_subgroups`. Only widths of up to 64 are used, so that the ballot covers every lane. --- src/oneAPIKernels.jl | 74 ++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 68 insertions(+), 6 deletions(-) diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index 8bc28dbdc..fb9d64e1e 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -33,8 +33,18 @@ KI.supports_float64(::oneAPIBackend) = device_limits().supports_float64 KI.supports_unified(::oneAPIBackend) = true KI.supports_atomics(::oneAPIBackend) = true KI.supports_subgroups(::oneAPIBackend) = device_limits().sub_group_size > 0 -function KI.supports_shuffle(::oneAPIBackend, ::Type{T}) where {T} - T in SPIRVIntrinsics.gentypes || return false +# Intel's GPU compiler forms sub-groups from consecutive work-items (which the KernelInterface +# testsuite checks), and they execute independently +KI.supports_linear_subgroups(backend::oneAPIBackend) = KI.supports_subgroups(backend) +KI.supports_independent_subgroups(backend::oneAPIBackend) = KI.supports_subgroups(backend) + +# The types `sub_group_shuffle` supports; other types are shuffled as words or field by +# field. `Float16` and `Float64` are reported as unsupported on devices without them, rather +# than shuffled as words: the overrides below are selected by type, not by device. +const ShuffleTypes = Union{SPIRVIntrinsics.gentypes...} + +function KI.supports_shuffle(backend::oneAPIBackend, ::Type{T}) where {T <: ShuffleTypes} + KI.supports_subgroups(backend) || return false T === Float64 && return device_limits().supports_float64 T === Float16 && return device_limits().supports_float16 return true @@ -175,8 +185,9 @@ function device_limits(dev::oneAPI.oneL0.ZeDevice = device()) props = oneAPI.oneL0.compute_properties(dev) module_props = oneAPI.oneL0.module_properties(dev) # the sub-group width that `kernel_function` compiles for: 32, like a CUDA warp, if - # the device supports it, and 0 if the device has no sub-groups - sg_sizes = props.subGroupSizes + # the device supports it, and 0 if the device has no sub-groups. Only widths of up to + # 64, so that `KI.sub_group_ballot` (a `UInt64` mask) covers every lane. + sg_sizes = filter(<=(64), props.subGroupSizes) sub_group_size = 32 in sg_sizes ? 32 : maximum(sg_sizes; init = 0) (; max_work_group_size = props.maxTotalGroupSize, max_work_group_dims = (props.maxGroupSizeX, props.maxGroupSizeY, props.maxGroupSizeZ), @@ -253,8 +264,59 @@ end sub_group_barrier(SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE) end -@device_override function KI.shfl_down(val::T, offset::Integer) where T - sub_group_shuffle(val, get_sub_group_local_id() + offset) +@device_override KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} = + sub_group_shuffle(val, lane) + +# past the sub-group width, `shfl_down` and `shfl_up` return the work-item's own value, which +# `sub_group_shuffle` (like SPIR-V's `OpGroupNonUniformShuffleDown`) leaves undefined +@device_override function KI.shfl_down(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = get_sub_group_local_id() + # compared before adding, so that large offsets don't overflow + inside = offset <= get_max_sub_group_size() - lane + return sub_group_shuffle(val, ifelse(inside, lane + offset, lane)) +end + +@device_override function KI.shfl_up(val::T, offset::Integer) where {T <: ShuffleTypes} + lane = get_sub_group_local_id() + return sub_group_shuffle(val, ifelse(lane > offset, lane - offset, lane)) +end + +@device_override KI.shfl_xor(val::T, mask::Integer) where {T <: ShuffleTypes} = + sub_group_shuffle_xor(val, mask) + +@device_override KI.sub_group_any(pred::Bool) = SPIRVIntrinsics.sub_group_any(pred) + +@device_override KI.sub_group_all(pred::Bool) = SPIRVIntrinsics.sub_group_all(pred) + +@device_override function KI.sub_group_ballot(pred::Bool) + mask = SPIRVIntrinsics.sub_group_ballot(pred) + return UInt64(mask[1].value) | (UInt64(mask[2].value) << 32) +end + +# The native reductions and scans, for `+` on 32- and 64-bit integers and floats, and +# `min`/`max` on integers (OpenCL's `min` and `max` treat NaN and the sign of zero differently +# from Julia's). +const CollectiveIntTypes = Union{Int32, UInt32, Int64, UInt64} +const CollectiveTypes = Union{CollectiveIntTypes, Float16, Float32, Float64} + +@device_override KI.sub_group_reduce(::typeof(+), val::CollectiveTypes) = + SPIRVIntrinsics.sub_group_reduce_add(val) +@device_override KI.sub_group_reduce(::typeof(min), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_reduce_min(val) +@device_override KI.sub_group_reduce(::typeof(max), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_reduce_max(val) + +@device_override KI.sub_group_scan(::typeof(+), val::CollectiveTypes) = + SPIRVIntrinsics.sub_group_scan_inclusive_add(val) +@device_override KI.sub_group_scan(::typeof(min), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_scan_inclusive_min(val) +@device_override KI.sub_group_scan(::typeof(max), val::CollectiveIntTypes) = + SPIRVIntrinsics.sub_group_scan_inclusive_max(val) + +# the native exclusive scans start from the identity, not from `init` +@device_override function KI.sub_group_exclusive_scan(::typeof(+), val::T, init::T) where {T <: CollectiveTypes} + prefix = SPIRVIntrinsics.sub_group_scan_exclusive_add(val) + return ifelse(get_sub_group_local_id() == 1, init, init + prefix) end @device_override @inline function KI._print(args...) From 930aff595cefb8ddb1022d28ac72446fad562983 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Sun, 11 Oct 2026 11:21:32 +0200 Subject: [PATCH 12/12] Report partitionable sub-groups Report `KI.supports_partitionable_subgroups`: an Intel GPU sub-group is the SIMD lanes of one EU thread, and the shuffles and ballot that the width functions use are `OpGroupNonUniform*` instructions over the active lanes. A sub-group `OpControlBarrier` has to be executed by the whole sub-group, so `KI.sub_group_barrier(width)` for a segment narrower than the sub-group is a sub-group-scope memory barrier on local and global memory instead. --- src/oneAPIKernels.jl | 25 +++++++++++++++++++++---- 1 file changed, 21 insertions(+), 4 deletions(-) diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index fb9d64e1e..664aed0ac 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -37,10 +37,15 @@ KI.supports_subgroups(::oneAPIBackend) = device_limits().sub_group_size > 0 # testsuite checks), and they execute independently KI.supports_linear_subgroups(backend::oneAPIBackend) = KI.supports_subgroups(backend) KI.supports_independent_subgroups(backend::oneAPIBackend) = KI.supports_subgroups(backend) - -# The types `sub_group_shuffle` supports; other types are shuffled as words or field by -# field. `Float16` and `Float64` are reported as unsupported on devices without them, rather -# than shuffled as words: the overrides below are selected by type, not by device. +# A sub-group is the SIMD lanes of one EU thread, so a segment of it executes the functions +# with a `width` together without the rest of the sub-group: the shuffles and the ballot are +# `OpGroupNonUniform*` instructions over the active lanes, and `KI.sub_group_barrier(width)` +# only has to order memory (see below). +KI.supports_partitionable_subgroups(backend::oneAPIBackend) = KI.supports_subgroups(backend) + +# The types `sub_group_shuffle` supports; other types are shuffled as words. `Float16` and +# `Float64` are reported as unsupported on devices without them, rather than shuffled as +# words: the overrides below are selected by type, not by device. const ShuffleTypes = Union{SPIRVIntrinsics.gentypes...} function KI.supports_shuffle(backend::oneAPIBackend, ::Type{T}) where {T <: ShuffleTypes} @@ -264,6 +269,18 @@ end sub_group_barrier(SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE) end +# A sub-group `OpControlBarrier` has to be executed by the whole sub-group, so a narrower +# segment, whose lanes execute together anyway, only fences its memory accesses +@device_override @inline function KI.sub_group_barrier(width::Integer) + if width == get_max_sub_group_size() + KI.sub_group_barrier() + else + SPIRVIntrinsics.atomic_work_item_fence( + SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE, + SPIRVIntrinsics.memory_order_acq_rel, SPIRVIntrinsics.memory_scope_sub_group) + end +end + @device_override KI.shfl(val::T, lane::Integer) where {T <: ShuffleTypes} = sub_group_shuffle(val, lane)