[Bug][Relax] BundleModelParams output (tuple param with symbolic-shape fields) segfaults VMShapeLower
- Dominant language
- Python
- Stars
- 13.7k
- Forks
- 4k
- Avg merge
- 2d 1h
- Merged PRs (30d)
- 112
Description
Expected behavior
A Relax module that the compiler cannot build should be rejected with a clean diagnostic. In particular, relax.transform.BundleModelParams followed by relax.build (default pipeline) — or by
relax.transform.VMShapeLower alone — must not crash the compiler process.
Actual behavior
For a module with a R.Prim parameter and a symbolic-shape tensor parameter (where num_input is smaller than the number of parameters), BundleModelParams bundles them into a tuple parameter whose field types
lose their shapes (R.Tensor(dtype="float32", ndim=1) — shape is unknown), and synthesizes a match_cast to re-bind the extent and the weight. A subsequent relax.build (default pipeline) segfaults inside
VMShapeLowerMutator::VisitType_(TensorTypeNode const*, ...) while walking the tuple-typed parameter. The crash is deterministic (3/3 runs) and is compile-time (exec_mode irrelevant).
Key frames:
tvm::relax::VMShapeLowerMutator::VisitType_(tvm::relax::TensorTypeNode const*, ...)
tvm::relax::TypeFunctor<...>::VisitType(...)
tvm::relax::VMShapeLowerMutator::VisitType_(tvm::relax::TupleTypeNode const*, ...)
tvm::relax::VMShapeLowerMutator::Rewrite(tvm::GlobalVar, tvm::relax::Function)
tvm::relax::VMShapeLowerMutator::Lower(tvm::IRModule, bool)
Severity note: the same module without BundleModelParams is rejected cleanly by the VM codegen (Prim-typed parameters are not supported), so this bug downgrades a clean diagnostic into a process-level
segfault. The module itself is well-formed Relax (it parses and the passes accept it); a pass emitting IR that crashes a later pass — instead of raising — is the defect.
Environment
OS: Linux x86_64
Target: llvm
TVM version: 0.26.dev0
TVM commit: 2a2b293 (current main at time of writing)
Steps to reproduce
import tvm
from tvm import relax
from tvm.script import ir as I, relax as R
MOD = """
@I.ir_module
class Module:
@R.function
def main(x: R.Tensor([32], "float32"), extent: R.Prim("int64"),
weight: R.Tensor(["extent"], "float32")):
R.func_attr({"num_input": 1})
out = R.add(x, weight)
return out
"""
from tvm.script import from_source
def
# Trigger: BundleModelParams first -> deterministic SIGSEGV in VMShapeLower
build(MOD, bundle=True)
# Single-pass reproduction: crash is inside VMShapeLower itself
mod = from_source(MOD)
mod = relax.transform.BundleModelParams()(mod)
mod = relax.transform.VMShapeLower()(mod) # segfaults without any other pass
After BundleModelParams the module looks like this (note the tuple field lost its shape, and the synthesized match_cast):
@I.ir_module
class Module:
@R.function
def main(x: R.Tensor((32,), dtype="float32"),
model_params: R.Tuple(T.int64, R.Tensor(dtype="float32", ndim=1))
) -> R.Tensor(dtype="float32", ndim=1):
R.func_attr({"num_input": 1})
extent: T.int64 = model_params[0]
weight: R.Tensor((extent,), dtype="float32") = \
R.match_cast(model_params[1], R.Tensor((extent,), dtype="float32"))
out: R.Tensor(dtype="float32", ndim=1) = R.add(x, weight)
return out
Controls (each removes exactly one trigger condition)
- Trigger: Prim param + symbolic-shape weight + num_input=1
without BMP: clean error (rc=1)
with BMP: SIGSEGV (rc=-11), 3/3 runs
- Same + match_cast in the body (the originally-found pattern)
without BMP: clean error
with BMP: SIGSEGV (match_cast in user code is NOT required, BMP synthesizes it)
- Static-shape weight (R.Tensor([32])), rest identical
without BMP: clean error
with BMP: clean error (symbolic shape in the bundled field is required)
- No Prim param (all-static, plain bundling)
without BMP: builds OK
with BMP: builds OK
- No num_input attribute (BMP is a no-op)
without BMP: clean error
with BMP: clean error
Notes on triage
- Searched upstream for VMShapeLower (16 hits) and BundleModelParams (8 hits): no matching report. #17876 / #18491 are opt_level=1 InternalErrors from the default pipeline — different failure and different
path. BundleModelParams issues on file are 2023–24 feature/PR threads.
- Suggested labels: needs-triage, type: bug, relax.
Contributor guide
No contributing guide indexed for this repository
Research direction
Run the provided Python reproducer, then inspect relax.transform.BundleModelParams and VMShapeLower, especially VMShapeLowerMutator::VisitType_ for TensorTypeNode and TupleTypeNode. Verify the symbolic-shape bundled-parameter case is rejected with a clean diagnostic rather than segfaulting, while the listed control cases retain their behavior.
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
- 55/100