Skip to content

[FIX][Relax] Preserve symbolic shape dependencies in FuseOps - #20333

Open
Junius-Wynn wants to merge 1 commit into
apache:mainfrom
Junius-Wynn:fix/issue-20177-vm-shape-lower
Open

Junius-Wynn wants to merge 1 commit into
apache:mainfrom
Junius-Wynn:fix/issue-20177-vm-shape-lower

Conversation

@Junius-Wynn

@Junius-Wynn Junius-Wynn commented Sep 13, 2026

Copy link
Copy Markdown
Contributor

Summary

Fixes #20177.

FuseOps can schedule a call_tir before the match_cast that defines its output shape when another branch fuses with a later consumer. The resulting allocation references a shape symbol before it is defined, causing VMShapeLower to fail on the ONNX Size/Slice/Reshape model in the issue.

Track the first definition of each shape symbol and include its uses in both fusion partitioning and group scheduling. Preserve the ordinary input dependencies of match_cast to prevent cyclic fusion through these shape dependencies, and detect cycles when ordering groups. The collector uses StructuralVisit and StructuralWalk, treats scalar parameters as already defined, isolates definitions in each If branch, and respects parameter-type definitions, repeated match_cast checks, and symbols local to callable signatures.

Add Relax regressions for definition-before-use ordering, valid fusion, branch-local symbol scope, and scalar parameters in pattern fusion. Add two ONNX execution cases using the original pipeline with FuseOps(fuse_opt_level=2), comparing VM outputs against ONNX Runtime and the expected result. Explicitly set the test models to ONNX IR version 8 so newer ONNX packages do not generate models that CI's ONNX Runtime cannot load. The VMShapeLower assertion remains in place.

Validation

Local validation on macOS arm64 with Python 3.12 and LLVM 15, using this checkout's TVM and a tvm-ffi build matching the checked-out submodule:

  • TVM build: passed.
  • Targeted tests: 134 passed — FuseOps (45), FuseTIR (35), VMShapeLower (10), FuseOpsByPattern (30), MergeCompositeFunctions (12), and the two ONNX regression cases.
  • The If-scope and scalar-parameter regressions both pass structural equality checks against the expected IR. Both ONNX cases compile and execute with matching numerical results.
  • Additional before/after checks with ONNX IR 8 and IR 10 both reproduce the original VMShapeLower assertion using the parent commit's FuseOps implementation, and pass compilation, execution, and numerical comparisons using the fix.
  • Both ONNX regressions also passed with the default model IR forced to 13, while the test models remained explicitly pinned to IR 8.
  • Changed-file pre-commit hooks in an isolated environment: passed, including Ruff checks, Ruff formatting, and clang-format.
  • git diff --check: passed.

Regression command, run from the repository root in the configured build/test environment with this checkout's TVM and matching tvm-ffi:

python -m pytest \
  tests/python/relax/test_transform_fuse_ops.py \
  tests/python/relax/test_transform_fuse_tir.py \
  tests/python/relax/test_backend_transform_shape_lower.py \
  tests/python/relax/test_transform_fuse_ops_by_pattern.py \
  tests/python/relax/test_transform_merge_composite_functions.py \
  tests/python/relax/test_frontend_onnx.py::test_size_slice_reshape_with_fusion -q

@Junius-Wynn
Junius-Wynn force-pushed the fix/issue-20177-vm-shape-lower branch from 5416119 to 192c254 Compare September 15, 2026 04:53
FuseOps can move a call_tir before the match_cast that defines its output
shape when another branch fuses with a later consumer. This leaves the
allocation shape undefined and causes VMShapeLower to fail on the
Size/Slice/Reshape model in apache#20177.

Track the first definition of each shape symbol in both graph partitioning
and group scheduling. Preserve match_cast input edges and detect cycles
in the fusion group dependency graph. Isolate definitions in each If
branch, recognize scalar parameters as existing definitions, and respect
repeated match_cast checks and callable-type scopes.

Add Relax regressions for shape dependencies, branch-local definitions,
and scalar parameters in pattern fusion. Add two ONNX execution cases
using the original pipeline at fusion level 2, checking VM results against
ONNX Runtime while preserving valid fusion. Pin the test models to ONNX
IR 8 so newer ONNX packages remain compatible with CI's ONNX Runtime.

Validation:
- TVM build passed on macOS arm64 with LLVM 15.
- 134 tests passed across FuseOps, FuseTIR, VMShapeLower, FuseOpsByPattern,
  MergeCompositeFunctions, and the ONNX regression.
- ONNX regressions passed with the default model IR forced to 13.
- Changed-file pre-commit hooks and git diff --check passed.

Fixes apache#20177
@Junius-Wynn
Junius-Wynn force-pushed the fix/issue-20177-vm-shape-lower branch from 192c254 to a1734c0 Compare September 15, 2026 04:59
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] Static-shape Size/Slice/Reshape hits an uncomputed shape slot

1 participant