[Nvidia][TritonGPU] Gate RLC phases via compiler.py env switches

Author: zeroherolin

CMakeLists.txt

@@ -35,7 +35,6 @@ endif()
 if(NOT FLAGTREE_BACKEND)
   add_definitions(-D__NVIDIA__)
   add_definitions(-D__AMD__)
-  add_definitions(-D__FLAGTREE_RLC_ENHANCE__)
 elseif(FLAGTREE_BACKEND STREQUAL "iluvatar")
   add_definitions(-D__ILUVATAR__)
   set(FLAGTREE_TLE OFF)

include/triton/Dialect/TritonGPU/Transforms/Passes.td

@@ -273,6 +273,23 @@ def TritonGPURemoveLayoutConversions : Pass<"tritongpu-remove-layout-conversions
   let dependentDialects = ["mlir::triton::gpu::TritonGPUDialect",
                            "mlir::triton::TritonDialect"];
 
+  let options = [
+    Option<"enableCostBasedResolution", "enable-cost-based-resolution",
+           "bool", /*default=*/"false",
+           "Use cost-based conflict resolution instead of heuristic rules">,
+    Option<"enableBackwardPropagation", "enable-backward-propagation",
+           "bool", /*default=*/"false",
+           "Enable backward layout propagation from convert_layout ops">,
+    Option<"enableSmallComponentSolving", "enable-small-component-solving",
+           "bool", /*default=*/"false",
+           "Enable small local layout component solving">,
+    Option<"enableStoreLayoutRematerialization",
+           "enable-store-layout-rematerialization",
+           "bool", /*default=*/"false",
+           "Drop store/atomic writeback convert_layout ops by rematerializing "
+           "their address/mask chains in the value's layout">
+  ];
+
 }
 
 def TritonGPUOptimizeThreadLocality : Pass<"tritongpu-optimize-thread-locality", "mlir::ModuleOp"> {

lib/Dialect/TritonGPU/Transforms/RemoveLayoutConversions.cpp

@@ -47,22 +47,6 @@ namespace mlir::triton::gpu {
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
-// Master build-time switch for the enhancement phases. The NVIDIA backend
-// defines __FLAGTREE_RLC_ENHANCE__ from CMake (on by default); every other
-// backend leaves it undefined, keeping the pass at its original behavior.
-#ifdef __FLAGTREE_RLC_ENHANCE__
-static constexpr bool kFlagtreeRlcEnhance = true;
-#else
-static constexpr bool kFlagtreeRlcEnhance = false;
-#endif
-
-// Per-phase switches, on by default; each is AND-ed with the master switch at
-// its use site. Flip one to false here to disable that single phase.
-static constexpr bool kEnableCostBasedResolution = true;
-static constexpr bool kEnableBackwardPropagation = true;
-static constexpr bool kEnableSmallComponentSolving = true;
-static constexpr bool kEnableStoreLayoutRematerialization = true;
-
 namespace {
 
 #ifdef __TLE__
@@ -4489,10 +4473,11 @@ class TritonGPURemoveLayoutConversionsPass
     MLIRContext *context = &getContext();
     ModuleOp m = getOperation();
 
-    bool costBased = kFlagtreeRlcEnhance && kEnableCostBasedResolution;
-    bool backwardProp = kFlagtreeRlcEnhance && kEnableBackwardPropagation;
-    bool smallComponentSolving =
-        kFlagtreeRlcEnhance && kEnableSmallComponentSolving;
+    bool costBased = enableCostBasedResolution;
+    bool backwardProp = enableBackwardPropagation;
+    bool smallComponentSolving = enableSmallComponentSolving;
+    bool storeLayoutRemat = enableStoreLayoutRematerialization;
+
     if (!costBased) {
       backwardProp = false;
       smallComponentSolving = false;
@@ -4527,8 +4512,6 @@ class TritonGPURemoveLayoutConversionsPass
 
     cleanupConvertOps();
 
-    bool storeLayoutRemat =
-        kFlagtreeRlcEnhance && kEnableStoreLayoutRematerialization;
     bool changed = false;
     do {
       changed = false;

python/src/passes.cc

@@ -76,8 +76,18 @@ void init_triton_passes_ttgpuir(py::module &&m) {
   ADD_PASS_OPTION_WRAPPER_1("add_f32_dot_tc", createTritonGPUF32DotTC, bool);
   ADD_PASS_OPTION_WRAPPER_1("add_optimize_dot_operands",
                             createTritonGPUOptimizeDotOperands, bool);
-  ADD_PASS_WRAPPER_0("add_remove_layout_conversions",
-                     createTritonGPURemoveLayoutConversions);
+  m.def(
+      "add_remove_layout_conversions",
+      [](mlir::PassManager &pm, bool costBased, bool backwardProp,
+         bool smallComponentSolving, bool storeLayoutRemat) {
+        pm.addPass(createTritonGPURemoveLayoutConversions(
+            {costBased, backwardProp, smallComponentSolving,
+             storeLayoutRemat}));
+      },
+      py::arg("pm"), py::arg("enable_cost_based_resolution") = false,
+      py::arg("enable_backward_propagation") = false,
+      py::arg("enable_small_component_solving") = false,
+      py::arg("enable_store_layout_rematerialization") = false);
   ADD_PASS_WRAPPER_0("add_reduce_data_duplication",
                      createTritonGPUReduceDataDuplication);
   ADD_PASS_WRAPPER_0("add_allocate_warp_groups",

third_party/nvidia/backend/compiler.py

@@ -37,6 +37,25 @@ def get_ptxas(arch: int) -> knobs.NvidiaTool:
     return knobs.nvidia.ptxas_blackwell if arch >= 100 else knobs.nvidia.ptxas
 
 
+def _env_bool(name: str, default: bool) -> bool:
+    value = os.environ.get(name)
+    if value is None:
+        return default
+    return value.strip().lower() not in {"0", "false", "off", "no"}
+
+
+# flagtree: enable RemoveLayoutConversions phases for NVIDIA (env-configurable).
+def _add_remove_layout_conversions(pm):
+    master = _env_bool("FLAGTREE_RLC_ENHANCE", True)
+    passes.ttgpuir.add_remove_layout_conversions(
+        pm,
+        master and _env_bool("FLAGTREE_RLC_COST_BASED_RESOLUTION", True),
+        master and _env_bool("FLAGTREE_RLC_BACKWARD_PROPAGATION", True),
+        master and _env_bool("FLAGTREE_RLC_SMALL_COMPONENT_SOLVING", True),
+        master and _env_bool("FLAGTREE_RLC_STORE_LAYOUT_REMATERIALIZATION", True),
+    )
+
+
 @functools.lru_cache()
 def get_ptxas_version(arch: int = 80):
     mock_ver = knobs.nvidia.mock_ptx_version
@@ -299,7 +318,7 @@ def make_ttgir(mod, metadata, opt, capability):
         passes.ttgpuir.add_f32_dot_tc(pm, emuTF32)
         # TODO(Qingyi): Move PlanCTAPass to the front of CoalescePass
         nvidia.passes.ttnvgpuir.add_plan_cta(pm)
-        passes.ttgpuir.add_remove_layout_conversions(pm)
+        _add_remove_layout_conversions(pm)
         passes.ttgpuir.add_optimize_thread_locality(pm)
         tle.passes.add_early_assign_memory_space(pm)
         # begin flagtree tle
@@ -310,7 +329,7 @@ def make_ttgir(mod, metadata, opt, capability):
         # end flagtree tle
         passes.ttgpuir.add_accelerate_matmul(pm)
         tle.passes.add_lower_wgmma(pm)
-        passes.ttgpuir.add_remove_layout_conversions(pm)
+        _add_remove_layout_conversions(pm)
         passes.ttgpuir.add_optimize_dot_operands(pm, capability >= 80)
         tle.passes.add_promote_local_store_staging(pm)
         nvidia.passes.ttnvgpuir.add_optimize_descriptor_encoding(pm)
@@ -361,7 +380,7 @@ def make_ttgir(mod, metadata, opt, capability):
             tle.passes.add_lower_tma_copy(pm)
             tle.passes.add_schedule_tma_store_sync(pm)
             nvidia.passes.ttnvgpuir.add_tma_lowering(pm)
-        passes.ttgpuir.add_remove_layout_conversions(pm)
+        _add_remove_layout_conversions(pm)
         nvidia.passes.ttnvgpuir.add_interleave_tmem(pm)
         passes.ttgpuir.add_reduce_data_duplication(pm)
         passes.ttgpuir.add_reorder_instructions(pm)