[Gluon] Move gluon-specific passes into gluon dialect (#7489)

This renames `--tritongpu-canonicalize` -> `--gluon-canonicalize` and
also makes the gluon inliner pass into a real mlir pass so it can be
called from `triton-opt` as `--gluon-inline`.

Author: peterbell10

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

@@ -11,4 +11,28 @@ def GluonResolveAutoEncodingsPass : Pass<"gluon-resolve-auto-encodings", "mlir::
 
 }
 
+def GluonCanonicalize: Pass<"gluon-canonicalize"> {
+  let summary = "reduced set of simplifications for TTGIR";
+
+  let description = [{
+    The `gluon-canonicalize` pass applies a reduced set of simplification
+    and canonicalization patterns to the module.
+  }];
+  let dependentDialects = [
+    "mlir::arith::ArithDialect",
+    "mlir::cf::ControlFlowDialect",
+    "mlir::scf::SCFDialect",
+  ];
+}
+
+def GluonInline: Pass<"gluon-inline"> {
+  let summary = "reduced set of simplifications for TTGIR";
+
+  let description = [{
+    The `gluon-inline` pass applies a reduced set of simplification
+    and canonicalization patterns to the module.
+  }];
+  let dependentDialects = [];
+}
+
 #endif

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

@@ -360,18 +360,4 @@ def TritonGPUCoalesceAsyncCopy: Pass<"tritongpu-coalesce-async-copy", "mlir::Mod
                            "mlir::triton::TritonDialect"];
 }
 
-def TritonGPUCanonicalize: Pass<"tritongpu-canonicalize"> {
-  let summary = "reduced set of simplifications for TTGIR";
-
-  let description = [{
-    The `tritongpu-canonicalize` pass applies a reduced set of simplification
-    and canonicalization patterns to the module.
-  }];
-  let dependentDialects = [
-    "mlir::arith::ArithDialect",
-    "mlir::cf::ControlFlowDialect",
-    "mlir::scf::SCFDialect",
-  ];
-}
-
 #endif

lib/Dialect/Gluon/Transforms/CMakeLists.txt

