[mlir][VectorOps] Introduce a `vector.fma` op that works on n-D vectors and lowers to `llvm.intrin.fmuladd`

[nicolasvasilache]

Summary

The vector.fma operation is portable enough across targets that we do not want
to keep it wrapped under vector.outerproduct and llvm.intrin.fmuladd.
This revision lifts the op into the vector dialect and implements the lowering to LLVM by using two patterns:

1. a pattern that lowers from n-D to (n-1)-D by unrolling when n > 2
2. a pattern that converts from 1-D to the proper LLVM representation

[nicolasvasilache]

Avoid deleting all the tests.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[aartbik]

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

382

typo: that operate (plural)

but more in general, can you describe the semantics in a bit more detail than this? In particular, the lowering part to llvm could be mentioned at one point as motivation to have this, but it seems a bit strange to mention that in the very first sentence already.

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

618

perhaps use the %[[s0:.*]] = .... to capture the values and make sure they are used where expected?

[dcaballe]

Thanks, Nicolas! LGTM. Just minor changes.

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

413

Adding doc about what this pattern is matching and what is not would be great

526

I think adding a small example to the doc would help better understand what this patter is doing.

1013

Rename VectorFMAOpRewritePattern and VectorFMAOpConversion to be more aligned with what they match?
You'll probably find something better but something along the lines of SingleDimVectorFMAOpRewritePattern and MultiDimVectorFMAOpConversion

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

618

+1

[nicolasvasilache]

ddress review comments.

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

382

Made the description more general.
There is a notion of guaranteed fmuladd in the LLVM case that I kept.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[aartbik]

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

383

much better, and the mention of llvm later is fine of course

I would still say something about that all shapes need to match exactly (rank and dimensions), it is implied by the syntax, and of course enforced by the traits, but why not be a bit more explicit in the doc :-)

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

614

note that something like

// CHECK-SAME: %[[A:arg[0-9]+]]:

can be used if you want to match the argument inputs as well, see e.g. insert_strided_slice3 (probably need to rename a/b into arg0/arg1 to be sure, although I think that renaming happens anyway)

[aartbik]

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

534

this is not a check pattern, so why not fill out the %a, %b, %c and ssa vars for readability?

[fhahn]

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

384

IIRC llvm.fmuladd does not specify the rounding behaviour (mul and add may or may not be fused).

The way I read the description here it seems like fusion is ecpexted for vector.fma unless I am missing something. It might be worth clarifying the rounding behaviour

[nicolasvasilache]

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

384

Thanks @fhahn ! I overlooked the semantics section that clearly states to use fma.
What I want is indeed fma, will adapt the revision accordingly.

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

614

I do not see this as necessary, there is a 1-1 mapping that is already captured by the type.

[aartbik]

Is there a way to remove your blocking review but without approving yet? I don't see that option here.
Anyway, please look at the other comments that are still open.

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

614

It was of course just a suggestion :-)

Regardless, I find %[[arg0]] easier to parse by eye than {{.*}} and it is the direction we are generally taking with new tests. But I won't insist.

[nicolasvasilache]

Is there a way to remove your blocking review but without approving yet? I don't see that option here.
Anyway, please look at the other comments that are still open.

@aartbik that's fine, I still need to expose the fma intrinsic from LLVM and use it in this revision, fmuladd is not the right thing as @fhahn 's comment made me realize.

[nicolasvasilache]

Rebase on top of llvm.intr.fma addition.
Address review comments.

[nicolasvasilache]

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

614

ok, I am sensitive to general simplicity, maybe erring on the side of terseness is counter-productive, I'll just update it :)

mlir/include/mlir/Dialect/VectorOps/VectorOps.td

let extraClassDeclaration = [{
     static StringRef getSizesAttrName() { return "sizes"; }
     static StringRef getStridesAttrName() { return "strides"; }
   }];
   let assemblyFormat = [{
     $vector `,` $sizes `,` $strides attr-dict `:` type($vector) `into`
     type(results)
   }];
 }
 
