[Bug] [Relax]Symbolic Flatten/Slice/Cast triggers a duplicate remap assertion
- 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. ONNX checker, ONNX Runtime session construction, and `from_onnx` succeed. Compiling it with the shown Relax pipeline fails in `FuseTIR`:
```text
Traceback (most recent call last):
...
File ".../src/s_tir/transform/renew_defs.cc", line 167, in AddDefRemap
TVM_FFI_ICHECK(remap_.count(source) == 0)
tvm.error.InternalError: Check failed: (remap_.count(source) == 0) is false:
```
### 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 FuseTIR/RenewDefs Flatten-Slice-Cast 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, ["n", 8])
out = helper.make_tensor_value_info("out", TensorProto.INT64, [1, 1])
initializers = [
helper.make_tensor("starts", TensorProto.INT64, [1], [0]),
helper.make_tensor("ends", TensorProto.INT64, [1], [1]),
helper.make_tensor("axes", TensorProto.INT64, [1], [1]),
helper.make_tensor("steps", TensorProto.INT64, [1], [1]),
]
nodes = [
helper.make_node("Flatten", ["x"], ["flat"], axis=0),
helper.make_node(
"Slice",
["flat", "starts", "ends", "axes", "steps"],
["sliced"],
),
helper.make_node("Cast", ["sliced"], ["out"], to=TensorProto.INT64),
]
graph = helper.make_graph(
nodes,
"renew_defs_flatten_slice_cast",
[x],
[out],
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 frontend expands the float-to-int `Cast` into an `isfinite`/`where`/`astype` sequence. With the symbolic input dimension `n`, the normal fusion pipeline reaches `RenewDefs` and fails on a duplicate remap entry. Replacing `n` with a concrete dimension avoids this failure on this checkout.
### Triage
* needs-triage
Contributor guide
No contributing guide indexed for this repository
Research direction
Start by running the self-contained reproduction through the listed Relax pipeline and compare symbolic n with a concrete dimension. Then inspect src/s_tir/transform/renew_defs.cc around AddDefRemap and follow how FuseTIR reaches RenewDefs for the frontend's expanded Cast sequence. Done means the symbolic Flatten/Slice/Cast model compiles without the duplicate remap assertion.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- compilers, machine-learning
- Issue type
- Bug
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 48/100