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 fallback, as packed UInt32 words.
• 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.
• KI.supports_linear_subgroups and KI.supports_independent_subgroups are true: wavefronts are formed from consecutive linear work-item ids, x fastest (KernelInterface's layout tests check this), and are independent. KernelInterface no longer implies either from supports_subgroups.
• 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.
• KI.supports_partitionable_subgroups is true: the functions with a width only need the lanes of the caller's segment. ds_bpermute only involves the lanes it reads from, llvm.amdgcn.ballot only reports the lanes that execute it, and the sub-group barrier is wavefront-scope fences around llvm.amdgcn.wave.barrier, which HIP's tile synchronization also uses in divergent code. KernelAbstractions can then pack several narrow tiles into a wavefront.

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 points KernelAbstractions and KernelInterface at vc/ki-subgroup-ops. SPIRVIntrinsics 1.4 is registered, so the earlier pin to OpenCL.jl's vc/subgroup-votes is gone. Drop the commit once KernelAbstractions.jl#831 is merged.

Local testing

gfx1036 iGPU (wave32), Julia 1.12, on ka-0.10 rebased onto main, KernelAbstractions at the current #831: KernelInterface testsuite 6089/6089, KernelAbstractions testsuite 2639 passed, 4 broken, in three runs without hangs. Offsets up to typemax(Int64) give the own value past the wavefront.