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 Metal.jl to that model in one step. It replaces #917, which added KernelInterface next to the existing KernelAbstractions 0.9 back end, and the kaintr branch, which ported an earlier version of it to KernelAbstractions 0.10.

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

using Metal
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 = Metal.rand(Float32, 1000); b = Metal.rand(Float32, 1000); c = similar(a)
KI.@launch MetalBackend() 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 stack allocation behind @private, and the rule for moving an MtlArray to the CPU with adapt. Metal's own copy of the KernelAbstractions launch path goes away: partitioning the ndrange, building the kernel's context, sizing the threadgroups, calling the kernel, and the overrides of KernelAbstractions' index functions. KI.copyto! now also accepts contiguous views of host arrays, and adapt(MetalBackend(), x) moves any array to the GPU, as adapt(MtlArray, x) does, not only Arrays. Threadgroups are still as large as the pipeline's maxTotalThreadsPerThreadgroup allows, which is what KernelInterface's default launch_configuration picks, so Metal doesn't override it.

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 divisions to compute its index, which the test suite checks for the generated LLVM IR. On Metal this matters a lot, since 64-bit integer division is slow on Apple GPUs. With a KA kernel that copies a Float32 array, on an M1 (GPU time of the command buffer):

	main (KA 0.9)	this PR
256×256×256, @index(Global, Cartesian)	12.5 ms	2.3 ms
256×256×256, @index(Global, Linear)	11.4 ms	2.4 ms
255×257×129, @index(Global, Cartesian)	6.4 ms	1.1 ms
255×257×129, @index(Global, Linear)	5.9 ms	1.1 ms
launch of an empty kernel	2.33 µs, 768 bytes	2.58 µs, 768 bytes
launch of a kernel with 28 arguments	9.6 µs, 4480 bytes	3.1 µs, 2112 bytes

The kernel times vary by about 10% between runs of this PR. The launch benchmark uses 28 arguments rather than the 40 of CUDA.jl's, because a Metal kernel takes at most 31 buffers, and every argument, the kernel's context and Metal's kernel state each take one.

