apache / apache/tvm

Bug][Relax] Unchecked indices in gather/scatter ops: silent OOB reads and memory-corrupting OOB writes (reachable through the ONNX importer)

Open
#20,297 2 comments 0 reactions 0 assignees View on GitHub
needs-triage type: bug
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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.