+def Vector_FMAOp :
+  Op<Vector_Dialect, "fma", [NoSideEffect,
+                             AllTypesMatch<["lhs", "rhs", "acc", "result"]>]>,
+    Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyVector:$acc)>,
+    Results<(outs AnyVector:$result)> {
+  let summary = "vector fused multiply-add";
+  let description = [{
+    Multiply-add expressions that operates on n-D f32 or f64 vectors and lower
+    to the llvm.fmuladd.* intrinsic.
+
+    Example:
+    
+    ```
+      %3 = vector.fma %0, %1, %2: vector<8x16xf32>
+    ```
+  }];
+  // Fully specified by traits.
+  let verifier = ?;
+  let assemblyFormat = "$lhs `,` $rhs `,` $acc attr-dict `:` type($lhs)";
+  let builders = [OpBuilder<
+    "Builder *b, OperationState &result, Value lhs, Value rhs, Value acc",
+    "build(b, result, lhs.getType(), lhs, rhs, acc);">];
+  let extraClassDeclaration = [{
+    VectorType getVectorType() { return lhs().getType().cast<VectorType>(); }
+  }];
+}
+
 def Vector_InsertElementOp :
   Vector_Op<"insertelement", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      PredOpTrait<"dest operand and result have same type",
                  TCresIsSameAsOpBase<0, 1>>]>,
      Arguments<(ins AnyType:$source, AnyVector:$dest, AnyInteger:$position)>,
      Results<(outs AnyVector)> {

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position);
     extracted =
         rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant);
     rewriter.replaceOp(op, extracted);
 
     return matchSuccess();
   }
 };
 
+class VectorFMAOpConversion : public LLVMOpLowering {
+public:
+  explicit VectorFMAOpConversion(MLIRContext *context,
+                                 LLVMTypeConverter &typeConverter)
+      : LLVMOpLowering(vector::FMAOp::getOperationName(), context,
+                       typeConverter) {}
+
+  PatternMatchResult
+  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
+                  ConversionPatternRewriter &rewriter) const override {
+    auto adaptor = vector::FMAOpOperandAdaptor(operands);
+    vector::FMAOp fmaOp = cast<vector::FMAOp>(op);
+    VectorType vType = fmaOp.getVectorType();
+    if (vType.getRank() != 1)
+      return matchFailure();
+    rewriter.replaceOpWithNewOp<LLVM::FMulAddOp>(op, adaptor.lhs(),
+                                                 adaptor.rhs(), adaptor.acc());
+    return matchSuccess();
+  }
+};
+
 class VectorInsertElementOpConversion : public LLVMOpLowering {
 public:
   explicit VectorInsertElementOpConversion(MLIRContext *context,
                                            LLVMTypeConverter &typeConverter)
       : LLVMOpLowering(vector::InsertElementOp::getOperationName(), context,
                        typeConverter) {}
 
   PatternMatchResult
matchAndRewrite(Operation *op, ArrayRef<Value> operands,
     if (positionAttrs.size() > 1) {
       oneDVectorType = reducedVectorTypeBack(destVectorType);
       auto nMinusOnePositionAttrs =
           ArrayAttr::get(positionAttrs.drop_back(), context);
       extracted = rewriter.create<LLVM::ExtractValueOp>(
           loc, lowering.convertType(oneDVectorType), extracted,
           nMinusOnePositionAttrs);
     }
 
     // Insertion of an element into a 1-D LLVM vector.
     auto i64Type = LLVM::LLVMType::getInt64Ty(lowering.getDialect());
     auto constant = rewriter.create<LLVM::ConstantOp>(loc, i64Type, position);
     Value inserted = rewriter.create<LLVM::InsertElementOp>(
         loc, lowering.convertType(oneDVectorType), extracted, adaptor.source(),
         constant);
 
     // Potential insertion of resulting 1-D vector into array.
     if (positionAttrs.size() > 1) {
       auto nMinusOnePositionAttrs =
           ArrayAttr::get(positionAttrs.drop_back(), context);
       inserted = rewriter.create<LLVM::InsertValueOp>(loc, llvmResultType,
                                                       adaptor.dest(), inserted,
                                                       nMinusOnePositionAttrs);
     }
 
     rewriter.replaceOp(op, inserted);
     return matchSuccess();
   }
 };
 
+// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1.
+class VectorFMAOpRewritePattern : public OpRewritePattern<FMAOp> {
+public:
+  using OpRewritePattern<FMAOp>::OpRewritePattern;
+
+  PatternMatchResult matchAndRewrite(FMAOp op,
+                                     PatternRewriter &rewriter) const override {
+    auto vType = op.getVectorType();
+    if (vType.getRank() < 2)
+      return matchFailure();
+
+    auto loc = op.getLoc();
+    auto elemType = vType.getElementType();
+    Value zero = rewriter.create<ConstantOp>(loc, elemType,
+                                             rewriter.getZeroAttr(elemType));
+    Value desc = rewriter.create<SplatOp>(loc, vType, zero);
+    for (int64_t i = 0, e = vType.getShape().front(); i != e; ++i) {
+      Value extrLHS = rewriter.create<ExtractOp>(loc, op.lhs(), i);
+      Value extrRHS = rewriter.create<ExtractOp>(loc, op.rhs(), i);
+      Value extrACC = rewriter.create<ExtractOp>(loc, op.acc(), i);
+      Value fma = rewriter.create<FMAOp>(loc, extrLHS, extrRHS, extrACC);
+      desc = rewriter.create<InsertOp>(loc, fma, desc, i);
+    }
+    rewriter.replaceOp(op, desc);
+    return matchSuccess();
+  }
+};
+
 // When ranks are different, InsertStridedSlice needs to extract a properly
 // ranked vector from the destination vector into which to insert. This pattern
 // only takes care of this part and forwards the rest of the conversion to
 // another pattern that converts InsertStridedSlice for operands of the same
 // rank.
 //
 // RewritePattern for InsertStridedSliceOp where source and destination vectors
 // have different ranks. In this case:

 ///   2. extract + optional strided_slice + insert for the n-D case.
 class VectorStridedSliceOpConversion : public OpRewritePattern<StridedSliceOp> {
 public:
   using OpRewritePattern<StridedSliceOp>::OpRewritePattern;
 
   PatternMatchResult matchAndRewrite(StridedSliceOp op,
                                      PatternRewriter &rewriter) const override {
     auto dstType = op.getResult().getType().cast<VectorType>();
 
     assert(!op.offsets().getValue().empty() && "Unexpected empty offsets");
 
     int64_t offset =
         op.offsets().getValue().front().cast<IntegerAttr>().getInt();
     int64_t size = op.sizes().getValue().front().cast<IntegerAttr>().getInt();
     int64_t stride =
         op.strides().getValue().front().cast<IntegerAttr>().getInt();
 
     auto loc = op.getLoc();
     auto elemType = dstType.getElementType();
     assert(elemType.isIntOrIndexOrFloat());
     Value zero = rewriter.create<ConstantOp>(loc, elemType,
                                              rewriter.getZeroAttr(elemType));

 };
 
 } // namespace
 
 /// Populate the given list with patterns that convert from Vector to LLVM.
 void mlir::populateVectorToLLVMConversionPatterns(
     LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
   MLIRContext *ctx = converter.getDialect()->getContext();
-  patterns.insert<VectorInsertStridedSliceOpDifferentRankRewritePattern,
+  patterns.insert<VectorFMAOpRewritePattern,
+                  VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
                   VectorStridedSliceOpConversion>(ctx);
   patterns.insert<VectorBroadcastOpConversion, VectorShuffleOpConversion,
                   VectorExtractElementOpConversion, VectorExtractOpConversion,
-                  VectorInsertElementOpConversion, VectorInsertOpConversion,
-                  VectorOuterProductOpConversion, VectorTypeCastOpConversion,
-                  VectorPrintOpConversion>(ctx, converter);
+                  VectorFMAOpConversion, VectorInsertElementOpConversion,
+                  VectorInsertOpConversion, VectorOuterProductOpConversion,
+                  VectorTypeCastOpConversion, VectorPrintOpConversion>(
+      ctx, converter);
 }
 
 namespace {
 struct LowerVectorToLLVMPass : public ModulePass<LowerVectorToLLVMPass> {
   void runOnModule() override;
 };
 } // namespace
 

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

 //      CHECK: %[[s0:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1x1xf32>) : !llvm<"[1 x <1 x float>]">
 //      CHECK: %[[s1:.*]] = llvm.extractvalue %[[A]][2] : !llvm<"[3 x <3 x float>]">
 //      CHECK: %[[s3:.*]] = llvm.mlir.constant(dense<0.000000e+00> : vector<1xf32>) : !llvm<"<1 x float>">
 //      CHECK: %[[s4:.*]] = llvm.mlir.constant(2 : index) : !llvm.i64
 //      CHECK: %[[s5:.*]] = llvm.extractelement %[[s1]][%[[s4]] : !llvm.i64] : !llvm<"<3 x float>">
 //      CHECK: %[[s6:.*]] = llvm.mlir.constant(0 : index) : !llvm.i64
 //      CHECK: %[[s7:.*]] = llvm.insertelement %[[s5]], %[[s3]][%[[s6]] : !llvm.i64] : !llvm<"<1 x float>">
 //      CHECK: %[[s8:.*]] = llvm.insertvalue %[[s7]], %[[s0]][0] : !llvm<"[1 x <1 x float>]">
 //      CHECK: llvm.return %[[s8]] : !llvm<"[1 x <1 x float>]">
