Repository navigation
KernelInterface: sub-group shuffles, votes, reductions and scans - #831
Merged
Merged
Conversation
Contributor
Benchmark ResultsShow table
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #831 +/- ##
==========================================
+ Coverage 80.82% 81.31% +0.48%
==========================================
Files 26 27 +1
Lines 2227 2280 +53
==========================================
+ Hits 1800 1854 +54
+ Misses 427 426 -1 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
vchuravy
added this pull request to stack #834
October 4, 2026 06:46
vchuravy
commented
Oct 4, 2026
This was referenced Oct 4, 2026
vchuravy
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 4, 2026
KernelInterface (JuliaGPU/KernelAbstractions.jl#831) now shuffles primitive types that a back-end doesn't support natively as `UInt32` words, so the Metal-specific split into halves isn't needed anymore. Assisted-by: Claude Code (Opus 5.5)
vchuravy
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 6, 2026 12:29
c3e0b7a to
64ac668
Compare
Member
Author
|
Rebased onto main and addressed the review:
The full test suite passes locally on Julia 1.10 and 1.12. Still open from the review: Float16/Float64 on PoCL devices without fp16/fp64 are reported as unsupported for shuffles, even though the #830 is based on this branch and needs a rebase. |
vchuravy
commented
Oct 10, 2026
vchuravy
commented
Oct 10, 2026
vchuravy
commented
Oct 10, 2026
vchuravy
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 10, 2026 13:31
757eafd to
d53946a
Compare
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 10, 2026 16:59
d53946a to
0d37f80
Compare
maleadt
pushed a commit
to JuliaGPU/CUDA.jl
that referenced
this pull request
Oct 10, 2026
…roup-ops branch Test against JuliaGPU/KernelAbstractions.jl#831. Drop this commit (going back to KernelAbstractions' main branch) once that is merged. Assisted-by: Claude Code (Opus 5.5)
maleadt
changed the title
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 10, 2026 18:31
0d37f80 to
b730e69
Compare
maleadt
pushed a commit
to JuliaGPU/CUDA.jl
that referenced
this pull request
Oct 10, 2026
…roup-ops branch Test against JuliaGPU/KernelAbstractions.jl#831. Drop this commit (going back to KernelAbstractions' main branch) once that is merged. Assisted-by: Claude Code (Opus 5.5)
maleadt
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 10, 2026
Get KernelAbstractions and KernelInterface from KernelAbstractions.jl's `vc/ki-subgroup-ops` branch instead of `main`. Julia 1.10 ignores [sources], so Buildkite checks out and develops them there. Drop this commit once JuliaGPU/KernelAbstractions.jl#831 is merged. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 10, 2026
Implement the sub-group communication of KernelInterface (JuliaGPU/KernelAbstractions.jl#831): - `shfl`, `shfl_down`, `shfl_up` and `shfl_xor` map to the SIMD-group shuffles, only for the primitive types Metal supports natively, so that other types reach KernelInterface's fallbacks (64-bit integers as `UInt32` words, structs field by field). Offsets past the SIMD-group width return the work-item's own value, also ones that don't fit in the intrinsics' `Int16`. - `sub_group_ballot` is `simd_ballot` masked to the lanes of the SIMD-group, since MSL leaves the upper bits undefined, and `sub_group_any`/`sub_group_all` are derived from it. - `get_max_sub_group_size` returns the constant `SIMD_WIDTH`, which `kernel_function` checks the pipeline's execution width against. - Sub-groups are reported as linear and independent: Metal doesn't specify the layout, but Apple GPUs form SIMD-groups from consecutive threads, which the testsuite checks. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 10, 2026
Get KernelAbstractions and KernelInterface from KernelAbstractions.jl's `vc/ki-subgroup-ops` branch instead of `main`. Julia 1.10 ignores [sources], so Buildkite checks out and develops them there. Drop this commit once JuliaGPU/KernelAbstractions.jl#831 is merged. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 10, 2026
Implement the sub-group communication of KernelInterface (JuliaGPU/KernelAbstractions.jl#831): - `shfl`, `shfl_down`, `shfl_up` and `shfl_xor` map to the SIMD-group shuffles, only for the primitive types Metal supports natively, so that other types reach KernelInterface's fallbacks (64-bit integers as `UInt32` words, structs field by field). Offsets past the SIMD-group width return the work-item's own value, also ones that don't fit in the intrinsics' `Int16`. - `sub_group_ballot` is `simd_ballot` masked to the lanes of the SIMD-group, since MSL leaves the upper bits undefined, and `sub_group_any`/`sub_group_all` are derived from it. - `get_max_sub_group_size` returns the constant `SIMD_WIDTH`, which `kernel_function` checks the pipeline's execution width against. - Sub-groups are reported as linear and independent: Metal doesn't specify the layout, but Apple GPUs form SIMD-groups from consecutive threads, which the testsuite checks. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
added a commit
to JuliaGPU/oneAPI.jl
that referenced
this pull request
Oct 10, 2026
Get KernelAbstractions and KernelInterface from JuliaGPU/KernelAbstractions.jl#831, which adds the sub-group operations implemented here. Drop this commit when it is merged.
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 11, 2026 10:01
b730e69 to
e65d2c8
Compare
Open
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 12, 2026 07:52
e65d2c8 to
5d12c31
Compare
maleadt
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 12, 2026
Implement the sub-group communication of KernelInterface (JuliaGPU/KernelAbstractions.jl#831): - `shfl`, `shfl_down`, `shfl_up` and `shfl_xor` map to the SIMD-group shuffles, only for the primitive types Metal supports natively, so that other types reach KernelInterface's fallbacks (64-bit integers as `UInt32` words, structs field by field). Offsets past the SIMD-group width return the work-item's own value, also ones that don't fit in the intrinsics' `Int16`. - `sub_group_ballot` is `simd_ballot` masked to the lanes of the SIMD-group, since MSL leaves the upper bits undefined, and `sub_group_any`/`sub_group_all` are derived from it. - `get_max_sub_group_size` returns the constant `SIMD_WIDTH`, which `kernel_function` checks the pipeline's execution width against. - Sub-groups are reported as linear and independent: Metal doesn't specify the layout, but Apple GPUs form SIMD-groups from consecutive threads, which the testsuite checks. Co-authored-by: Tim Besard <tim@juliahub.com>
maleadt
pushed a commit
to JuliaGPU/CUDA.jl
that referenced
this pull request
Oct 12, 2026
…roup-ops branch Test against JuliaGPU/KernelAbstractions.jl#831. Drop this commit (going back to KernelAbstractions' main branch) once that is merged. Assisted-by: Claude Code (Opus 5.5)
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
3 times, most recently
from
October 12, 2026 09:43
d2339f4 to
c46beae
Compare
vchuravy
commented
Oct 12, 2026
| `GPUCompiler.method_tables` instead. | ||
| - `foreach_index(f, A)` runs `f` once per index of the array `A` without writing a kernel out, | ||
| and `foreach_index(f, backend, indices)` once per index in a range or `CartesianIndices`. | ||
| - KernelInterface 0.4 adds sub-groups: ids, shuffles, votes, reductions and scans, also |
Member
Author
There was a problem hiding this comment.
Suggested change
| - KernelInterface 0.4 adds sub-groups: ids, shuffles, votes, reductions and scans, also | |
| - KernelInterface 0.4.1 adds sub-groups: ids, shuffles, votes, reductions and scans, also |
Add `shfl`, `shfl_up` and `shfl_xor` next to `shfl_down`, their variants within segments of a `width`, and the votes `sub_group_any`, `sub_group_all` and `sub_group_ballot`. Back ends implement the shuffles for their native primitive types; other primitive types are shuffled as `UInt32` or `UInt64` words and `isbits` structs field by field, with `supports_shuffle` derived accordingly. Which work-items form a sub-group stays unspecified. Back ends that form them from consecutive work-items report `supports_linear_subgroups`, and those whose sub-groups can communicate in control flow that differs between the sub-groups of a work-group report `supports_independent_subgroups`. The testsuite computes the expected results from the sub-groups that formed, so that it runs without either guarantee. The POCL back end implements the shuffles and votes with SPIRVIntrinsics, for widths of up to 64 so that the ballot covers every lane. PoCL forms sub-groups from consecutive work-items, but executes sub-group operations as work-group barriers, so it reports linear but not independent sub-groups. Co-authored-by: Tim Besard <tim@juliahub.com>
Add `sub_group_reduce`, `sub_group_scan` and `sub_group_exclusive_scan`, in the order of the lanes, for associative operators that needn't be commutative. The fallbacks are built on the shuffles. Back ends may implement them natively per operator and type, as long as the result is the documented one, except that floating-point `+` and `*` may be reassociated. The POCL back end uses PoCL's native reductions and scans where they have Julia's semantics: `+` on integers and floats, and `min`/`max` on integers. Co-authored-by: Tim Besard <tim@juliahub.com>
Structs and tuples were shuffled field by field, a shuffle per primitive
field, so that e.g. an `NTuple{2, Float16}` took two. Pack their fields
into `UInt32` words instead, with fields of less than 4 bytes sharing a
word, and shuffle the words through `shfl_word_tuple`. A backend can
implement that to exchange all words at once: on PoCL, where every
shuffle is a work-group barrier, that would help, once its CPU device
accepts vector shuffles from SPIR-V.
The generated fallbacks also weren't inlined: `@inline` on a generated
function doesn't apply to a body that the generator returns with an
explicit `return`, so shuffling an `SVector` in a larger kernel left a
call, with the value passed through memory. They now mark their bodies
for inlining.
Tiles narrower than a sub-group have to share it, and to be useful they have to be able to diverge from each other, like CUDA's `tiled_partition<N>`. Add the votes, collectives and barrier within segments of `width` lanes (the shuffles already had a `width`), and `supports_partitionable_subgroups`: whether a sub-group can be partitioned into segments that diverge independently, so that only the work-items of a segment have to execute these together. It extends `supports_independent_subgroups` to segments, and implies it. The fallbacks build on the functions without a `width`, so that they are correct wherever the whole sub-group executes them. Backends that report partitionable sub-groups have to make them work for a segment alone: CUDA needs a mask of the segment's lanes, other backends a ballot that only reports active lanes and a barrier that a segment can execute. OpenCL and Metal have no barrier for part of a sub-group, so there it takes hardware that executes a sub-group in lock-step, like Intel's and Apple's GPUs, and a fence. PoCL never reports it, as its sub-group functions are work-group barriers. The shuffles with a `width` of other types now shuffle `UInt32` words with a `width`, so that a backend's shuffles with a `width` handle them.
Every shuffle on PoCL is a work-group barrier, so a struct packed into words cost a barrier per word: tile kernels shuffling an `SVector` were several times slower than ones using local memory. Exchange the words with the vector shuffles of `SPV_INTEL_subgroups`, a barrier per up to 16 words, which PoCL's CPU device accepts with `cl_intel_subgroups` since pocl/pocl#2394. Four shuffles of an `SVector{4, Float32}` over 4M work-items take 1.7 ms instead of 4.6 ms. KernelInterface gets `shuffle_source`, the lane a shuffle reads from, for backends that exchange the words with a shuffle from a lane.
maleadt
force-pushed
the
vc/ki-subgroup-ops
branch
from
October 12, 2026 16:46
c46beae to
7615cb3
Compare
vchuravy
pushed a commit
that referenced
this pull request
Oct 12, 2026
* Fix the import conflicts warned about when precompiling on Julia 1.10 `using .POCL` brought in POCL's `@device_code_*` macros, which clash with KernelAbstractions' own, and POCL's `using LLVM.IR` clashed with its `context`. Import only the names that are used. * Reword the 0.10 release notes on KernelInterface sub-groups KernelInterface 0.3 already had the sub-group ids, `sub_group_barrier` and `shfl_down`; say what 0.4 adds, and that the layout guarantee moved behind `supports_linear_subgroups`. Also fix the `cpu=` option's deprecation version, which said 0.11, and a few typos. * Raise the GPUToolbox floor, and tag KernelInterface releases GPUToolbox 3.3.2 requires LLVM 9 or earlier, so with LLVM 10 the effective lower bound was already 3.3.3. TagBot only tagged the root package; add a job for lib/KernelInterface, whose 0.2 and 0.3 releases went untagged. * CI: test the back ends' main branches The back ends' KernelInterface 0.4 support is merged into their default branches, so stop testing the sub-group branches (added temporarily in #831) and the ka-0.10 branches before them.
maleadt
added a commit
to JuliaGPU/Metal.jl
that referenced
this pull request
Oct 13, 2026
Implement the sub-group communication of KernelInterface (JuliaGPU/KernelAbstractions.jl#831): - `shfl`, `shfl_down`, `shfl_up` and `shfl_xor` map to the SIMD-group shuffles, only for the primitive types Metal supports natively, so that other types reach KernelInterface's fallbacks (64-bit integers as `UInt32` words, structs field by field). Offsets past the SIMD-group width return the work-item's own value, also ones that don't fit in the intrinsics' `Int16`. - `sub_group_ballot` is `simd_ballot` masked to the lanes of the SIMD-group, since MSL leaves the upper bits undefined, and `sub_group_any`/`sub_group_all` are derived from it. - `get_max_sub_group_size` returns the constant `SIMD_WIDTH`, which `kernel_function` checks the pipeline's execution width against. - Sub-groups are reported as linear and independent: Metal doesn't specify the layout, but Apple GPUs form SIMD-groups from consecutive threads, which the testsuite checks. Co-authored-by: Tim Besard <tim@juliahub.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Sub-group communication for KernelInterface, for kernels written against KI (e.g. Molly's
force kernels) and for kernel languages built on it.
Contract
is guaranteed:
supports_linear_subgroups(backend): 1-D work-groups, or ones whose x extent is amultiple of
W, form sub-groups from consecutive work-items, x fastest. CUDA, AMDGPU,Metal and PoCL report it; OpenCL runtimes generally don't promise it.
supports_independent_subgroups(backend): communication in control flow that is uniformper sub-group but not per work-group. PoCL reports
false(its sub-group operations arework-group barriers).
supports_partitionable_subgroups(backend): a sub-group can be partitioned intosegments that diverge independently: the functions with a
widthonly need the caller'ssegment of
widthlanes (narrow tiles sharing a sub-group). Impliessupports_independent_subgroups. CUDA (segment masks), AMDGPU, Intel GPUs and Metal report it. Metal and OpenCLhave no barrier for part of a sub-group, so on Intel and Apple GPUs the segment barrier is
a sub-group-scope fence that relies on the hardware executing a sub-group in lock-step
(validated on M1 and M3 with a stress test of diverging segments). PoCL never reports it:
its sub-group operations are work-group barriers.
KI 0.3's "a 1-D work-group of at most
Wwork-items is one sub-group" moves undersupports_linear_subgroups.sub_group_barrierand the votes undersupports_subgroups; shufflesper type under
supports_shuffle; the collectives need their type's shuffles or a nativeimplementation.
get_max_sub_group_size()SHOULD be a constant of the generated code; it is never knownduring inference (pass
sub_group_size(backend)from the host).Device functions
shfl,shfl_down,shfl_up,shfl_xorUInt32; past the widthshfl_down/shfl_upreturn the caller's value;shfl_xorneedsmask < Wshfl(val, x, width)etc.shfl; CUDA's width semantics; only the segment participates withsupports_partitionable_subgroupssub_group_ballot/any/all(pred, width),sub_group_reduce/scan/exclusive_scan(..., width),sub_group_barrier(width)width == W), the ballot and the width shuffles; back ends with partitionable sub-groups make them work for a segment alonesub_group_any/all,sub_group_balloti-1iff laneiexists and its predicate holdssub_group_reduce,sub_group_scan,sub_group_exclusive_scan(op, val, init)op; native overrides must give the documented result, except float+/*may be reassociatedOther primitive types are shuffled as
UInt32/UInt64words,isbitsstructs packed intoUInt32words (small fields share one), through an internal hook a back end can implement toexchange several words at once.
supports_shufflereports transport, not arithmetic.Not included (no consumer yet, can be added without breaking anything):
sub_group_match_any,sub_group_broadcast.POCL
Shuffles and votes through SPIRVIntrinsics, native
+reductions/scans (and integermin/max), widths ≤ 64.supports_linear_subgroups = true,supports_independent_subgroups = supports_partitionable_subgroups = false. Packed structs areexchanged with the vector shuffles of
SPV_INTEL_subgroups, a work-group barrier per up to 16words instead of per word (4.6 → 1.7 ms for four
SVector{4, Float32}shuffles over 4Mwork-items). That needs pocl/pocl#2394 in pocl_jll.
Documentation
The KernelInterface manual's Sub-groups section describes the contract and the per-back-end capabilities, Implementing a backend says what a back end implements and when it may claim each capability, and the 0.10 release notes have an entry.
Testing
The KI testsuite records each work-item's sub-group and lane and computes expected results
per observed sub-group, so it runs on back ends without the linear layout; layout and
independence tests are gated on the queries. Float64 is only moved through
UInt64storagewhere the device lacks fp64.
The last commit is
[TEMP]: Buildkite tests the back ends'vc/ki-subgroup-ops/tb/ki-subgroup-opsbranches (JuliaGPU/CUDA.jl#3331, JuliaGPU/Metal.jl#998,JuliaGPU/AMDGPU.jl#1133, JuliaGPU/OpenCL.jl#554, JuliaGPU/oneAPI.jl#680).