[Bug][Relax] FoldConstant crashes on shape_to_tensor with symbolic shapes (AssertionError) and folds it into a Constant that breaks reshape
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
`shape_to_tensor` on a tensor with symbolic dims is a core part of the dynamic-reshape idiom (reshape to a runtime shape). `FoldConstant` should skip folding it when the shape values are not statically known,
and both official build pipelines should compile valid dynamic-shape modules.
### Actual behavior
Two failure modes on the current `main`:
1. `shape_of -> shape_to_tensor -> tensor_to_shape -> reshape` with a symbolic dim: the cpu_generic pipeline (`relax.get_default_pipeline`) crashes with an empty `AssertionError` inside `FoldConstant`.
2. `shape_of -> shape_to_tensor -> reshape` (without `tensor_to_shape`): `FoldConstant` folds the `shape_to_tensor` call into a `Constant` **tensor**, which `reshape` then rejects with `TypeError: Reshape
requires the input new shape to be Shape`.
### Environment
```text
OS: Linux x86_64
Target: llvm
TVM commit: 2a2b293c02269f4d9f3526c5b03a7548578e78e8 (current main)
```
### Steps to reproduce
```python
import tvm
from tvm import relax
from tvm import tirx as tir
from tvm.relax import transform
bb = relax.BlockBuilder()
x = relax.Var("x", relax.TensorType([tir.Var("m", "int64")], "float32"))
with bb.function("main", params=[x]):
with bb.dataflow():
s = bb.emit(relax.op.shape_of(x), "s")
t = bb.emit(relax.op.shape_to_tensor(s), "t")
b = bb.emit(relax.op.tensor_to_shape(t), "b")
y = bb.emit(relax.op.reshape(x, b), "y")
gv = bb.emit_output(y)
bb.emit_func_output(gv)
mod = bb.get()
with tvm.target.Target("llvm"):
transform.FoldConstant()(mod) # AssertionError
```
Behavior breakdown:
- cpu_generic pipeline: `AssertionError` in `FoldConstant`
- default pipeline (no `FoldConstant`): builds and runs correctly
- static control (`m` replaced by literal `8`): passes — a symbolic dim is the necessary condition
### Root cause
`src/relax/transform/fold_constant.cc`, the `relax.shape_to_tensor` special case:
```cpp
for (size_t i = 0; i < values.size(); i++) {
PrimExpr val = values[i];
arr.push_back(val.as()->value); // unchecked deref
is_known &= val.ty().MatchesElementType(DLDataTypeCode::kDLInt, 64);
}
```
When a shape value is a symbolic variable, the dereference happens **before** the `is_known` guard can skip it. Moving the IntImm check ahead of the `push_back` (treating non-IntImm values as `!is_known`)
looks sufficient for failure mode 1.
For failure mode 2, the folded result is a runtime `Constant` tensor while `reshape`'s type contract expects a `Shape`-typed operand, so the fold itself changes the operand kind.
### Suggested fix
Check `val->IsInstance()` (or use the checked form of `as<>()`) before dereferencing, and only fold when all shape values are concrete IntImms. For mode 2, avoid folding `shape_to_tensor` into a
`Constant` when any consumer requires a `Shape`-typed operand.
### Related
- Initially reported in #20193 (no minimal repro / root cause there; this issue adds both). Still reproduces on current `main` as of commit `2a2b293`.
This is the canonical dynamic-reshape idiom, so dynamic-shape models (dynamic batch / sequence length) hit it whenever the fusion pipeline runs.
Contributor guide
No contributing guide indexed for this repository
Research direction
Start in src/relax/transform/fold_constant.cc at the relax.shape_to_tensor special case, then run the provided symbolic-shape reproducer through FoldConstant and the cpu_generic pipeline. Verify that non-IntImm shape values are not dereferenced or folded, and that reshape still receives a Shape-typed operand in both failure modes.
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
- 68/100