+
+// CHECK-LABEL: llvm.func @vector_fma
+func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>)
+  -> (vector<8xf32>, vector<2x4xf32>)
+{
+  //         CHECK: llvm.intr.fmuladd{{.*}}: (!llvm<"<8 x float>">, !llvm<"<8 x float>">, !llvm<"<8 x float>">) -> !llvm<"<8 x float>">
+  %0 = vector.fma %a, %a, %a : vector<8xf32>
+  // CHECK-COUNT-3: llvm.extractvalue {{.*}}[0] : !llvm<"[2 x <4 x float>]">
+  //         CHECK: llvm.intr.fmuladd{{.*}} : (!llvm<"<4 x float>">, !llvm<"<4 x float>">, !llvm<"<4 x float>">) -> !llvm<"<4 x float>">
+  //         CHECK: llvm.insertvalue {{.*}}, {{.*}}[0] : !llvm<"[2 x <4 x float>]">
+  // CHECK-COUNT-3: llvm.extractvalue {{.*}}[1] : !llvm<"[2 x <4 x float>]">
+  //         CHECK: llvm.intr.fmuladd{{.*}} : (!llvm<"<4 x float>">, !llvm<"<4 x float>">, !llvm<"<4 x float>">) -> !llvm<"<4 x float>">
+  //         CHECK: llvm.insertvalue {{.*}}, {{.*}}[1] : !llvm<"[2 x <4 x float>]">
+  %1 = vector.fma %b, %b, %b : vector<2x4xf32>
+  return %0, %1: vector<8xf32>, vector<2x4xf32>
+}

mlir/test/Dialect/VectorOps/ops.mlir

func @reshape(%arg0 : vector<3x2x4xf32>) -> (vector<2x3x4xf32>) {
   %c6 = constant 6 : index
   // CHECK:      %[[C9:.*]] = constant 9 : index
   %c9 = constant 9 : index
   // CHECK: vector.reshape %{{.*}}, [%[[C3]], %[[C6]]], [%[[C2]], %[[C9]]], [4] : vector<3x2x4xf32> to vector<2x3x4xf32>
   %1 = vector.reshape %arg0, [%c3, %c6], [%c2, %c9], [4]
     : vector<3x2x4xf32> to vector<2x3x4xf32>
 
   return %1 : vector<2x3x4xf32>
 }
+
+// CHECK-LABEL: @vector_fma
+func @vector_fma(%a: vector<8xf32>, %b: vector<8x4xf32>) {
+  // CHECK: vector.fma %{{.*}} : vector<8xf32>
+  vector.fma %a, %a, %a : vector<8xf32>
+  // CHECK: vector.fma %{{.*}} : vector<8x4xf32>
+  vector.fma %b, %b, %b : vector<8x4xf32>
+  return
+}