[mlir][vector] Implement unrolling of ReductionOp

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

def Vector_ContractionOp :
   let hasCanonicalizer = 1;
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def Vector_ReductionOp :
   Vector_Op<"reduction", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
-                 TCresVTEtIsSameAsOpBase<0, 0>>]>,
+                 TCresVTEtIsSameAsOpBase<0, 0>>,
+     DeclareOpInterfaceMethods<VectorUnrollOpInterface,
+                               ["getShapeForUnroll"]>]>,
     Arguments<(ins Vector_CombiningKindAttr:$kind, AnyVector:$vector,
                    Optional<AnyType>:$acc)>,
     Results<(outs AnyType:$dest)> {
   let summary = "reduction operation";
   let description = [{
     Reduces an 1-D vector "horizontally" into a scalar using the given
     operation (add/mul/min/max for int/fp and and/or/xor for int only).
     Some reductions (add/mul for fp) also allow an optional fused

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
   // TODO: Add remaining reduction operations.
   default:
     (void)emitOptionalError(loc, "Reduction operation type not supported");
     break;
   }
   return nullptr;
 }
 
+Optional<SmallVector<int64_t, 4>> ReductionOp::getShapeForUnroll() {
+  return llvm::to_vector<4>(getVectorType().getShape());
+}
+
 //===----------------------------------------------------------------------===//
 // ContractionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
                                   ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
                                   ArrayRef<StringRef> iteratorTypes) {

mlir/lib/Dialect/Vector/Transforms/VectorUnrollDistribute.cpp

LogicalResult matchAndRewrite(vector::TransferWriteOp write,
     rewriter.create<vector::TransferWriteOp>(
         loc, insert.vector(), write.source(), indices,
         write.permutation_mapAttr(), write.in_boundsAttr());
     rewriter.eraseOp(write);
     return success();
   }
 };
 
+struct UnrollReductionPattern : public OpRewritePattern<vector::ReductionOp> {
+  UnrollReductionPattern(MLIRContext *context,
+                         const vector::UnrollVectorOptions &options)
+      : OpRewritePattern<vector::ReductionOp>(context, /*benefit=*/1),
+        options(options) {}
+
+  LogicalResult matchAndRewrite(vector::ReductionOp reductionOp,
+                                PatternRewriter &rewriter) const override {
+    Optional<SmallVector<int64_t, 4>> targetShape =
+        getTargetShape(options, reductionOp);
+    if (!targetShape)
+      return failure();
+    SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
+    int64_t ratio = (*shapeRatio(originalSize, *targetShape))[0];
+
+    // Create unrolled vector reduction.
+    Location loc = reductionOp.getLoc();
+    Value accumulator = nullptr;
+    for (int64_t i = 0; i < ratio; ++i) {
+      SmallVector<int64_t> offsets =
+          getVectorOffset(originalSize, *targetShape, i);
+      SmallVector<int64_t> strides(offsets.size(), 1);
+      Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
+          loc, reductionOp.vector(), offsets, *targetShape, strides);
+      Operation *newOp = cloneOpWithOperandsAndTypes(
+          rewriter, loc, reductionOp, slicedOperand, reductionOp.getType());
+      Value result = newOp->getResult(0);
+
+      if (!accumulator) {
+        // This is the first reduction.
+        accumulator = result;
+      } else {
+        // On subsequent reduction, combine with the accumulator.
+        accumulator = makeArithReduction(rewriter, loc, reductionOp.kind(),
+                                         accumulator, result);
+      }
+    }
+
+    rewriter.replaceOp(reductionOp, accumulator);
+    return success();
+  }
+
+private:
+  const vector::UnrollVectorOptions options;
+};
+
 } // namespace
 
 void mlir::vector::populateVectorUnrollPatterns(
     RewritePatternSet &patterns, const UnrollVectorOptions &options) {
   patterns.add<UnrollTransferReadPattern, UnrollTransferWritePattern,
                UnrollContractionPattern, UnrollElementwisePattern,
-               UnrollMultiReductionPattern>(patterns.getContext(), options);
+               UnrollReductionPattern, UnrollMultiReductionPattern>(
+      patterns.getContext(), options);
 }
 
 void mlir::vector::populatePropagateVectorDistributionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<PointwiseExtractPattern, ContractExtractPattern,
                TransferReadExtractPattern, TransferWriteInsertPattern>(
       patterns.getContext());
 }

