[b-fg]

On a backend with a 3-d grid, an ndrange with more than three dimensions is launched as a LinearLaunch: every work-item decomposes its linear local and group ids into Cartesian positions. With a tuned workgroup size those are divisions by runtime values, and on PoCL such kernels are 30 to 100 times slower than the same kernel over a 3-d ndrange. WaterLily.jl hit this in its loops over vector fields, whose component is a fourth index (a time step about 2.5 times slower than on KA 0.9), and NNlib in imrotate (#913).

This launches them as an NDLaunch instead when the workgroups have one work-item along the dimensions past the third. The work-items keep the direct x/y/z mapping of a 3-d launch; the third and later dimensions share the grid's z-dimension, so only the group's z-id is decomposed along them, which is the same for every work-item of a workgroup. A static workgroup size folds when it has one work-item past the third dimension and the folded grid fits max_num_groups.

A tuned workgroup size is then chosen with work-items along the first three dimensions only, which gives small workgroups when those are small. On GPUs that loses against a linear launch, so tuned launches only fold when the first three dimensions hold at least 64 work-items. Folded over linear, for the scale! kernel of #913 on an RTX 5080 (the contributor's RTX 4060 gives the same picture):

ndrange	first three dims	CUDA	PoCL, 8 threads
(128, 128, 16, 8)	262144	1.70	33.8
(8, 8, 8, 1536)	512	1.62	7.8
(4, 4, 4, 8192)	64	1.31	4.3
(4, 4, 3, 10923)	48	1.03	4.7
(4, 4, 2, 16384)	32	0.71	6.1
(2, 3, 4, 50000)	24	0.50	3.9
(1, 1, 1, 60000)	1	0.07	1.2

PoCL would gain from folding below 64 as well; that would need a per-backend threshold, which can come later. Shapes that stay linear are no slower than before.

Median time per launch on PoCL (1 thread), tuned workgroup size:

ndrange	main	this PR	same work-items as 3-d
(128, 128, 16, 8)	12.4 ms	117 µs	127 µs
(64, 64, 64, 3)	4.6 ms	51 µs
(16, 16, 4, 4, 32)	1.5 ms	9 µs

The back ends' launch configuration tests assert which launch kind is selected for a 4-d range, so they fail with this change. JuliaGPU/CUDA.jl#3458, JuliaGPU/AMDGPU.jl#1183, JuliaGPU/Metal.jl#1107, JuliaGPU/oneAPI.jl#734 and JuliaGPU/OpenCL.jl#582 replace those assertions with behavioural checks, and should be merged first.

Tested with the KA test suite on PoCL, which now also launches a tuned 5-d range that folds, and with the KernelAbstractions tests of CUDA.jl (RTX 5080) and OpenCL.jl (PoCL).

Fixes #913

[vchuravy]

Looks like this might also be #913

[maleadt]

This indeed fixes #913 too. Reproduced on PoCL: (128,128,16,8) goes from 6.4 ms to 78 µs, and the test suite passes.

One change Claude suggests: the tuned-size condition (prod(min.(extent[1:3], max_dims[1:3])) >= max_items) keeps the slow path for shapes like (8,8,8,1536) or (256,2,2,768). PoCL's max_work_group_size is 4096, and forcing the fold there gives 2.4 ms → 27 µs and 2.4 ms → 19 µs. Even (2,3,4,50000), with 24-item workgroups, is 8× faster folded. Could you drop the condition, or use a small fixed threshold if it measurably helps on GPU?

The CI failures are the back ends' own launch configuration tests asserting LinearLaunch for 4-d.

[b-fg]

I checked, and a small threshold does measurably help on GPU. Measured on an RTX 4060 (with CUDA.jl main) and on PoCL, forcing each launch kind for 4-d ndranges with a tuned workgroup size, this is the speedup of the fold over the LinearLaunch (time of the linear launch over time of the folded one, the scale! kernel of #913, median per launch):

ndrange	first three dims	CUDA	PoCL, 2 threads
(128, 128, 16, 8)	262144	2.01	36.2
(8, 8, 8, 1536)	512	1.82	27.1
(256, 2, 2, 768)	1024	1.82	25.3
(4, 4, 4, 8192)	64	1.55	22.6
(4, 4, 3, 10923)	48	1.29	21.3
(4, 4, 2, 16384)	32	0.85	16.5
(2, 3, 4, 50000)	24	0.62	6.40
(2, 2, 2, 8192)	8	0.51	4.82
(1, 1, 1, 60000)	1	0.09	1.42

On PoCL the fold advantage is consistent, but on the GPU it loses when the first three dimensions hold fewer than about 48 work-items (tuned blocks are then smaller than two warps). So a fixed threshold of 64 work-items (013ab13): (8, 8, 8, 1536) and (256, 2, 2, 768) fold now, and shapes like (2, 3, 4, 50000) stay linear. That leaves the 6x on PoCL for those, but folding them costs 1.6x on CUDA, and up to 11x for (1, 1, 1, 60000). Favouring PoCL there would need a per-backend threshold but KI.sub_group_size is 32 on both so it cannot tell them apart.

With the threshold, the back ends' launch configuration tests need two changes: (4, 4, 4, 4) now folds too, while (2, 2, 2, 2) stays a LinearLaunch{Int32}, and (1, 1, 1, 2^31 + 1) instead of (2^11 + 1, 2^10, 2^10, 1) stays a LinearLaunch{Int} on every back end. Both shapes give the same result on main, so the tests can be updated before this PR is merged.

[maleadt]

Rebased and slightly reworked. I'll open back-end PRs that remove the now-problematic test.

[maleadt]

Thanks @b-fg! I've reopened this as #919, from a branch in this repository, so that #918 (which makes the 64-item threshold per back end) can be stacked on it; GitHub doesn't stack pull requests from forks. The commit is still yours.

Since you last looked, I squashed your two commits and rebased them on main, and:

• select_launch only checks the folded grid when it actually folds;
• the index-layout tests add (9, 7, 3, 5, 2), so a tuned launch that folds gets tested;
• the docs describe the fold.

The threshold of 64 holds on an RTX 5080 as on your RTX 4060. Rather than swapping the shapes in the back ends' tests, JuliaGPU/CUDA.jl#3458, JuliaGPU/AMDGPU.jl#1183, JuliaGPU/Metal.jl#1107, JuliaGPU/oneAPI.jl#734 and JuliaGPU/OpenCL.jl#582 drop their assertions on the launch kind; those go in first.

[b-fg]

Noted, thanks!