[Bug][Relax] split with out-of-range indices: type layer clamps but topi computes a negative extent ? bad_alloc / uint error / silent oversized (OOB-read) output depending on config
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
`split` with an out-of-range split point is handled inconsistently across compiler layers, producing three different outcomes for the same module:
1. **default pipeline, static shapes** → `std::bad_alloc` (the negative extent `-2` is consumed as a huge unsigned allocation size),
2. **fused (cpu_generic) pipeline, static shapes** → `InternalError: cannot make uint from negative value -2` (the same `-2` hits an IntImm→unsigned conversion),
3. **fused pipeline, symbolic shapes** → **compiles and runs silently, returning a tensor of shape (10,) from an input of shape (8,)** — the two extra elements are read out of bounds and materialized into the
user-visible output.
The root cause is a layer disagreement:
- `InferTypeSplit` (`src/relax/op/tensor/manipulate.cc`) **clamps** indices into `[0, dim)` and computes `split_dim = max(right - left, 0)` — the type layer promises a valid shape;
- `split_indices_array` (`include/tvm/topi/transform.h`) validates that indices are **sorted** but not that they are **in range**, and computes the output extent with plain subtraction: `src_axis_size -
begin_ids[i]` = `8 - 10 = -2`, which flows into the output tensor shape unchecked.
## Expected behavior
Indices outside `[0, axis_length)` should be rejected when the op is built (or at legalization), consistently across pipelines. ONNX `Split` requires indices within the axis extent (and non-decreasing);
out-of-range indices are an input error, not a reason to allocate 2⁶⁴ bytes or to return a larger-than-input tensor.
Notably, the **unsorted-but-in-range** case is already correctly rejected by an ICHECK in `split_indices_array` — only the range check is missing.
## Repro (`main` @ `2a2b293`, llvm, CPU)
```python
import numpy as np, tvm
from tvm import relax
from tvm import tirx as tir
def build(indices, sym):
d = tir.Var("d", "int64") if sym else 8
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType([d], "float32"))
with bb.function("main", params=[x]):
with bb.dataflow():
y = bb.emit(relax.TupleGetItem(relax.op.split(x, indices, axis=0), 0))
gv = bb.emit_output(y)
bb.emit_func_output(gv)
return bb.get()
def run(mod, fused):
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(np.ones(8, "float32"), tvm.cpu())).numpy()
for indices in ([10], [5, 3], [3, 5]):
for sym in (False, True):
outs = {}
for fused in (False, True):
try:
outs["fused" if fused else "default"] = str(run(build(indices, sym), fused).shape)
except Exception as e:
outs["fused" if fused else "default"] = f"ERR:{str(e)[:50]}"
print(f"indices={indices} sym={sym}: default={outs['default']} fused={outs['fused']}")
```
Output:
```
indices=[10] sym=False: default=ERR:std::bad_alloc fused=ERR:cannot make uint from negative value -2
indices=[10] sym=True: default=ERR:std::bad_alloc fused=(10,) # silent: larger than input, OOB read
indices=[5,3] sym=False: default=ERR:Check failed: idx_node->value > ... fused=ERR:... # unsorted in-range: correctly rejected
indices=[3,5] sym=False: default=(3,) fused=(3,) # in-range: correct
```
The `fused + symbolic` cell is the severe one: the model compiles and runs, and `main` returns a 10-element tensor built from an 8-element input — a silent wrong-code with an out-of-bounds read baked into the
output.
## Root cause
| layer | location | behavior |
|---|---|---|
| type inference | `src/relax/op/tensor/manipulate.cc`, `InferTypeSplit` | clamps: `min(max(idx,0),dim)`, `max(right-left, 0)` |
| legalization | `python/tvm/relax/transform/legalize_ops/manipulate.py`, `_split` | passes indices through to `topi.split` |
| topi | `include/tvm/topi/transform.h`, `split_indices_array` | sorted check only; extent = `src_axis_size - begin_ids[i]` (no range clamp) → `-2` enters the output shape |
| consumers | e.g. `include/tvm/tirx/op.h:973` | negative IntImm → unsigned conversion error; allocation path → `bad_alloc`; symbolic path → oversized runtime extent |
## Suggested fix
Mirror the type layer's clamping (or better, reject) in `split_indices_array`: ICHECK each index into `[0, src_axis_size]` at legalization time, so all pipelines reject the module consistently instead of
disagreeing downstream. The unsorted check is already there; this adds the missing range check.
## Environment
- TVM built from source, `main` @ `2a2b293`; target `llvm`, `exec_mode="compiled"`; Python 3.10; Ubuntu 22.04.
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the supplied reproducer first, then compare InferTypeSplit in src/relax/op/tensor/manipulate.cc with _split in python/tvm/relax/transform/legalize_ops/manipulate.py and split_indices_array in include/tvm/topi/transform.h. Trace how out-of-range indices reach the output extent and consumers such as include/tvm/tirx/op.h:973. Done means out-of-range indices are rejected consistently across default, fused, static, and symbolic pipelines while valid and unsorted cases retain their stated behavior.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- cpp, python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 64/100