[mlir][Vector] Support 0-D vectors in FMAOp

[nicolasvasilache]

picking this up to land and avoid letting it rot.

[nicolasvasilache]

Rebase.

[michalt]

In D115742#3745986, @nicolasvasilache wrote:

picking this up to land and avoid letting it rot.

Thanks a lot! :D I somehow thought that this already got merged... Sorry for that!

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

def Vector_ExtractMapOp :
   let hasVerifier = 1;
 }
 
 def Vector_FMAOp :
   Op<Vector_Dialect, "fma", [
        NoSideEffect, AllTypesMatch<["lhs", "rhs", "acc", "result"]>,
        DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
      ] # ElementwiseMappable.traits>,
-    Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyVector:$acc)>,
-    Results<(outs AnyVector:$result)> {
+    Arguments<(ins AnyVectorOfAnyRank:$lhs,
+                   AnyVectorOfAnyRank:$rhs,
+                   AnyVectorOfAnyRank:$acc)>,
+    Results<(outs AnyVectorOfAnyRank:$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

mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp

 class VectorFMAOp1DConversion : public ConvertOpToLLVMPattern<vector::FMAOp> {
 public:
   using ConvertOpToLLVMPattern<vector::FMAOp>::ConvertOpToLLVMPattern;
 
   LogicalResult
   matchAndRewrite(vector::FMAOp fmaOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     VectorType vType = fmaOp.getVectorType();
-    if (vType.getRank() != 1)
+    if (vType.getRank() > 1)
       return failure();
     rewriter.replaceOpWithNewOp<LLVM::FMulAddOp>(
         fmaOp, adaptor.getLhs(), adaptor.getRhs(), adaptor.getAcc());
     return success();
   }
 };
 
 class VectorInsertElementOpConversion

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

 //       CHECK:    %[[V4_1:.*]] = llvm.extractvalue {{.*}}[1] : !llvm.array<2 x vector<4xf32>>
 //       CHECK:    %[[V4_0_1:.*]] = llvm.extractvalue {{.*}}[0, 1] : !llvm.array<16 x array<4 x vector<8xf32>>>
 //       CHECK:    %[[R8_2:.*]] = llvm.shufflevector %[[V4_1]], %[[V4_1]] [0, 1, 2, 3, 0, 0, 0, 0] : vector<4xf32>
 //       CHECK:    %[[R8_3:.*]] = llvm.shufflevector %[[R8_2]], %[[V4_0_1]] [8, 9, 0, 1, 2, 3, 14, 15] : vector<8xf32>
 //       CHECK:    llvm.insertvalue %[[R8_3]], {{.*}}[1] : !llvm.array<4 x vector<8xf32>>
 
 // -----
 
-func.func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>, %c: vector<1x1x1xf32>) -> (vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>) {
+func.func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>, %c: vector<1x1x1xf32>, %d: vector<f32>) -> (vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>, vector<f32>) {
   // CHECK-LABEL: @vector_fma
   //  CHECK-SAME: %[[A:.*]]: vector<8xf32>
   //  CHECK-SAME: %[[B:.*]]: vector<2x4xf32>
   //  CHECK-SAME: %[[C:.*]]: vector<1x1x1xf32>
   //       CHECK: %[[BL:.*]] = builtin.unrealized_conversion_cast %[[B]] : vector<2x4xf32> to !llvm.array<2 x vector<4xf32>>
   //       CHECK: "llvm.intr.fmuladd"
   //  CHECK-SAME:   (vector<8xf32>, vector<8xf32>, vector<8xf32>) -> vector<8xf32>
   %0 = vector.fma %a, %a, %a : vector<8xf32>
func.func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>, %c: vector<1x1x1xf32>, %d: vector<f32>) -> (vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>, vector<f32>) {
   //  CHECK-SAME: (vector<4xf32>, vector<4xf32>, vector<4xf32>) -> vector<4xf32>
   //       CHECK: llvm.insertvalue %[[B1]], {{.*}}[1] : !llvm.array<2 x vector<4xf32>>
   %1 = vector.fma %b, %b, %b : vector<2x4xf32>
 
   //       CHECK: %[[C0:.*]] = "llvm.intr.fmuladd"
   //  CHECK-SAME:   (vector<1xf32>, vector<1xf32>, vector<1xf32>) -> vector<1xf32>
   %2 = vector.fma %c, %c, %c : vector<1x1x1xf32>
 
-  return %0, %1, %2: vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>
+  //       CHECK: %[[D0:.*]] = "llvm.intr.fmuladd"
+  //  CHECK-SAME:   (vector<1xf32>, vector<1xf32>, vector<1xf32>) -> vector<1xf32>
+  %3 = vector.fma %d, %d, %d : vector<f32>
+
+  return %0, %1, %2, %3: vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>, vector<f32>
 }
 
 // -----
 
 func.func @reduce_0d_f32(%arg0: vector<f32>) -> f32 {
   %0 = vector.reduction <add>, %arg0 : vector<f32> into f32
   return %0 : f32
 }

mlir/test/Dialect/Vector/ops.mlir

func.func @bitcast(%arg0 : vector<5x1x3x2xf32>,
 
   // CHECK: vector.bitcast %{{.*}} : vector<f32> to vector<i32>
   %8 = vector.bitcast %arg4 : vector<f32> to vector<i32>
 
   return %0, %1, %2, %3, %4, %5, %6, %7, %8 : vector<5x1x3x4xf16>, vector<5x1x3x8xi8>, vector<8x4xi8>, vector<8x1xf32>, vector<16x1x2xi32>, vector<16x1x4xi16>, vector<16x1x1xindex>, vector<8x2x2xf32>, vector<i32>
 }
 
 // CHECK-LABEL: @vector_fma
-func.func @vector_fma(%a: vector<8xf32>, %b: vector<8x4xf32>) {
+func.func @vector_fma(%a: vector<8xf32>, %b: vector<8x4xf32>, %c: vector<f32>) {
   // CHECK: vector.fma %{{.*}} : vector<8xf32>
   vector.fma %a, %a, %a : vector<8xf32>
   // CHECK: vector.fma %{{.*}} : vector<8x4xf32>
   vector.fma %b, %b, %b : vector<8x4xf32>
+  // CHECK: vector.fma %{{.*}} : vector<f32>
+  vector.fma %c, %c, %c : vector<f32>
   return
 }
 
 // CHECK-LABEL: @reduce_fp
 func.func @reduce_fp(%arg0: vector<16xf32>, %arg1: f32) -> f32 {
   // CHECK:    vector.reduction <add>, %{{.*}} : vector<16xf32> into f32
   vector.reduction <add>, %arg0 : vector<16xf32> into f32
   // CHECK:    vector.reduction <add>, %{{.*}}, %{{.*}} : vector<16xf32> into f32

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

func.func @create_mask_0d(%zero : index, %one : index) {
 
   %one_mask = vector.create_mask %one : vector<i1>
   // CHECK: ( 1 )
   vector.print %one_mask : vector<i1>
 
   return
 }
 
-func.func @reduce_add(%arg0: vector<f32>) -> f32 {
+func.func @reduce_add(%arg0: vector<f32>) {
   %0 = vector.reduction <add>, %arg0 : vector<f32> into f32    
-  return %0 : f32
+  vector.print %0 : f32
+  // CHECK: 5
+  return
+}
+
+func.func @fma_0d(%four: vector<f32>) {
+  %0 = vector.fma %four, %four, %four : vector<f32>
+  // 4 * 4 + 4 = 20
+  // CHECK: ( 20 )
+  vector.print %0: vector<f32>
+  return
 }
 
 func.func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
 
func.func @entry() {
   %bigger = arith.constant dense<4242> : vector<i32>
   call  @arith_cmpi_0d(%smaller, %bigger) : (vector<i32>, vector<i32>) -> ()
 
   %zero_idx = arith.constant 0 : index
   %one_idx = arith.constant 1 : index
   call  @create_mask_0d(%zero_idx, %one_idx) : (index, index) -> ()
 
   %red_array = arith.constant dense<5.0> : vector<f32>
-  %red_res = call  @reduce_add(%red_array) : (vector<f32>) -> (f32)  
-  vector.print %red_res : f32
-  // CHECK: 5
+  call  @reduce_add(%red_array) : (vector<f32>) -> ()
+
+  %5 = arith.constant dense<4.0> : vector<f32>
+  call  @fma_0d(%5) : (vector<f32>) -> ()
 
   return
 }