[mlir][vector] Add support for unrolling vector.fma

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

def Vector_ExtractMapOp :
   let assemblyFormat = [{
     $vector `[` $ids `]` attr-dict `:` type($vector) `to` type(results)
   }];
 
   let hasFolder = 1;
 }
 
 def Vector_FMAOp :
-  Op<Vector_Dialect, "fma", [NoSideEffect,
-                             AllTypesMatch<["lhs", "rhs", "acc", "result"]>]>,
+  Op<Vector_Dialect, "fma", [
+       NoSideEffect, AllTypesMatch<["lhs", "rhs", "acc", "result"]>,
+       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
+     ]>,
     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`

mlir/lib/Dialect/Vector/VectorOps.cpp

AffineMap calculateImplicitMap(MapOp op) {
   auto map = AffineMap::get(op.getSourceVectorType().getRank(), 0, perm,
                             op.getContext());
   return map;
 }
 
 AffineMap ExtractMapOp::map() { return calculateImplicitMap(*this); }
 
 //===----------------------------------------------------------------------===//
+// FmaOp
+//===----------------------------------------------------------------------===//
+
+Optional<SmallVector<int64_t, 4>> FMAOp::getShapeForUnroll() {
+  return llvm::to_vector<4>(getVectorType().getShape());
+}
+
+//===----------------------------------------------------------------------===//
 // BroadcastOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(BroadcastOp op) {
   VectorType srcVectorType = op.getSourceType().dyn_cast<VectorType>();
   VectorType dstVectorType = op.getVectorType();
   // Scalar to vector broadcast is always valid. A vector
   // to vector broadcast needs some additional checking.
OpFoldResult TransferReadOp::fold(ArrayRef<Attribute>) {
   if (succeeded(foldTransferMaskAttribute(*this)))
     return getResult();
   if (succeeded(foldMemRefCast(*this)))
     return getResult();
   return OpFoldResult();
 }
 
 Optional<SmallVector<int64_t, 4>> TransferReadOp::getShapeForUnroll() {
-  auto s = getVectorType().getShape();
-  return SmallVector<int64_t, 4>{s.begin(), s.end()};
+  return llvm::to_vector<4>(getVectorType().getShape());
 }
 
 void TransferReadOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   if (getShapedType().isa<MemRefType>())
     effects.emplace_back(MemoryEffects::Read::get(), source(),
                          SideEffects::DefaultResource::get());

mlir/test/Dialect/Vector/vector-unroll-options.mlir

 //  CHECK-SAME:     vector<4x4xf16>, vector<4x4xf16> into vector<4x4xf16>
 //       CHECK:   vector.contract {
 //  CHECK-SAME:     vector<4x4xf16>, vector<4x4xf16> into vector<4x4xf16>
 //       CHECK:   vector.contract {
 //  CHECK-SAME:     vector<4x4xf16>, vector<4x4xf16> into vector<4x4xf16>
 //       CHECK:   vector.contract {
 //  CHECK-SAME:     vector<4x4xf16>, vector<4x4xf16> into vector<4x4xf16>
 //       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>

mlir/test/lib/Transforms/TestVectorTransforms.cpp

struct TestVectorUnrollingPatterns
   TestVectorUnrollingPatterns() = default;
   TestVectorUnrollingPatterns(const TestVectorUnrollingPatterns &pass) {}
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;
     patterns.insert<UnrollVectorPattern>(
         ctx, UnrollVectorOptions()
                  .setNativeShape(ArrayRef<int64_t>{2, 2})
-                 .setFilterConstraint(
-                     [](Operation *op) { return success(isa<AddFOp>(op)); }));
+                 .setFilterConstraint([](Operation *op) {
+                   return success(isa<AddFOp, vector::FMAOp>(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()
                                  .getElementType()