gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

make numStages an option in PipelinePass

Browse Source
This commit is contained in:
Yan Da
2022-05-23 12:47:55 +08:00
parent 39b1235082
commit 36c45ec687
6 changed files with 23 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/src/triton.cc
View File

@@ -1341,8 +1341,8 @@ 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());
.def("add_tritongpu_pipeline_pass", [](mlir::PassManager &self, int numStages) {
self.addPass(mlir::createTritonGPUPipelinePass(numStages));
})
.def("add_triton_gpu_combine_pass", [](mlir::PassManager &self) {
self.addPass(mlir::triton::gpu::createCombineOpsPass());