Accept index ranges in ndrange - #771
Merged
Merged
Conversation
giordano
added this pull request to stack #773
September 11, 2026 22:35
giordano
force-pushed
the
mg/ndrange-offsets
branch
from
September 11, 2026 22:35
cc39448 to
b6d37b5
Compare
Contributor
Benchmark ResultsShow table
Benchmark PlotsA plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR. |
giordano
force-pushed
the
mg/ndrange-offsets
branch
from
September 12, 2026 17:30
b6d37b5 to
ec7f3c7
Compare
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #771 +/- ##
==========================================
+ Coverage 63.73% 64.71% +0.98%
==========================================
Files 23 23
Lines 1941 1995 +54
==========================================
+ Hits 1237 1291 +54
Misses 704 704 ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Kernels can now iterate over a region whose indices do not start at 1.
Each entry of `ndrange` may be a range instead of an extent, and the
whole `ndrange` may be a single range or a `CartesianIndices`, both
statically (`kernel(backend, workgroupsize, (-2:N+3, 0:M+1))`) and at
launch (`ndrange=(-2:N+3, 0:M+1)`). `@index(Global, Cartesian)` and
`@index(Global, NTuple)` return the shifted indices, `@index(Global,
Linear)` counts the region from 1 in column-major order, `@ndrange()`
returns the extents.
Downstream packages such as Oceananigans implement this today by
pirating `partition`, `expand`, `__ndrange` and `__groupsize` with a
custom `_Size` subtype smuggled into `NDRange`'s dynamic-workitems type
parameter, which is fragile and broke with the compiled CPU backend.
Implementation:
- `StaticSize` stores `UnitRange{Int}` axes next to `Int` extents;
`StaticSize(ranges)` and `StaticSize(::CartesianIndices)` normalise
their input, `Base.OneTo` axes become plain extents.
- `NDRange` gains a `mapping` field holding a `StaticOffset` or
`DynamicOffset` (or `nothing`), which `expand` adds to the blocked
index. Plain-size launches keep `mapping === nothing`, so their
`NDRange` types are unchanged.
- `partition` normalises `ndrange`/`workgroupsize` (integer, range,
tuple, `CartesianIndices`) and compares static and launch ndranges by
extents and offsets, so `(1:128,)` matches a static `(128,)`.
- `CompilerMetadata` keeps an offset `CartesianIndices` as `ndrange`
(`CartesianIndices(::CartesianIndices)` would drop the offsets), so
the existing `I in __ndrange(ctx)` bounds check is correct.
- The global linear index is computed by `linear_index`, since
`LinearIndices` only supports 1-based axes.
- The POCL autotune path uses `extents(ndrange)`.
Offsets need no backend changes as long as the workgroup size is static
or given at launch; autotuning a dynamic workgroup size from a range
`ndrange` requires backends to call `extents` where they use
`prod(ndrange)`.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHciC8x39gm97sABrSvBkt
`linear_index` built a `LinearIndices` from the extents of the `ndrange`. Each `Base.OneTo` axis constructed that way clamps its length at zero, which left two `llvm.smax` per work item in the device code of kernels using `@index(Global, Linear)` on a multi-dimensional `ndrange`, compared to the code generated before index ranges were supported. Folding the column-major index directly from the extents and the axis starts produces the same instruction sequence as before. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01VHciC8x39gm97sABrSvBkt
vchuravy
force-pushed
the
mg/ndrange-offsets
branch
from
September 13, 2026 12:21
ec7f3c7 to
1f7ed89
Compare
1 task
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.
Redo of #764 in this repo, ref #757. CC @simone-silvestri.