[Relax][Backend] Fix TVM crashes with default relax pipeline when opt_level=1: InternalError: Check failed: (slot->value_computed) is false

python/tvm/relax/backend/cpu_generic/pipeline.py

@@ -53,6 +53,7 @@ def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argume
         relax.transform.KillAfterLastUse(),
         relax.transform.LowerRuntimeBuiltin(),
         relax.transform.ComputePrimValue(),
+        relax.transform.CanonicalizeShapeExpr(),
         relax.transform.VMShapeLower(),
         relax.transform.AttachGlobalSymbol(),
     ]

python/tvm/relax/backend/cuda/pipeline.py

@@ -65,6 +65,7 @@ def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argume
         relax.transform.KillAfterLastUse(),
         relax.transform.LowerRuntimeBuiltin(),
         relax.transform.ComputePrimValue(),
+        relax.transform.CanonicalizeShapeExpr(),
         relax.transform.VMShapeLower(),
         relax.transform.AttachGlobalSymbol(),
     ]

python/tvm/relax/backend/gpu_generic/pipeline.py

@@ -64,6 +64,7 @@ def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argume
         relax.transform.KillAfterLastUse(),
         relax.transform.LowerRuntimeBuiltin(),
         relax.transform.ComputePrimValue(),
+        relax.transform.CanonicalizeShapeExpr(),
         relax.transform.VMShapeLower(),
         relax.transform.AttachGlobalSymbol(),
     ]

python/tvm/relax/backend/rocm/pipeline.py

@@ -64,6 +64,7 @@ def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argume
         relax.transform.KillAfterLastUse(),
         relax.transform.LowerRuntimeBuiltin(),
         relax.transform.ComputePrimValue(),
+        relax.transform.CanonicalizeShapeExpr(),
         relax.transform.VMShapeLower(),
         relax.transform.AttachGlobalSymbol(),
     ]

python/tvm/relax/transform/__init__.py

@@ -28,6 +28,7 @@
     BundleModelParams,
     CallTIRRewrite,
     CanonicalizeBindings,
+    CanonicalizeShapeExpr,
     CombineParallelMatmul,
     ComputePrimValue,
     ConvertLayout,

python/tvm/relax/transform/transform.py

@@ -735,6 +735,26 @@ def FoldConstant() -> tvm.ir.transform.Pass:
     return _ffi_api.FoldConstant()  # type: ignore
 
 
+def CanonicalizeShapeExpr() -> tvm.ir.transform.Pass:
+    """Canonicalize ShapeExpr by lifting compound PrimExpr into separate bindings.
+
+    VMShapeLower can only handle ShapeExpr where each dimension is either:
+    - IntImm (concrete integer constant)
+    - tir::Var (symbolic variable)
+
+    This pass lifts compound PrimExpr (e.g., n+1, 4*n*m, etc.) into separate  shape bindings
+    with MatchCast to extract symbolic variables, ensuring VMShapeLower receives only
+    canonical shape expressions.
+
+    This pass should be applied after ComputePrimValue and before VMShapeLower.
+
+    Returns
+    -------
+    ret: tvm.ir.transform.Pass
+    """
+    return _ffi_api.CanonicalizeShapeExpr()  # type: ignore
+
+
 def ExpandTupleArguments() -> tvm.ir.transform.Pass:
     """Expand tuple arguments to internal functions