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

[vchuravy]

Implements the sub-group contract of KernelInterface from JuliaGPU/KernelAbstractions.jl#831 for ROCBackend: the shuffles shfl, shfl_down, shfl_up and shfl_xor, the votes sub_group_any, sub_group_all and sub_group_ballot, and a get_max_sub_group_size that is a constant of the generated code.

What's implemented

KernelInterface	AMDGPU
shfl(val, lane)	ds_bpermute (Device.bpermute) from hardware lane lane - 1
shfl_down(val, offset)	ds_bpermute from hardware lane mbcnt + offset, or mbcnt past the wavefront
shfl_up(val, offset)	ds_bpermute from hardware lane mbcnt - offset, or mbcnt if that is negative
shfl_xor(val, mask)	ds_bpermute from hardware lane mbcnt ⊻ mask
sub_group_ballot(pred)	Device.ballot(pred) (llvm.amdgcn.ballot.i32/.i64, widened to UInt64)
sub_group_any(pred)	Device.ballot(pred) != 0
sub_group_all(pred)	Device.ballot(!pred) == 0
get_max_sub_group_size()	Device.wavefrontsize(), folded to a constant (unchanged)

• The shuffles use the hardware lane (mbcnt), like get_sub_group_local_id, not activelane. They don't use Device.shfl*, which clamp out-of-range offsets to the own lane and work on activelane. ds_bpermute only uses the low bits of the address, so a lane out of range reads an unspecified value rather than trapping. Lanes and offsets are truncated with % Cint, so there is no checked conversion: the IR of a kernel using all shuffles and votes with runtime Int lanes/offsets has no throw/trap/unreachable.
• The shuffles and supports_shuffle are now only defined for const ShuffleTypes = Union{Bool, Base.BitInteger, Base.IEEEFloat} (what Device._shfl decomposes into 32-bit bpermutes). The previous generic where {T} shfl_down override and the supports_shuffle that listed Complex are gone, so Complex and other isbits structs/tuples go through KernelInterface's field-by-field fallback.
• shfl_down/shfl_up return the work-item's own value where the source lane is past the wavefront (as KernelInterface now requires, like CUDA): one compare and select against the constant wavefront size before the ds_bpermute. shfl_xor needs no change, as KernelInterface requires mask below the width.
• The shuffles with a width (KI.shfl(val, lane, width) etc.) use KernelInterface's fallbacks: one ds_bpermute from a lane computed with a few integer operations, the same as Device.shfl etc. DPP/ds_swizzle could be cheaper for constant offsets, but isn't used here.
• Sub-groups are formed from consecutive linear work-item ids, x fastest, as KernelInterface now guarantees (its new layout tests pass).
• The votes rely on ballot setting only the bits of active lanes, i.e. of the work-items of the (possibly partial) sub-group. The result is uniform.

Constant width

This needed no new code. KI.kernel_function compiles for the device's wavefront size (and rejects a conflicting wavefrontsize64), and fold_wavefrontsize! in finish_module! already replaces llvm.amdgcn.wavefrontsize with the compiled-for size before optimization. I checked this with @device_code_llvm on a kernel that stores KI.get_max_sub_group_size(): it compiles to store i64 32 with wavefrontsize64=false and store i64 64 with wavefrontsize64=true, and the call to llvm.amdgcn.wavefrontsize is gone. I added a comment saying so.

[TEMP] commit

[TEMP] Test against KernelAbstractions' vc/ki-subgroup-ops branch sits on top of the existing [TEMP] commit. It:

• points KernelAbstractions and KernelInterface at vc/ki-subgroup-ops
• takes SPIRVIntrinsics from OpenCL.jl vc/subgroup-votes (lib/intrinsics, SPIRVIntrinsics: sub-group votes and collectives, unchecked shuffle lanes OpenCL.jl#526). KA gets it through its own [sources], which dependents don't pick up, so it is added to [extras]/[sources] here and developed on Julia 1.10 in the pipeline.
• allows AcceleratedKernels 0.5, which AK's main now is.

Drop it once #831 is merged.

Note: KernelAbstractions main, and so #831, require LLVM.jl 10 and GPUCompiler 2.10, but ka-0.10 is still on LLVM.jl 9. This branch (like ka-0.10 against KA main) won't resolve until ka-0.10 picks up the LLVM.jl 10 port in #1132.

Local testing

Update (after the shuffle/layout contract changes in JuliaGPU/KernelAbstractions.jl#831): KernelInterface testsuite: all pass but 4: sub_group_reduce/sub_group_scan of max on Float32 with a NaN (32 and 29 work-items). That is not this PR: AMDGPU's device override of Base.max for floats lowers to llvm.maxnum, which drops NaNs, unlike Julia's max. kernelabstractions_tests: 2564 passed, 4 broken. wave32 and wave64 compile; shfl_down/shfl_up are a v_cmp/v_cndmask and a ds_bpermute each.

Earlier results:

Tested on an AMD Radeon RX 6800 XT (gfx1030, wave32), ROCm in /opt/rocm, Julia 1.12.7. To get past the LLVM.jl 10 constraint, I used a local-only branch (not pushed): this branch merged with origin/main, with the 5 commits of tb/llvm10 (#1132) cherry-picked on top. The conflicts were resolved in favour of ka-0.10's ROCKernels.jl, local-memory initializer and wavefront fence.

• julia --project=test test/runtests.jl --jobs=2 kernelinterface kernelabstractions: 3152 passed, 4 broken, 0 failed (kernelinterface_tests and kernelabstractions_tests).
• KernelInterface testsuite on its own (Testsuite.testsuite(ROCBackend(), ROCArray)): 569/569 passed. This includes the new shuffle tests (lanes, out-of-range, shfl_up/shfl_xor, structs), votes and the width.
• @device_code_llvm/@device_code_gcn with wavefrontsize64=false and true: everything compiles, and the width is constant (see above). Only wave32 could be run on this GPU.

Not tested locally: Julia 1.10/1.11, and running in wave64 mode.

🤖 Generated with Claude Code