[XPU] Conditionally add -tritonintelgpu-optimize-reduction-locality to pipeline (#2553)

Add the `-tritonintelgpu-optimize-reduction-locality` pass to the
pipeline if the `TRITON_INTEL_OPTIMIZE_REDUCTION_LOCALITY` is set to 1.

As shown in #2266, this pass gives quite promising results, although
there is still room for improvement. Conditionally enabling it will
greatly help performance investigation.

Signed-off-by: victor-eds <[email protected]>

Author: victor-eds

third_party/intel/backend/compiler.py

@@ -245,6 +245,8 @@ def make_ttgir(mod, metadata, opt, properties):
         passes.common.add_cse(pm)
         passes.ttgpuir.add_prefetch(pm)
         passes.ttgpuir.add_optimize_dot_operands(pm, True)
+        if os.getenv("TRITON_INTEL_OPTIMIZE_REDUCTION_LOCALITY", "0") == 1:
+            intel.passes.ttgpuir.add_optimize_reduction_locality(pm)
         intel.passes.ttgpuir.add_remove_layout_conversions(pm)
         intel.passes.ttgpuir.add_reduce_data_duplication(pm)
         passes.ttgpuir.add_reorder_instructions(pm)

third_party/intel/triton_xpu.cc

@@ -99,6 +99,8 @@ void init_triton_intel_passes_ttgpuir(py::module &&m) {
                      gpu::intel::createTritonIntelGPUReduceDataDuplication);
   ADD_PASS_WRAPPER_0("add_materialize_block_pointer",
                      gpu::intel::createTritonIntelGPUMaterializeBlockPointer);
+  ADD_PASS_WRAPPER_0("add_optimize_reduction_locality",
+                     gpu::intel::createTritonIntelGPUOptimizeReductionLocality);
 }
 
 void init_triton_intel(py::module &&m) {