[Backend] Refactor Transform Pipeline to support different backends (#2189)

* [Backend] Refactor Pipeline to support different backends

Introduce Pipeline abstraction in backend/pipeline.py with per-backend
registration. Each backend (cuda, hip, c, llvm) now registers its own
compilation pass pipeline. engine/lower.py resolves the pipeline from
the target instead of hardcoding phase imports.

* enhance the structure

* fuse two stage

* fix after rebase

* update template pass pipeline

* pipeline

* fix test

* fix test

* fix cuda pipeline prologue

Author: SiriusNEO

testing/python/issue/test_tilelang_issue_2123.py

@@ -4,7 +4,7 @@
 from tilelang import tvm
 from tvm import tirx
 from tvm.tirx import op
-from tilelang.engine.phase import LowerAndLegalize
+from tilelang.backend.cuda.pipeline import CUDAPassPipelineBodyPrologue
 from tilelang.transform import LowerAccessPtr
 
 
@@ -65,7 +65,7 @@ def test_issue_2123_atomic_load_lower_access_ptr_pipeline():
     func = issue_2123_atomic_load_repro(4).with_attr("global_symbol", "main")
     mod = tvm.IRModule.from_expr(func)
 
-    lowered = LowerAndLegalize(mod, target)
+    lowered = CUDAPassPipelineBodyPrologue(mod, target)
 
     _assert_access_ptr_lowered(lowered)
 

testing/python/transform/test_tilelang_transform_inject_tcgen05_fence.py

@@ -2,7 +2,7 @@
 from tilelang import tvm as tvm
 import tilelang as tl
 import tilelang.language as T
-from tilelang.engine.phase import LowerAndLegalize
+from tilelang.backend.cuda.pipeline import CUDAPassPipelineBodyPrologue
 from tvm import tirx
 
 
@@ -118,7 +118,7 @@ def func(X: T.Tensor((256, 256), T.float16), Y: T.Tensor((256, 256), T.float16))
 
     mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
     with sm100_target:
-        mod = LowerAndLegalize(mod, sm100_target)
+        mod = CUDAPassPipelineBodyPrologue(mod, sm100_target)
         mod = tl.transform.LowerSharedTmem()(mod)
 
     body = mod["main"].body
@@ -166,7 +166,7 @@ def func(X: T.Tensor((256, 256), T.bfloat16)):
 
     mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
     with sm100_target:
-        mod = LowerAndLegalize(mod, sm100_target)
+        mod = CUDAPassPipelineBodyPrologue(mod, sm100_target)
         mod = tl.transform.LowerSharedTmem()(mod)
 
     body = mod["main"].body

testing/python/transform/test_tilelang_transform_lexical_alloc_scope.py

@@ -12,7 +12,7 @@
 import tilelang as tl
 import tilelang.language as T
 from tilelang import tvm
-from tilelang.engine.phase import LowerAndLegalize
+from tilelang.backend.cuda.pipeline import CUDAPassPipelineBodyPrologue
 from tvm.tirx.stmt_functor import post_order_visit
 import tilelang.testing
 
@@ -51,7 +51,7 @@ def _apply_lower_opaque_pipeline(func, target, pass_configs=None):
     mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
     pass_configs = pass_configs or {}
     with target, tvm.transform.PassContext(config=pass_configs):
-        mod = LowerAndLegalize(mod, target)
+        mod = CUDAPassPipelineBodyPrologue(mod, target)
         mod = tl.transform.LowerSharedTmem()(mod)
         mod = tl.transform.IfStmtBinding()(mod)
         mod = tl.transform.PlanAndUpdateBufferAllocationLocation()(mod)

testing/python/transform/test_tilelang_transform_lower_shared_barrier.py

@@ -4,7 +4,7 @@
 from tilelang.utils.target import determine_target
 import tilelang.language as T
 import tilelang.testing
-from tilelang.engine.phase import LowerAndLegalize
+from tilelang.backend.cuda.pipeline import CUDAPassPipelineBodyPrologue
 from tvm import tirx
 
 auto_target = tvm.target.Target(determine_target("auto"))
@@ -158,7 +158,7 @@ def func(
     target = tvm.target.Target({"kind": "cuda", "arch": "sm_100"})
     with tvm.transform.PassContext(config=pass_configs), target:
         mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
-        mod = LowerAndLegalize(mod, target)
+        mod = CUDAPassPipelineBodyPrologue(mod, target)
         mod = tl.transform.LowerSharedTmem()(mod)
         mod = tl.transform.IfStmtBinding()(mod)
         mod = tl.transform.PlanAndUpdateBufferAllocationLocation()(mod)

testing/python/transform/test_tilelang_transform_plan_update_buffer_allocation_location.py

@@ -2,14 +2,14 @@
 import tilelang.language as T
 import tilelang.testing
 from tilelang import tvm
-from tilelang.engine.phase import LowerAndLegalize
+from tilelang.backend.cuda.pipeline import CUDAPassPipelineBodyPrologue
 
 
 def _apply_plan_update(func: tvm.tirx.PrimFunc) -> tvm.IRModule:
     target = tvm.target.Target("cuda")
     mod = tvm.IRModule.from_expr(func.with_attr("global_symbol", "main"))
     with target:
-        mod = LowerAndLegalize(mod, target)
+        mod = CUDAPassPipelineBodyPrologue(mod, target)
         mod = tl.transform.LowerSharedTmem()(mod)
         mod = tl.transform.IfStmtBinding()(mod)
         mod = tl.transform.PlanAndUpdateBufferAllocationLocation()(mod)

tilelang/backend/__init__.py

@@ -1 +1,6 @@
-# Backend packages for Metal (other backends live in cpu/, cuda/, rocm/).
+# Import built-in backend packages so their pipelines register.
+from . import cpu as _cpu  # noqa: F401,E402
+from . import common as _common  # noqa: F401,E402
+from . import cuda as _cuda  # noqa: F401,E402
+from . import metal as _metal  # noqa: F401,E402
+from . import rocm as _rocm  # noqa: F401,E402

tilelang/backend/common.py

@@ -0,0 +1,7 @@
+from __future__ import annotations
+
+from tilelang.backend.pipeline import Pipeline, register_pipeline
+from tilelang.backend.cpu.pipeline import CPUPassPipelineBody
+
+
+register_pipeline(Pipeline("webgpu", CPUPassPipelineBody))

tilelang/backend/cpu/__init__.py

@@ -0,0 +1 @@
+from . import pipeline  # noqa: F401

tilelang/backend/cpu/pipeline.py

@@ -0,0 +1,88 @@
+from __future__ import annotations
+
+from tvm import IRModule, s_tir, tirx
+from tvm.target import Target
+
+import tilelang
+from tilelang.backend.pipeline import Pipeline, register_pipeline
+from tilelang.backend.pipeline_utils import (
+    LayoutVisual,
+    allow_vectorize,
+    should_disable_shared_memory_reuse,
+    should_enable_aggressive_merge,
+    should_enable_race_check,
+    should_force_let_inline,
+)
+
+
+def CPUPassPipelineBody(mod: IRModule, target: Target) -> IRModule:
+    mod = tirx.transform.BindTarget(target)(mod)
+    pass_ctx = tilelang.transform.get_pass_context()
+
+    if should_force_let_inline():
+        mod = tilelang.transform.LetInline()(mod)
+    mod = tilelang.transform.AddWrapperForSingleBufStore()(mod)
+    mod = tilelang.transform.LegalizeNegativeIndex()(mod)
+    if should_enable_race_check():
+        mod = tilelang.transform.VerifyParallelLoop()(mod)
+    mod = tilelang.transform.InjectAssumes()(mod)
+    mod = tilelang.transform.Simplify()(mod)
+    mod = tilelang.transform.LayoutReducer()(mod)
+
+    mod = tilelang.transform.IfStmtBinding()(mod)
+    mod = tilelang.transform.PipelinePlanning()(mod)
+    mod = tilelang.transform.InjectSoftwarePipeline()(mod)
+    mod = tilelang.transform.Simplify()(mod)
+
+    mod = tilelang.transform.LayoutInference()(mod)
+    LayoutVisual(mod)
+    mod = tilelang.transform.LowerTileOp()(mod)
+
+    mod = tilelang.transform.DecoupleTypeCast()(mod)
+    mod = tilelang.transform.LegalizeVectorizedLoop()(mod)
+    mod = tilelang.transform.LegalizeSafeMemoryAccess()(mod)
+    mod = tilelang.transform.LowerAccessPtr()(mod)
+    mod = tilelang.transform.Simplify()(mod)
+    mod = tilelang.transform.HoistNonRestrictParams()(mod)
+
+    mod = tilelang.transform.PlanAndUpdateBufferAllocationLocation()(mod)
+    mod = tilelang.transform.HoistGlobalBufferAllocations()(mod)
+    mod = tilelang.transform.LowerOpaqueBlock()(mod)
+    mod = tilelang.transform.Simplify()(mod)
+    mod = tirx.transform.NarrowDataType(32)(mod)
+    mod = tilelang.transform.FlattenBuffer()(mod)
+    mod = tilelang.transform.ConfigIndexBitwidth()(mod)
+    mod = tirx.transform.Simplify()(mod)
+    mod = tilelang.transform.VectorizeLoop(enable_vectorize=allow_vectorize(pass_ctx=pass_ctx))(mod)
+    mod = tilelang.transform.StorageRewrite()(mod)
+    mod = tilelang.transform.LoopUnswitching()(mod)
+    mod = tilelang.transform.UnrollLoop()(mod)
+    mod = s_tir.transform.RenormalizeSplitPattern()(mod)
+    mod = tirx.transform.Simplify()(mod)
+    mod = tirx.transform.RemoveNoOp()(mod)
+    mod = s_tir.transform.HoistIfThenElse()(mod)
+
+    mod = tirx.transform.VerifyMemory()(mod)
+    mod = tirx.transform.AnnotateEntryFunc()(mod)
+    mod = s_tir.transform.InferFragment()(mod)
+    mod = tilelang.transform.LowerThreadAllreduce()(mod)
+
+    mod = tilelang.transform.AnnotateDeviceRegions()(mod)
+    mod = tilelang.transform.SplitHostDevice()(mod)
+    mod = tilelang.transform.AnnotateReadOnlyParams()(mod)
+
+    enable_aggressive_merge = should_enable_aggressive_merge(pass_ctx=pass_ctx, target=target)
+    disable_reuse = should_disable_shared_memory_reuse(pass_ctx=pass_ctx)
+    mod = tilelang.transform.MergeSharedMemoryAllocations(enable_aggressive_merge=enable_aggressive_merge, disable_reuse=disable_reuse)(mod)
+
+    mod = tilelang.transform.ThreadSync("shared")(mod)
+    mod = tilelang.transform.ThreadSync("shared.dyn")(mod)
+    mod = tilelang.transform.MergeIfStmt()(mod)
+    mod = tilelang.transform.MakePackedAPI()(mod)
+    mod = tilelang.transform.Simplify()(mod)
+    mod = tilelang.transform.LowerDeviceKernelLaunch()(mod)
+    return mod
+
+
+for _kind in ("c", "llvm"):
+    register_pipeline(Pipeline(_kind, CPUPassPipelineBody))

tilelang/backend/cuda/__init__.py

@@ -0,0 +1 @@
+from . import pipeline  # noqa: F401