Bug][Relax] Unchecked indices in gather/scatter ops: silent OOB reads and memory-corrupting OOB writes (reachable through the ONNX importer)
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
The `gather_nd` / `gather_elements` and `scatter_elements` / `scatter_nd` operators
perform **no bounds checking on index values** at the op/topi level. Consequences,
all reproducible on current `main` (which includes the #20219 fix):
1. `scatter_elements` / `scatter_nd` with an out-of-range index **write outside the
output buffer** — with a far out-of-range index this is a deterministic
**SIGSEGV**, and it is reachable from a hand-written ONNX model through the
official `from_onnx` → `relax.build` → VM pipeline.
2. A one-past-the-end scatter index writes outside the buffer **silently**, with no
diagnostic (we also observed delayed `malloc(): invalid next size` aborts after
several such writes).
3. `gather_nd` / `gather_elements` with an in-range **negative** index or an
out-of-range index silently read out-of-bounds memory: the same index spelled
four equivalent ways in one process (static/symbolic shapes × constant/
runtime-computed indices) returns **different wrong values** (zeros or garbage
on the order of 1e5–1e30, varying run to run).
## Expected behavior
- In-range negative indices (`[-s, s-1]` for an axis of size `s`) should count from
the end — the semantics ONNX defines and that #20219 already implements in the
ONNX importer, and that `topi.scatter_elements` already implements internally
(`shifted = k + (k < 0) * axis_range`).
- Out-of-range indices should raise a clear error (as ONNX Runtime does), or at
minimum be documented as rejected — never silently read/write outside the tensor.
## Actual behavior
No diagnostic of any kind; OOB reads return adjacent memory (zeros or garbage);
OOB writes corrupt the heap or segfault.
### Repro 1 — deterministic segfault, native relax
```python
import numpy as np, tvm
from tvm import relax
def scatter_nd(v):
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType([8], "float32"))
with bb.function("main", params=[x]):
with bb.dataflow():
y = bb.emit(relax.op.scatter_nd(
x, relax.const([[v]], "int64"), relax.const([9.0], "float32")))
gv = bb.emit_output(y)
bb.emit_func_output(gv)
return bb.get()
X = np.clip(np.random.RandomState(23).randn(8), -1, 1).astype("float32")
exe = tvm.relax.build(scatter_nd(10_000_000),
target=tvm.target.Target("llvm"), exec_mode="compiled")
print(relax.VirtualMachine(exe, tvm.cpu())["main"](
tvm.runtime.tensor(X, tvm.cpu())).numpy()) # → SIGSEGV
```
Same for `relax.op.scatter_elements(..., axis=0)` with index `10_000_000`
(`returncode -11`). With index `8` (axis length is 8) the program "succeeds":
the update is silently written outside the buffer and the returned output is the
input unchanged — no error. After several such boundary writes we observed
`malloc(): invalid next size (unsorted)` + `SIGABRT` (delayed corruption).
### Repro 2 — same through the official ONNX path
```python
import numpy as np, onnx, tvm
from onnx import helper, TensorProto
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
def scatternd_onnx(idx_val):
data = helper.make_tensor_value_info("data", TensorProto.FLOAT, [8])
out = helper.make_tensor_value_info("y", TensorProto.FLOAT, [8])
init = [helper.make_tensor("indices", TensorProto.INT64, [1, 1], [idx_val]),
helper.make_tensor("updates", TensorProto.FLOAT, [1], [9.0])]
node = helper.make_node("ScatterND", ["data", "indices", "updates"], ["y"])
return helper.make_model(
helper.make_graph([node], "g", [data], [out], initializer=init),
opset_imports=[helper.make_opsetid("", 17)])
mod = from_onnx(scatternd_onnx(10_000_000))
exe = tvm.relax.build(mod, target=tvm.target.Target("llvm"), exec_mode="compiled")
X = np.clip(np.random.RandomState(23).randn(8), -1, 1).astype("float32")
print(relax.VirtualMachine(exe, tvm.cpu())["main"](
tvm.runtime.tensor(X, tvm.cpu())).numpy()) # → SIGSEGV (139, core dumped)
```
With `idx_val = 8` the model runs to completion and returns the input unchanged —
the write landed outside the tensor, silently. So the memory-safety issue is
reachable from ONNX models (malformed, corrupted, or fuzzed weights/indices),
not only from hand-written Relax.
### Repro 3 — gather: same index, four equivalent spellings, four different answers
Same value `-4`, same input, same process — differing only in how the index /
shape is spelled (all four must produce identical results by definition):
| index | expected (ONNX/numpy) | static/const | symbolic/const | static/computed | symbolic/computed |
|---|---|---|---|---|---|
| `-4` | `X[4] = 0.7017` | `+302242.5` | `+0.0` | `+343455.5` | `+0.0` |
| `8` (one past end) | error | `+339527.0` | `+3.5e+29` | `+0.0` | `+0.0` |
"computed" = index produced at runtime via `shape_to_tensor(ShapeExpr([expr(d]))`
(the pattern produced by shape arithmetic); "symbolic" = symbolic input shape.
No error is raised in any cell; the garbage values differ run to run, i.e. the
result depends on adjacent memory contents.
## Root cause
| location | issue |
|---|---|
| `include/tvm/topi/transform.h:1572` (`gather_nd`) | `data(real_indices)` with raw index values — no clamp, no check |
| `include/tvm/topi/transform.h:1517` (`gather`) | same |
| `python/tvm/topi/scatter_elements.py:121-124` | negative indices are intentionally shifted (`k + (k<0)*axis_range`) but the shifted index is never range-checked; positive OOB passes through raw |
| `python/tvm/topi/scatter.py` (`gen_ir`) | `scatter_nd` has no negative shift at all and no bounds check — inconsistent with its sibling `scatter_elements` |
| `src/relax/op/tensor/manipulate.cc:2285+` | `InferTypeGatherND` validates structure only (dtype/rank/batch_dims); no index value contract, and the op docstrings don't state one either |
## Relation to #20219
#20219 (fix 3aa0eb1, included in the tree used here) adds
`_normalize_negative_indices` — `where(idx < 0, idx + extent, idx)` — in the
**ONNX importer** for Gather/GatherElements/GatherND/ScatterND/OneHot. That
handles runtime-computed negative indices correctly, but:
1. it covers only models entering through the ONNX frontend — native Relax
programs and other frontends are unprotected (Repro 3 reads OOB on the very
commit that contains the fix);
2. it normalizes negatives only — **positive out-of-range indices are unchecked
everywhere**, which is what makes the OOB *write* in Repro 2 reachable
through the already-fixed path;
3. the underlying ops still have no index contract, so every frontend has to
re-implement the same patch.
## Suggested fix
Sink the handling into the op layer (legalization or topi), once, for all
frontends:
```
idx = where(idx < 0, idx + axis_extent, idx) # already the #20219 semantics
error/clamp if idx not in [0, axis_extent) # currently missing everywhere
```
This also resolves the sibling inconsistency (`scatter_elements` shifts
negatives, `scatter_nd` doesn't) and would let the per-frontend normalization
in the ONNX importer retire.
## Environment
- TVM built from source, `main` @ `2a2b293` (includes #20219 fix `3aa0eb1`)
- Target `llvm`, `exec_mode="compiled"`; CPU; Python 3.10; Ubuntu 22.04 (Linux 6.8)
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the three reproductions first, then inspect include/tvm/topi/transform.h, python/tvm/topi/scatter_elements.py, python/tvm/topi/scatter.py, and src/relax/op/tensor/manipulate.cc. Trace how native Relax and the ONNX importer reach these operators. Done means negative indices follow the stated semantics and every gather/scatter path rejects positive and negative out-of-range indices without memory corruption.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 5/5
- Estimated time
- Over a week
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 35/100