Implement KernelInterface's sub-group shuffles, votes and collectives

[maleadt]

Implements KernelInterface's sub-group operations (JuliaGPU/KernelAbstractions.jl#831) on oneAPI. Stacked on #653.

• shfl, shfl_up, shfl_down and shfl_xor for the types SPIR-V shuffles natively; other types use KernelInterface's fallbacks, as UInt32 words. Past the sub-group width, shfl_down/shfl_up return the work-item's own value, which SPIR-V leaves undefined.
• sub_group_any/all, and sub_group_ballot with the native ballot.
• Native reductions and scans where they have Julia's semantics: + on 32/64-bit integers and floats, min/max on integers. Float min/max and other operators use KernelInterface's ordered fallbacks.
• supports_linear_subgroups and supports_independent_subgroups are both true: IGC forms sub-groups from consecutive work-items, and they execute independently. The KernelInterface testsuite checks the layout and independent control flow.
• supports_partitionable_subgroups is true: an Intel GPU sub-group is the SIMD lanes of one EU thread, and the shuffles and ballot the width functions use are OpGroupNonUniform* instructions over the active lanes. A sub-group OpControlBarrier needs 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. KernelAbstractions can then pack several narrow tiles into a sub-group.
• Kernels only use sub-group widths of up to 64, so that the ballot covers every lane.

The [TEMP] commit takes KernelAbstractions and KernelInterface from KernelAbstractions.jl's vc/ki-subgroup-ops branch; drop it once #831 is merged.

#653 was rebased onto main (SPIRVIntrinsics 1.4, the atomics descriptor; AcceleratedKernels is no longer pinned, as the registered 0.5 supports KernelAbstractions 0.10).

Tested on an Iris Xe (Julia 1.12, KernelAbstractions.jl#831 at 0d37f80a): KernelInterface testsuite 5711/5711 at the default width 32, and 5018/5018 and 5346/5346 with kernels compiled for widths 8 and 16 (layout tests included); KernelAbstractions testsuite 2533 passed, 6 broken.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/level-zero/event.jl b/lib/level-zero/event.jl
index d805cf1..be21064 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; signal=0, wait=0)
-        desc_ref = Ref(ze_event_desc_t(; index=index-1, signal, wait))
+    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/synchronization.jl b/lib/level-zero/synchronization.jl
index 5d52351..40e3996 100644
--- a/lib/level-zero/synchronization.jl
+++ b/lib/level-zero/synchronization.jl
@@ -36,7 +36,8 @@ gcsafe_synchronize(queue::ZeCommandQueue) =
         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
