microsoft / microsoft/onnxruntime
CPU ScatterElements always copies its data input, even when the producer's buffer could be reused in place
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
### Describe the issue
`ScatterData` copies the entire `data` input to the output before applying updates, and skips that copy only when the runtime has already aliased the two buffers:
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/providers/cpu/tensor/scatter.cc#L298-L309
```cpp
// We allow runtime to re-use input for output. If input/output Tensor* are the same
// we do not copy
if (src_base != dst_base) {
...
memcpy(static_cast(dst_base), static_cast(src_base), total_input_bytes);
}
```
All five `Scatter`/`ScatterElements` registrations declare `.MayInplace(0, 0)`, so the intent is clearly that the planner may hand the kernel an aliased buffer. **In practice I could not find any configuration where it does.** I instrumented the kernel to print `src_base`, `dst_base` and whether they were equal, and ran four graphs through a real `InferenceSession` with `ORT_ENABLE_ALL`:
| graph | aliased? |
|---|---|
| all-zero initializer → `ScatterElements` | no — full copy |
| `ConstantOfShape(Shape(x))` → `ScatterElements`, static shapes | no — full copy |
| `ConstantOfShape` → `ScatterElements` → `Relu` (scatter output not a graph output) | no — full copy |
| `ConstantOfShape` → `ScatterElements` → `Relu`, **symbolic** dims | no — full copy |
The last row is the one that matters most, since it is the shape real exports have. `torch.zeros(...).scatter_add_(...)` and ORT's own gradient builder both emit `ConstantOfShape(Shape(x))` rather than a stored initializer, so the zeros tensor is produced at runtime, consumed only by the scatter, and then never used again — exactly the case `MayInplace` exists for.
The cost is a full extra pass over the data tensor. For the `ConstantOfShape` pattern the node ends up writing N elements to materialize the zeros, then reading N and writing N to copy them, where N writes would do.
I have not confirmed the mechanism, only the effect. Reading `allocation_planner.cc`, the `MayInplace` path requires `SameSize(*p_input_arg, *p_output_arg)`, and `SameSize` returns false whenever either shape is unresolved:
https://github.com/microsoft/onnxruntime/blob/main/onnxruntime/core/framework/allocation_planner.cc#L548-L555
```cpp
// If the shapes are unknown, we conservatively assume they may be of different size.
if ((nullptr == p_shape1) || (nullptr == p_shape2)) return false;
```
That would explain the symbolic-shape rows but not the static ones, so there is probably more to it. Treat the table as the finding and the rest as a starting point.
If this is intended — for instance if in-place reuse is deliberately conservative for reasons I have not found — it might be worth a comment at the `MayInplace(0, 0)` call sites, since the kernel comment currently reads as though the aliased case is expected to be common.
### To reproduce
Build a graph of `Shape → ConstantOfShape(value=0) → ScatterElements(axis=0, reduction='add') → Relu` with symbolic leading dims, run it on the CPU EP, and observe in `ScatterData` that `data_input->DataRaw() != data_output->MutableDataRaw()`, so the `memcpy` runs. I can share the generator script if useful.
### Urgency
Low — this is a missed optimization, not a correctness problem. Filing it because it is invisible from the outside: the copy just shows up as part of the node's time.
### Platform
Mac
### Execution Provider
Default CPU
### Note
Related but separate: #32059 makes that copy multi-threaded, which helps when it does happen. It does not address whether the copy needs to happen at all, which is what this issue is about.
Contributor guide
Assessment
This issue has not been assessed yet.