Skip to content

[FIX][Relax] Guard unknown tensor shapes in VMShapeLower - #20345

Open
ruiling-smartbear wants to merge 2 commits into
apache:mainfrom
ruiling-smartbear:fix/vm-shape-unknown-tuple
Open

ruiling-smartbear wants to merge 2 commits into
apache:mainfrom
ruiling-smartbear:fix/vm-shape-unknown-tuple

Conversation

@ruiling-smartbear

@ruiling-smartbear ruiling-smartbear commented Sep 15, 2026 •

Copy link
Copy Markdown

Fixes #20335.

The static-shape shortcut dereferenced a missing ShapeExpr for unknown-shape tensor fields. Guard that shortcut and keep the existing runtime checks. Add coverage for unknown-shape tuple parameters and the reported BundleModelParams → VMShapeLower sequence.

Validation on Linux CPU / LLVM 16, based on cc0f9f07 and its pinned TVM FFI:

  • Original relevant tests: 19 passed. New unknown-shape reproducer and full build reproducer: both SIGSEGV (-11).
  • After the fix: shape-lowering and parameter-bundling suites, 21 passed.
  • The full build reproducer now raises an ordinary InternalError at VM codegen instead of crashing; this does not add Prim-parameter execution support.
  • Changed-file pre-commit checks passed. No GPU tests run.

Only apply the static-shape shortcut when a ShapeExpr exists. Keep unknown tensor shapes on the existing runtime-check path and cover bundled symbolic parameters.
@ruiling-smartbear
ruiling-smartbear force-pushed the fix/vm-shape-unknown-tuple branch from 3233f40 to e8a968f Compare September 15, 2026 21:57
@wwoosshh

wwoosshh commented Oct 7, 2026 •

Copy link
Copy Markdown

I checked this on top of current main (d55759e). The C++ change merges and builds cleanly, and test_check_weights_with_unknown_shape passes with it and crashes without it.

test_lower_bundled_symbolic_shape fails on main with AttributeError: module 'tvm.relax.script.ir_builder' has no attribute 'Prim', because R.Prim was removed with the new TVMScript parser in #20423. The current way to write this is to bind the prim parameter and reuse it in the weight's shape, as tests/python/relax/script/test_relax_script_dynamic_shape.py does:

            extent: T.int64,
            weight: R.Tensor([extent], "float32"),  # noqa: F821

This needs from __future__ import annotations at the top of the file, so that Python does not evaluate the annotation itself. With that, BundleModelParams produces the same R.Tuple(T.int64, R.Tensor(dtype="float32", ndim=1)) parameter as in #20335, the test segfaults without the C++ change and passes with it, all 12 tests in the file pass, and pre-commit is clean.

@wwoosshh wwoosshh left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the update. I re-checked f38615db7 merged onto current main (d55759e): all 12 tests in tests/python/relax/test_backend_transform_shape_lower.py pass, both new tests crash when vm_shape_lower.cc is reverted to main, and pre-commit passes on the changed files.

LGTM.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

3 participants