Parallelize the ScatterElements input-to-output copy - #32059
Open
Alexander Novikov (novikov-alexander) wants to merge 4 commits into
Open
Parallelize the ScatterElements input-to-output copy#32059Alexander Novikov (novikov-alexander) wants to merge 4 commits into
Alexander Novikov (novikov-alexander) wants to merge 4 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>
ScatterData copies the whole data input to the output before applying any
updates, and does it with a single memcpy while the operator thread pool sits
idle. For a 16 MB data tensor that is around a millisecond on one core.
Split the copy across the pool when the tensor is large enough to be worth it.
Small tensors keep the plain memcpy, where sharding would cost more than it
saves, and the string path is untouched.
Measured on an M3 Pro, Release, 12 threads, timing the whole node through an
InferenceSession, axis=0 float reduction='add':
before after speedup
N=4096 C=1024 2.25 ms 1.61 ms 1.40x
N=4096 C=256 0.23 ms 0.20 ms 1.16x
N=8192 C=64 0.15 ms 0.13 ms 1.16x
N=65536 C=32 1.00 ms 0.99 ms 1.02x
Against main on its own the same measurement shows no change beyond noise
(11.72 ms vs 11.47 ms on the first shape). That is expected: the strided
scatter loop dominates there and the copy is only about a tenth of the node.
It is worth landing once the loop cost is out of the way.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Alexander Novikov (novikov-alexander)
force-pushed
the
scatter-parallel-copy-stacked
branch
from
August 13, 2026 11:10
1ca4fc3 to
961affd
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
ScatterDatacopies the entire data input to the output before applying any updates, and does it with a singlememcpywhile the operator thread pool sits idle. For a 16 MB data tensor that's roughly a millisecond on one core.This splits the copy across the pool when the tensor is large enough to be worth it. Small tensors keep the plain
memcpy— sharding them costs more than it saves — and thestd::stringpath is untouched.Motivation and Context
Measured on an M3 Pro, Release, 12 threads, timing the whole node through an
InferenceSession(axis=0, float,reduction='add'), A/B back to back:Being straight about the sequencing: measured against
mainon its own, this change shows no improvement beyond noise — 11.72 ms vs 11.47 ms on the first shape. That's expected rather than surprising. Onmainthe strided scatter loop dominates the node and the copy is only about a tenth of it, so making the copy faster is invisible. Once #32026 and #32055 remove that loop overhead, the copy becomes the largest remaining component and parallelizing it is worth 1.4x.So this is only worth landing after those two. I'm raising it now so it isn't lost, not to jump the queue — happy for it to sit until they're in.
One thing I could not measure: this should help more on typical x86 servers than on Apple Silicon, where a single core already reaches a large fraction of total memory bandwidth. I have no x86 machine to check that on, so treat it as a hypothesis rather than a claim.
Tests
No behaviour change — the copy produces identical bytes whether it runs on one thread or twelve, and the existing suite covers the copy on every
Scatter/ScatterElementspath. Fullonnxruntime_provider_testsweep: 5642 tests, 5454 passed, 0 failures, unchanged from baseline.