[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

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

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

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

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

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

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

@@ -63,6 +63,7 @@ def finalize_passes(target: tvm.target.Target):  # pylint: disable=unused-argume
         relax.transform.LowerAllocTensor(),
         relax.transform.KillAfterLastUse(),
         relax.transform.LowerRuntimeBuiltin(),
+        relax.transform.CanonicalizeShapeExpr(),
         relax.transform.ComputePrimValue(),
         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,32 @@ def FoldConstant() -> tvm.ir.transform.Pass:
     return _ffi_api.FoldConstant()  # type: ignore
 
 
+def CanonicalizeShapeExpr() -> tvm.ir.transform.Pass:
+    """Canonicalize ShapeExpr by replacing compound PrimExpr with fresh symbolic variables.
+
+    VMShapeLower can only handle ShapeExpr where each dimension is either:
+    - IntImm (concrete integer constant)
+    - tir::Var (symbolic variable from function parameters or match_cast)
+
+    This pass transforms compound PrimExpr (e.g., n+1, 4*n*m) by:
+    1. Creating a fresh tir::Var for each compound expression
+    2. Emitting a MatchCast that binds the fresh var to a PrimValue computing the expression
+    3. Replacing the compound expression in ShapeExpr with teh fresh var
+
+    Example transformation:
+        Before: y = R.zeros(R.shape([n + 1]), dtype="float32")
+        After:  _s0_pv: R.Prim(value=_s0) = R.match_cast(R.prim_value(n+1), R.Prim(value=_s0))
+                y = R.zeros(R.shape([_s0]), dtype="float32")
+
+    This pass should be applied before 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