[mlir][vector] Fold one-element reduction into extract

[antiagainst]

Summary

If there is only one single element in the vector, then we can
just extract the element as the result for the reduction.

[mravishankar]

Just a small question.

mlir/test/Dialect/Vector/canonicalize.mlir

1581

Couldnt you just make this vector.extract -> addf ?

[antiagainst]

Address comments

[antiagainst]

mlir/test/Dialect/Vector/canonicalize.mlir

1581

Sure. Implemented this case too.

[mravishankar]

Awesome!

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

let builders = [
     OpBuilder<(ins "CombiningKind":$kind, "Value":$vector, "Value":$acc)>,
     // Builder that infers the type of `dest` and has no accumulator.
     OpBuilder<(ins "CombiningKind":$kind, "Value":$vector)>
   ];
 
   // TODO: Migrate to assemblyFormat once `AllTypesMatch` supports optional
   // operands.
   let hasCustomAssemblyFormat = 1;
+  let hasCanonicalizer = 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>,

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

Value mlir::vector::getVectorReductionOp(arith::AtomicRMWKind op,
   }
   return nullptr;
 }
 
 Optional<SmallVector<int64_t, 4>> ReductionOp::getShapeForUnroll() {
   return llvm::to_vector<4>(getVectorType().getShape());
 }
 
+namespace {
+struct ElideSingleElementReduction : public OpRewritePattern<ReductionOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ReductionOp reductionOp,
+                                PatternRewriter &rewriter) const override {
+    if (reductionOp.getVectorType().getDimSize(0) != 1)
+      return failure();
+
+    Location loc = reductionOp.getLoc();
+    Value result = rewriter.create<ExtractOp>(loc, reductionOp.getType(),
+                                              reductionOp.getVector(),
+                                              rewriter.getI64ArrayAttr(0));
+
+    if (Value acc = reductionOp.getAcc()) {
+      assert(reductionOp.getType().isa<FloatType>());
+      switch (reductionOp.getKind()) {
+      case CombiningKind::ADD:
+        result = rewriter.create<arith::AddFOp>(loc, result, acc);
+        break;
+      case CombiningKind::MUL:
+        result = rewriter.create<arith::MulFOp>(loc, result, acc);
+        break;
+      default:
+        assert(false && "invalid op!");
+      }
+    }
+
+    rewriter.replaceOp(reductionOp, result);
+    return success();
+  }
+};
+} // namespace
+
+void ReductionOp::getCanonicalizationPatterns(RewritePatternSet &results,
+                                              MLIRContext *context) {
+  results.add<ElideSingleElementReduction>(context);
+}
+
 //===----------------------------------------------------------------------===//
 // ContractionOp
 //===----------------------------------------------------------------------===//
 
 void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
                                   ArrayRef<ArrayRef<AffineExpr>> indexingExprs,
                                   ArrayRef<StringRef> iteratorTypes) {

mlir/test/Dialect/Vector/canonicalize.mlir

 //  CHECK-SAME: (%[[ARG:.+]]: i32)
 //       CHECK:   return %[[ARG]]
 func.func @extract_element_splat_fold(%a : i32) -> i32 {
   %v = vector.splat %a : vector<4xi32>
   %i = arith.constant 2 : i32
   %1 = vector.extractelement %v[%i : i32] : vector<4xi32>
   return %1 : i32
 }
+
+// -----
+
+// CHECK-LABEL: func @reduce_one_element_vector_extract
+//  CHECK-SAME: (%[[V:.+]]: vector<1xf32>)
+//       CHECK:   %[[S:.+]] = vector.extract %[[V]][0] : vector<1xf32>
+//       CHECK:   return %[[S]] : f32
+func @reduce_one_element_vector_extract(%a : vector<1xf32>) -> f32 {
+  %s = vector.reduction <add>, %a : vector<1xf32> into f32
+  return %s : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduce_one_element_vector_addf
+//  CHECK-SAME: (%[[V:.+]]: vector<1xf32>, %[[B:.+]]: f32)
+//       CHECK:   %[[A:.+]] = vector.extract %[[V]][0] : vector<1xf32>
+//       CHECK:   %[[S:.+]] = arith.addf %[[A]], %arg1 : f32
+//       CHECK:   return %[[S]]
+func @reduce_one_element_vector_addf(%a : vector<1xf32>, %b: f32) -> f32 {
+  %s = vector.reduction <add>, %a, %b : vector<1xf32> into f32
+  return %s : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @reduce_one_element_vector_mulf
+//  CHECK-SAME: (%[[V:.+]]: vector<1xf32>, %[[B:.+]]: f32)
+//       CHECK:   %[[A:.+]] = vector.extract %[[V]][0] : vector<1xf32>
+//       CHECK:   %[[S:.+]] = arith.mulf %[[A]], %arg1 : f32
+//       CHECK:   return %[[S]]
+func @reduce_one_element_vector_mulf(%a : vector<1xf32>, %b: f32) -> f32 {
+  %s = vector.reduction <mul>, %a, %b : vector<1xf32> into f32
+  return %s : f32
+}
+
+// -----
+
+// CHECK-LABEL: func @dont_reduce_one_element_vector
+//       CHECK: vector.reduction
+func @dont_reduce_one_element_vector(%a : vector<4xf32>) -> f32 {
+  %s = vector.reduction <add>, %a : vector<4xf32> into f32
+  return %s : f32
+}