Integrate

Author: victor-eds

third_party/intel/backend/compiler.py

@@ -245,6 +245,7 @@ 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)
+        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/lib/TritonIntelGPUToLLVM/ConvertLayoutOpToLLVM.cpp

@@ -486,8 +486,20 @@ struct ConvertLayoutOpUsingLinearLayoutsConversion
     constexpr std::size_t laneIndex = 0;
     constexpr std::size_t registerIndex = 1;
     int32_t size = conversion->getInDimSize(kLane);
+    std::vector<std::vector<int32_t>> registerBases =
+      buildBasis(size, registerIndex);
+    {
+      // Populate register bases for N > 8.
+      std::vector<int32_t> base(2);
+      for (int32_t i = registerBases.back()[registerIndex] * 2,
+	     n = conversion->getInDimSize(kRegister); i < n;
+           i *= 2) {
+        base.front() = i;
+        registerBases.push_back(base);
+      }
+    }
     std::array<std::pair<StringAttr, std::vector<std::vector<int32_t>>>, 2>
-        bases{{{kRegister, buildBasis(size, registerIndex)},
+      bases{{{kRegister, std::move(registerBases)},
                {kLane, buildBasis(size, laneIndex)}}};
     std::array<StringAttr, 2> outDimNames{kRegister, kLane};
     return conversion == LinearLayout(bases, outDimNames);

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) {