More on the pipeline pass

python/src/triton.cc

@@ -16,6 +16,7 @@
 #include "triton/Dialect/Triton/IR/Dialect.h"
 #include "triton/Dialect/Triton/IR/Types.h"
 #include "triton/Dialect/Triton/Transforms/Passes.h"
+#include "triton/Dialect/TritonGPU/Transforms/Passes.h"
 
 #include "llvm/IR/Module.h"
 #include "llvm/IR/LegacyPassManager.h"
@@ -1340,6 +1341,9 @@ void init_triton_ir(py::module &&m) {
       .def("add_convert_triton_to_tritongpu_pass", [](mlir::PassManager &self) {
         self.addPass(mlir::triton::createConvertTritonToTritonGPUPass());
       })
+      .def("add_tritongpu_pipeline_pass", [](mlir::PassManager &self) {
+        self.addPass(mlir::createTritonGPUPipelinePass());
+      })
       ;
 }
 

python/triton/code_gen.py

@@ -1314,6 +1314,7 @@ class JITFunction:
         pm.add_triton_combine_pass()
         pm.add_canonicalizer_pass()
         pm.add_convert_triton_to_tritongpu_pass()
+        pm.add_tritongpu_pipeline_pass()
         pm.run(mod)
         return mod