Add a contiguous fast path to the CPU ScatterElements kernel

[Alexander Novikov (novikov-alexander)]

Description

The generic scatter loop gives 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 == 1, there are only C work units, and each 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 that strided walk becomes a sequential one the compiler can vectorize. This adds that path, detected at runtime from the indices tensor.

Detection bails on the first mismatch, so ordinary indices pay for one slice's worth of reads.

Results are unchanged, bit for bit. Each work unit still owns a distinct (outer, inner-range) region of the output, so no two touch the same element; and within a unit the axis is walked in ascending order, which is the order the generic loop applies updates in. A destination receiving several updates therefore accumulates them in exactly the same sequence — which matters for float add. Splitting the contiguous run into blocks, rather than parallelizing over update rows, is what preserves that.

Two guards are load-bearing:

• Trailing dims must agree. The run is contiguous only if the updates and the data agree on every dimension after the axis. ScatterElements permits the updates to be smaller there, in which case an inner coordinate maps to different offsets in the two tensors. There is a test for this.
• Short runs are skipped. Below roughly a cache line the generic loop's stride is already local, so taking the fast path would trade away parallelism for nothing.

Motivation and Context

Row-broadcast indices are not an exotic case — they are how a row scatter is spelled in PyTorch, because Tensor.scatter_ requires index and src to have the same rank as the destination:

x.scatter_(0, idx.unsqueeze(-1).expand(-1, C), src)

That exports as Expand → ScatterElements, producing exactly indices[i][j] == indices[i][0]. The same shape shows up in scatter-add pooling, MoE dispatch/combine, KV-cache reorder, and in GatherElementsGrad whenever the forward was the torch.gather row idiom.

Because detection is at runtime rather than at graph level, it fires whether or not the Expand was constant-folded, and whether the row ids came from an initializer or from a runtime op.

Measured on an M3 Pro, Release, 12 threads, timing the whole node through an InferenceSession (axis=0, float, reduction='add', row-broadcast indices):

shape	before	after	speedup
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

To be straight about where the remaining time goes: the scatter loop itself speeds up considerably more than these node-level numbers suggest. GetIndices materializes a normalized int64_t for every element up front, and after this change that dominates. Narrowing it looks like a worthwhile follow-up, and it would compound with this.

Tests

• RowBroadcastIndicesAdd — the fast path for reduction='add', with two update rows landing on one destination row.
• RowBroadcastIndicesNoneKeepsLastUpdate — duplicate destinations under reduction='none'. Fails if the axis is visited in any order other than ascending, which is the property the bit-exactness argument rests on.
• RowBroadcastIndicesRank3 — rank 3, confirming the run spans all dimensions after the axis.
• RowBroadcastIndicesNarrowerTrailingDim — updates narrower than the data in the last dimension. Must stay on the generic path; taking the fast path here would write the second half of each slice at the wrong offset.

I confirmed by temporary instrumentation that the first three tests actually reach the fast path and the fourth does not — the existing suite would not have exercised it, since its RunTest helper fills indices per element with std::rand().

Full onnxruntime_provider_test sweep: 5640 tests, 5452 passed, 0 failures, unchanged from baseline. That includes GatherElements, which shares ScatterData.

[azure-pipelines]

Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

[Copilot]

Pull request overview

This PR adds a CPU ScatterElements optimization that detects row-broadcast (inner-constant) indices at runtime and, when safe, switches from a strided per-element loop to a contiguous-run loop that is friendlier to vectorization—while preserving bit-exact update order semantics (including for reduction="add" and reduction="none" cases with duplicates).

Changes:

• Added a runtime detector for “indices constant along inner dimensions” and a contiguous fast path in the CPU ScatterElements kernel when trailing dims match.
• Introduced new CPU provider tests covering the fast path, order sensitivity for reduction="none", higher-rank behavior, and a guard case where trailing dims differ.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

File	Description
onnxruntime/core/providers/cpu/tensor/scatter.cc	Adds indices-pattern detection and a contiguous-blocked parallel fast path when safe.
onnxruntime/test/providers/cpu/tensor/scatter_op_test.cc	Adds targeted tests to exercise the fast path and its safety/ordering guards.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Alexander Novikov (novikov-alexander)]

Thanks — that's a fair catch on the comment, and CI turned up two real problems with my tests. Both are now fixed in 693c978.

The comment. You're right that it over-promised. With TryParallelFor the other workers keep going until they reach their next slice boundary and observe the flag, so the work isn't bounded to one slice and the mismatch that stops the scan isn't necessarily the first one. Reworded to say that workers stop at their next slice boundary, so non-broadcast indices are rejected after a small amount of work rather than a full pass — which is the property the fast path actually relies on.

The CI failures, since they're worth being explicit about: all six were my tests, not the kernel. No CPU pipeline failed.

• RowBroadcastIndicesNoneKeepsLastUpdate failed on the CUDA plugin EP and on WebGPU. That test asserts the later update wins when several target one element — but ONNX leaves that unspecified, so it's a property of the CPU kernel, not of the operator, and providers that apply updates concurrently pick an arbitrary winner. My mistake was letting a kernel-ordering test run on every provider. It now runs on CPU only, which is what it was always testing. This also matches how the existing RunTest helper in this file sidesteps the issue, by making duplicate destinations receive identical values so ordering can't matter.
• RowBroadcastIndicesNarrowerTrailingDim failed on QNN, which returned 0 where 2 was expected. That test exists to confirm the CPU fast path declines a shape where an update slice isn't a contiguous run of the output, so it's likewise CPU-internal and now runs on CPU only.

On the second one: QNN produced a wrong value rather than declining the node, on data {4,2,8} / updates {2,2,4} with axis=0. That looks like it may be a genuine QNN gap, but it's unrelated to this PR and I haven't investigated — happy to file it separately if that's useful.

The two dependent PRs (#32055, #32059) have been rebased onto this.