KernelAbstractions backend: warp-level primitives on KernelInterface sub-groups

[vchuravy]

Draft: this depends on unreleased KernelAbstractions changes, JuliaGPU/KernelAbstractions.jl#830 and #831, which add portable sub-group (warp) operations to KernelAbstractions 0.10 through its KernelInterface. Until those are released, [sources] points at their branch, and at JuliaGPU/OpenCL.jl#526 for the CPU backend's SPIR-V intrinsics.

So far, ParallelStencil's warp macros threw "not supported" on the KernelAbstractions backend. This implements them on KernelInterface:

macro	KernelInterface
@warpsize()	KI.get_max_sub_group_size() (a compile-time constant)
@laneid()	KI.get_sub_group_local_id() (1-based, as in ParallelStencil)
@active_mask()	the mask of the lanes that exist in the sub-group
@shfl_sync(m, v, lane[, w])	KI.shfl(v, lane, w); without w, the sub-group width
@shfl_up_sync / @shfl_down_sync / @shfl_xor_sync	KI.shfl_up / KI.shfl_down / KI.shfl_xor with a width (CUDA's semantics: lanes without a partner keep their own value)
@vote_any_sync / @vote_all_sync / @vote_ballot_sync	KI.sub_group_any / KI.sub_group_all / KI.sub_group_ballot (UInt64, as on the AMDGPU and CPU backends)

Semantics:

• Masks: KernelInterface has no partial participation. All lanes of the sub-group have to take part, so masks are ignored at run time, and a literal partial mask (e.g. 0x0000ffff) is rejected when the macro expands.
• @active_mask: reports the lanes that exist in the sub-group, not the lanes active in a divergent branch.
• Padding lanes: ParallelStencil launches whole workgroups and checks the ranges itself, so KernelAbstractions' padding work-items aren't involved. As on CUDA, warp operations must only be used where the whole warp is inside the ranges.
• KernelAbstractions 0.9: the warp macros throw an error saying KernelAbstractions 0.10 is needed.

Also in this PR:

• @sharedMem with KernelAbstractions 0.10: KernelAbstractions' @localmem interpolates a module object, which breaks @kernel's deepcopy of the pre-expanded kernel. The fix emits KernelInterface.localmemory by name. The @sharedMem tests now also run on the KernelAbstractions backend.
• Bug fix: @ArgumentError in the code generated by @parallel with more than one positional argument wasn't qualified, giving an UndefVarError in the caller.

Testing (Julia 1.12, KernelAbstractions' CPU backend (POCL), 4 threads):

• ParallelKernel/test_kernel_language.jl passes 67/67 on KernelAbstractions. That includes every primitive with and without a width, a tuple shuffle, a butterfly reduction and the ballot.
• The other KernelAbstractions-backend test files pass, e.g. test_parallel 113/113, FiniteDifferences3D 108/108, test_allocators 428/428.
• With KernelAbstractions 0.9, the new tests are skipped.
• A diffusion stencil, the shared-memory example and a warp reduction in 2-D blocks match host references on the CPU backend and on CUDA (Quadro RTX 4000) through KernelAbstractions.

Known issues:

• Pkg.test() doesn't resolve with KernelAbstractions 0.10 yet: the test target includes Enzyme, which is incompatible with GPUCompiler 2.10 that KernelAbstractions 0.10 needs. I ran the test files with a driver that leaves Enzyme out.
• Existing differences found while testing, not changed here:
  • The CPU and AMDGPU backends subtract 1 from the lane mask of @shfl_xor_sync, unlike CUDA and the docs.
  • The AMDGPU @laneid uses activelane(), which is wrong under divergence.

🤖 Generated with Claude Code

[omlins]

@vchuravy : Thanks for opening the PR! Looking at the CI. the issue with Julia1.11 appears to be compatibility requirements with GPUToolbox visible in the last line of the output here ("KernelAbstractions [63c18a36] is fixed to version 0.10.0-dev"):

   ├─restricted by compatibility requirements with GPUToolbox [096a3bc2] to versions: [0.1.0 - 1.2.3, 2.6.0 - 2.8.0] or uninstalled, leaving only versions: [1.0.0 - 1.2.3, 2.6.0 - 2.8.0]
   │ └─GPUToolbox [096a3bc2] log:
   │   ├─possible versions are: 0.1.0 - 3.3.3 or uninstalled
   │   ├─restricted to versions 3.3.2 - 3 by KernelAbstractions [63c18a36], leaving only versions: 3.3.2 - 3.3.3
   │   │ └─KernelAbstractions [63c18a36] log:
   │   │   ├─possible versions are: 0.10.0 or uninstalled
   │   │   ├─restricted to versions 0.9 - 0.10 by project [2e668c77], leaving only versions: 0.10.0
   │   │   │ └─project [2e668c77] log: see above
   │   │   └─KernelAbstractions [63c18a36] is fixed to version 0.10.0-dev