Added canonicalization for vector.multi_reduction

[vmurali]

Summary

If there are reductions only along unit dimensions, then they are folded

[dcaballe]

Thanks for adding this canonicalizer pattern! LGTM overall. Just a couple of comments.

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

313

nit: Elide -> Fold? (I'm not the best one to provide this feedback but elide sounds more like skip to me)?

Can you also please add doc to this class about the cases that are folded and the ones that are not?

321

I understand this will only fold multi reductions when all the reduction dims are 1. Would it make sense to also support partial folding (e.g., [0, 1], vector<1x64>)?

mlir/test/Dialect/Vector/canonicalize.mlir

1360

I would add a test for the partial folding case even if we don't support it for now.

[vmurali]

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

313

I used elide because of https://github.com/vmurali/llvm-project/blob/murali_llvm/mlir/lib/Dialect/Vector/IR/VectorOps.cpp#L452 (it's literally doing the same thing).

I will see about adding the documentation. I didn't bother earlier because it's literally mimicking ReductionOp's canonicalizer.

321

I don't think I understand what partial folding is. Can you give an example?

[vmurali]

Addressed the review comments

[vmurali]

Explicitly state that this transform happens only if all the reduction dimensions are unit dimensions.

[dcaballe]

LGTM, thanks!

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

317

// -> /// for doc.

321

Sure, it would be something like:

%0 = vector.multi_reduction <mul>, %source, %acc [3] : vector<5x1x4x1x20xf32> to vector<5x4x20xf32>

Two unit sizes, only one reduced.

[aartbik]

I see this failure in some of our tests after this change. Any pointers?

assert.h assertion failed at third_party/llvm/llvm-project/mlir/include/mlir/IR/Value.h:434 in mlir::detail::TypedValue<mlir::VectorType>::TypedValue(Value) [Ty = mlir::VectorType]: !val || val.getType().template isa<Ty>()

[vmurali]

In D134996#3838312, @aartbik wrote:

I see this failure in some of our tests after this change. Any pointers?

assert.h assertion failed at third_party/llvm/llvm-project/mlir/include/mlir/IR/Value.h:434 in mlir::detail::TypedValue<mlir::VectorType>::TypedValue(Value) [Ty = mlir::VectorType]: !val || val.getType().template isa<Ty>()

On it.

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

static SmallVector<bool> getReductionMask(
       for (auto idx : reductionDims)
         res[idx] = true;
       return res;
     }
   }];
   let assemblyFormat =
     "$kind `,` $source `,` $acc attr-dict $reduction_dims `:` type($source) `to` type($dest)";
   let hasFolder = 1;
+  let hasCanonicalizer = 1;
   let hasVerifier = 1;
 }
 
 def Vector_BroadcastOp :
   Vector_Op<"broadcast", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>]>,
     Arguments<(ins AnyType:$source)>,

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

LogicalResult MultiDimReductionOp::verify() {
   if (getType() != inferredReturnType)
     return emitOpError() << "destination type " << getType()
                          << " is incompatible with source type "
                          << getSourceVectorType();
 
   return success();
 }
 
+namespace {
+// Only unit dimensions that are being reduced are folded. If the dimension is
+// unit, but not reduced, it is not folded, thereby keeping the output type the
+// same. If not all dimensions which are reduced are of unit dimension, this
+// transformation does nothing. This is just a generalization of
+// ElideSingleElementReduction for ReduceOp.
+struct ElideUnitDimsInMultiDimReduction
+    : public OpRewritePattern<MultiDimReductionOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(MultiDimReductionOp reductionOp,
+                                PatternRewriter &rewriter) const override {
+    ArrayRef<int64_t> shape = reductionOp.getSourceVectorType().getShape();
+    for (auto dim : enumerate(shape)) {
+      if (reductionOp.isReducedDim(dim.index()) && dim.value() != 1)
+        return failure();
+    }
+    Location loc = reductionOp.getLoc();
+    Value acc = reductionOp.getAcc();
+    Value cast = rewriter.create<vector::ShapeCastOp>(
+        loc, reductionOp.getDestType(), reductionOp.getSource());
+    Value result = vector::makeArithReduction(rewriter, loc,
+                                              reductionOp.getKind(), acc, cast);
+    rewriter.replaceOp(reductionOp, result);
+    return success();
+  }
+};
+} // namespace
+
+void MultiDimReductionOp::getCanonicalizationPatterns(
+    RewritePatternSet &results, MLIRContext *context) {
+  results.add<ElideUnitDimsInMultiDimReduction>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ReductionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
                                 CombiningKind kind, Value vector) {
   build(builder, result, kind, vector, /*acc=*/Value());
 }

