SPIRVIntrinsics: sub-group votes and collectives, unchecked shuffle lanes

[vchuravy]

Upstreams two workarounds from KernelAbstractions' POCL back-end (JuliaGPU/KernelAbstractions.jl#831, see this review thread), where they implement KernelInterface's sub-group operations.

Votes

New sub_group_any(::Bool) and sub_group_all(::Bool) (cl_khr_subgroups), and sub_group_ballot(::Bool)::NTuple{4, VecElement{UInt32}} (cl_khr_subgroup_ballot, OpenCL's uint4 mask).

These call the OpenCL C built-ins, which the SPIR-V back-end lowers, rather than __spirv_GroupNonUniformAny/All/Ballot: @builtin_ccall can't mangle their bool arguments. Defining them here also registers the names in known_intrinsics properly. From another package (KA's POCL back-end), @builtin_ccall registers them at macro-expansion time into this package's array, and that registration doesn't survive precompilation.

Shuffle lanes

sub_group_shuffle converted the lane with UInt32(i - 1), which throws an InexactError (and leaves an error branch in the kernel) for an out-of-range lane, e.g. when implementing a shuffle-up as sub_group_shuffle(x, lane - offset). The lane is now passed modulo UInt32, and so is the mask of sub_group_shuffle_xor. An out-of-range lane gives an undefined value, as with the OpenCL C built-ins.

Collectives

The cl_khr_subgroups collectives, which the SPIR-V back-end lowers to OpGroup* instructions, for Int32, UInt32, Int64, UInt64, Float16, Float32 and Float64:

• sub_group_reduce_{add,min,max}
• sub_group_scan_inclusive_{add,min,max}
• sub_group_scan_exclusive_{add,min,max}
• sub_group_broadcast(x, lane), with a 1-based lane like sub_group_shuffle

KernelAbstractions uses these to implement KernelInterface.sub_group_reduce/sub_group_scan natively.

The non-uniform arithmetic of cl_khr_subgroup_non_uniform_arithmetic (mul, and, or, xor, …) and the char/short types of cl_khr_subgroup_extended_types aren't included, because PoCL doesn't support them, so they couldn't be tested here.

Convergence

All collectives (shuffles, votes, ballot, reductions, scans, broadcast) are now declared convergent. @builtin_ccall gets a convergent = true option for this. Without it, LLVM was free to make the calls control-dependent on more values. For r = sub_group_reduce_add(i <= m ? x[i] : 0), the optimizer sank the call into both arms of the branch, so the work-items of a sub-group executed two different calls, which is undefined behavior. On PoCL, which implements the collectives with work-group barriers, this returned 0.

A bounds-checked variant still fails on PoCL 7.2. There, the (never taken) early exit of the bounds check makes WorkitemLoops give the peeled first work-item its own scratch memory. PoCL fixed this on main in pocl/pocl@8fa5d2732 (#2239); Julia's PoCL builds carry it as a patch with JuliaPackaging/Yggdrasil#15001. The test marks that variant as broken on PoCL 7.2.

Tests

• New tests in test/intrinsics.jl:
  • a shuffle from an out-of-range lane
  • any/all for several predicate patterns
  • ballot, when the device has cl_khr_subgroup_ballot
  • the reductions, scans and broadcast
• Local results: pocl/intrinsics and intel/intrinsics pass (2238 tests). Both devices support sub-groups and ballot.

KernelAbstractions can drop its copies once this is released, so a version bump of SPIRVIntrinsics (1.3.0, new features) would be needed.

🤖 Generated with Claude Code