Refactor: CanonicalizeShapeExpr

Author: cchung100m

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

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

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

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

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

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

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

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

python/tvm/relax/transform/transform.py

@@ -736,17 +736,23 @@ def FoldConstant() -> tvm.ir.transform.Pass:
 
 
 def CanonicalizeShapeExpr() -> tvm.ir.transform.Pass:
-    """Canonicalize ShapeExpr by lifting compound PrimExpr into separate bindings.
+    """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)
+    - tir::Var (symbolic variable from function parameters or match_cast)
 
-    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 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
 
-    This pass should be applied after ComputePrimValue and before VMShapeLower.
+    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
     -------