Read ScatterElements indices in place instead of materializing them - #32055
Open
Alexander Novikov (novikov-alexander) wants to merge 3 commits into
Open
Read ScatterElements indices in place instead of materializing them#32055Alexander Novikov (novikov-alexander) wants to merge 3 commits into
Alexander Novikov (novikov-alexander) wants to merge 3 commits into
Conversation
The generic loop hands each work unit a single inner index, so a work unit walks one column of the updates with a stride of inner_size. For the very common axis=0 rank-2 case that means outer_size is 1, there are only C work units, and each one traverses a row-major buffer column-wise. When the indices repeat across the dimensions after the axis, each update slice lands on a single contiguous run of the output, so the strided walk becomes a sequential one the compiler can vectorize. Those indices are what an Expand of an [N, 1] index tensor to [N, C] produces, which is how a row scatter is spelled in PyTorch (Tensor.scatter_ requires index and src to have the same rank), and it is also what GatherElementsGrad sees. Detection is at runtime from the indices tensor, so it does not depend on the Expand having been constant-folded. The scan bails on the first mismatch, which is the first slice for ordinary indices. Results are unchanged, bit for bit. Each work unit still owns a distinct (outer, inner-range) region of the output, and within a unit the axis is walked in ascending order, so a destination receiving several updates accumulates them in the same sequence as before. Splitting the run into blocks rather than parallelizing over update rows is what preserves that. Two guards matter. The run is contiguous only when the updates and the data agree on every dimension after the axis; ScatterElements permits the updates to be smaller there, and an inner coordinate then maps to different offsets in the two tensors. And runs shorter than about a cache line are skipped, since the generic loop's stride is already local and the fast path would trade away parallelism for nothing. Measured on an M3 Pro, Release, 12 threads, timing the whole node through an InferenceSession, axis=0 float reduction='add' with row-broadcast indices: N=4096 C=1024 13.42 ms -> 3.07 ms 4.4x N=4096 C=256 0.94 ms -> 0.38 ms 2.5x N=8192 C=64 0.37 ms -> 0.22 ms 1.7x N=65536 C=32 4.97 ms -> 1.95 ms 2.5x The scatter loop itself speeds up considerably more than that; index validation, which materializes a normalized int64 per element, dominates what is left. Tests cover the fast path for 'add' and rank 3, a duplicate-destination case under reduction='none' that fails if the axis is visited in any other order, and updates narrower than the data in the last dimension, which must stay on the generic path. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
|
Azure Pipelines: There may be pipelines that require an authorized user to comment /azp run to run. |
CI showed two of the new tests failing on providers they were never meant to say anything about. RowBroadcastIndicesNoneKeepsLastUpdate asserts that the later update wins when several target one element. ONNX leaves that unspecified, so it is a property of this kernel rather than of the operator, and providers that apply updates concurrently pick an arbitrary winner. It failed on the CUDA plugin EP and on WebGPU for that reason. RowBroadcastIndicesNarrowerTrailingDim checks which internal path the CPU kernel picks for a shape the fast path has to decline; it failed on QNN, which returns zeros for that shape. Both now run on CPU only, which is what they were always testing. Also reworded the comment on the row-broadcast detector. It claimed the scan bails out on the first mismatch and costs one slice, which overstates what a parallel loop does: other workers stop at their next slice boundary, so the bound is small but not one slice. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
GetIndices normalized every index into a std::vector<int64_t> before any
scatter work began. For [4096, 1024] indices that is a 33 MB allocation and a
full write pass, and the scatter then reads those 8 bytes per element back
rather than the 4 an int32 indices tensor would have cost.
Validate the indices in place instead, and let the scatter read the indices
tensor directly, normalizing negatives where they are used. Validation still
happens up front so an out-of-range index is still reported before the output
is touched.
The index element type is carried by a small reader rather than a template
parameter on ScatterData. Templating would have been marginally faster but
instantiates the whole scatter, including the contiguous fast path, once per
index type: measured on this file that is 2.20 MB of object code against
1.12 MB for the reader, and the branch it replaces is loop-invariant in
practice. The 1.08 MB seemed like the wrong thing to spend for roughly 7% on
the largest shape.
Measured on an M3 Pro, Release, 12 threads, timing the whole node through an
InferenceSession, axis=0 float reduction='add' with row-broadcast indices:
before after speedup
N=4096 C=1024 2.52 ms 1.91 ms 1.3x
N=4096 C=256 0.35 ms 0.27 ms 1.3x
N=8192 C=64 0.20 ms 0.13 ms 1.6x
N=65536 C=32 1.42 ms 1.26 ms 1.1x
Adds tests for negative indices on the contiguous path, and for a slice that
mixes the two spellings of one row, which is not constant by raw value and so
has to fall back to the generic path.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Alexander Novikov (novikov-alexander)
force-pushed
the
scatter-index-no-materialize
branch
from
August 13, 2026 11:04
9d9124c to
7126841
Compare
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
1 participant
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.
Description
GetIndicesnormalized every index into astd::vector<int64_t>before any scatter work began. For a[4096, 1024]indices tensor that is a 33 MB allocation plus a full write pass, and the scatter then reads 8 bytes per element back rather than the 4 anint32indices tensor would have cost.This validates the indices in place and lets the scatter read the indices tensor directly, normalizing negatives at the point of use. Validation still runs up front, so an out-of-range index is still reported before the output is touched — the error path and its message are unchanged.
Why a reader rather than a template parameter
The obvious implementation templates
ScatterDataon the index type. I built that first and measured it: it instantiates the entire scatter — including the contiguous fast path — once per index type.scatter.ccThe templated version was about 7% faster on the largest shape. Spending 1.08 MB of object code for that seemed like the wrong trade for a runtime that ships to mobile and has a minimal-build mode, so the element type lives in a small
ScatterIndicesstruct instead. The branch it introduces is loop-invariant in practice, and on the fast path only one index is read per row anyway.Happy to switch to the templated version if maintainers weigh binary size differently — it's a small change either way.
Motivation and Context
After #32026 the scatter loop itself is no longer the bottleneck for row-broadcast indices; the up-front index materialization is. This removes it.
Measured on an M3 Pro, Release, 12 threads, timing the whole node through an
InferenceSession(axis=0, float,reduction='add', row-broadcast indices):Cumulatively against current
main, measured in the same session so the numbers are comparable:The allocation removal also matters on its own: peak memory for the node drops by
8 bytes x indices count, which for the first row above is 33 MB that no longer has to be reserved, written and read back.Tests
Existing coverage already exercises negative indices heavily — the
RunTesthelper deliberately generates them — and that now goes through the new normalize-on-read path.Added two cases for the interaction between negative indices and the contiguous path, which is new:
RowBroadcastIndicesNegative— a row of all-3addressing the same row as all1, on the fast path.RowBroadcastIndicesMixedSignFallsBack— a slice mixing-3and1. These are the same destination row but differ by raw value, so the constancy check rejects the slice and it must fall back to the generic path and still land correctly.GatherElementsGradshares this code and its caller is updated accordingly; note it is behindENABLE_TRAINING_OPS, which my local build does not compile, so that path is compile-checked by review and CI rather than by me.Full
onnxruntime_provider_testsweep: 5642 tests, 5454 passed, 0 failures, unchanged from baseline.