[Bug] [Relax] Static-shape Size/Slice/Reshape hits an uncomputed shape slot
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
### Expected behavior
The Relax LLVM compilation pipeline should handle this valid ONNX model without an internal assertion.
### Actual behavior
The script constructs a valid ONNX model with static input and output shapes. ONNX checker, ONNX Runtime session construction, and `from_onnx` succeed. Compiling it with the shown Relax pipeline fails in `VMShapeLower`:
```text
Traceback (most recent call last):
...
File ".../src/relax/backend/vm/vm_shape_lower.cc", line 505, in MakeSymbolicShapeArg
TVM_FFI_ICHECK(slot->value_computed)
tvm.error.InternalError: Check failed: (slot->value_computed) is false: PrimExpr T.int64(4) * s in function I.GlobalVar("main") has not been computed
```
### Environment
* OS: macOS 15.6 (Darwin 24.6.0, arm64)
* Python: 3.12.2
* TVM: c7b458e946bc4266915da582457476bdcd9705ae (tag v0.26.0; package reports 0.26.dev0)
* ONNX: 1.17.0
* ONNX Runtime: 1.21.1
* Target: `llvm`
* Frontend: `tvm.relax.frontend.onnx.from_onnx`
### Steps to reproduce
The following self-contained script constructs and compiles the model:
```python
#!/usr/bin/env python3
"""Reproduce the VMShapeLower Size/Slice/Reshape failure."""
import onnx
import onnxruntime as ort
import tvm
from onnx import TensorProto, helper
from tvm import relax
from tvm.relax.frontend.onnx import from_onnx
def main() -> None:
x = helper.make_tensor_value_info("x", TensorProto.FLOAT, [1, 4, 5, 5])
y = helper.make_tensor_value_info("y", TensorProto.FLOAT, [1, 4, 5, 5])
initializers = [
helper.make_tensor("flat_shape", TensorProto.INT64, [1], [-1]),
helper.make_tensor("size_shape", TensorProto.INT64, [1], [1]),
helper.make_tensor("slice_starts", TensorProto.INT64, [1], [0]),
helper.make_tensor("slice_axes", TensorProto.INT64, [1], [0]),
helper.make_tensor("slice_steps", TensorProto.INT64, [1], [1]),
]
nodes = [
helper.make_node("Relu", ["x"], ["h0"]),
helper.make_node("Relu", ["h0"], ["positive"]),
helper.make_node("Reshape", ["x", "flat_shape"], ["flat"]),
helper.make_node("Size", ["flat"], ["size"]),
helper.make_node("Reshape", ["size", "size_shape"], ["size_1d"]),
helper.make_node(
"Slice",
["flat", "slice_starts", "size_1d", "slice_axes", "slice_steps"],
["projected"],
),
helper.make_node("Shape", ["x"], ["x_shape"]),
helper.make_node("Reshape", ["projected", "x_shape"], ["splice"]),
helper.make_node("Add", ["positive", "splice"], ["y"]),
]
graph = helper.make_graph(
nodes,
"vm_shape_lower_size_slice",
[x],
[y],
initializers,
)
model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 18)])
onnx.checker.check_model(model)
ort.InferenceSession(model.SerializeToString(), providers=["CPUExecutionProvider"])
mod = from_onnx(model, opset=18, keep_params_in_input=True)
pipeline = tvm.transform.Sequential(
[
relax.backend.DispatchSampling(),
relax.backend.DispatchSortScan(),
relax.transform.LegalizeOps(),
relax.transform.AnnotateTIROpPattern(),
relax.transform.FoldConstant(),
relax.transform.FuseOps(fuse_opt_level=2),
relax.transform.FuseTIR(),
relax.transform.RewriteDataflowReshape(),
relax.transform.ToNonDataflow(),
relax.transform.RemovePurityChecking(),
relax.transform.CallTIRRewrite(),
relax.transform.StaticPlanBlockMemory(),
relax.transform.LowerAllocTensor(),
relax.transform.KillAfterLastUse(),
relax.transform.LowerRuntimeBuiltin(),
relax.transform.ComputePrimValue(),
relax.transform.VMShapeLower(emit_err_ctx=True),
relax.transform.AttachGlobalSymbol(),
]
)
tvm.compile(mod, target="llvm", relax_pipeline=pipeline)
if __name__ == "__main__":
main()
```
### Analysis
The model has static input and output tensor shapes, but uses `Size(flat)` as a `Slice` end before a later `Reshape`. In `VMShapeLower`, this path reaches `T.int64(4) * s` with a slot whose `value_computed` flag is false. `FuseOps(fuse_opt_level=0)` avoids the failure on this checkout.
### Triage
* needs-triage
Contributor guide
No contributing guide indexed for this repository
Research direction
Reproduce the self-contained script and inspect src/relax/backend/vm/vm_shape_lower.cc at MakeSymbolicShapeArg, especially the value_computed assertion around line 505. Trace how ComputePrimValue and VMShapeLower handle the Size, Slice, and later Reshape after FuseOps; compare with the fuse_opt_level=0 workaround. Done means the shown pipeline compiles the model without the assertion and regression coverage protects the case.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100