Add CHECKS

Author: nbpatel

mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp

@@ -830,16 +830,17 @@ struct WgToSgMultiDimReductionOp
   matchAndRewrite(vector::MultiDimReductionOp op, OneToNOpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
-    // Only support reduction with layout and on a single dimension for now.
     VectorType srcType = dyn_cast<VectorType>(op.getSource().getType());
     VectorType accType = dyn_cast<VectorType>(op.getAcc().getType());
     VectorType resType = dyn_cast<VectorType>(op.getResult().getType());
     Type elemTy = srcType.getElementType();
     if (!srcType || !accType || !resType)
       return failure();
 
+    // Support only 2D vectors
+    if (srcType.getShape().size() != 2 && resType.getShape().size() != 1)
+      return failure();
     ArrayRef<int64_t> wgShape = resType.getShape();
-    // Handle both LayoutAttr and SliceAttr for the op result.
     auto layoutName = xegpu::getLayoutName(op->getResult(0));
     auto sliceAttr = op->getAttrOfType<xegpu::SliceAttr>(layoutName);
     if (!sliceAttr || sliceAttr.getRank() != 1)
@@ -871,7 +872,6 @@ struct WgToSgMultiDimReductionOp
           VectorType::get(shapeCastShape, srcType.getElementType());
       auto shapeCast = rewriter.create<vector::ShapeCastOp>(
           op.getLoc(), shapeCastTy, sgReduce.getResult());
-      // TODO: Change it to shapeCast
       newReductions.push_back(shapeCast.getResult());
     }
 
@@ -889,23 +889,23 @@ struct WgToSgMultiDimReductionOp
     auto slmTy = MemRefType::get(slmSize, rewriter.getI8Type(), {}, 3);
     auto slm = rewriter.create<memref::AllocaOp>(loc, slmTy);
 
-    // Create a view for the SLM buffer using xegpu.create_mem_desc
-    SmallVector<int64_t> viewShape;
+    // Create a SLM buffer using xegpu.create_mem_desc
+    SmallVector<int64_t> memDescShape;
     auto srcVecType = dyn_cast<VectorType>(adaptor.getSource()[0].getType());
     ArrayRef<int64_t> srcShape =
         srcVecType ? srcVecType.getShape() : ArrayRef<int64_t>();
     for (size_t i = 0; i < srcShape.size(); ++i) {
       if (static_cast<int64_t>(i) == reduceDim) {
         // For the reduced dimension, use sgLayoutParent[i]
-        viewShape.push_back(sgLayoutParent[i]);
+        memDescShape.push_back(sgLayoutParent[i]);
       } else {
         // For other dimensions, multiply sgLayoutParent[i] by sgShape[i]
-        viewShape.push_back(sgLayoutParent[i] * srcShape[i]);
+        memDescShape.push_back(sgLayoutParent[i] * srcShape[i]);
       }
     }
 
-    auto memDescType = xegpu::MemDescType::get(rewriter.getContext(), viewShape,
-                                               elemTy, nullptr);
+    auto memDescType = xegpu::MemDescType::get(rewriter.getContext(),
+                                               memDescShape, elemTy, nullptr);
     auto memDesc =
         rewriter.create<xegpu::CreateMemDescOp>(loc, memDescType, slm);
 
@@ -951,16 +951,16 @@ struct WgToSgMultiDimReductionOp
     // Step 3: Load from SLM for the second reduction
     SmallVector<int64_t> slmLoadShape;
 