+    event::ze_event_handle_t, typemax(UInt64)::UInt64
+)::ze_result_t
 
 """
     nonblocking_synchronize(list_or_queue)
diff --git a/src/compiler/compilation.jl b/src/compiler/compilation.jl
index f8aa0cb..9bdf27d 100644
--- a/src/compiler/compilation.jl
+++ b/src/compiler/compilation.jl
@@ -1,7 +1,7 @@
 ## gpucompiler interface implementation
 
 Base.@kwdef struct oneAPICompilerParams <: AbstractCompilerParams
-    sub_group_size::Union{Nothing,Int} = nothing
+    sub_group_size::Union{Nothing, Int} = nothing
 end
 
 const oneAPICompilerConfig = CompilerConfig{SPIRVCompilerTarget, oneAPICompilerParams}
@@ -298,8 +298,10 @@ function _atomics_capabilities(dev, properties; supports_fp16, supports_fp64)
     )
 end
 
-@noinline function _compiler_config(dev; kernel=true, name=nothing, always_inline=false,
-                                    sub_group_size=nothing, 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
diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl
index 664aed0..151eabf 100644
--- a/src/oneAPIKernels.jl
+++ b/src/oneAPIKernels.jl
@@ -113,7 +113,7 @@ struct oneAPIKernel{F, H <: oneAPI.HostKernel}
     host::H
 end
 
-function KI.kernel_function(backend::oneAPIBackend, f::F, tt::TT=Tuple{}; name = nothing, kwargs...) where {F,TT}
+function KI.kernel_function(backend::oneAPIBackend, f::F, tt::TT = Tuple{}; name = nothing, kwargs...) where {F, TT}
     # 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
@@ -125,12 +125,14 @@ function KI.kernel_function(backend::oneAPIBackend, f::F, tt::TT=Tuple{}; name =
         zefunction(kernel_convert(f), tt; name, backend.always_inline, kwargs...)
     end
     kern = oneAPIKernel(f, host)
-    KI.Kernel{oneAPIBackend, typeof(kern)}(backend, kern)
+    return KI.Kernel{oneAPIBackend, typeof(kern)}(backend, kern)
 end
 
 # 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...)
+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`"))
@@ -186,7 +188,7 @@ 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
+    return 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
@@ -194,12 +196,14 @@ function device_limits(dev::oneAPI.oneL0.ZeDevice = device())
         # 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),
-           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)
+        (;
+            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
@@ -277,7 +281,8 @@ end
     else
         SPIRVIntrinsics.atomic_work_item_fence(
             SPIRVIntrinsics.LOCAL_MEM_FENCE | SPIRVIntrinsics.GLOBAL_MEM_FENCE,
-            SPIRVIntrinsics.memory_order_acq_rel, SPIRVIntrinsics.memory_scope_sub_group)
+            SPIRVIntrinsics.memory_order_acq_rel, SPIRVIntrinsics.memory_scope_sub_group
+        )
     end
 end
 
@@ -361,11 +366,15 @@ 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)
+    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)
diff --git a/src/utils.jl b/src/utils.jl
index 05bceb8..ce3dcac 100644
--- a/src/utils.jl
+++ b/src/utils.jl
@@ -31,13 +31,15 @@ function versioninfo(io::IO=stdout)
     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))
+        return (pkg[2], get(Base.loaded_modules, id, nothing))
     end
 
     println(io, "Julia packages:")
     println(io, "- oneAPI.jl: $(Base.pkgversion(oneAPI))")
-    for pkg in [:GPUArrays, :GPUCompiler, ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"),
-                 :KernelInterface, :LLVM, :SPIRVIntrinsics]
+    for pkg in [
+            :GPUArrays, :GPUCompiler, ("63c18a36-062a-441e-b654-da1e3ab1ce7c", "KernelAbstractions"),
+            :KernelInterface, :LLVM, :SPIRVIntrinsics,
+        ]
         name, mod = get_module(pkg)
         isnothing(mod) || println(io, "- $(name): $(Base.pkgversion(mod))")
     end
diff --git a/test/execution.jl b/test/execution.jl
index 19d71ae..3369cac 100644
--- a/test/execution.jl
+++ b/test/execution.jl
@@ -41,20 +41,20 @@ end
 end
 
 
-function store_max_sub_group_size(a)
-    @inbounds a[1] = get_max_sub_group_size()
-    return
-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
+    @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
-    @test_throws ArgumentError @oneapi sub_group_size=maximum(sizes)+1 dummy()
-end
 
 
 @testset "inference" begin
diff --git a/test/kernelabstractions.jl b/test/kernelabstractions.jl
index e3f3b56..68fa9dd 100644
--- a/test/kernelabstractions.jl
+++ b/test/kernelabstractions.jl
@@ -6,12 +6,12 @@ include(joinpath(dirname(pathof(KernelAbstractions)), "..", "test", "testsuite.j
 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",
+        "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)
@@ -28,7 +28,7 @@ end
 
 @testset "launch configuration" begin
     backend = oneAPIBackend()
-    function select(kernel, ndrange, workgroupsize=nothing)
+    function select(kernel, ndrange, workgroupsize = nothing)
         ndrange, workgroupsize, iterspace, _ = KA.launch_config(kernel, ndrange, workgroupsize)
         KA.select_launch(kernel, workgroupsize, iterspace)
     end
@@ -40,20 +40,22 @@ end
 
     # 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)))
+    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)
+    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}()))
+    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)
+        kernel(A; ndrange = dims)
         @test Array(A) == [prod(dims), dims[2]]
     end
 end
@@ -69,14 +71,14 @@ 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
+    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
+    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 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
index 27c0194..c42c9c9 100644
--- a/test/kernelinterface.jl
+++ b/test/kernelinterface.jl
@@ -30,15 +30,15 @@ end
     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)
+    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]
+    @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)
+    @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)
@@ -51,15 +51,15 @@ end
 
 @testset "launch keywords" begin
     A = oneAPI.zeros(Int, 4)
-    kernel = KI.@launch oneAPIBackend() launch=false ki_fill!(A)
+    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()))
+    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)
+    @test_throws ArgumentError kernel(A; ndrange = 4, items = 8)
+    @test_throws ArgumentError kernel(A; ndrange = 4, groups = 2)
 end
 
 @testset "versioninfo" begin

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 79.58%. Comparing base (90c6de9) to head (930aff5).

Additional details and impacted files

@@             Coverage Diff             @@
##           ka-0.10     #680      +/-   ##
===========================================
+ Coverage    79.56%   79.58%   +0.01%     
===========================================
  Files           51       51              
  Lines         3950     3953       +3     
===========================================
+ Hits          3143     3146       +3     
  Misses         807      807              

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.