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

2022-06-10 17:31:48 +08:00
parent 117a402c1b
commit 0ee6e486f8
2 changed files with 8 additions and 3 deletions
--- a/python/triton/code_gen.py
+++ b/python/triton/code_gen.py
@@ -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)