Skip to content

KernelInterface: sub-group shuffles, votes and a constant width - #3331

Closed
vchuravy wants to merge 11 commits into
ka-0.10from
vc/ki-subgroup-ops
Closed

vchuravy wants to merge 11 commits into
ka-0.10from
vc/ki-subgroup-ops

Conversation

@vchuravy

@vchuravy vchuravy commented Oct 4, 2026 •

Copy link
Copy Markdown
Member

Implements KernelInterface's sub-group contract from JuliaGPU/KernelAbstractions.jl#831 in
the CUDA back-end (CUDACore/src/CUDAKernels.jl). Stacked on #3314 (ka-0.10).

  • Shuffles KI.shfl, shfl_down, shfl_up, shfl_xor and their width variants, each
    a single shfl.sync with the full warp mask. They and KI.supports_shuffle only match the
    primitive types CUDA.jl's warp shuffles handle (Bool, Base.BitInteger,
    Base.IEEEFloat); KernelInterface shuffles other isbits types (Complex, structs,
    tuples) as packed UInt32 words. shfl.sync only uses the low 5 bits of an offset, so offsets of
    32 and more become 0, which gives the thread's own value as the contract requires. Offsets
    are compared before they are narrowed, so wide ones (2^32 + 1, 2^64 + 1, a regression
    test) do too.
  • Votes KI.sub_group_any, sub_group_all, sub_group_ballot (bit i - 1 for lane i).
  • Capabilities: KI.supports_linear_subgroups and KI.supports_independent_subgroups
    are true: warps are formed from consecutive linear thread indices, and are independent.
  • Constant width: the warp size is the constant 32 instead of a read of %WARP_SZ, so
    get_max_sub_group_size is a constant of the generated code; the other queries use
    unsigned 32-bit shifts.
  • redux.sync: on sm_80+, KI.sub_group_reduce of + min max & | ⊻ on Int32/UInt32
    is a single redux.sync over the full warp mask (not activemask(), which threads that
    aren't converged could read differently). The operators are commutative, so the lane order the contract
    specifies doesn't matter. Everything else uses KernelInterface's ordered fallback.
  • Partitionable sub-groups: KI.supports_partitionable_subgroups is true. The functions with a width
    pass the member mask of the caller's segment to the .sync intrinsics instead of the full
    mask: the width shuffles, the votes (masking other segments out of the ballot and shifting
    it to the segment), the redux.sync reductions, and KI.sub_group_barrier(width), which
    becomes bar.warp.sync on the segment. A width of 32 still compiles to the full mask. This
    lets KernelAbstractions put several narrow tiles in a warp: a DG-like kernel with tiles of
    16 runs 2.4x faster with 16 tiles per block than with a block per tile.

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

Tested locally (RTX 5080): KernelInterface testsuite 6091/6091, KernelAbstractions testsuite 3103 pass,
17 broken.

@codecov

codecov Bot commented Oct 10, 2026 •

Copy link
Copy Markdown

Codecov Report

❌ Patch coverage is 90.47619% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 38.34%. Comparing base (34e43e0) to head (41b9f97).

Files with missing lines Patch % Lines
CUDACore/src/CUDAKernels.jl 90.47% 4 Missing ⚠️
Additional details and impacted files
@@             Coverage Diff             @@
##           ka-0.10    #3331      +/-   ##
===========================================
+ Coverage    37.01%   38.34%   +1.33%     
===========================================
  Files          191      191              
  Lines        19296    18698     -598     
===========================================
+ Hits          7142     7170      +28     
+ Misses       12154    11528     -626     

☔ 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 force-pushed the vc/ki-subgroup-ops branch from 613fd94 to cff88d4 Compare October 10, 2026 18:44
maleadt and others added 10 commits October 12, 2026 09:30
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 CUDA's copy of
that launch path: partitioning the ndrange, building the kernel's context, and
tuning the workgroup size.

`CUDABackend` now implements KernelInterface in CUDACore, which 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 the `maxthreads` hint for a static workgroup size.
`prefer_blocks` now applies in `KI.launch_configuration`, which receives the
number of work-items.

`KI.launch` passes the kernel arguments on as a tuple, so kernels with many
arguments stay cheap to launch, as #3309 made them for the old launch path. It
rejects `threads` and `blocks`, which would override the launch geometry that
KernelInterface validated. `KI.copyto!` accepts dense arrays and contiguous
views of them, as the old `KA.copyto!` did.

KernelAbstractions converts the arguments twice, to determine the types to
compile for and again when launching, where CUDA's launch converted them once.
`cudaconvert` has to be pure, so this only shows with conversions that have
side effects, as in two tests that count them.

Co-authored-by: Christian Guinard <28689358+christiangnrd@users.noreply.github.com>
Sub-groups are warps: implement the sub-group queries, `sub_group_barrier` and
`shfl_down`, and report their support to KernelInterface. KernelInterface
leaves unspecified how work-items are grouped into sub-groups; CUDA forms warps
from consecutive linear thread indices, so the last warp of a block can be
partial, which a test checks.

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

KernelInterface's testsuite checks them once `record_event` returns an event.
`KI.versioninfo(CUDABackend())` prints `CUDA.versioninfo()`.
… main branch

Neither is registered yet. Julia 1.10 and 1.11 don't pick up the test
project's [sources], so Buildkite develops both explicitly there, and the
GPU-less and Enzyme jobs move to Julia 1.12.

Drop this commit once both are registered.
KernelAbstractions now backs @Private with a stack-allocated PrivateArray,
provided as an overlay in GPUCompiler's shared method table, which CUDA's
method table already stacks on top of. Drop our MArray-based Scratchpad
override so that it takes effect, and require GPUCompiler 2.12.

That leaves StaticArrays only for the dimension_mismatch_fail device
override, so make it a weak dependency with a regular extension.
KI.max_work_group_size and KI.launch_configuration are specialized on the
kernel, so their bodies were compiled for every kernel launched through
KernelAbstractions. Forward them to functions of the CuFunction, which are
compiled once.
…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)
Implement the sub-group communication of KernelInterface's updated contract: `shfl`,
`shfl_up` and `shfl_xor` next to `shfl_down`, their variants within segments of a `width`
(a single `shfl.sync` each), and the votes `sub_group_any`, `sub_group_all` and
`sub_group_ballot`, using the warp intrinsics with the full mask. Warps are formed from
consecutive linear thread indices and are independent, so report
`supports_linear_subgroups` and `supports_independent_subgroups`.

