apache / apache/tvm

[Bug][Relax] FoldConstant crashes on shape_to_tensor with symbolic shapes (AssertionError) and folds it into a Constant that breaks reshape

Open
#20,260 0 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

### 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.