mlir/test/Dialect/Vector/vector-unroll-options.mlir

 //       CHECK:   %[[R4:.*]] = vector.multi_reduction <add>, %[[E4]] [1] : vector<2x2xf32> to vector<2xf32>
 //       CHECK:   %[[A2:.*]] = arith.addf %[[R4]], %[[R3]] : vector<2xf32>
 //       CHECK:   %[[E5:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 4], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
 //       CHECK:   %[[R5:.*]] = vector.multi_reduction <add>, %[[E5]] [1] : vector<2x2xf32> to vector<2xf32>
 //       CHECK:   %[[A3:.*]] = arith.addf %[[R5]], %[[A2]] : vector<2xf32>
 //       CHECK:   %[[V1:.*]] = vector.insert_strided_slice %[[A1]], %[[V0]] {offsets = [0], strides = [1]} : vector<2xf32> into vector<4xf32>
 //       CHECK:   %[[V2:.*]] = vector.insert_strided_slice %[[A3]], %[[V1]] {offsets = [2], strides = [1]} : vector<2xf32> into vector<4xf32>
 //       CHECK:   return %[[V2]] : vector<4xf32>
+
+// CHECK-LABEL: func @vector_reduction(
+//  CHECK-SAME:     %[[v:.*]]: vector<8xf32>
+//       CHECK:   %[[s0:.*]] = vector.extract_strided_slice %[[v]] {offsets = [0], sizes = [2]
+//       CHECK:   %[[r0:.*]] = vector.reduction <add>, %[[s0]]
+//       CHECK:   %[[s1:.*]] = vector.extract_strided_slice %[[v]] {offsets = [2], sizes = [2]
+//       CHECK:   %[[r1:.*]] = vector.reduction <add>, %[[s1]]
+//       CHECK:   %[[add1:.*]] = arith.addf %[[r0]], %[[r1]]
+//       CHECK:   %[[s2:.*]] = vector.extract_strided_slice %[[v]] {offsets = [4], sizes = [2]
+//       CHECK:   %[[r2:.*]] = vector.reduction <add>, %[[s2]]
+//       CHECK:   %[[add2:.*]] = arith.addf %[[add1]], %[[r2]]
+//       CHECK:   %[[s3:.*]] = vector.extract_strided_slice %[[v]] {offsets = [6], sizes = [2]
+//       CHECK:   %[[r3:.*]] = vector.reduction <add>, %[[s3]]
+//       CHECK:   %[[add3:.*]] = arith.addf %[[add2]], %[[r3]]
+//       CHECK:   return %[[add3]]
+func @vector_reduction(%v : vector<8xf32>) -> f32 {
+  %0 = vector.reduction <add>, %v : vector<8xf32> into f32
+  return %0 : f32
+}
+

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

void runOnOperation() override {
     RewritePatternSet patterns(ctx);
     populateVectorUnrollPatterns(
         patterns, UnrollVectorOptions()
                       .setNativeShape(ArrayRef<int64_t>{2, 2})
                       .setFilterConstraint([](Operation *op) {
                         return success(isa<arith::AddFOp, vector::FMAOp,
                                            vector::MultiDimReductionOp>(op));
                       }));
+    populateVectorUnrollPatterns(
+        patterns, UnrollVectorOptions()
+                      .setNativeShape(ArrayRef<int64_t>{2})
+                      .setFilterConstraint([](Operation *op) {
+                        return success(isa<vector::ReductionOp>(op));
+                      }));
 
     if (unrollBasedOnType) {
       UnrollVectorOptions::NativeShapeFnType nativeShapeFn =
           [](Operation *op) -> Optional<SmallVector<int64_t, 4>> {
         vector::ContractionOp contractOp = cast<vector::ContractionOp>(op);
         SmallVector<int64_t, 4> nativeShape = {4, 4, 2};
         if (auto floatType = contractOp.getLhsType()
                                  .getElementType()