Add reduction_size option for xetile.reduction op. (#1058)

Add canonicalization pattern for new variant.
Add test case for canonicalization.

Author: silee2

include/imex/Dialect/XeTile/IR/XeTileOps.td

@@ -540,7 +540,8 @@ def XeTile_ReductionOp: XeTile_Op<"reduction", []> {
 
     let arguments = (ins Vector_CombiningKindAttr: $kind,
                          XeTile_2DVector: $source,
-                         DenseI64ArrayAttr: $reduction_dims);
+                         DenseI64ArrayAttr: $reduction_dims,
+                         DefaultValuedAttr<I64Attr, "0">: $reduction_size);
     let results = (outs XeTile_2DVector: $result);
     let assemblyFormat = [{
         $kind `,` $source $reduction_dims attr-dict `:` type($source) `->` type($result)

lib/Dialect/XeTile/IR/XeTileOps.cpp

@@ -322,9 +322,28 @@ mlir::LogicalResult TransposeOp::verify() {
 mlir::LogicalResult ReductionOp::verify() {
   auto dims = getReductionDims();
   auto resShape = getResult().getType().getShape();
-  for (auto i : dims)
-    if (resShape[i] != 1)
-      return emitOpError("reduction dimension of result must have size 1");
+  if (getReductionSize() > 0) {
+    if (dims.size() > 1)
+      // When reduction size is specified,
+      // only a single dimension can be reduced.
+      return emitOpError(
+          "when reduction size is specified, only a single reduction "
+          "dimension is allowed.");
+    auto srcTy = getSource().getType();
+    if (srcTy.getRank() != 2)
+      return emitOpError(
+          "when reduction size is specified, source must be a 2D vector.");
+    auto redDim = dims.front();
+    if (resShape[redDim] !=
+        srcTy.getShape()[redDim] / static_cast<int64_t>(getReductionSize()))
+      return emitOpError(
+          "reduction size does not match the expected size of the result "
+          "dimension after reduction.");
+  } else {
+    for (auto i : dims)
+      if (resShape[i] != 1)
+        return emitOpError("reduction dimension of result must have size 1");
+  }
   return mlir::success();
 }
 

lib/Dialect/XeTile/Transforms/Canonicalization.cpp

@@ -365,6 +365,57 @@ struct RemoveRedundantTransposeOpPattern
   }
 };
 
+// Remove XeTile reduction's reduction_size attribute and replace it with shape
+// casts around it.
+struct RemoveReductionSizePattern
+    : public mlir::OpRewritePattern<imex::xetile::ReductionOp> {
+  using mlir::OpRewritePattern<imex::xetile::ReductionOp>::OpRewritePattern;
+  mlir::LogicalResult
+  matchAndRewrite(imex::xetile::ReductionOp op,
+                  mlir::PatternRewriter &rewriter) const override {
+    if (op.getReductionSize() == 0)
+      return mlir::failure();
+
+    auto loc = op.getLoc();
+    auto srcTy = op.getSource().getType();
+    auto elemTy = srcTy.getElementType();
+    auto reductionSize = static_cast<int64_t>(op.getReductionSize());
+    // Get reduction dimension "i"
+    // Op validation ensures that only a single reduction dimension is
+    // present if reduction size is set to non zero.
+    // Also source rank is restricted to 2D.
+    auto redDim = op.getReductionDims().front();
+    auto resultTy = op.getType();
+    auto numRes = resultTy.getNumElements();
+    llvm::SmallVector<int64_t> newShape;
+    if (redDim == 0) {
+      newShape.push_back(reductionSize);
+      newShape.push_back(numRes);
+    } else {
+      newShape.push_back(numRes);
+      newShape.push_back(reductionSize);
+    }
+    llvm::SmallVector<int64_t> newRedShape;
+    if (redDim == 0) {
+      newRedShape.push_back(1);
+      newRedShape.push_back(numRes);
+    } else {
+      newRedShape.push_back(numRes);
+      newRedShape.push_back(1);
+    }
+    mlir::Value newCast = rewriter.create<mlir::vector::ShapeCastOp>(
+        loc, mlir::VectorType::get(newShape, elemTy), op.getSource());
+    mlir::Value newReductionOp = rewriter.create<imex::xetile::ReductionOp>(
+        loc, mlir::VectorType::get(newRedShape, elemTy), op.getKind(), newCast,
+        op.getReductionDims());
+
+    auto shapeCastOp = rewriter.create<mlir::vector::ShapeCastOp>(
+        op.getLoc(), op.getType(), newReductionOp);
+    rewriter.replaceOp(op, shapeCastOp);
+    return mlir::success();
+  }
+};
+
 struct XeTileCanonicalizationPass final
     : public imex::impl::XeTileCanonicalizationBase<
           XeTileCanonicalizationPass> {
@@ -430,6 +481,7 @@ struct XeTileCanonicalizationPass final
 
       target.addLegalOp<mlir::memref::ReinterpretCastOp>();
       target.addLegalOp<imex::xetile::TransposeOp>();
+      target.addLegalOp<mlir::vector::ShapeCastOp>();
       // Col-major tile creattion is not allowed.
       target.addDynamicallyLegalOp<imex::xetile::InitTileOp>(
           [&](imex::xetile::InitTileOp op) {
@@ -450,6 +502,11 @@ struct XeTileCanonicalizationPass final
           [&](imex::xetile::LoadTileOp op) {
             return isValidTile(op.getTileType());
           });
+      // ReductionOp is legal if reduction size is 0.
+      target.addDynamicallyLegalOp<imex::xetile::ReductionOp>(
+          [&](imex::xetile::ReductionOp op) {
+            return op.getReductionSize() == 0;
+          });
       // If any iterArg of the forOp is a col-major tile, it is illegal.
       target.addDynamicallyLegalOp<mlir::scf::ForOp>([&](mlir::scf::ForOp op) {
         for (auto arg : op.getRegionIterArgs()) {
@@ -475,6 +532,7 @@ struct XeTileCanonicalizationPass final
           .add<InitTileOpPattern, LoadTileOpPattern, UpdateTileOffsetOpPattern,
                PrefetchTilePattern, ScfForOpPattern, ScfYieldOpPattern>(
               typeConverter, context);
+      patterns.add<RemoveReductionSizePattern>(context);
 
       if (mlir::failed(mlir::applyPartialConversion(getOperation(), target,
                                                     std::move(patterns))))

test/Dialect/XeTile/IR/ops.mlir

@@ -330,6 +330,12 @@ func.func @test_reduce(%source: vector<8x16xf16>) {
   return
 }
 
+func.func @test_reduce_reduction_size(%source: vector<8x16xf16>) {
+  // CHECK: xetile.reduction {{.*}} [0] {reduction_size = 8 : i64} : vector<8x16xf16> -> vector<1x16xf16>
+  %1 = xetile.reduction <add>, %source [0] { reduction_size = 8 : i64 } : vector<8x16xf16> -> vector<1x16xf16>
+  return
+}
+
 func.func @test_reduce_map(%source: vector<256x128xf16>) {
   // CHECK: xetile.reduction {{.*}} [1] {map1 = #xetile.wg_map<sg_layout = [32, 1], sg_data = [8, 128]>, map2 = #xetile.wg_map<sg_layout = [32, 1], sg_data = [8, 1]>} : vector<256x128xf16> -> vector<256x1xf16>
   %1 = xetile.reduction <add>, %source [1] {map1 = #wg_map_a, map2 = #wg_map_a2} : vector<256x128xf16> -> vector<256x1xf16>

test/Dialect/XeTile/Transforms/canonicalization.mlir

@@ -344,3 +344,33 @@ gpu.module @test_module {
 //CHECK: %[[transpose:.*]] = memref.transpose %[[arg0]] (d0, d1) -> (d1, d0) : memref<512x128xf16, 3> to memref<128x512xf16, strided<[1, 128]>, 3>
 //CHECK: %[[r2:.*]] = xetile.init_tile %[[transpose]][16, 32] : memref<128x512xf16, strided<[1, 128]>, 3> -> !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0, 1], memory_space = 3 : i64>>
 //CHECK: xetile.store_tile %[[r1]],  %[[r2]] : vector<16x32xf16>, !xetile.tile<16x32xf16, #xetile.tile_attr<order = [0, 1], memory_space = 3 : i64>>
+
+// -----
+gpu.module @test_module {
+  // CHECK-LABEL: gpu.func @test_reduction_size
+  // CHECK-SAME: (%[[ARG0:.*]]: vector<64x256xf32>
+  gpu.func @test_reduction_size(%arg0: vector<64x256xf32>, %arg1: memref<64x8xf32>) {
+    %cst0 = arith.constant 0 : index
+    // CHECK: %[[V0:.*]] = vector.shape_cast %[[ARG0]] : vector<64x256xf32> to vector<512x32xf32>
+    // CHECK: %[[V1:.*]] = xetile.reduction <add>, %[[V0]] [1] : vector<512x32xf32> -> vector<512x1xf32>
+    // CHECK: %[[V2:.*]] = vector.shape_cast %[[V1]] : vector<512x1xf32> to vector<64x8xf32>
+    %reduced = xetile.reduction <add>, %arg0 [1] { reduction_size = 32 : i64 } : vector<64x256xf32> -> vector<64x8xf32>
+    vector.store %reduced, %arg1[%cst0, %cst0] : memref<64x8xf32>, vector<64x8xf32>
+    gpu.return
+  }
+}
+
+// -----
+gpu.module @test_module {
+  // CHECK-LABEL: gpu.func @test_reduction_size
+  // CHECK-SAME: (%[[ARG0:.*]]: vector<64x256xf32>
+  gpu.func @test_reduction_size(%arg0: vector<64x256xf32>, %arg1: memref<2x256xf32>) {
+    %cst0 = arith.constant 0 : index
+    // CHECK: %[[V0:.*]] = vector.shape_cast %[[ARG0]] : vector<64x256xf32> to vector<32x512xf32>
+    // CHECK: %[[V1:.*]] = xetile.reduction <add>, %[[V0]] [0] : vector<32x512xf32> -> vector<1x512xf32>
+    // CHECK: %[[V2:.*]] = vector.shape_cast %[[V1]] : vector<1x512xf32> to vector<2x256xf32>
+    %reduced = xetile.reduction <add>, %arg0 [0] { reduction_size = 32 : i64 } : vector<64x256xf32> -> vector<2x256xf32>
+    vector.store %reduced, %arg1[%cst0, %cst0] : memref<2x256xf32>, vector<2x256xf32>
+    gpu.return
+  }
+}