KernelInterface 0.4 passes the arguments to the back end as a tuple (JuliaGPU/KernelAbstractions.jl#811). Metal has no public launch function that takes a tuple, but calling a HostKernel forwards its arguments as one to an internal function, and KI.launch calls that function directly, so nothing is splatted. KI.launch rejects threads and groups, which would override the launch geometry KernelInterface has validated, but passes Metal's other launch options on:

kernel = KI.@launch MetalBackend() launch=false vadd(c, a, b)
kernel(c, a, b; ndrange=length(c), queue=global_queue(device()))   # fine
kernel(c, a, b; ndrange=length(c), threads=256)                     # ArgumentError

The first commit is the port. The next ones implement parts of KernelInterface that are optional: sub-groups (SIMD-groups, including a partial last one), ordering work across tasks with Metal events for KernelAbstractions.@spawn, where the waiting task's command buffer waits on the GPU instead of blocking the host, and KI.versioninfo.

KernelAbstractions' conversion tests, skipped until now, pass with 32-bit indices. Its test that launches a kernel with 41 arguments is skipped, since a Metal kernel takes at most 31 buffers, one per argument. That limit isn't new, but it now shows in KernelAbstractions' test suite. Unlike CUDA.jl, Metal already converted the arguments twice (once for the types to compile for, and again while encoding them), so KernelAbstractions' generic launch doesn't add a conversion.

This needs a breaking release, since it requires KernelAbstractions 0.10, which isn't registered yet. Until then, the last commit takes KernelAbstractions and KernelInterface from their development branch through [sources]. Julia 1.10 ignores [sources] and Julia 1.11 doesn't support workspaces, so on Buildkite those two set up the test environment by hand and run the tests in it directly. That commit is dropped before merging.

[github-actions]

Metal Benchmarks

Details

Benchmark suite	Current: 5d61e33	Previous: d209fe6	Ratio
array/accumulate/Float32/1d	389500 ns	388291 ns	1.00
array/accumulate/Float32/dims=1	364708 ns	350375 ns	1.04
array/accumulate/Float32/dims=1L	8766209 ns	8763333 ns	1.00
array/accumulate/Float32/dims=2	426333 ns	428458 ns	1.00
array/accumulate/Float32/dims=2L	2513625 ns	2526958 ns	0.99
array/accumulate/Int64/1d	839666 ns	846375 ns	0.99
array/accumulate/Int64/dims=1	911042 ns	904042 ns	1.01
array/accumulate/Int64/dims=1L	9440542 ns	9422041 ns	1.00
array/accumulate/Int64/dims=2	1211167 ns	1211708 ns	1.00
array/accumulate/Int64/dims=2L	6427834 ns	6424250 ns	1.00
array/broadcast	240834 ns	228459 ns	1.05
array/construct	2291 ns	2291 ns	1
array/permutedims/2d	338333 ns	452666 ns	0.75
array/permutedims/3d	463291 ns	1027541 ns	0.45
array/permutedims/4d	616041 ns	1098416 ns	0.56
array/private/copy	221959 ns	226750 ns	0.98
array/private/copyto!/cpu_to_gpu	218167 ns	211583 ns	1.03
array/private/copyto!/gpu_to_cpu	219875 ns	208917 ns	1.05
array/private/copyto!/gpu_to_gpu	216666 ns	220375 ns	0.98
array/private/iteration/findall/bool	1056584 ns	1048542 ns	1.01
array/private/iteration/findall/int	1213084 ns	1224750 ns	0.99
array/private/iteration/findfirst/bool	1152916 ns	1153792 ns	1.00
array/private/iteration/findfirst/int	1151333 ns	1156375 ns	1.00
array/private/iteration/findmin/1d	1189208 ns	1218292 ns	0.98
array/private/iteration/findmin/2d	1028292 ns	1024250 ns	1.00
array/private/iteration/logical	1653167 ns	1661084 ns	1.00
array/private/iteration/scalar	1382375 ns	1381000 ns	1.00
array/random/rand/Float32	438584 ns	414334 ns	1.06
array/random/rand/Int64	519042 ns	494209 ns	1.05
array/random/rand!/Float32	428083 ns	402500 ns	1.06
array/random/rand!/Int64	443792 ns	431167 ns	1.03
array/random/randn/Float32	364625 ns	388208 ns	0.94
array/random/randn!/Float32	355334 ns	367166 ns	0.97
array/reductions/mapreduce/Float32/1d	447084 ns	450000 ns	0.99
array/reductions/mapreduce/Float32/dims=1	353959 ns	349250 ns	1.01
array/reductions/mapreduce/Float32/dims=1L	623666 ns	611750 ns	1.02
array/reductions/mapreduce/Float32/dims=2	352875 ns	355334 ns	0.99
array/reductions/mapreduce/Float32/dims=2L	1236000 ns	1220750 ns	1.01
array/reductions/mapreduce/Int64/1d	632709 ns	628041 ns	1.01
array/reductions/mapreduce/Int64/dims=1	634458 ns	631541 ns	1.00
array/reductions/mapreduce/Int64/dims=1L	1007958 ns	1006459 ns	1.00
array/reductions/mapreduce/Int64/dims=2	789209 ns	788791 ns	1.00
array/reductions/mapreduce/Int64/dims=2L	2194542 ns	2200959 ns	1.00
array/reductions/reduce/Float32/1d	437542 ns	449833 ns	0.97
array/reductions/reduce/Float32/dims=1	357542 ns	356375 ns	1.00
array/reductions/reduce/Float32/dims=1L	615750 ns	607625 ns	1.01
array/reductions/reduce/Float32/dims=2	245417 ns	245042 ns	1.00
array/reductions/reduce/Float32/dims=2L	484875 ns	482709 ns	1.00
array/reductions/reduce/Int64/1d	632917 ns	628250 ns	1.01
array/reductions/reduce/Int64/dims=1	637000 ns	637625 ns	1.00
array/reductions/reduce/Int64/dims=1L	1006333 ns	1004292 ns	1.00
array/reductions/reduce/Int64/dims=2	270250 ns	269292 ns	1.00
array/reductions/reduce/Int64/dims=2L	667583 ns	665459 ns	1.00
array/shared/copy	130333 ns	128833 ns	1.01
array/shared/copyto!/cpu_to_gpu	37625 ns	37834 ns	0.99
array/shared/copyto!/gpu_to_cpu	36959 ns	37542 ns	0.98
array/shared/copyto!/gpu_to_gpu	37458 ns	37583 ns	1.00
array/shared/iteration/findall/bool	1049125 ns	1046041 ns	1.00
array/shared/iteration/findall/int	1221125 ns	1209208 ns	1.01
array/shared/iteration/findfirst/bool	975167 ns	980292 ns	0.99
array/shared/iteration/findfirst/int	988042 ns	988542 ns	1.00
array/shared/iteration/findmin/1d	1052667 ns	1055334 ns	1.00
array/shared/iteration/findmin/2d	1027083 ns	1032833 ns	0.99
array/shared/iteration/logical	1528750 ns	1528625 ns	1.00
array/shared/iteration/scalar	3765.625 ns	3671.875 ns	1.03
array/sorting/1d	2096000 ns	2083750 ns	1.01
array/sorting/2d	8420542 ns	8302083 ns	1.01
integration/byval/reference	1116542 ns	1117666 ns	1.00
integration/byval/slices=1	1119833 ns	1119833 ns	1
integration/byval/slices=2	2029041 ns	2010584 ns	1.01
integration/byval/slices=3	6540208 ns	6549916 ns	1.00
integration/metaldevrt	396958 ns	393166 ns	1.01
kernel/indexing	214125 ns	211500 ns	1.01
kernel/indexing_checked	395500 ns	394709 ns	1.00
kernel/launch	2051 ns	2050.8888888888887 ns	1.00
kernel/rand	405083 ns	408041 ns	0.99
latency/import	1796217750 ns	1797481542 ns	1.00
latency/precompile	32871589291 ns	32287978750 ns	1.02
latency/ttfp	2291828875 ns	2252476334 ns	1.02
metal/synchronization/context	537.2592592592592 ns	536.4052631578948 ns	1.00
metal/synchronization/stream	339.58181818181816 ns	343.70506912442397 ns	0.99

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