Repository navigation
Conversation
maleadt
added this pull request to stack #920
October 23, 2026 14:42
This was referenced Oct 23, 2026
Merged
Merged
Merged
Merged
Contributor
|
Documentation preview: https://juliagpu.github.io/KernelAbstractions.jl/previews/PR919/ |
Contributor
Benchmark ResultsShow table
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
On a backend with a 3-d grid, an ndrange with more than three dimensions was always launched as a `LinearLaunch`, which decomposes the linear ids of every work-item into Cartesian positions. With a tuned workgroup size, these are divisions by runtime values, which made such kernels 30 to 100 times slower than the same kernel over a 3-d ndrange on PoCL (#913). Launch it as an `NDLaunch` instead when its workgroups have one work-item along the dimensions past the third: the work-items keep the direct x/y/z mapping, the later dimensions share the grid's z-dimension with the third, and only the group's z-id is decomposed along them, not the ids of the work-items. A tuned workgroup then only has work-items along the first three dimensions, so it only folds if these hold at least 64 work-items. Below that, the folded launch is slower on GPUs (0.7x at 32 work-items and 0.07x for (1, 1, 1, 60000) on an RTX 5080), while from 64 on it is 1.3 to 1.7 times faster there, and many times faster on PoCL. Co-authored-by: Tim Besard <tim.besard@gmail.com> Assisted-By: Claude Opus 5.5 <noreply@anthropic.com>
maleadt
force-pushed
the
tb/ndlaunch-fold
branch
from
October 25, 2026 14:14
bfbf67d to
78b2c4e
Compare
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 inimrotate(#913).This launches them as an
NDLaunchinstead 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 fitsmax_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):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:
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).
This is #912 by @b-fg, reopened from a branch in this repository so that follow-ups can be stacked on it.
Fixes #913