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

[Nishant Patel]

================
@@ -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) {
+      Value dimVal = sgIds[i];
+      int64_t stride =
+          (llvm::is_contained(reductionDims, i)) ? 0 : sgSrcShape[i];
+      Value strideVal = arith::ConstantIndexOp::create(rewriter, loc, stride);
+      Value offsetVal = arith::MulIOp::create(rewriter, loc, dimVal, strideVal);
+      slmLoadOffsets.push_back(offsetVal);
+    }
 
-    auto loadOp = xegpu::LoadMatrixOp::create(
-        rewriter, loc, loadType2D, memDesc.getResult(), loadOffsets2D,
+    VectorType slmLoadType = VectorType::get(slmLoadDataShape, elemTy);
+    auto slmLoadOp = xegpu::LoadMatrixOp::create(
+        rewriter, loc, slmLoadType, memDesc.getResult(), slmLoadOffsets,
         /*layout=*/nullptr);
 
     // Step 6: Perform final reduction with ZERO accumulator
-    SmallVector<int64_t> finalReductionDims = {0};
-    SmallVector<int64_t> finalResultShape = {localElements};
-    VectorType finalResultType = VectorType::get(finalResultShape, elemTy);
-
     auto neutralFinalAcc =
-        createAccumulator(rewriter, loc, finalResultType, op.getKind());
+        createAccumulator(rewriter, loc, newDstType, op.getKind());
 
     auto finalReduce = vector::MultiDimReductionOp::create(
-        rewriter, loc, finalResultType, op.getKind(), loadOp.getResult(),
-        neutralFinalAcc, finalReductionDims);
+        rewriter, loc, newDstType, op.getKind(), slmLoadOp.getResult(),
+        neutralFinalAcc, reductionDims);
 
     // Step 7: Add the original accumulator at the end
     Value originalAcc = adaptor.getAcc()[0];
     Value accToAdd = originalAcc;
 
-    // Handle shape mismatch by shape casting
-    if (originalAcc.getType() != finalReduce.getResult().getType()) {
-      auto originalAccType = cast<VectorType>(originalAcc.getType());
-      auto finalResultType =
-          cast<VectorType>(finalReduce.getResult().getType());
-
-      // If they have the same number of elements, just shape cast
-      if (originalAccType.getNumElements() ==
-          finalResultType.getNumElements()) {
-        auto shapeCast = vector::ShapeCastOp::create(
-            rewriter, loc, finalResultType, originalAcc);
-        accToAdd = shapeCast.getResult();
-      }
-    }
-
     auto finalResult = vector::makeArithReduction(
         rewriter, loc, op.getKind(), finalReduce.getResult(), accToAdd);
----------------
nbpatel wrote:

nit: can directly use adaptor.getAcc()[0] here instead of accToAdd and delete the two lines above

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