[NFC] Remove uses of deprecated GEN_PASS_CLASSES for Triton/Transforms (#6785)

Continuation of triton-lang/triton#3971

Signed-off-by: Anatoly Myachev <[email protected]>

Author: anmyachev

bin/RegisterTritonDialects.h

@@ -42,7 +42,7 @@ void registerTestTritonAMDGPURangeAnalysis();
 
 inline void registerTritonDialects(mlir::DialectRegistry &registry) {
   mlir::registerAllPasses();
-  mlir::registerTritonPasses();
+  mlir::triton::registerTritonPasses();
   mlir::triton::gpu::registerTritonGPUPasses();
   mlir::registerTritonNvidiaGPUPasses();
   mlir::test::registerTestAliasPass();

include/triton/Dialect/Triton/Transforms/Passes.h

@@ -6,18 +6,14 @@
 namespace mlir {
 namespace triton {
 
-std::unique_ptr<Pass> createCombineOpsPass();
-
-std::unique_ptr<Pass> createLoopInvariantCodeMotionPass();
-std::unique_ptr<Pass> createReorderBroadcastPass();
-std::unique_ptr<Pass> createRewriteTensorPointerPass();
-std::unique_ptr<Pass> createLoopUnrollPass();
-
-} // namespace triton
+// Generate the pass class declarations.
+#define GEN_PASS_DECL
+#include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
 #define GEN_PASS_REGISTRATION
 #include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
+} // namespace triton
 } // namespace mlir
 
 #endif

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

@@ -19,8 +19,6 @@ def TritonCombineOps : Pass</*cli-arg*/"triton-combine", /*Op*/"mlir::ModuleOp">
        => dot(x,y,splat(0))`
   }];
 
-  let constructor = "mlir::triton::createCombineOpsPass()";
-
   let dependentDialects = ["mlir::arith::ArithDialect"];
 }
 
@@ -33,7 +31,7 @@ def TritonReorderBroadcast : Pass</*cli-arg*/"triton-reorder-broadcast", /*Op*/"
     In the event of a match, the broadcast (or splat) operation is delayed
     and performed after the ElementWise operation.
   }];
-  let constructor = "mlir::triton::createReorderBroadcastPass()";
+
   let dependentDialects = ["mlir::triton::TritonDialect"];
 }
 
@@ -45,8 +43,6 @@ def TritonRewriteTensorPointer : Pass</*cli-arg*/"triton-rewrite-tensor-pointer"
     the pointer/mask/other for each load/store.
   }];
 
-  let constructor = "mlir::triton::createRewriteTensorPointerPass()";
-
   let dependentDialects = ["mlir::triton::TritonDialect"];
 }
 
@@ -56,7 +52,7 @@ def TritonLoopUnroll : Pass</*cli-arg*/"triton-loop-unroll", /*Op*/"mlir::Module
     The pass unrolls a scf loop with tt.loop_unroll_factor attribute. The attribute specialises how many iterations
     the loop should be unrolled.
   }];
-  let constructor = "mlir::triton::createLoopUnrollPass()";
+
   let dependentDialects = ["mlir::triton::TritonDialect"];
 }
 
@@ -68,7 +64,7 @@ def TritonLoopInvariantCodeMotion : Pass</*cli-arg*/"triton-licm", /*Op*/"mlir::
     generates a trip-count check. For scf.while loops, it clones the condition
     from the before body.
   }];
-  let constructor = "mlir::triton::createLoopInvariantCodeMotionPass()";
+
   let dependentDialects = ["mlir::triton::TritonDialect"];
 }
 

lib/Dialect/Triton/Transforms/Combine.cpp

@@ -1,5 +1,3 @@
-#include <memory>
-
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
@@ -10,10 +8,11 @@
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 
-#define GEN_PASS_CLASSES
+namespace mlir::triton {
+
+#define GEN_PASS_DEF_TRITONCOMBINEOPS
 #include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
-namespace mlir::triton {
 namespace {
 
 bool isZero(Value val) {
@@ -240,7 +239,9 @@ class RankedReduceDescriptorLoads : public mlir::OpRewritePattern<ReshapeOp> {
   }
 };
 
-class CombineOpsPass : public TritonCombineOpsBase<CombineOpsPass> {
+} // anonymous namespace
+
+class CombineOpsPass : public impl::TritonCombineOpsBase<CombineOpsPass> {
 public:
   void runOnOperation() override {
     MLIRContext *context = &getContext();
@@ -264,10 +265,4 @@ class CombineOpsPass : public TritonCombineOpsBase<CombineOpsPass> {
   }
 };
 
-} // anonymous namespace
-
-std::unique_ptr<mlir::Pass> createCombineOpsPass() {
-  return std::make_unique<CombineOpsPass>();
-}
-
 } // namespace mlir::triton

lib/Dialect/Triton/Transforms/LoopInvariantCodeMotion.cpp

@@ -4,19 +4,18 @@
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 #include "llvm/Support/Debug.h"
 
-#define GEN_PASS_CLASSES
+namespace mlir::triton {
+
+#define GEN_PASS_DEF_TRITONLOOPINVARIANTCODEMOTION
 #include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
 #define DEBUG_TYPE "triton-licm"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
-namespace mlir::triton {
-
-namespace {
-
 class LoopInvariantCodeMotionPass
-    : public TritonLoopInvariantCodeMotionBase<LoopInvariantCodeMotionPass> {
+    : public impl::TritonLoopInvariantCodeMotionBase<
+          LoopInvariantCodeMotionPass> {
 
   DenseMap<LoopLikeOpInterface, bool> isLoopMemoryEffectFreeOrOnlyRead;
 
@@ -81,10 +80,4 @@ class LoopInvariantCodeMotionPass
   }
 };
 
-} // anonymous namespace
-
-std::unique_ptr<mlir::Pass> createLoopInvariantCodeMotionPass() {
-  return std::make_unique<LoopInvariantCodeMotionPass>();
-}
-
 } // namespace mlir::triton

lib/Dialect/Triton/Transforms/LoopUnroll.cpp

@@ -1,5 +1,3 @@
-#include <memory>
-
 #include "mlir/Dialect/SCF/Utils/Utils.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/Matchers.h"
@@ -12,18 +10,16 @@
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 #include "llvm/Support/Debug.h"
 
-#define GEN_PASS_CLASSES
+namespace mlir::triton {
+
+#define GEN_PASS_DEF_TRITONLOOPUNROLL
 #include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
 #define DEBUG_TYPE "triton-loop-unroll"
 #define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 #define LDBG(X) LLVM_DEBUG(DBGS() << X << "\n")
 
-namespace mlir::triton {
-
-namespace {
-
-class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
+class LoopUnrollPass : public impl::TritonLoopUnrollBase<LoopUnrollPass> {
 
   int getUnrollFactorOrDefault(scf::ForOp forOp) {
     // Use the attribute attached to the loop if it exists otherwise set the
@@ -38,8 +34,6 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
   const char *pipelineStagesAttrName = "tt.num_stages";
 
 public:
-  LoopUnrollPass() = default;
-  LoopUnrollPass(const LoopUnrollPass &) {}
   void runOnOperation() override {
     LDBG("Loop unroll pass");
     SmallVector<scf::ForOp, 4> loops;
@@ -65,10 +59,4 @@ class LoopUnrollPass : public TritonLoopUnrollBase<LoopUnrollPass> {
   }
 };
 
-} // anonymous namespace
-
-std::unique_ptr<mlir::Pass> createLoopUnrollPass() {
-  return std::make_unique<LoopUnrollPass>();
-}
-
 } // namespace mlir::triton

lib/Dialect/Triton/Transforms/ReorderBroadcast.cpp

@@ -10,12 +10,11 @@
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 
-// TODO(jlebar): Move this and all other generatede code into namespace
-// mlir::triton.
+namespace mlir::triton {
+
 #define GEN_PASS_DEF_TRITONREORDERBROADCAST
 #include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
-namespace mlir::triton {
 namespace {
 
 Operation *cloneWithNewArgsAndResultTypes(PatternRewriter &rewriter,
@@ -208,8 +207,10 @@ struct MoveBroadcastAfterElementwisePattern
   }
 };
 
+} // namespace
+
 class ReorderBroadcastPass
-    : public ::impl::TritonReorderBroadcastBase<ReorderBroadcastPass> {
+    : public impl::TritonReorderBroadcastBase<ReorderBroadcastPass> {
 public:
   void runOnOperation() override {
     MLIRContext *context = &getContext();
@@ -228,10 +229,4 @@ class ReorderBroadcastPass
   }
 };
 
-} // namespace
-
-std::unique_ptr<mlir::Pass> createReorderBroadcastPass() {
-  return std::make_unique<ReorderBroadcastPass>();
-}
-
 } // namespace mlir::triton

lib/Dialect/Triton/Transforms/RewriteTensorPointer.cpp

@@ -1,4 +1,3 @@
-#include <memory>
 #include <stack>
 
 #include "mlir/Dialect/ControlFlow/IR/ControlFlow.h"
@@ -10,9 +9,9 @@
 #include "triton/Dialect/Triton/IR/Utility.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
 
-using namespace mlir;
+namespace mlir::triton {
 
-#define GEN_PASS_CLASSES
+#define GEN_PASS_DEF_TRITONREWRITETENSORPOINTER
 #include "triton/Dialect/Triton/Transforms/Passes.h.inc"
 
 namespace {
@@ -196,7 +195,7 @@ struct RewritedInfo {
 // very fragile and to solve we should expose convert Ptr of tensor to a
 // structure containins all values and not only offsets.
 class RewriteTensorPointerPass
-    : public TritonRewriteTensorPointerBase<RewriteTensorPointerPass> {
+    : public impl::TritonRewriteTensorPointerBase<RewriteTensorPointerPass> {
 private:
   DenseMap<Value, RewritedInfo> rewritedInfo;
 
@@ -560,6 +559,4 @@ class RewriteTensorPointerPass
   }
 };
 
-std::unique_ptr<Pass> triton::createRewriteTensorPointerPass() {
-  return std::make_unique<RewriteTensorPointerPass>();
-}
+} // namespace mlir::triton

python/src/passes.cc

@@ -36,12 +36,12 @@ void init_triton_passes_common(py::module &&m) {
 
 void init_triton_passes_ttir(py::module &&m) {
   using namespace mlir::triton;
-  ADD_PASS_WRAPPER_0("add_combine", createCombineOpsPass);
-  ADD_PASS_WRAPPER_0("add_reorder_broadcast", createReorderBroadcastPass);
+  ADD_PASS_WRAPPER_0("add_combine", createTritonCombineOps);
+  ADD_PASS_WRAPPER_0("add_reorder_broadcast", createTritonReorderBroadcast);
   ADD_PASS_WRAPPER_0("add_rewrite_tensor_pointer",
-                     createRewriteTensorPointerPass);
-  ADD_PASS_WRAPPER_0("add_loop_unroll", createLoopUnrollPass);
-  ADD_PASS_WRAPPER_0("add_triton_licm", createLoopInvariantCodeMotionPass);
+                     createTritonRewriteTensorPointer);
+  ADD_PASS_WRAPPER_0("add_loop_unroll", createTritonLoopUnroll);
+  ADD_PASS_WRAPPER_0("add_triton_licm", createTritonLoopInvariantCodeMotion);
   ADD_PASS_OPTION_WRAPPER_4("add_convert_to_ttgpuir",
                             createConvertTritonToTritonGPU, const std::string &,
                             int, int, int);