apache / apache/tvm

[Bug][Relax] BundleModelParams output (tuple param with symbolic-shape fields) segfaults VMShapeLower

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

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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.