[mlir][vector] Add unrolling pattern for multidim_reduce op

[springerm]

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

332

nit: I'd spell out the type here.

339

Why is this needed?

366

typo

367

keeps

[springerm]

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

344

printf

[ThomasRaoux]

Address review comments and rebase

[ThomasRaoux]

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

339

It is because there is no default DenseMapInfo for SmallVector type, so I need to give OffsetMapInfo. Let me know if you see a better data structure.

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

def Vector_ReductionOp :
   let hasCustomAssemblyFormat = 1;
   let hasVerifier = 1;
 }
 
 def Vector_MultiDimReductionOp :
   Vector_Op<"multi_reduction", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
-     DeclareOpInterfaceMethods<InferTypeOpInterface>]>,
+     DeclareOpInterfaceMethods<InferTypeOpInterface>,
+     DeclareOpInterfaceMethods<VectorUnrollOpInterface,
+                               ["getShapeForUnroll"]>]>,
     Arguments<(ins Vector_CombiningKindAttr:$kind,
                    AnyVector:$source,
                    I64ArrayAttr:$reduction_dims)>,
     Results<(outs AnyType:$dest)> {
   let summary = "Multi-dimensional reduction operation";
   let description = [{
     Reduces an n-D vector into an (n-k)-D vector (or a scalar when k == n)
     using the given operation (add/mul/min/max for int/fp and and/or/xor for

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

 
 OpFoldResult MultiDimReductionOp::fold(ArrayRef<Attribute> operands) {
   // Single parallel dim, this is a noop.
   if (getSourceVectorType().getRank() == 1 && !isReducedDim(0))
     return source();
   return {};
 }
 
+Optional<SmallVector<int64_t, 4>> MultiDimReductionOp::getShapeForUnroll() {
+  return llvm::to_vector<4>(getSourceVectorType().getShape());
+}
+
 //===----------------------------------------------------------------------===//
 // ReductionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
                                 CombiningKind kind, Value vector) {
   build(builder, result, kind, vector, /*acc=*/Value());
 }

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

else
       rewriter.eraseOp(writeOp);
     return success();
   }
 
 private:
   vector::UnrollVectorOptions options;
 };
 
-struct UnrollContractionPattern
-    : public OpRewritePattern<vector::ContractionOp> {
-  struct OffsetMapInfo {
-    static SmallVector<int64_t> getEmptyKey() { return {int64_t(-1)}; }
 
-    static SmallVector<int64_t> getTombstoneKey() { return {int64_t(-2)}; }
 
-    static unsigned getHashValue(const SmallVector<int64_t> &v) {
-      return static_cast<unsigned>(
-          llvm::hash_combine_range(v.begin(), v.end()));
-    }
+    return static_cast<unsigned>(llvm::hash_combine_range(v.begin(), v.end()));
 
-    static bool isEqual(const SmallVector<int64_t> &lhs,
-                        const SmallVector<int64_t> &rhs) {
-      return lhs == rhs;
-    }
-  };
+
+struct UnrollContractionPattern
+    : public OpRewritePattern<vector::ContractionOp> {
   UnrollContractionPattern(MLIRContext *context,
                            const vector::UnrollVectorOptions &options)
       : OpRewritePattern<vector::ContractionOp>(context, /*benefit=*/1),
         options(options) {}
 
   LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                 PatternRewriter &rewriter) const override {
     auto targetShape = getTargetShape(options, contractOp);
LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
     rewriter.replaceOp(contractOp, result);
     return success();
   }
 
 private:
   vector::UnrollVectorOptions options;
 };
 
+struct UnrollMultiReductionPattern
+    : public OpRewritePattern<vector::MultiDimReductionOp> {
+  UnrollMultiReductionPattern(MLIRContext *context,
+                              const vector::UnrollVectorOptions &options)
+      : OpRewritePattern<vector::MultiDimReductionOp>(context, /*benefit=*/1),
+        options(options) {}
+
+  LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
+                                PatternRewriter &rewriter) const override {
+    Optional<SmallVector<int64_t, 4>> targetShape =
+        getTargetShape(options, reductionOp);
+    if (!targetShape)
+      return failure();
+    SmallVector<int64_t, 4> originalSize = *reductionOp.getShapeForUnroll();
+    SmallVector<int64_t, 4> ratio = *shapeRatio(originalSize, *targetShape);
+    llvm::MapVector<
+        SmallVector<int64_t>, Value,
+        llvm::DenseMap<SmallVector<int64_t>, unsigned, OffsetMapInfo>>
+        accCache;
+    // Compute shape ratio of 'shape' and 'sizes'.
+    int64_t sliceCount = computeMaxLinearIndex(ratio);
+    Location loc = reductionOp.getLoc();
+    for (int64_t i = 0; i < sliceCount; i++) {
+      SmallVector<int64_t, 4> offsets =
+          getVectorOffset(originalSize, *targetShape, i);
+
+      SmallVector<int64_t, 4> operandStrides(offsets.size(), 1);
+      Value slicedOperand = rewriter.create<vector::ExtractStridedSliceOp>(
+          loc, reductionOp.getOperand(), offsets, *targetShape, operandStrides);
+
+      SmallVector<int64_t> dstShape;
+      SmallVector<int64_t> destOffset;
+      for (size_t i : llvm::seq(size_t(0), targetShape->size())) {
+        if (!reductionOp.isReducedDim(i)) {
+          destOffset.push_back(offsets[i]);
+          dstShape.push_back((*targetShape)[i]);
+        }
+      }
+      auto targetType = VectorType::get(
+          dstShape, reductionOp.getSourceVectorType().getElementType());
+      Operation *newOp = cloneOpWithOperandsAndTypes(rewriter, loc, reductionOp,
+                                                     slicedOperand, targetType);
+      Value result = newOp->getResult(0);
+      // Save the accumulated value until all the loops are unrolled since
+      // reduction loop keeps updating the accumulator.
+      auto accIt = accCache.find(destOffset);
+      if (accIt != accCache.end())
+        result = makeArithReduction(rewriter, loc, reductionOp.kind(), result,
+                                    accIt->second);
+      accCache[destOffset] = result;
+    }
+    // Assemble back the accumulator into a single vector.
+    Value result = rewriter.create<arith::ConstantOp>(
+        loc, reductionOp.getDestType(),
+        rewriter.getZeroAttr(reductionOp.getDestType()));
+    for (const auto &it : accCache) {
+      SmallVector<int64_t> dstStrides(it.first.size(), 1);
+      result = rewriter.create<vector::InsertStridedSliceOp>(
+          loc, it.second, result, it.first, dstStrides);
+    }
+    rewriter.replaceOp(reductionOp, result);
+    return success();
+  }
+
+private:
+  vector::UnrollVectorOptions options;
+};
+
 struct UnrollElementwisePattern : public RewritePattern {
   UnrollElementwisePattern(MLIRContext *context,
                            const vector::UnrollVectorOptions &options)
       : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context),
         options(options) {}
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
     if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
