[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

398

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

645

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

540

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

1061–1064

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

645

+1

[nicolasvasilache]

ddress review comments.

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

398

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

399

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

641

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

548

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

400

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

400

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

641

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

641

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

641

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 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 operate on n-D vectors and compute a fused
+    pointwise multiply-and-accumulate: `$result = `$lhs * $rhs + $acc`.
+    All operands and result have the same vector type. The semantics
+    of the operation correspond to those of the `llvm.fma`
+    [intrinsic](https://llvm.org/docs/LangRef.html#int-fma). In the
+    particular case of lowering to LLVM, this is guaranteed to lower
+    to the `llvm.fma.*` 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,
     extracted =
         rewriter.create<LLVM::ExtractElementOp>(loc, extracted, constant);
     rewriter.replaceOp(op, extracted);
 
     return matchSuccess();
   }
 };
 
+/// Conversion pattern that turns a vector.fma on a 1-D vector
+/// into an llvm.intr.fmuladd. This is a trivial 1-1 conversion.
+/// This does not match vectors of n >= 2 rank.
+///
+/// Example:
+/// ```
+///  vector.fma %a, %a, %a : vector<8xf32>
+/// ```
+/// is converted to:
+/// ```
+///  llvm.intr.fma %va, %va, %va:
+///    (!llvm<"<8 x float>">, !llvm<"<8 x float>">, !llvm<"<8 x float>">)
+///    -> !llvm<"<8 x float>">
+/// ```
+class VectorFMAOp1DConversion : public LLVMOpLowering {
+public:
+  explicit VectorFMAOp1DConversion(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::FMAOp>(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
if (positionAttrs.size() > 1) {
                                                       nMinusOnePositionAttrs);
     }
 
     rewriter.replaceOp(op, inserted);
     return matchSuccess();
   }
 };
 
+/// Rank reducing rewrite for n-D FMA into (n-1)-D FMA where n > 1.
+///
+/// Example:
+/// ```
+///   %d = vector.fma %a, %b, %c : vector<2x4xf32>
+/// ```
+/// is rewritten into:
+/// ```
+///  %r = splat %f0: vector<2x4xf32>
+///  %va = vector.extractvalue %a[0] : vector<2x4xf32>
+///  %vb = vector.extractvalue %b[0] : vector<2x4xf32>
+///  %vc = vector.extractvalue %c[0] : vector<2x4xf32>
+///  %vd = vector.fma %va, %vb, %vc : vector<4xf32>
+///  %r2 = vector.insertvalue %vd, %r[0] : vector<4xf32> into vector<2x4xf32>
+///  %va2 = vector.extractvalue %a2[1] : vector<2x4xf32>
+///  %vb2 = vector.extractvalue %b2[1] : vector<2x4xf32>
+///  %vc2 = vector.extractvalue %c2[1] : vector<2x4xf32>
+///  %vd2 = vector.fma %va2, %vb2, %vc2 : vector<4xf32>
+///  %r3 = vector.insertvalue %vd2, %r2[1] : vector<4xf32> into vector<2x4xf32>
+///  // %r3 holds the final value.
+/// ```
+class VectorFMAOpNDRewritePattern : 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:

 };
 
 } // 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<VectorFMAOpNDRewritePattern,
+                  VectorInsertStridedSliceOpDifferentRankRewritePattern,
                   VectorInsertStridedSliceOpSameRankRewritePattern,
                   VectorStridedSliceOpConversion>(ctx);
   patterns.insert<VectorBroadcastOpConversion, VectorShuffleOpConversion,
                   VectorExtractElementOpConversion, VectorExtractOpConversion,
-                  VectorInsertElementOpConversion, VectorInsertOpConversion,
-                  VectorOuterProductOpConversion, VectorTypeCastOpConversion,
-                  VectorPrintOpConversion>(ctx, converter);
+                  VectorFMAOp1DConversion, 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: %[[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(
+//  CHECK-SAME: %[[A:.*]]: !llvm<"<8 x float>">, %[[B:.*]]: !llvm<"[2 x <4 x float>]">)
+//  CHECK-SAME: -> !llvm<"{ <8 x float>, [2 x <4 x float>] }"> {
+func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>) -> (vector<8xf32>, vector<2x4xf32>) {
+  //         CHECK: "llvm.intr.fma"(%[[A]], %[[A]], %[[A]]) :
+  //    CHECK-SAME:   (!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: %[[b00:.*]] = llvm.extractvalue %[[B]][0] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[b01:.*]] = llvm.extractvalue %[[B]][0] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[b02:.*]] = llvm.extractvalue %[[B]][0] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[B0:.*]] = "llvm.intr.fma"(%[[b00]], %[[b01]], %[[b02]]) :
+  //  CHECK-SAME: (!llvm<"<4 x float>">, !llvm<"<4 x float>">, !llvm<"<4 x float>">) -> !llvm<"<4 x float>">
+  //       CHECK: llvm.insertvalue %[[B0]], {{.*}}[0] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[b10:.*]] = llvm.extractvalue %[[B]][1] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[b11:.*]] = llvm.extractvalue %[[B]][1] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[b12:.*]] = llvm.extractvalue %[[B]][1] : !llvm<"[2 x <4 x float>]">
+  //       CHECK: %[[B1:.*]] = "llvm.intr.fma"(%[[b10]], %[[b11]], %[[b12]]) :
+  //  CHECK-SAME: (!llvm<"<4 x float>">, !llvm<"<4 x float>">, !llvm<"<4 x float>">) -> !llvm<"<4 x float>">
+  //       CHECK: llvm.insertvalue %[[B1]], {{.*}}[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 @shape_cast(%arg0 : vector<5x1x3x2xf32>,
   %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<15x2xf32>
 
   // CHECK-NEXT: vector.shape_cast %{{.*}} : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to tuple<vector<20x2xf32>, vector<12x2xf32>>
   %1 = vector.shape_cast %arg1 : tuple<vector<5x4x2xf32>, vector<3x4x2xf32>> to
                                  tuple<vector<20x2xf32>, vector<12x2xf32>>
 
   return %0, %1 : vector<15x2xf32>, tuple<vector<20x2xf32>, vector<12x2xf32>>
 }
+
+// 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
+}