Implement KernelInterface, and support KernelAbstractions 0.10

[maleadt]

KernelAbstractions 0.10 moves its back-end API into a separate package, KernelInterface. A back end implements KernelInterface: allocating and copying memory, selecting devices, compiling and launching kernels, and the device-side intrinsics. KernelAbstractions then launches @kernel kernels the same way on every back end (JuliaGPU/KernelAbstractions.jl#801). This PR ports AMDGPU.jl to that model in one step, as JuliaGPU/CUDA.jl#3314 does for CUDA. It supersedes #1047, which added KernelInterface next to the KernelAbstractions 0.9 back end, and Christian's intrinsics branch, which ported an earlier version of it to KernelAbstractions 0.10.

ROCBackend now implements KernelInterface, and AMDGPU depends on KernelInterface instead of KernelAbstractions. That makes the KernelInterface layer usable on its own, without KernelAbstractions' macros:

using AMDGPU
import KernelInterface as KI

function vadd(c, a, b)
    i = KI.get_global_id().x
    if i <= length(c)
        @inbounds c[i] = a[i] + b[i]
    end
    return
end

a = AMDGPU.rand(1000); b = AMDGPU.rand(1000); c = similar(a)
KI.@launch ROCBackend() ndrange=length(c) vadd(c, a, b)

KernelAbstractions becomes a weak dependency. Its extension only provides what KernelAbstractions still needs from a back end: the MArray behind @private, the Adapt rule behind @Const (which still reads through the constant address space), and the maxthreads hint for kernels with a static workgroup size. AMDGPU's own copy of the KernelAbstractions launch path goes away: partitioning the ndrange, building the kernel's context, tuning the workgroup size with the occupancy API, and calling the kernel. In practice KernelAbstractions is still always loaded, because AcceleratedKernels depends on it, so the extension is about where the code lives rather than about load time.

Kernels behave the same, but they are launched differently. KernelAbstractions now launches them on a 3-D grid, and computes @index in 32 bits when the iteration space fits (JuliaGPU/KernelAbstractions.jl#797). So a kernel over a 3-D ndrange doesn't need integer divisions to compute its index anymore, which matter more on AMD GPUs than elsewhere because they have no hardware integer division. The test suite checks that for the generated LLVM IR. I have no AMD GPU to measure the effect on, so unlike #3314 there are no timings here.

Some smaller changes came along:

• Each call site of KI.localmemory (and thus @localmem) needs its own memory, and KernelInterface doesn't pass an id to tell them apart. alloc_special now defines local memory with internal linkage, so that allocations don't alias because they share a name.
• AMDGPU.sync_wavefront() is a new barrier for the lanes of a wavefront, fenced at wavefront scope. It backs KI.sub_group_barrier.
• HIP.max_workgroup_dims(dev) reports the per-dimension workgroup limits, cached in the device, since KernelInterface queries them on every launch whose workgroup size it picks.
• shfl_down of a Complex{Int32} compiles now.

KI.launch rejects groupsize and gridsize, which would override the launch geometry KernelInterface has validated, but passes AMDGPU's other launch options on:

kernel = KI.@launch ROCBackend() launch=false vadd(c, a, b)
kernel(c, a, b; ndrange=length(c), stream=AMDGPU.stream())   # fine
kernel(c, a, b; ndrange=length(c), groupsize=256)            # ArgumentError

KI.kernel_function receives the kernel's callable unconverted. It compiles the converted form, which only holds pointers to the arrays a closure captures, so the returned kernel keeps the original callable alive. Every launch converts it again for the launch's stream, like the arguments, so that the arrays it captures are handed to that stream the same way.

KernelInterface 0.4 passes the kernel arguments to the back end as one tuple (JuliaGPU/KernelAbstractions.jl#811), so that kernels with many arguments stay cheap to launch. A HIPKernel can't be launched with a tuple yet, so KI.launch splats it for now, and launching a kernel with more than about 32 arguments still goes through dynamic dispatch, as it does on master. #1115 adds the tuple launch; once it is in, KI.launch should forward the tuple to it.

The first two commits prepare local memory and the workgroup limits, and the third is the port. The next ones implement parts of KernelInterface that are optional: sync_wavefront and sub-groups (wavefronts, including a partial last wavefront and lane ids that don't change in divergent code), ordering work across tasks with HIP events for KernelAbstractions.@spawn, and KI.versioninfo.

What changes for users: ROCBackend is a KernelInterface.Backend rather than a KernelAbstractions.GPU, which no longer exists. Code that uses KernelAbstractions.get_backend, allocate, synchronize and so on keeps working, since KernelAbstractions re-exports them from KernelInterface. @print in a kernel still does nothing on AMDGPU.

I could not run any of this on AMD hardware. Locally, without a GPU, I checked that AMDGPU and its KernelAbstractions extension precompile and load on Julia 1.10, 1.11, 1.12 and 1.13, that ROCBackend has methods for everything KernelInterface requires, and, compiling for gfx1030 and gfx90a through GPUCompiler, that KernelAbstractions kernels using local memory, @private, @Const, atomics and @print, the sub-group queries, and shfl_down for every type KI.supports_shuffle claims all compile. Whether they run correctly depends on AMDGPU's CI.

This needs a breaking release, since it requires KernelAbstractions 0.10, which isn't registered yet, as well as an AcceleratedKernels release that allows it (the registered 0.4.3 caps KernelAbstractions at 0.9). Until then, the last commit takes KernelAbstractions and KernelInterface from their development branch and AcceleratedKernels from its main branch, through [sources]. Julia 1.10 doesn't support [sources], and without workspaces Julia 1.11 only picks up KernelAbstractions from the test project's, so Buildkite runs the tests in the test project on those versions, and the Enzyme and GPU-less jobs move to Julia 1.12. That commit is dropped before merging, and the AcceleratedKernels compat bound should then require the release that supports KernelAbstractions 0.10.

[gbaraldi]

Tested on MI300A + MI250 (ROCm 7.2.4):

1. CI failure. KA.Scratchpad builds a StaticArraysCore.MArray(undef), whose constructor and indexing live in StaticArrays, and KA 0.10 no longer loads that. GPUArrays' matmul (@private) then fails with InvalidIRError (jl_f_throw_methoderror): 524 failures and 67 errors across 9 files on both GPUs, while main passes. Wrapping GPUCompiler.alloca(T, Val(prod(dims)), Val(AS.Private)) in a ROCDeviceArray (like POCL) fixes it; depending on StaticArrays would too.
2. Device switch. Calling a KI.Kernel after device! to another GPU segfaults in hipFuncGetAttribute (via KI.max_work_group_size(kernel)). Store the device in ROCKernel and re-resolve or throw.
3. [sources]: AcceleratedKernels main is now 0.5.0, outside compat, so instantiate fails. AK 0.5 also makes backend a kwarg, which breaks accumulate/cumsum/cumprod in src/kernels/accumulate.jl.
4. Nit: each launch does 2–3 uncached attribute queries (~0.4–0.5 µs each); they could be cached like max_workgroup_dims.

Perf on MI300A, main → PR: async launch 11.8 → 4.4 µs, launch+sync 28.1 → 17.8 µs, 3-D scale 0.334 → 0.222 ms. div/rem are gone from 3-D indexing (ISA 188 → 36 instructions).

[github-actions]

AMDGPU.jl Benchmarks

Details

Benchmark suite	Current: 0a11881	Previous: 67b2a7d	Ratio
amdgpu/synchronization/context/device	1777.5 ns	580 ns	3.06
amdgpu/synchronization/stream/blocking	230 ns	260 ns	0.88
amdgpu/synchronization/stream/nonblocking	355 ns	350 ns	1.01
applications/bitonic_sort	1147917.75 ns
applications/convolution	73168.5 ns
applications/floyd_warshall	9068925 ns
applications/histogram	805095.75 ns
applications/prefix_sum	192682.5 ns
array/accumulate/Float32/1d	100741.25 ns	84511 ns	1.19
array/accumulate/Float32/dims=1	280233.75 ns	249484 ns	1.12
array/accumulate/Float32/dims=1L	109804 ns	133142 ns	0.82
array/accumulate/Float32/dims=2	104211.5 ns	105042 ns	0.99
array/accumulate/Float32/dims=2L	2182838.75 ns	3113817 ns	0.70
array/accumulate/Int64/1d	117076.5 ns	99832 ns	1.17
array/accumulate/Int64/dims=1	253358.25 ns	246424 ns	1.03
array/accumulate/Int64/dims=1L	114309 ns	161882 ns	0.71
array/accumulate/Int64/dims=2	106686.5 ns	132652 ns	0.80
array/accumulate/Int64/dims=2L	2522810.75 ns	3327530 ns	0.76
array/broadcast	57470.75 ns	144662 ns	0.40
array/construct	2192.5 ns	1690 ns	1.30
array/copy	39855.5 ns	39090 ns	1.02
array/copyto!/cpu_to_gpu	85836 ns	115172 ns	0.75
array/copyto!/gpu_to_cpu	85953.5 ns	115222 ns	0.75
array/copyto!/gpu_to_gpu	59063.25 ns	43831 ns	1.35
array/iteration/findall/bool	96433.75 ns	174192 ns	0.55
array/iteration/findall/int	106759 ns	181862 ns	0.59
array/iteration/findfirst/bool	135414.5 ns	173572 ns	0.78
array/iteration/findfirst/int	135214.25 ns	153812 ns	0.88
array/iteration/findmin/1d	106508.75 ns	169412 ns	0.63
array/iteration/findmin/2d	94158.75 ns	155622 ns	0.61
array/iteration/logical	96866.25 ns	317275 ns	0.31
array/iteration/scalar	306299 ns	296015 ns	1.03
array/permutedims/2d	49883 ns	62141 ns	0.80
array/permutedims/3d	59275.75 ns	72921 ns	0.81
array/permutedims/4d	51005.75 ns	80191 ns	0.64
array/random/rand/Float32	46518 ns	49610 ns	0.94
array/random/rand/Int64	46258.25 ns	57351 ns	0.81
array/random/rand!/Float32	51143 ns	144202 ns	0.35
array/random/rand!/Int64	63610.75 ns	144743 ns	0.44
array/random/randn/Float32	70778.25 ns	78391 ns	0.90
array/random/randn!/Float32	84478.5 ns	102532 ns	0.82
array/reductions/mapreduce/Float32/1d	80328.5 ns	132932 ns	0.60
array/reductions/mapreduce/Float32/dims=1	61701 ns	95342 ns	0.65
array/reductions/mapreduce/Float32/dims=1L	69585.75 ns	772841 ns	0.09003889545197524
array/reductions/mapreduce/Float32/dims=2	70231 ns	96711 ns	0.73
array/reductions/mapreduce/Float32/dims=2L	82586 ns	296274 ns	0.28
array/reductions/mapreduce/Int64/1d	70846 ns	133642 ns	0.53
array/reductions/mapreduce/Int64/dims=1	61878.25 ns	77651 ns	0.80
array/reductions/mapreduce/Int64/dims=1L	80823.5 ns	783182 ns	0.10
array/reductions/mapreduce/Int64/dims=2	69351 ns	96432 ns	0.72
array/reductions/mapreduce/Int64/dims=2L	86503.75 ns	298044 ns	0.29
array/reductions/reduce/Float32/1d	80046 ns	132832 ns	0.60
array/reductions/reduce/Float32/dims=1	61320.75 ns	95191 ns	0.64
array/reductions/reduce/Float32/dims=1L	69936 ns	773712 ns	0.09039022271853092
array/reductions/reduce/Float32/dims=2	71008.5 ns	96761 ns	0.73
array/reductions/reduce/Float32/dims=2L	83178.5 ns	294874 ns	0.28
array/reductions/reduce/Int64/1d	81456 ns	133652 ns	0.61
array/reductions/reduce/Int64/dims=1	60673.25 ns	77301 ns	0.78
array/reductions/reduce/Int64/dims=1L	80828.5 ns	779802 ns	0.10
array/reductions/reduce/Int64/dims=2	60828.25 ns	96382 ns	0.63
array/reductions/reduce/Int64/dims=2L	86661 ns	294464 ns	0.29
array/reverse/1d	48863 ns	43691 ns	1.12
array/reverse/1dL	68548.5 ns	74571 ns	0.92
array/reverse/1dL_inplace	79793.75 ns	73091 ns	1.09
array/reverse/1d_inplace	47898 ns	89692 ns	0.53
array/reverse/2d	48668.25 ns	54800 ns	0.89
array/reverse/2dL	73321 ns	101512 ns	0.72
array/reverse/2dL_inplace	95168.75 ns	164353 ns	0.58
array/reverse/2d_inplace	50183.25 ns	136732 ns	0.37
array/sorting/1d	336692 ns	345405 ns	0.97
gemm/tiled	1923390.5 ns
gemm/tiled_unbounded	1923590.25 ns
integration/byval/reference	40330 ns	39220 ns	1.03
integration/byval/slices=1	39601 ns	41250 ns	0.96
integration/byval/slices=2	157272 ns	141373 ns	1.11
integration/byval/slices=3	228653 ns	233903 ns	0.98
integration/volumerhs	4842354 ns	4943464 ns	0.98
kernel/indexing	35328 ns	127662 ns	0.28
kernel/indexing_checked	42818 ns	131583 ns	0.33
kernel/launch	1205 ns	1370 ns	0.88
kernel/rand	48108 ns	191753 ns	0.25
latency/import	1927304333 ns	1724918584 ns	1.12
latency/precompile	34262706605 ns	36524319792 ns	0.94
latency/ttfp	2454903779 ns	2278376161 ns	1.08
stencil/diffusion3d	1575973.25 ns
stencil/diffusion3d_checked	1636521.5 ns

This comment was automatically generated by workflow using github-action-benchmark.

[maleadt]

CI failure is JuliaGPU/KernelAbstractions.jl#888