[Mlir-commits] [mlir] [MLIR][XeGPU] Add 2D `vector.multi_reduction` optimization (PR #171154)

[Jianhui Li]

================
@@ -416,12 +416,102 @@ class VectorExtractOpPattern final
   }
 };
 
+class MultiRed2dOp : public OpConversionPattern<vector::MultiDimReductionOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(vector::MultiDimReductionOp reductionOp, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto sourceVecType = reductionOp.getSourceVectorType();
+    if (reductionOp.getReductionDims().size() != 2 ||
+        sourceVecType.getRank() != 2)
+      return rewriter.notifyMatchFailure(
+          reductionOp, "Expected 2D multi reduction of a 2D source");
+    auto resLayout = xegpu::getDistributeLayoutAttr(reductionOp.getResult());
+    // Retrieve and order dims for 1D decomposition (prefer intra-lane first).
+    auto dims = llvm::to_vector(reductionOp.getReductionDims());
+    auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, resLayout);
+    // Order does not matter
+    if (intraLaneDim == -1 || crossLaneDim == -1) {
+      intraLaneDim = dims[0];
+      crossLaneDim = dims[1];
+    }
+    auto loc = reductionOp.getLoc();
+    auto acc = reductionOp.getAcc();
+
+    // The first reduction's dist attribute does not have the cross lane dim.
+    auto resSliceLayoutAttr = cast<xegpu::SliceAttr>(resLayout);
----------------
Jianhui-Li wrote:

Sorry I didn't make it clear. I meant dropSliceDims(user-specified-dims) for sliceAttr (not DistributeLayoutAttr). 

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