The shuffles, and `supports_shuffle`, are restricted to the primitive types CUDA's warp
shuffles support, so that KernelInterface shuffles structs and tuples (including `Complex`)
field by field. `shfl.sync` only uses the low 5 bits of an offset, so offsets past the warp
(including ones that don't fit in 32 bits) become 0, giving the thread's own value as
KernelInterface requires.

The warp size is the constant 32 rather than a read of `%WARP_SZ`, so that
`get_max_sub_group_size` is a constant of the generated code, and the sub-group queries are
computed with unsigned 32-bit shifts.

Co-authored-by: Tim Besard <tim@juliahub.com>
On sm_80 and later, implement `sub_group_reduce` of `+`, `min`, `max`, `&`, `|` and `⊻` on
32-bit integers with a single `redux.sync`. These operators are commutative, so the order of
the lanes, which KernelInterface's reduction follows, doesn't change the result. Other
operators, types and devices use KernelInterface's fallback.

Co-authored-by: Tim Besard <tim@juliahub.com>
@maleadt
maleadt force-pushed the vc/ki-subgroup-ops branch from 6b496b8 to b2dcea4 Compare October 12, 2026 07:52
Report `supports_partitionable_subgroups`: the functions with a `width` now pass the member
mask of the caller's segment to the `.sync` warp intrinsics instead of the full mask, so that
only the threads of that segment have to execute them together. This covers the width shuffles,
the votes (masking other segments out of the ballot and shifting it to the segment), the
32-bit integer reductions with `redux.sync`, and the width `sub_group_barrier`, which becomes
`bar.warp.sync` on the segment. Other reductions and scans use KernelInterface's fallbacks on
top of the width shuffles.

The width is a constant, so a width of 32 still compiles to the full mask. This lets
KernelAbstractions put several narrow tiles in one warp: a DG-like kernel with tiles of 16
runs 2.4x faster with 16 tiles per block than with a block per tile.
@maleadt
maleadt force-pushed the vc/ki-subgroup-ops branch from b2dcea4 to 41b9f97 Compare October 12, 2026 09:08
@maleadt maleadt closed this Oct 13, 2026
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.

3 participants