Skip to content

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

Closed
maleadt wants to merge 12 commits into
ka-0.10from
tb/ki-subgroup-ops
Closed

maleadt wants to merge 12 commits into
ka-0.10from
tb/ki-subgroup-ops

Conversation

@maleadt

@maleadt maleadt commented Oct 10, 2026 •

Copy link
Copy Markdown
Member

Implements KernelInterface's sub-group API (JuliaGPU/KernelAbstractions.jl#831) for OpenCL.jl. Stacked on #519.

  • Shuffles (shfl, shfl_up, shfl_down, shfl_xor) for the types sub_group_shuffle supports, with cl_khr_subgroup_shuffle. They're overridden only for those types, so other types reach KernelInterface's fallbacks, which shuffle them as UInt32 words. shfl_down/shfl_up return the work-item's own value past the width, also for very large offsets.
  • Float16/Float64 aren't reported as shuffleable on devices without cl_khr_fp16/cl_khr_fp64. A kernel can't even move such values there: GPUCompiler rejects the IR after LLVM turns an integer load that is reinterpreted as a float into a float load.
  • Votes. The ballot uses cl_khr_subgroup_ballot where available, and otherwise two 32-bit sub_group_reduce_adds.
  • Native sub_group_reduce/sub_group_scan/sub_group_exclusive_scan for + (integers and floats), and min/max on integers only, because OpenCL's float min/max don't propagate NaN or order signed zeros the way Julia does.
  • Sub-groups are still only reported with cl_intel_required_subgroup_size, which kernels need to get a fixed width. They're now also detected when the device reports them as an OpenCL 3.0 core feature, and only widths of at most 64 are reported, so the ballot has a bit for every lane. Requiring the extension is an initial limitation; other ways to fix the width can be added later.
  • supports_linear_subgroups: PoCL and Intel's CPU and GPU runtimes, all of which pass KernelInterface's layout tests.
  • supports_independent_subgroups: false on PoCL, whose sub-group operations are work-group barriers. It's also false on Intel's CPU runtime, which loses the writes after a sub_group_barrier that other sub-groups returned before ([OpenCL CPU Runtime] Writes after a sub-group OpControlBarrier are lost when other sub-groups returned early intel/llvm#23449). Intel's GPU runtime passes the independence test.
  • supports_partitionable_subgroups: only on Intel's GPU runtime with cl_khr_subgroup_ballot, where a sub-group is the SIMD lanes of one EU thread and the shuffles and ballot of the width functions only involve the active lanes. OpenCL has no barrier for part of a sub-group, so a new :partitionable_sub_groups device feature makes KI.sub_group_barrier(width) a sub-group-scope memory barrier there for segments narrower than the sub-group. Elsewhere (PoCL, Intel's CPU runtime, unknown runtimes) it's false, and KernelAbstractions' tiles aren't available on PoCL and Intel's CPU runtime, whose sub-groups aren't independent.
  • cl.sub_group_size returns 0 for devices without sub-groups (fixes cl.sub_group_size doesn't return 0 for devices without sub-groups #553).

The TEMP commit gets KernelAbstractions and KernelInterface from KA's vc/ki-subgroup-ops branch.

Tested locally, KI and KA testsuites, Julia 1.12 (and 1.10 on PoCL and Intel's CPU runtime):

@maleadt
maleadt changed the base branch from ka-0.10 to tb/private-array October 10, 2026 17:36
@maleadt
maleadt added this pull request to stack #528 October 10, 2026 17:36
@maleadt
maleadt removed this pull request from stack #528 October 10, 2026 17:36
@maleadt
maleadt changed the base branch from tb/private-array to ka-0.10 October 10, 2026 17:36
@maleadt
maleadt added this pull request to stack #555 October 10, 2026 17:36
@maleadt
maleadt force-pushed the tb/ki-subgroup-ops branch from 885c9fc to b06173c Compare October 10, 2026 18:05
@maleadt
maleadt force-pushed the tb/ki-subgroup-ops branch 2 times, most recently from 429bfe2 to b06173c Compare October 10, 2026 18:39
@maleadt
maleadt force-pushed the tb/ki-subgroup-ops branch from b06173c to 67e20e3 Compare October 10, 2026 20:33
maleadt and others added 11 commits October 12, 2026 09:30
Each layer of the launch path took the kernel arguments as varargs and splatted
them into the next: the kernel object's call, `launch_with_exception_mailbox`,
`cl.call` and `cl.set_args!`. Julia doesn't turn a splat of more than 32
elements into a direct call, and a method with both varargs and keyword
arguments splats them into its body, so launching a kernel with many arguments
went through `Core._apply_iterate` several times.

Pass the arguments along as one tuple instead, and generate the per-argument
code. The kernel object keeps its call syntax, with its keyword method defined
explicitly, and `OpenCL.launch_tuple` takes the tuple directly. `cl.call` now
takes the arguments as a tuple; `cl.set_args!` and `clcall` keep their varargs
signatures.

On PoCL, launching a kernel with 40 arguments with `@opencl` goes from 10.3 µs
and 21.5 kB to 5.2 µs and 976 bytes, which is what a kernel with one argument
allocates (it takes 4.3 µs).
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 OpenCL.jl's copy
of that launch path: partitioning the ndrange, building the kernel's context,
and tuning the workgroup size.

`OpenCLBackend` now implements KernelInterface, and OpenCL.jl depends on it
instead of on KernelAbstractions. What KernelAbstractions still needs from a
back end moves to an extension: the `MArray` behind `@private`, and the Adapt
rules for `@Const` and for moving arrays to the CPU.

The platform is part of an `OpenCLBackend`'s configuration, but it was only
checked, with a warning, before allocating and on every launch. Instead, work
for a backend whose platform isn't active now activates the platform's default
device, as `KI.device!` would, so that the arrays it creates can be used
afterwards, and host queries answer for that device. `get_backend` returns a
backend for the array's platform. Kernel queries answer for the device the
kernel was compiled for, and launching a kernel on another device throws an
error instead of failing in the driver.

`KI.launch` passes the kernel arguments on as a tuple, so kernels with many
arguments stay cheap to launch. It rejects `global_size` and `local_size`,
which would override the launch geometry that KernelInterface validated, but
passes OpenCL's other launch options on. The device properties that launches
query are cached per device.

`KI.kernel_function` compiles the converted callable, and keeps the original:
it is converted again at every launch, like the arguments, as `@opencl` already
does, so that the arrays a closure captures stay alive.

`KI.synchronize` still blocks the thread while waiting for the device, rather
than waiting cooperatively as KernelInterface asks.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Implement the sub-group queries, `sub_group_barrier` and `shfl_down` with
OpenCL's sub-group builtins, whose semantics already match KernelInterface's:
the last sub-group of a work-group can be partial, and `get_sub_group_size`
counts the work-items that are present.

KernelInterface requires kernels to execute with the sub-group width that
`KI.sub_group_size` reports. `clfunction` requests that width on devices with
`cl_intel_required_subgroup_size`, so only those report sub-group support, and
shuffles for the types that `cl_khr_subgroup_shuffle` covers on the device. A
`sub_group_size` compiler option asking for another width is rejected.

How work-items form sub-groups differs between devices: Intel's CPU runtime
forms them per row of a multi-dimensional work-group. KernelInterface doesn't
promise more than that (JuliaGPU/KernelAbstractions.jl#815).

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Implement `KI.record_event` with a marker on the task's queue, and
`KI.wait_event` with a barrier on the waiting task's queue.
`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.

A queue can only wait for events of its own context, and OpenCL.jl creates a
context per device, so waiting for an event of another device waits on the
host.
`KI.versioninfo(OpenCLBackend())` prints `OpenCL.versioninfo()`.
… main branch

Neither is registered yet. Julia 1.10 and 1.11 don't pick up the test
project's [sources], so CI develops both explicitly there. Pkg < 1.12 can only
develop the weak dependency on KernelAbstractions if it's also listed in
[extras].

Drop this commit once both are registered.
They ended the `julia -e '...'` argument early.
Documenter failed on the OpenCLBackend docstring missing from the manual.
KernelAbstractions now backs @Private with a stack allocation through an
overlay in GPUCompiler's shared method table, so drop our MArray-based
Scratchpad override along with the StaticArrays dependency. This relies
on OpenCL consulting the shared table, which OpenCL.jl#530 does by
declaring our method tables through method_tables.
Override the shuffles for the types `sub_group_shuffle` supports (with
`cl_khr_subgroup_shuffle`), and only those, so that other types reach KernelInterface's word
and field fallbacks; `shfl_down` and `shfl_up` return the own value past the width. Add the
votes, with the ballot from `cl_khr_subgroup_ballot` or, without it, from 32-bit
reductions, and the native reductions and scans where they have Julia's semantics.

Sub-groups still need `cl_intel_required_subgroup_size` for a fixed width, now also when
reported as an OpenCL 3.0 core feature, and of at most 64 lanes for the ballot. The linear
layout is reported for PoCL and Intel's CPU and GPU runtimes, which the layout tests confirm;
independent sub-groups everywhere but PoCL (whose sub-group operations are work-group
barriers) and Intel's CPU runtime (which loses the writes after a `sub_group_barrier` that
other sub-groups returned before, intel/llvm#23449).

OpenCL C compilation through spirv2clc needs 0.2.5, which translates the sub-group
operations.
@maleadt
maleadt force-pushed the tb/ki-subgroup-ops branch from 892959f to 0ac38d8 Compare October 12, 2026 07:52
Report `KI.supports_partitionable_subgroups` only on Intel's GPU runtime with
`cl_khr_subgroup_ballot`, where a sub-group is the SIMD lanes of one EU thread
and the width functions' shuffles and ballot only involve the active lanes.
OpenCL has no barrier for part of a sub-group, so a new
`:partitionable_sub_groups` device feature lets `KI.sub_group_barrier(width)`
use a sub-group-scope memory barrier there for segments narrower than the
sub-group; elsewhere, including PoCL (whose sub-group functions are work-group
barriers) and Intel's CPU runtime, the whole sub-group keeps executing the
width functions together.
@maleadt
maleadt force-pushed the tb/ki-subgroup-ops branch from 0ac38d8 to 9c807cf Compare October 12, 2026 09:08
@codecov

codecov Bot commented Oct 12, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 94.64286% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 87.06%. Comparing base (e867580) to head (9c807cf).

Files with missing lines Patch % Lines
src/OpenCLKernels.jl 94.44% 3 Missing ⚠️
Additional details and impacted files
@@             Coverage Diff             @@
##           ka-0.10     #554      +/-   ##
===========================================
- Coverage    87.13%   87.06%   -0.08%     
===========================================
  Files           18       18              
  Lines         1664     1716      +52     
===========================================
+ Hits          1450     1494      +44     
- Misses         214      222       +8     

☔ 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.

@maleadt maleadt closed this Oct 13, 2026
An error occurred while trying to automatically change base from ka-0.10 to main October 13, 2026 07:56
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

cl.sub_group_size doesn't return 0 for devices without sub-groups

2 participants