mlir/test/Dialect/Vector/canonicalize.mlir

func.func @vector_multi_reduction_single_parallel(%arg0: vector<2xf32>, %acc: vector<2xf32>) -> vector<2xf32> {
     %0 = vector.multi_reduction <mul>, %arg0, %acc [] : vector<2xf32> to vector<2xf32>
 
 //       CHECK:     return %[[v]] : vector<2xf32>
     return %0 : vector<2xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @vector_multi_reduction_unit_dimensions(
+//  CHECK-SAME: %[[SOURCE:.+]]: vector<5x1x4x1x20xf32>, %[[ACC:.+]]: vector<5x4x20xf32>
+func.func @vector_multi_reduction_unit_dimensions(%source: vector<5x1x4x1x20xf32>, %acc: vector<5x4x20xf32>) -> vector<5x4x20xf32> {
+//       CHECK:   %[[CAST:.+]] = vector.shape_cast  %[[SOURCE]] : vector<5x1x4x1x20xf32> to vector<5x4x20xf32>
+//       CHECK:   %[[RESULT:.+]] = arith.mulf  %[[ACC]], %[[CAST]] : vector<5x4x20xf32>
+    %0 = vector.multi_reduction <mul>, %source, %acc [1, 3] : vector<5x1x4x1x20xf32> to vector<5x4x20xf32>
+
+//       CHECK:     return %[[RESULT]] : vector<5x4x20xf32>
+    return %0 : vector<5x4x20xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @vector_multi_reduction_unit_dimensions_fail(
+//  CHECK-SAME: %[[SRC:.+]]: vector<5x1x4x1x20xf32>, %[[ACCUM:.+]]: vector<5x1x1x20xf32>
+func.func @vector_multi_reduction_unit_dimensions_fail(%source: vector<5x1x4x1x20xf32>, %acc: vector<5x1x1x20xf32>) -> vector<5x1x1x20xf32> {
+//       CHECK:   %[[RES:.+]] = vector.multi_reduction  <mul>, %[[SRC]], %[[ACCUM]] [2] : vector<5x1x4x1x20xf32> to vector<5x1x1x20xf32>
+    %0 = vector.multi_reduction <mul>, %source, %acc [2] : vector<5x1x4x1x20xf32> to vector<5x1x1x20xf32>
+
+//       CHECK:     return %[[RES]] : vector<5x1x1x20xf32>
+    return %0 : vector<5x1x1x20xf32>
+}
+
+// -----
+
 // CHECK-LABEL: func @insert_strided_slice_full_range
 //  CHECK-SAME: %[[SOURCE:.+]]: vector<16x16xf16>, %{{.+}}: vector<16x16xf16>
 func.func @insert_strided_slice_full_range(%source: vector<16x16xf16>, %dest: vector<16x16xf16>) -> vector<16x16xf16> {
   %0 = vector.insert_strided_slice %source, %dest {offsets = [0, 0], strides = [1, 1]} : vector<16x16xf16> into vector<16x16xf16>
   // CHECK: return %[[SOURCE]]
   return %0: vector<16x16xf16>
 }