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

[lrcyyds1]

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 [Fix][Relax] Normalize negative indices in Gather/Scatter/OneHot Ops #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

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

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 [Fix][Relax] Normalize negative indices in Gather/Scatter/OneHot Ops #20219 fix 3aa0eb1)
• Target llvm, exec_mode="compiled"; CPU; Python 3.10; Ubuntu 22.04 (Linux 6.8)

[lrcyyds1]

Follow-up: the same unchecked indices also fire at compile time — FoldConstant evaluates OOB gathers inside the compiler and bakes the result into the compiled artifact.

While fuzzing around this issue with a pattern generator, we found a third manifestation of the missing bounds check, this time during compilation rather than execution.

Mechanism. When a gather has an all-constant input, FoldConstant evaluates it at compile time. With an out-of-range index, that evaluation performs the OOB read inside the compiler process — reading
whatever lies adjacent to the constant in memory — and writes the result into the compiled module as if it were a legitimate folded constant. The wrong value is then frozen into the artifact: every run of the
compiled model returns it deterministically.

Minimal repro (main @ 2a2b293, llvm, CPU):

import numpy as np, tvm
from tvm import relax

C = np.clip(np.random.RandomState(27).randn(8), -1, 1).astype("float32")  # 8 elems, valid idx 0..7

def build(v):
    bb = relax.BlockBuilder()
    x = relax.Var("x", relax.TensorType([8], "float32"))
    with bb.function("main", params=[x]):
        with bb.dataflow():
            c = bb.emit(relax.const(C.tolist(), "float32"))
            k = bb.emit(relax.op.gather_elements(c, relax.const([v], "int64"), axis=0))
            y = bb.emit(relax.op.add(x, k))
            gv = bb.emit_output(y)
        bb.emit_func_output(gv)
    return bb.get()

def run(mod, fused, X):
    kw = {"relax_pipeline": relax.get_default_pipeline(tvm.target.Target("llvm"))} if fused else {}
    exe = tvm.relax.build(mod, target=tvm.target.Target("llvm"), exec_mode="compiled", **kw)
    return relax.VirtualMachine(exe, tvm.cpu())["main"](tvm.runtime.tensor(X, tvm.cpu())).numpy()

X = np.clip(np.random.RandomState(23).randn(8), -1, 1).astype("float32")
for v in (7, 8, 9):
    d, f = run(build(v), False, X), run(build(v), True, X)
    print(v, "default k =", d[0]-X[0], " fused(folded) k =", f[0]-X[0])

Output — the folded constant k for v=8/9 is a compile-time OOB read; the two official configurations disagree:

idx=7: default k=+0.2428   fused k=+0.2428      # in-bounds: consistent, correct
idx=8: default k=+1.1e13   fused k=+0.0000      # one-past-end: garbage vs baked-in OOB read
idx=9: default k=-1.4e11   fused k=+0.0000

Across a generated pattern sweep (gather/gather_nd on constant tensors × boundary indices, 3 official configs each), we measured 29 divergent patterns from this template alone. Notably, 14 of them diverge by
0.38–1.7 — i.e. normal-looking values, because the compile-time read lands in the compiler's constant pool (neighboring constants), not in zeroed heap. That makes the wrong-code silent: no crash, no NaN,
values in the expected range, frozen into the artifact, and different across build configurations.

We also observed the write-side variant: with a constant scatter_elements and a far out-of-range index, evaluating during folding corrupts the compiler process's heap (subprocess killed with SIGABRT,
malloc(): invalid next size).

So alongside the runtime effects already reported here (OOB reads via twin spellings, OOB writes/segfaults through the ONNX path), the same root cause also gives:

1. an OOB read inside the compiler while folding — compiling an untrusted model is itself a memory-safety surface, and
2. a silent wrong value permanently baked into the compiled artifact.

The op-level fix discussed above (normalize negative indices + range-check at legalization) would cover this path as well, since the folding evaluator ultimately calls the same op implementations.

[linsen458-spec]

Great report — I can confirm all three symptoms on current main, and I'd like to work on this.

Looking at the code, the fix seems to span a few layers, so I'd like to agree on the approach first:

1. Negative index normalization in topi.gather_nd / gather_elements (the C++ gather_nd in include/tvm/topi/transform.h passes indices through raw; scatter_elements already does k + (k < 0) * axis_range). This part is uncontroversial and matches the ONNX spec.
2. Compile-time validation of constant indices — the repro indices are relax.const, so validating them where they're folded would turn the SIGSEGV into a clear error. This would also cover the FoldConstant path mentioned above.
3. Runtime indices — a per-element bounds check in generated code has a perf cost, so I'd like to hear maintainer preference: an assert guarded by a debug flag, clamping, or documenting out-of-range as undefined?

Happy to start with (1) + (2) and leave (3) to whatever direction you prefer.