[Mlir-commits] [mlir] [MLIR][XeGPU] Improve workgroup to subgroup distribution pattern for mulit-reduction op (PR #182178)

[Artem Kroviakov]

================
@@ -1484,78 +1440,54 @@ struct WgToSgMultiDimReductionOp
       return failure();
     SmallVector<Value> sgIds = *sgIdsResult;
 
-    // Row offset: linearize reduction dimension indices
-    Value rowOffsetStore = linearizeSubgroupIndices(
-        rewriter, loc, sgIds, crossSgReductionDims, sgLayout);
-
-    // Column offset: linearize non-reduction dimension indices
-    SmallVector<int64_t> nonReductionDims;
-    for (size_t i = 0; i < sgLayout.size(); ++i) {
-      if (!llvm::is_contained(reductionDims, static_cast<int64_t>(i))) {
-        nonReductionDims.push_back(static_cast<int64_t>(i));
-      }
+    SmallVector<OpFoldResult> slmStoreOffsets;
+    for (int i = 0; i < srcVecRank; ++i) {
+      Value dimVal = sgIds[i];
+      int64_t stride =
+          (llvm::is_contained(reductionDims, i)) ? 1 : sgSrcShape[i];
+      Value strideVal = arith::ConstantIndexOp::create(rewriter, loc, stride);
+      Value offsetVal = arith::MulIOp::create(rewriter, loc, dimVal, strideVal);
+      slmStoreOffsets.push_back(offsetVal);
     }
 
-    Value colOffset = linearizeSubgroupIndices(rewriter, loc, sgIds,
-                                               nonReductionDims, sgLayout);
-
-    Value localElementsVal =
-        arith::ConstantIndexOp::create(rewriter, loc, localElements);
-    colOffset =
-        arith::MulIOp::create(rewriter, loc, colOffset, localElementsVal);
-
-    SmallVector<OpFoldResult> storeOffsets2D = {rowOffsetStore, colOffset};
-
-    xegpu::StoreMatrixOp::create(rewriter, loc, storeData, memDesc.getResult(),
-                                 storeOffsets2D, /*layout=*/nullptr);
+    xegpu::StoreMatrixOp::create(rewriter, loc, slmStoreData,
+                                 memDesc.getResult(), slmStoreOffsets,
+                                 /*layout=*/nullptr);
 
     gpu::BarrierOp::create(rewriter, loc);
 
     // Step 5: Load from SLM for final reduction
-    SmallVector<int64_t> loadShape2D = {totalReductionSubgroups, localElements};
-    VectorType loadType2D = VectorType::get(loadShape2D, elemTy);
-
-    // Load offsets - each subgroup loads its column based on non-reduction
-    // position
-    Value rowOffsetLoad = arith::ConstantIndexOp::create(rewriter, loc, 0);
-
-    SmallVector<OpFoldResult> loadOffsets2D = {rowOffsetLoad, colOffset};
+    SmallVector<int64_t> slmLoadDataShape(sgSrcShape.begin(), sgSrcShape.end());
+    for (int64_t dim : reductionDims)
+      slmLoadDataShape[dim] = slmShape[dim];
+
+    SmallVector<OpFoldResult> slmLoadOffsets;
+    for (int i = 0; i < srcVecRank; ++i) {
----------------
akroviakov wrote:

A lambda might be a better option to avoid logic duplication for load/store.

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