Conversation
giordano
added this pull request to stack #773
September 11, 2026 22:35
giordano
force-pushed
the
mg/ndrange-index-map
branch
from
September 11, 2026 22:35
c0b180e to
c4daa40
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-index-map
branch
from
September 12, 2026 17:30
c4daa40 to
aa7d1ca
Compare
Member
|
Reposting @vchuravy's comment here:
Originally posted by @vchuravy in #765 (comment) |
vchuravy
force-pushed
the
mg/ndrange-index-map
branch
from
September 13, 2026 12:21
aa7d1ca to
13f4f01
Compare
`ndrange` may now be a device vector of `CartesianIndex{N}` or
`NTuple{N, <:Integer}` elements. Each work item handles one listed index:
`@index(Global, Cartesian)` and `@index(Global, NTuple)` return it,
`@index(Global, Linear)` its position in the vector, `@ndrange()` the
vector's length. The kernel must have a dynamic `ndrange` and a 1-D
workgroup size, static or given at launch; the partial last workgroup is
bounds-checked against the vector's length.
This is how Oceananigans iterates over the active cells of an immersed
boundary grid. It currently wraps the kernel function, pirates
`Base.getproperty` on `Kernel`, `Adapt.adapt_structure` on `NDRange`
and `CompilerMetadata`, `partition`, `expand`, and adds one
`@device_override __validindex` per GPU backend, which the compiled CPU
backend's own `::Any` override shadows.
Implementation:
- `IndexMap{N}` wraps the vector and yields `CartesianIndex{N}` on
indexing; it is the `mapping` of a 1-D `NDRange` (`MappedNDRange`).
`expand` and `linear_index` on a `MappedNDRange` look the index up.
- `partition` dispatches vectors to `mapped_partition`, which validates
the kernel and workgroup size and partitions `(length(map),)`; the
dynamic-check flag is computed as for any other ndrange.
- Index validity moves into the generic `__validindex(ctx, groupidx,
idx)`, which dispatches on the iteration space (`I in ndrange` for
regular spaces, position `<= length(map)` for index maps). The 1-arg
`__validindex(ctx)` is now a plain method built on
`KernelInterface.get_group_id`/`get_local_id`, so the POCL override is
removed; `__index_Global_Linear` goes through `__global_linear` in the
same way. Backends that still override `__validindex(ctx)` keep
working for regular ndranges and should forward to the 3-arg method
to support index maps.
- `Adapt.adapt_structure` for `NDRange`, `IndexMap` and
`CompilerMetadata` moves the vector to the device with the kernel
arguments; `CompilerMetadata` treats a vector `ndrange` as
`CartesianIndices((length(v),))`.
Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01VHciC8x39gm97sABrSvBkt
vchuravy
force-pushed
the
mg/ndrange-index-map
branch
from
September 13, 2026 16:03
13f4f01 to
39488ea
Compare
Draft
Backends check the validity of a work item as `I in __ndrange(ctx)`, with `I` from `expand` or `__index_Global_Cartesian`. For a launch over an index map that compared an N-D index against a 1-D range, so every lane was inactive on OpenCL.jl. `expand` on a `MappedNDRange` now returns an `invalid_index` (`typemin(Int)` along every axis) for work items past the end of the map, and the `ndrange` of such a launch is a `MappedIndices`, which rejects only that index. Also fix the `Const` test: the CPU backend compiles with POCL, so the host `code_llvm` it used produced a kernel that only threw. Launch the kernel and check the result on every backend instead, keeping the IR checks for CUDA and ROCm. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CSb4tbVQi1dbKWvaP6R6EE
The docstrings of `MappedNDRange`, `MappedIndices` and `invalid_index` link to it, which made the docs build fail. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01CSb4tbVQi1dbKWvaP6R6EE
Collaborator
Author
|
Based on quick testing, we may not need this in Oceananigans after all, and #771 should be sufficient (it's hard to do full testing at the moment though, because of #734 (comment)). I'm leaving this open in case it's still useful for #779. |
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 #765, in this repo. Ref: #757. CC @simone-silvestri.