Launch KernelAbstractions kernels on N-d grids, indexing in 32 bits

[maleadt]

KernelAbstractions kernels (including GPUArrays' broadcast, permutedims! etc.) currently launch on a 1-D grid, and every thread recovers its Cartesian index from blockIdx/threadIdx with 64-bit divisions. With JuliaGPU/KernelAbstractions.jl#797, back ends can launch on a grid shaped like the iteration space instead. This PR does that for CUDA, on top of the KA 0.10 port (#3302):

• Iteration spaces of up to 3 dimensions launch on a matching 3-D grid, so @index needs no divisions. 4-D and larger ranges, or ranges needing more than 65535 blocks along y or z, keep the 1-D grid.
• Index arithmetic happens in Int32 whenever the padded iteration space fits (in Int otherwise), which also helps the 1-D fallback. @index still returns Ints.
• The launch is chosen before tuning the workgroup size, so a kernel is still compiled once.

Also included: the typed KernelInterface index queries now compute in the requested type (get_global_id(Int) used to wrap around at 2^31 threads), the KernelInterface per-dimension launch limits, and two fixes for #3302 (empty ndranges divided by zero when tuning; ranges with offsets couldn't be tuned).

Benchmarks

Kernel times on an RTX 5080 (median of 21×10 launches, best of two interleaved rounds), comparing #3302 with this PR. Both use KernelAbstractions with the first commit of JuliaGPU/KernelAbstractions.jl#797, which fixes a regression in KA 0.10 main (the untyped KernelInterface queries weren't inlined, making #3302 as-is up to 2.4× slower than CUDA.jl main):

operation	#3302	this PR	change
bcast 1D F32 a+b	954 µs	949 µs	-0%
bcast 2D F32 A+row	667 µs	657 µs	-2%
bcast 3D F32 A*v	163 µs	80 µs	-51%
bcast 3D F32 A+B	237 µs	238 µs	+0%
bcast 3D F32 small	4 µs	4 µs	-17%
bcast 3D Int8 A+v	633 µs	296 µs	-53%
bcast 3D F16 A+B+1	187 µs	121 µs	-36%
bcast 4D F32 A*v	255 µs	191 µs	-25%
bcast 2D view	677 µs	663 µs	-2%
permutedims 3D	275 µs	230 µs	-16%
KA 2D Cartesian	662 µs	659 µs	-0%
KA 3D Cartesian	179 µs	161 µs	-10%
KA 3D Linear	173 µs	162 µs	-6%
KA 3D Cartesian wgs=(32,8)	163 µs	164 µs	+0%
KA 3D stencil NTuple	195 µs	167 µs	-14%
KA 4D Cartesian	249 µs	192 µs	-23%

N-d broadcasts and dynamically-sized N-d kernels get considerably faster; 1-D and memory-bound 2-D kernels don't change. Neither does the kernel with a static (32, 8) workgroup, where decomposing the linear index only divides by constants.

Host overhead went down too: launching a tiny 3-D kernel takes 2.15 µs and allocates 688 B, vs. 2.35 µs and 1648 B with #3302.

Oceananigans.jl and Breeze.jl (on their KA 0.10 ports, CliMA/Oceananigans.jl#5799 and NumericalEarth/Breeze.jl#848), median ms/step over 200 steps, 3 interleaved rounds:

configuration	#3302	this PR	change
Oceananigans lat-lon F32 (F64 CATKE, as in CI)	37.53	37.02	−1.3%
Oceananigans lat-lon F32 (F32 CATKE)	21.00	20.32	−3.2%
Oceananigans tripolar F32 (F64 CATKE, as in CI)	40.60	40.22	−0.9%
Oceananigans tripolar F32 (F32 CATKE)	21.21	20.88	−1.6%
Oceananigans lat-lon / tripolar F64	200.62 / 158.45	200.49 / 158.22	−0.1%
Breeze CBL anelastic, 256×256×128 F32	18.32	18.10	−1.2%
Breeze CBL compressible	101.57	101.26	−0.3%
Breeze CBL anelastic + 1M microphysics	43.37	43.53	+0.4%

The largest per-kernel wins are Oceananigans' Gv (98 → 80 registers, −23%), update_hydrostatic_pressure (−33%) and broadcasts (−20%). A few of Breeze's momentum tendencies cross a register threshold instead (up to +20% for one kernel), hence the +0.4%. Oceananigans' results are bitwise identical; Breeze's anelastic ones differ at rounding level (FMA contraction).

Tests

KernelAbstractions' testsuite, which runs here too, now checks every @index flavour against the expected layout for 1–4-D, ragged, offset, static and empty ranges. The CUDA tests add launch selection, the absence of divisions in the PTX for a dynamic 3-D range, and ranges of more than 2^31 threads on both kinds of grid.

Temporary: [sources] and the Julia 1.10 CI step point KernelAbstractions and KernelInterface at the #797 branch (KernelInterface is repeated in every workspace project to work around JuliaLang/Pkg.jl#4831). To be reverted once KA 0.10 is registered.

[github-actions]

CUDA.jl Benchmarks

Details

Benchmark suite	Current: 3926ab5	Previous: f8606eb	Ratio
array/accumulate/Float32/1d	100579 ns	99613 ns	1.01
array/accumulate/Float32/dims=1	72763 ns	72400 ns	1.01
array/accumulate/Float32/dims=1L	1588905 ns	1589076 ns	1.00
array/accumulate/Float32/dims=2	138305 ns	137237 ns	1.01
array/accumulate/Float32/dims=2L	657360 ns	656147 ns	1.00
array/accumulate/Int64/1d	120070 ns	119011 ns	1.01
array/accumulate/Int64/dims=1	77782 ns	76083 ns	1.02
array/accumulate/Int64/dims=1L	1697910 ns	1698986 ns	1.00
array/accumulate/Int64/dims=2	151645 ns	150489 ns	1.01
array/accumulate/Int64/dims=2L	988862 ns	987960 ns	1.00
array/broadcast	10666 ns	16704 ns	0.64
array/broadcast launch	7681.25 ns	8249.666666666666 ns	0.93
array/construct	877.0545454545454 ns	914.7878787878788 ns	0.96
array/copy	16388 ns	16742 ns	0.98
array/copyto!/cpu_to_gpu	209576 ns	208127 ns	1.01
array/copyto!/gpu_to_cpu	241590 ns	239967 ns	1.01
array/copyto!/gpu_to_gpu	10173 ns	10240 ns	0.99
array/iteration/findall/bool	135814 ns	134305 ns	1.01
array/iteration/findall/int	146021 ns	145571 ns	1.00
array/iteration/findfirst/bool	71172 ns	69778 ns	1.02
array/iteration/findfirst/int	73871 ns	71546 ns	1.03
array/iteration/findmin/1d	67404 ns	63567 ns	1.06
array/iteration/findmin/2d	92386 ns	97878 ns	0.94
array/iteration/logical	190062 ns	188464 ns	1.01
array/iteration/scalar	61432 ns	61246 ns	1.00
array/permutedims/2d	45088 ns	47629 ns	0.95
array/permutedims/3d	43955 ns	48318 ns	0.91
array/permutedims/4d	46443 ns	50355 ns	0.92
array/random/rand/Float32	11503 ns	11938 ns	0.96
array/random/rand/Int64	20310 ns	20338 ns	1.00
array/random/rand!/Float32	7777.25 ns	7799.5 ns	1.00
array/random/rand!/Int64	18266 ns	16982 ns	1.08
array/random/randn/Float32	32400 ns	32630 ns	0.99
array/random/randn!/Float32	23485 ns	23378 ns	1.00
array/reductions/mapreduce/Float32/1d	34921 ns	32436 ns	1.08
array/reductions/mapreduce/Float32/dims=1	38121 ns	37579 ns	1.01
array/reductions/mapreduce/Float32/dims=1L	51527 ns	50850 ns	1.01
array/reductions/mapreduce/Float32/dims=2	55961 ns	55322 ns	1.01
array/reductions/mapreduce/Float32/dims=2L	68876 ns	67515 ns	1.02
array/reductions/mapreduce/Int64/1d	41719 ns	39642 ns	1.05
array/reductions/mapreduce/Int64/dims=1	41393 ns	40599 ns	1.02
array/reductions/mapreduce/Int64/dims=1L	89503 ns	88796 ns	1.01
array/reductions/mapreduce/Int64/dims=2	58770 ns	57822 ns	1.02
array/reductions/mapreduce/Int64/dims=2L	85308 ns	83743 ns	1.02
array/reductions/reduce/Float32/1d	34489 ns	32529 ns	1.06
array/reductions/reduce/Float32/dims=1	37946 ns	37411 ns	1.01
array/reductions/reduce/Float32/dims=1L	51268 ns	50621 ns	1.01
array/reductions/reduce/Float32/dims=2	55803 ns	55189 ns	1.01
array/reductions/reduce/Float32/dims=2L	69181 ns	67650 ns	1.02
array/reductions/reduce/Int64/1d	42070 ns	40254 ns	1.05
array/reductions/reduce/Int64/dims=1	41277 ns	40559 ns	1.02
array/reductions/reduce/Int64/dims=1L	89523 ns	88858 ns	1.01
array/reductions/reduce/Int64/dims=2	58850 ns	57570 ns	1.02
array/reductions/reduce/Int64/dims=2L	84941 ns	83538 ns	1.02
array/reverse/1d	17884 ns	16994 ns	1.05
array/reverse/1dL	70685 ns	69571 ns	1.02
array/reverse/1dL_inplace	67815 ns	67345 ns	1.01
array/reverse/1d_inplace	9157.666666666666 ns	8508.333333333334 ns	1.08
array/reverse/2d	20831 ns	20087 ns	1.04
array/reverse/2dL	74293 ns	73424 ns	1.01
array/reverse/2dL_inplace	67534 ns	67255 ns	1.00
array/reverse/2d_inplace	10305 ns	9913 ns	1.04
array/sorting/1d	2640524 ns	2645454 ns	1.00
array/sorting/2d	1018762 ns	1017223 ns	1.00
array/sorting/by	3176010 ns	3158337 ns	1.01
cuda/synchronization/context/auto	1002.1 ns	1019.4 ns	0.98
cuda/synchronization/context/blocking	788.8484848484849 ns	781.4509803921569 ns	1.01
cuda/synchronization/context/nonblocking	5684.666666666667 ns	5917 ns	0.96
cuda/synchronization/stream/auto	843.8636363636364 ns	870.5593220338983 ns	0.97
cuda/synchronization/stream/blocking	655.4382716049382 ns	662.6521739130435 ns	0.99
cuda/synchronization/stream/nonblocking	5765.833333333333 ns	5784.2 ns	1.00
integration/byval/reference	148251 ns	147901 ns	1.00
integration/byval/slices=1	149529 ns	149155 ns	1.00
integration/byval/slices=2	292418 ns	291823 ns	1.00
integration/byval/slices=3	435963 ns	434984 ns	1.00
integration/cudadevrt	105395 ns	105048 ns	1.00
integration/volumerhs	9154605 ns	9139761 ns	1.00
kernel/indexing	13309 ns	13334 ns	1.00
kernel/indexing_checked	13848 ns	13649 ns	1.01
kernel/launch	2582.6666666666665 ns	2124 ns	1.22
kernel/occupancy	1023.8 ns	692.9485294117648 ns	1.48
kernel/rand	14879 ns	16179 ns	0.92
latency/import	4019575880 ns	4227977216 ns	0.95
latency/precompile	4944914545 ns	4996503792 ns	0.99
latency/ttfp	4659000844 ns	4691818788 ns	0.99

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

[maleadt]

Superseded by #3314: KernelAbstractions now launches @kernel kernels on N-d grids and computes @index in 32 bits itself (JuliaGPU/KernelAbstractions.jl#797), and #3314 carries over the rest of this PR (typed index queries computed with % T, the per-dimension launch limits).