@@ -1,4 +1,6 @@
 add_triton_library(GluonTransforms
+  Canonicalize.cpp
+  Inline.cpp
   ResolveAutoEncodings.cpp
 
   DEPENDS

lib/Dialect/TritonGPU/Transforms/Canonicalize.cpp renamed to lib/Dialect/Gluon/Transforms/Canonicalize.cpp

@@ -1,25 +1,28 @@
+#include "triton/Dialect/Gluon/Transforms/Passes.h"
+
+#include "triton/Dialect/TritonGPU/IR/Dialect.h"
+#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
+#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
+
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
-#include "triton/Dialect/TritonGPU/IR/Dialect.h"
-#include "triton/Dialect/TritonGPU/Transforms/Utility.h"
-#include "triton/Dialect/TritonNvidiaGPU/IR/Dialect.h"
 
 using namespace mlir;
 using namespace triton;
 namespace ttg = triton::gpu;
 namespace ttng = triton::nvidia_gpu;
+namespace gluon = mlir::triton::gluon;
 
-namespace mlir::triton::gpu {
-#define GEN_PASS_DEF_TRITONGPUCANONICALIZE
-#include "triton/Dialect/TritonGPU/Transforms/Passes.h.inc"
-} // namespace mlir::triton::gpu
+namespace mlir::triton::gluon {
+#define GEN_PASS_DEF_GLUONCANONICALIZE
+#include "triton/Dialect/Gluon/Transforms/Passes.h.inc"
+} // namespace mlir::triton::gluon
 
 namespace {
-struct Canonicalize
-    : public ttg::impl::TritonGPUCanonicalizeBase<Canonicalize> {
+struct Canonicalize : public gluon::impl::GluonCanonicalizeBase<Canonicalize> {
   void runOnOperation() override;
 };
 } // namespace

lib/Dialect/Gluon/Transforms/Inline.cpp

@@ -0,0 +1,29 @@
+#include "triton/Dialect/Gluon/Transforms/Passes.h"
+
+#include "mlir/Pass/Pass.h"
+#include "mlir/Pass/PassManager.h"
+#include "mlir/Transforms/Passes.h"
+
+using namespace mlir;
+using namespace triton;
+namespace gluon = mlir::triton::gluon;
+
+namespace mlir::triton::gluon {
+#define GEN_PASS_DEF_GLUONINLINE
+#include "triton/Dialect/Gluon/Transforms/Passes.h.inc"
+} // namespace mlir::triton::gluon
+
+namespace {
+struct Inline : public gluon::impl::GluonInlineBase<Inline> {
+  void runOnOperation() override;
+};
+} // namespace
+
+void Inline::runOnOperation() {
+  mlir::PassManager pm(&getContext());
+  pm.addPass(createInlinerPass(/*opPipelines=*/{}, [](OpPassManager &pm) {
+    pm.addPass(gluon::createGluonCanonicalize());
+  }));
+  if (failed(pm.run(getOperation())))
+    return signalPassFailure();
+}

lib/Dialect/TritonGPU/Transforms/CMakeLists.txt

@@ -1,6 +1,5 @@
 add_triton_library(TritonGPUTransforms
   AccelerateMatmul.cpp
-  Canonicalize.cpp
   Coalesce.cpp
   F32DotTC.cpp
   FuseNestedLoops.cpp

python/src/passes.cc

@@ -89,12 +89,6 @@ void init_triton_passes_ttgpuir(py::module &&m) {
   ADD_PASS_WRAPPER_0("add_fuse_nested_loops", createTritonGPUFuseNestedLoops);
   ADD_PASS_WRAPPER_0("add_coalesce_async_copy",
                      createTritonGPUCoalesceAsyncCopy);
-  ADD_PASS_WRAPPER_0("add_canonicalizer", createTritonGPUCanonicalize);
-  ADD_PASS_WRAPPER_0("add_inliner", [] {
-    return createInlinerPass(/*opPipelines=*/{}, [](OpPassManager &pm) {
-      pm.addPass(createTritonGPUCanonicalize());
-    });
-  });
   ADD_PASS_WRAPPER_0("add_concurrency_sanitizer",
                      createTritonInstrumentConcurrencySanitizer);
 }
@@ -117,6 +111,8 @@ void init_gluon_passes(py::module &&m) {
   namespace gluon = mlir::triton::gluon;
   ADD_PASS_WRAPPER_0("add_resolve_auto_encodings",
                      gluon::createGluonResolveAutoEncodingsPass);
+  ADD_PASS_WRAPPER_0("add_canonicalizer", gluon::createGluonCanonicalize);
+  ADD_PASS_WRAPPER_0("add_inliner", gluon::createGluonInline);
 }
 
 void init_triton_passes(py::module &&m) {

third_party/amd/backend/compiler.py

@@ -270,11 +270,11 @@ def gluon_to_ttgir(src, metadata, options):
         pm = ir.pass_manager(mod.context)
         pm.enable_debug()
 
-        passes.ttgpuir.add_inliner(pm)
+        passes.gluon.add_inliner(pm)
         passes.gluon.add_resolve_auto_encodings(pm)
         passes.common.add_sccp(pm)
         passes.ttir.add_loop_aware_cse(pm)
-        passes.ttgpuir.add_canonicalizer(pm)
+        passes.gluon.add_canonicalizer(pm)
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
 
         pm.run(mod)

third_party/nvidia/backend/compiler.py

@@ -314,11 +314,11 @@ def gluon_to_ttgir(self, src, metadata, options, capability):
         pm = ir.pass_manager(mod.context)
         pm.enable_debug()
 
-        passes.ttgpuir.add_inliner(pm)
+        passes.gluon.add_inliner(pm)
         passes.gluon.add_resolve_auto_encodings(pm)
         passes.common.add_sccp(pm)
         passes.ttir.add_loop_aware_cse(pm)
-        passes.ttgpuir.add_canonicalizer(pm)
+        passes.gluon.add_canonicalizer(pm)
         passes.ttgpuir.add_combine_tensor_select_and_if(pm)
 
         pm.run(mod)