Repository navigation
Let a custom mapping define the linear index of its work-items - #848
Merged
Merged
Conversation
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #848 +/- ##
=======================================
Coverage 79.82% 79.83%
=======================================
Files 26 26
Lines 2151 2152 +1
=======================================
+ Hits 1717 1718 +1
Misses 434 434 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Benchmark ResultsShow table
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
vchuravy
marked this pull request as ready for review
October 6, 2026 14:55
giordano
added a commit
to CliMA/Oceananigans.jl
that referenced
this pull request
Oct 6, 2026
…ns' hook `@index(Global, Linear)` in a kernel launched over an active cells map returned the work item's position in the map through a method of `KernelAbstractions.__index_Global_Linear` for the context of mapped launches. That's an internal function, which KernelAbstractions documents as not to be extended, and the method only worked as long as `@index` lowered through it. KernelAbstractions now lets a custom `NDRange` mapping define the linear index of its work items with `NDIteration.linear_index(iterspace, ndrange, groupidx, idx)` (JuliaGPU/KernelAbstractions.jl#848), so implement that for `MappedNDRange` instead, and drop the then unused `MappedCompilerMetadata`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EtLmZUPY4g9hi8aJsM4Wrs
giordano
approved these changes
Oct 6, 2026
giordano
left a comment
giordano
left a comment
Collaborator
There was a problem hiding this comment.
Seems to do the job for Oceananigans!
`@index(Global, Linear)` for a custom `NDRange` mapping was the position of the expanded index in the context's `ndrange`. A mapping whose linear index isn't a function of the expanded index alone, such as a list of indices whose linear index is the position in the list, had to override the internal `__index_Global_Linear` instead. Add `NDIteration.linear_index(iterspace, ndrange, groupidx, idx)`, which defaults to the previous definition and which a mapping can specialize on its `NDRange` type. Test a launch over a list of indices in the launch testsuite. Assisted-by: Claude Code
vchuravy
force-pushed
the
vc/mapping-linear-index
branch
from
October 6, 2026 18:35
0235cae to
1656bf6
Compare
giordano
added a commit
to CliMA/Oceananigans.jl
that referenced
this pull request
Oct 6, 2026
…ns' hook `@index(Global, Linear)` in a kernel launched over an active cells map returned the work item's position in the map through a method of `KernelAbstractions.__index_Global_Linear` for the context of mapped launches. That's an internal function, which KernelAbstractions documents as not to be extended, and the method only worked as long as `@index` lowered through it. KernelAbstractions now lets a custom `NDRange` mapping define the linear index of its work items with `NDIteration.linear_index(iterspace, ndrange, groupidx, idx)` (JuliaGPU/KernelAbstractions.jl#848), so implement that for `MappedNDRange` instead, and drop the then unused `MappedCompilerMetadata`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EtLmZUPY4g9hi8aJsM4Wrs
giordano
added a commit
to CliMA/Oceananigans.jl
that referenced
this pull request
Oct 7, 2026
…ns' hook `@index(Global, Linear)` in a kernel launched over an active cells map returned the work item's position in the map through a method of `KernelAbstractions.__index_Global_Linear` for the context of mapped launches. That's an internal function, which KernelAbstractions documents as not to be extended, and the method only worked as long as `@index` lowered through it. KernelAbstractions now lets a custom `NDRange` mapping define the linear index of its work items with `NDIteration.linear_index(iterspace, ndrange, groupidx, idx)` (JuliaGPU/KernelAbstractions.jl#848), so implement that for `MappedNDRange` instead, and drop the then unused `MappedCompilerMetadata`. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01EtLmZUPY4g9hi8aJsM4Wrs
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
2 participants
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.
For a custom
NDRangemapping,@index(Global, Linear)islinear_index(__ndrange(ctx), expand(iterspace, group, item)): the position of the expanded index within the context'sndrange. That doesn't fit a mapping whose linear index isn't a function of the expanded index alone. The motivating case is Oceananigans launching kernels over a list of active cells (CliMA/Oceananigans.jl#5799), where the linear index of a work-item is its position in the list. A cell can't be mapped back to that position without a reverse lookup, so that PR overrides the internal__index_Global_Linearfor its context type instead. #781 documents that as something not to do, and it only works as long as@indexlowers through that function.This adds a hook on the iteration space:
It defaults to the previous definition,
linear_index(ndrange, expand(iterspace, groupidx, idx)), andglobal_linearcalls it for iteration spaces other than KernelAbstractions' own. A mapping specializes it on itsNDRangetype, e.g. for Oceananigans:KernelAbstractions' built-in iteration spaces (identity and offsets) still take the direct path, so their code doesn't change.
Open points
NDRangemappings #781 should list the hook next toexpandandcartesian. I've left Document the contract of customNDRangemappings #781 alone.Testing
@index(Global, Linear)is the position in the list. Without the hook the kernel doesn't compile (linear_index(::ListedIndices, ::CartesianIndex)has no method); with it the test passes on the CPU backend. The GPU back-ends will run it through the shared testsuite.IndexMaplaunch from CliMA/Oceananigans.jl@6852216 passes with the hook in place of the__index_Global_Linearoverride.🤖 Generated with Claude Code