More on the pipeline pass

2022-05-11 20:31:08 +08:00
parent 1a4fbed25b
commit d23d7b244c
4 changed files with 99 additions and 2 deletions
--- a/python/src/triton.cc
+++ b/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());
+      })
      ;
 }