struct TransferWriteInsertPattern
   }
 };
 
 } // namespace
 
 void mlir::vector::populateVectorUnrollPatterns(
     RewritePatternSet &patterns, const UnrollVectorOptions &options) {
   patterns.add<UnrollTransferReadPattern, UnrollTransferWritePattern,
-               UnrollContractionPattern, UnrollElementwisePattern>(
-      patterns.getContext(), options);
+               UnrollContractionPattern, UnrollElementwisePattern,
+               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:   return
 
 func @vector_fma(%a: vector<4x4xf32>, %b: vector<4x4xf32>, %c: vector<4x4xf32>) -> vector<4x4xf32> {
   %0 = vector.fma %a, %b, %c: vector<4x4xf32>
   return %0 : vector<4x4xf32>
 }
 //   CHECK-LABEL: func @vector_fma
 // CHECK-COUNT-4: vector.fma %{{.+}}, %{{.+}}, %{{.+}} : vector<2x2xf32>
+
+func @vector_multi_reduction(%v : vector<4x6xf32>) -> vector<4xf32> {
+  %0 = vector.multi_reduction #vector.kind<add>, %v [1] : vector<4x6xf32> to vector<4xf32>
+  return %0 : vector<4xf32>
+}
+// CHECK-LABEL: func @vector_multi_reduction
+//       CHECK:   %[[V0:.*]] = arith.constant dense<0.000000e+00> : vector<4xf32>
+//       CHECK:   %[[E0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
+//       CHECK:   %[[R0:.*]] = vector.multi_reduction <add>, %[[E0]] [1] : vector<2x2xf32> to vector<2xf32>
+//       CHECK:   %[[E1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
+//       CHECK:   %[[R1:.*]] = vector.multi_reduction <add>, %[[E1]] [1] : vector<2x2xf32> to vector<2xf32>
+//       CHECK:   %[[A0:.*]] = arith.addf %[[R1]], %[[R0]] : vector<2xf32>
+//       CHECK:   %[[E2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 4], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
+//       CHECK:   %[[R2:.*]] = vector.multi_reduction <add>, %5 [1] : vector<2x2xf32> to vector<2xf32>
+//       CHECK:   %[[A1:.*]] = arith.addf %[[R2]], %[[A0]] : vector<2xf32>
+//       CHECK:   %[[E3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
+//       CHECK:   %[[R3:.*]] = vector.multi_reduction <add>, %[[E3]] [1] : vector<2x2xf32> to vector<2xf32>
+//       CHECK:   %[[E4:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x6xf32> to vector<2x2xf32>
+//       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>

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

TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass)
       : PassWrapper(pass) {}
   void runOnOperation() override {
     MLIRContext *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     populateVectorUnrollPatterns(
         patterns, UnrollVectorOptions()
                       .setNativeShape(ArrayRef<int64_t>{2, 2})
                       .setFilterConstraint([](Operation *op) {
-                        return success(isa<arith::AddFOp, vector::FMAOp>(op));
+                        return success(isa<arith::AddFOp, vector::FMAOp,
+                                           vector::MultiDimReductionOp>(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()