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

add cse pass to the pipeline & pass num-warps as an argument

Browse Source
This commit is contained in:
Yan Da
2022-06-10 17:31:48 +08:00
parent 117a402c1b
commit 0ee6e486f8
2 changed files with 8 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
python/src/triton.cc
View File

@@ -1343,11 +1343,14 @@ void init_triton_ir(py::module &&m) {
.def("add_canonicalizer_pass", [](mlir::PassManager &self) {
self.addPass(mlir::createCanonicalizerPass());
})
.def("add_cse_pass", [](mlir::PassManager &self) {
self.addPass(mlir::createCSEPass());
})
.def("add_triton_combine_pass", [](mlir::PassManager &self) {
self.addPass(mlir::triton::createCombineOpsPass());
})
.def("add_convert_triton_to_tritongpu_pass", [](mlir::PassManager &self) {
self.addPass(mlir::triton::createConvertTritonToTritonGPUPass());
.def("add_convert_triton_to_tritongpu_pass", [](mlir::PassManager &self, int numWarps) {
self.addPass(mlir::triton::createConvertTritonToTritonGPUPass(numWarps));
})
.def("add_tritongpu_pipeline_pass", [](mlir::PassManager &self, int numStages) {
self.addPass(mlir::createTritonGPUPipelinePass(numStages));

4
python/triton/code_gen.py
View File

@@ -1316,9 +1316,11 @@ class JITFunction:
pm.add_inliner_pass()
pm.add_triton_combine_pass()
pm.add_canonicalizer_pass()
pm.add_convert_triton_to_tritongpu_pass()
pm.add_cse_pass()
pm.add_convert_triton_to_tritongpu_pass(num_warps)
pm.add_tritongpu_pipeline_pass(num_stages)
pm.add_canonicalizer_pass()
pm.add_cse_pass()
pm.add_triton_gpu_combine_pass()
pm.add_triton_gpu_verifier_pass()
return pm.run(mod)