-    for (size_t i = 0; i < viewShape.size(); ++i) {
+    for (size_t i = 0; i < memDescShape.size(); ++i) {
       if (static_cast<int64_t>(i) == reduceDim) {
-        slmLoadShape.push_back(viewShape[i]);
+        slmLoadShape.push_back(memDescShape[i]);
       } else {
         int64_t divisor = computeProduct(sgLayoutParent);
-        slmLoadShape.push_back(viewShape[i] / divisor);
+        slmLoadShape.push_back(memDescShape[i] / divisor);
       }
     }
 
-    // Calculate offsets for create_nd_desc
+    // Calculate offsets for load_matrix op
     SmallVector<OpFoldResult> slmLoadOffsets;
     for (size_t i = 0; i < sgLayoutParent.size(); ++i) {
       Value offset = rewriter.createOrFold<index::MulOp>(
@@ -975,8 +975,8 @@ struct WgToSgMultiDimReductionOp
         /*layout=*/nullptr);
 
     // Step 4: Create a constant accumulator for the second reduction
-    // with same vallue as adaptor.getAcc()[0] and shape set to
-    // the non reduce dimension of shapeCastLoad
+    // with same value as adaptor.getAcc()[0] and shape set to
+    // the non reduce dimension of load
     auto accShape = load.getType().getShape();
     SmallVector<int64_t> accShapeWithoutReduceDim;
     for (size_t i = 0; i < accShape.size(); ++i) {
@@ -1180,10 +1180,6 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
 
   target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
       [=](vector::MultiDimReductionOp op) -> bool {
-        // Only allow MultiDimReductionOp with a single reduction dimension
-        if (op.getReductionDims().size() != 1)
-          return true;
-
         // Check if the layout is legal
         return isLegal(xegpu::getDistributeLayoutAttr(op.getResult()));
       });

mlir/test/Dialect/XeGPU/xegpu-wg-to-sg-unify-ops.mlir

@@ -323,7 +323,33 @@ gpu.module @test_distribution {
   }
 
   //CHECK-LABEL: vector_reduce
+  // CHECK-SAME: %[[ARG_0:.*]]: memref<256x128xf32>
   gpu.func @vector_reduce(%src: memref<256x128xf32>) {
+    // CHECK: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<32xf32>
+    // CHECK: %[[TDESC:.*]] = xegpu.create_nd_tdesc %[[ARG_0]][{{%.*}}, {{%.*}}] : memref<256x128xf32> -> !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [8, 4], lane_data = [1, 1]>>
+    // CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC]] : !xegpu.tensor_desc<32x32xf32, #xegpu.layout<lane_layout = [8, 4], lane_data = [1, 1]>> -> vector<32x32xf32>
+    // CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, {{%.*}}, %[[CST]] [0] : vector<32x32xf32> to vector<32xf32>
+    // CHECK: %[[SHAPECAST:.*]] = vector.shape_cast %[[REDUCE]] : vector<32xf32> to vector<1x32xf32>
+    // CHECK: %[[ALLOCA:.*]] = memref.alloca() : memref<4096xi8, 3>
+    // CHECK: %[[MDESC:.*]] = xegpu.create_mem_desc %[[ALLOCA]] : memref<4096xi8, 3> -> !xegpu.mem_desc<8x128xf32>
+    // CHECK: %[[SGID:.*]] = gpu.subgroup_id : index
+    // CHECK: %[[C8:.*]] = arith.constant 8 : index
+    // CHECK: %[[C4:.*]] = arith.constant 4 : index
+    // CHECK: %[[C4_1:.*]] = arith.constant 4 : index
+    // CHECK: %[[ID_Y:.*]] = affine.apply #map()[%[[SGID]]]
+    // CHECK: %[[ID_X:.*]] = affine.apply #map1()[%[[SGID]]]
+    // CHECK: %[[C1:.*]] = arith.constant 1 : index
+    // CHECK: %[[C32:.*]] = arith.constant 32 : index
+    // CHECK: %[[L_OFF_X:.*]] = index.mul %[[ID_X]], %[[C32]]
+    // CHECK: xegpu.store_matrix {{.*}}, %[[MDESC]][%[[ID_Y]], %[[L_OFF_X]]] : vector<1x32xf32>, !xegpu.mem_desc<8x128xf32>, index, index
+    // CHECK: gpu.barrier
+    // CHECK: %[[C8_1:.*]] = arith.constant 8 : index
+    // CHECK: %[[OFF_Y:.*]] = index.mul %[[ID_Y]], %[[C8_1]]
+    // CHECK: %[[C4_2:.*]] = arith.constant 4 : index
+    // CHECK: %[[OFF_X:.*]] = index.mul %[[ID_X]], %[[C4_2]]
+    // CHECK: %[[LOAD:.*]] = xegpu.load_matrix %[[MDESC]][%[[OFF_Y]], %[[OFF_X]]] : !xegpu.mem_desc<8x128xf32>, index, index -> vector<8x4xf32>
+    // CHECK: %[[CST:.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
+    // CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[CST]] [0] : vector<8x4xf32> to vector<4xf32>
     %cst = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>, dims = [0]>} dense<1.0> : vector<128xf32>
     %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<256x128xf32>
       -> !xegpu.tensor_desc<256x128xf32, #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32], lane_layout = [8, 4], lane_data = [1, 1]>>