Conversation
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
force-pushed
the
vc/foreachindex
branch
from
September 13, 2026 14:01
c5e969d to
c033aef
Compare
vchuravy
force-pushed
the
mg/ndrange-index-map
branch
from
September 13, 2026 16:03
13f4f01 to
39488ea
Compare
A kernel whose body is a loop over the indices of an array needs none of the kernel language beyond the index itself, and writing it out is boilerplate that downstream packages repeat. Add `foreach_index(f, itr, backend = get_backend(itr); workgroupsize)`, which launches one work item per index of `eachindex(itr)` and calls `f` with the index a `for i in eachindex(itr)` loop would produce: a linear index for an `IndexLinear` array, a `CartesianIndex` otherwise. The index space is carried by the `ndrange`, so the kernel takes no argument besides the function, and an index space that does not start at 1 needs no special handling. The function is inlined into the kernel, as an out-of-line call to a closure can spill its captures to local memory. Ported from AcceleratedKernels.jl, without its CPU scheduling keywords (`max_tasks`, `min_elems`, `prefer_threads`): those pick a host-threaded loop, which 0.10 no longer has, and the workgroup size is left to the backend rather than defaulting to 256. Spelled with an underscore to keep the name distinct from the AcceleratedKernels function it does not behave identically to, so that importing both packages does not clash. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01DbbtGLjxXpt7kZj6h2yVdK
giordano
force-pushed
the
vc/foreachindex
branch
from
September 13, 2026 20:23
c033aef to
74d3e03
Compare
Collaborator
vchuravy
commented
Sep 13, 2026
Comment on lines
+10
to
+13
| @kernel function foreach_index_linear_kernel(f) | ||
| I = @index(Global, Cartesian) | ||
| @inline f(I.I[1]) | ||
| end |
Member
Author
There was a problem hiding this comment.
Suggested change
| @kernel function foreach_index_linear_kernel(f) | |
| I = @index(Global, Cartesian) | |
| @inline f(I.I[1]) | |
| end | |
| @kernel function foreach_index_linear_kernel(f) | |
| I = @index(Global, Linear) | |
| @inline f(I) | |
| end |
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.
Adds
foreach_index(f, itr, backend = get_backend(itr); workgroupsize = nothing), which runsfonce per index of
eachindex(itr)with one work item per index, so a loop over the indices of anarray needs no
@kerneldefinition:freceives the index afor i in eachindex(itr)loop would produce: a linear index for an arraywith
IndexLinearstyle, aCartesianIndexotherwise.Ported from AcceleratedKernels.jl, where
foreachindexis the primitive thatmap,reduce,sort,accumulate,findall,searchsortedandreverseare built on, so it has had a lot of use.Based on
mg/ndrange-index-mapBased on
mg/ndrange-index-maprather thanmain, because that branch is what makes this a thinwrapper: the index space is carried by the
ndrange, soeachindex(itr)as a kernel argument and indexes it with the global index, and@index(Global, Cartesian)already returns the shifted index. Covered by a test with an
IndexLineararray whoseaxes1is
-2:7.1f7ed89a("Compute the global linear index arithmetically") also removes theLinearIndiceslookup that pushed the AcceleratedKernels version to
unsafe_indices=true, so this uses theordinary bounds-checked path.
What the port drops
The AcceleratedKernels signature carries a CPU scheduling policy (
max_tasks,min_elems,prefer_threads) that selects aThreads.@spawnloop over the kernel. 0.10 has no host-threadedexecution path to select —
CPU === POCLBackendandget_backend(::Array) === POCLBackend()—and that policy is a backend's business anyway, which is where it lived in 0.9 (
CPU_GRAINSIZEand the
staticflag).block_size=256becomesworkgroupsize, KA's name for it, and defaultsto
nothingso that the backend's autotuning picks it.Open questions
AcceleratedKernels export, and so that the difference in behaviour (no scheduler keywords,
asynchronous, no
@inboundselision) does not hide behind an identical name. The cost is thatit no longer mirrors
eachindex, and that the ecosystem carries two near-identical names. Thealternative would be
foreachindexleft unexported.CPU(). Every distinct closure is a fresh POCL compile, so a one-shotforeach_indexover anArrayis far more expensive than the threaded loop it reads like. Finefor a loop that runs more than once, but worth saying out loud before users reach for this as a
@threadsreplacement.foraxes. Not included: with anndrangethat takes ranges,foreach_index(f, axes(A, 2))covers it (with the backend passed explicitly, since a range carries no backend).
reads like a
forloop that has finished. Documented, but it is a plausible footgun.Testing
test/foreach_index.jl, wired into the test suite: linear index spaces in 1-3 dimensions (resultplus each index visited exactly once), a
CartesianIndicesspace through a strided view, anoffset linear space, an index space with no device memory (
1:100), an explicitworkgroupsize,an explicit backend, an empty index space, and the error for an index space that is neither.
Passes on the CPU backend; not yet run on a GPU backend.
🤖 Generated with Claude Code
https://claude.ai/code/session_01DbbtGLjxXpt7kZj6h2yVdK