[REFACTOR][TIR][ARITH] Phase out ControlFlowGraph, NarrowPredicateExpression, and rename Simplify to StmtSimplify

include/tvm/tirx/transform.h

@@ -94,11 +94,11 @@ TVM_DLL Pass UnrollLoop();
 TVM_DLL Pass RemoveNoOp();
 
 /*!
- * \brief Run arithmetic simplifications on the statements and expressions.
+ * \brief Run statement-level arithmetic simplifications on the TIR PrimFunc.
  *
  * \return The pass.
  */
-TVM_DLL Pass Simplify();
+TVM_DLL Pass StmtSimplify();
 
 /*!
  * \brief Convert an IRModule to be SSA form.

python/tvm/s_tir/backend/adreno/pipeline.py

@@ -44,7 +44,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
             s_tir.transform.LowerAutoCopy(),
             s_tir.transform.UnifyThreadBinding(),
             s_tir.transform.LowerMatchBuffer(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
             s_tir.transform.InjectPermutedLayout(),
             s_tir.transform.AnnotateIrregularLoop(),
             s_tir.transform.InjectSoftwarePipeline(),
@@ -68,7 +68,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
                 s_tir.transform.HoistIfThenElse(),
                 tirx.transform.UnrollLoop(),
                 s_tir.transform.RenormalizeSplitPattern(),
-                tirx.transform.Simplify(),
+                tirx.transform.StmtSimplify(),
                 tirx.transform.RemoveNoOp(),
                 s_tir.transform.RewriteUnsafeSelect(),
             ]

python/tvm/s_tir/pipeline.py

@@ -45,7 +45,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
             s_tir.transform.LowerAutoCopy(),
             s_tir.transform.UnifyThreadBinding(),
             s_tir.transform.LowerMatchBuffer(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
             s_tir.transform.InjectPermutedLayout(),
             s_tir.transform.AnnotateIrregularLoop(),
             s_tir.transform.InjectSoftwarePipeline(),
@@ -68,7 +68,7 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
                 s_tir.transform.HoistIfThenElse(),
                 tirx.transform.UnrollLoop(),
                 s_tir.transform.RenormalizeSplitPattern(),
-                tirx.transform.Simplify(),
+                tirx.transform.StmtSimplify(),
                 tirx.transform.RemoveNoOp(),
                 s_tir.transform.RewriteUnsafeSelect(),
             ]
@@ -137,10 +137,10 @@ def finalize_host_passes():  # pylint: disable=unused-argument
 def finalize_device_passes():  # pylint: disable=unused-argument
     """The default finalization passes for TIR backend."""
     device_pass_list = [
-        tir.transform.LowerWarpMemory(),
-        tir.transform.Simplify(),
-        tir.transform.LowerCustomDatatypes(),
-        tir.transform.LowerIntrin(),
+        tirx.transform.LowerWarpMemory(),
+        tirx.transform.StmtSimplify(),
+        tirx.transform.LowerCustomDatatypes(),
+        tirx.transform.LowerIntrin(),
     ]
     return tvm.ir.transform.Sequential(device_pass_list)
 

python/tvm/testing/utils.py

@@ -2017,7 +2017,7 @@ class object that inherits from `Exception`.
     .. code-block:: python
 
         class TestRemoveIf(tvm.testing.CompareBeforeAfter):
-            transform = tvm.tirx.transform.Simplify()
+            transform = tvm.tirx.transform.StmtSimplify()
 
             def before(A: T.Buffer(1, "int32")):
                 if True:

python/tvm/tirx/compilation_pipeline.py

@@ -33,13 +33,13 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
         passes = [
             tirx.transform.LowerInitBlock(),
             tvm.s_tir.transform.UnifyThreadBinding(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
             tirx.transform.FlattenBuffer(),
             tirx.transform.BF16ComputeLegalize(),
             tirx.transform.NarrowDataType(32),
             tirx.transform.VectorizeLoop(not bool(config.get("tir.disable_vectorize", False))),
             tirx.transform.UnrollLoop(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
         ]
         if not bool(config.get("tir.disable_cse_tir", False)):
             passes.append(tirx.transform.CommonSubexprElim())
@@ -73,14 +73,14 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
         passes = [
             tirx.transform.LowerTIRx(),
             tvm.s_tir.transform.UnifyThreadBinding(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
             tirx.transform.LowerTIRxOpaque(),
             tirx.transform.FlattenBuffer(),
             tirx.transform.BF16ComputeLegalize(),
             tirx.transform.NarrowDataType(32),
             tirx.transform.VectorizeLoop(not bool(config.get("tir.disable_vectorize", False))),
             tirx.transform.UnrollLoop(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
         ]
         if not bool(config.get("tir.disable_cse_tir", False)):
             passes.append(tirx.transform.CommonSubexprElim())
@@ -115,11 +115,11 @@ def _pipeline(mod: tvm.ir.IRModule, _ctx: tvm.transform.PassContext) -> tvm.ir.I
             tirx.transform.trn.TrnNaiveAllocator(),
             tirx.transform.LowerTIRx(),
             tvm.s_tir.transform.DecorateDeviceScope(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
             tirx.transform.LowerTIRxOpaque(),
             tvm.s_tir.transform.LoopPartition(),
             tvm.s_tir.transform.HoistIfThenElse(),
-            tirx.transform.Simplify(),
+            tirx.transform.StmtSimplify(),
             tirx.transform.RemoveNoOp(),
             tirx.transform.AnnotateEntryFunc(),
             tirx.transform.AnnotateDeviceRegions(),
@@ -146,7 +146,7 @@ def finalize_device_passes():  # pylint: disable=unused-argument
     """The default finalization passes for TIR backend."""
     device_pass_list = [
         tirx.transform.LowerWarpMemory(),
-        tirx.transform.Simplify(),
+        tirx.transform.StmtSimplify(),
         tirx.transform.LowerCustomDatatypes(),
         tirx.transform.LowerIntrin(),
     ]
@@ -161,7 +161,7 @@ def finalize_device_passes_tirx():  # pylint: disable=unused-argument
 
 def finalize_device_passes_trn():  # pylint: disable=unused-argument
     """The default finalization passes for TRN backend."""
-    device_pass_list = [tirx.transform.Simplify()]
+    device_pass_list = [tirx.transform.StmtSimplify()]
     return tvm.ir.transform.Sequential(device_pass_list)
 
 

python/tvm/tirx/operator/tile_primitive/trn/compose_op/unary_reduce.py

@@ -118,7 +118,7 @@ def impl():
         import tvm
 
         mod = tvm.IRModule({"main": impl})
-        mod = tvm.tirx.transform.Simplify()(mod)
+        mod = tvm.tirx.transform.StmtSimplify()(mod)
         return mod["main"]
     else:
         # fmt: off

python/tvm/tirx/transform/transform.py

@@ -210,20 +210,20 @@ def CommonSubexprElim():
     return _ffi_api.CommonSubexprElim()  # type: ignore
 
 
-@_ffi.register_object("tirx.transform.SimplifyConfig")
-class SimplifyConfig(_ffi.Object):
-    """Config for simplify pass"""
+@_ffi.register_object("tirx.transform.StmtSimplifyConfig")
+class StmtSimplifyConfig(_ffi.Object):
+    """Config for stmt simplify pass"""
 
 
-def Simplify():
-    """Run arithmetic simplifications on the statements and expressions.
+def StmtSimplify():
+    """Run statement-level arithmetic simplifications on the TIR PrimFunc.
 
     Returns
     -------
     fpass : tvm.transform.Pass
         The result pass
     """
-    return _ffi_api.Simplify()  # type: ignore
+    return _ffi_api.StmtSimplify()  # type: ignore
 
 
 def ConvertSSA():