[Mlir-commits] [mlir] [MLIR][XeGPU] Improve multi-reduction layout inference (PR #224618)

[Artem Kroviakov]

================
@@ -2493,13 +2458,39 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
                  "source shape not divisible by sg_layout");
           sgData[i] = srcShape[i] / sgLayout[i];
           remainingSgCount /= sgLayout[i];
-          order[i] = remainOrder++;
         }
       }
-      DenseI32ArrayAttr resOrderAttr = DenseI32ArrayAttr::get(
-          context, SmallVector<int32_t>(order.begin(), order.end()));
-      if (!orderAttr || orderAttr.empty())
-        resOrderAttr = nullptr;
+
+      DenseI32ArrayAttr resOrderAttr = nullptr;
+      int numRetainedDims = srcRank - static_cast<int>(reductionDims.size());
+      if (orderAttr && !orderAttr.empty() &&
+          static_cast<int>(consumerOrder.size()) == numRetainedDims) {
+        // Fill known order
+        int retainedRank = 0;
+        for (int dim = 0; dim < srcRank; dim++)
+          if (!llvm::is_contained(reductionDims, dim))
+            order[dim] = consumerOrder[retainedRank++];
+        // Each new order is placed as close as possible to its neighbors,
+        // row-major. Re-index the order of higher walk orders.
+        for (int reducedDim = 0; reducedDim < srcRank; reducedDim++) {
+          if (!llvm::is_contained(reductionDims, reducedDim))
+            continue;
+          int64_t insertRank = 0;
+          // Add as an inner dim to the outer neighbor
+          if (reducedDim)
+            insertRank = order[reducedDim - 1];
+          // Add as an outer dim to the inner neighbor if no left neighbor
+          // exists.
+          else if (reducedDim + 1 < srcRank)
+            insertRank = order[reducedDim + 1] + 1;
----------------
akroviakov wrote:

reduceDim == 0 is implied, but the neighbor must be picked only when it exists (i.e., srcRank > 1), I simplified the else if condition to reflect this. 
In case of `reduceDim == 0` and `srcRank == 1`, none of the branches is visited and `int64_t insertRank = 0;` persists.

https://github.com/llvm/llvm-project/pull/224618