lrcyyds1 opened a new issue, #20260:
URL: https://github.com/apache/tvm/issues/20260
### 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<IntImmNode>()->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<IntImmNode>()` (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.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
