[mlir] [VectorOps] Replace zero fma with mult for vector.contract

[aartbik]

Summary

More efficient implementation of the multiply-reduce pair,
no need to add in a zero vector. Microbenchmarking on AVX2
yields the following difference in vector.contract speedup
(over strict-order scalar reduction).

SPEEDUP SIMD-fma SIMD-mul
4x4 1.45 2.00
8x8 1.40 1.90
32x32 5.32 5.80

mlir/lib/Dialect/Vector/VectorTransforms.cpp

LogicalResult matchAndRewrite(vector::OuterProductOp op,
     Value result = rewriter.create<ConstantOp>(loc, resType,
                                                rewriter.getZeroAttr(resType));
     for (int64_t d = 0, e = resType.getDimSize(0); d < e; ++d) {
       auto pos = rewriter.getI64ArrayAttr(d);
       Value x = rewriter.create<vector::ExtractOp>(loc, eltType, op.lhs(), pos);
       Value b = rewriter.create<vector::BroadcastOp>(loc, rhsType, x);
       Value m;
       if (acc) {
-        Value z = rewriter.create<vector::ExtractOp>(loc, rhsType, acc, pos);
-        m = rewriter.create<vector::FMAOp>(loc, b, op.rhs(), z);
+        Value e = rewriter.create<vector::ExtractOp>(loc, rhsType, acc, pos);
+        m = rewriter.create<vector::FMAOp>(loc, b, op.rhs(), e);
       } else {
         m = rewriter.create<MulFOp>(loc, b, op.rhs());
       }
       result = rewriter.create<vector::InsertOp>(loc, resType, m, result, pos);
     }
     rewriter.replaceOp(op, result);
     return success();
   }

 /// One:
 ///   %x = vector.contract with at least one free/batch dimension
 /// is replaced by:
 ///   %a = vector.contract with one less free/batch dimension
 ///   %b = vector.contract with one less free/batch dimension
 ///   ..
 ///   %x = combine %a %b ..
 /// until a pure contraction is reached (no free/batch dimensions),
-/// which is replaced by a fma/reduction op.
+/// which is replaced by a dot-product/reduction pair.
 ///
 /// TODO(ajcbik): break down into transpose/reshape/cast ops
 ///               when they become available to avoid code dup
 /// TODO(ajcbik): investigate lowering order impact on performance
 LogicalResult
 ContractionOpLowering::matchAndRewrite(vector::ContractionOp op,
                                        PatternRewriter &rewriter) const {
 
Value ContractionOpLowering::lowerReduction(vector::ContractionOp op,
   assert(lookupRhs.hasValue() && "missing RHS parallel index");
   int64_t lhsIndex = lookupLhs.getValue();
   int64_t rhsIndex = lookupRhs.getValue();
   int64_t dimSize = lhsType.getDimSize(lhsIndex);
   assert(dimSize == rhsType.getDimSize(rhsIndex) && "corrupt shape");
   // Base case.
   if (lhsType.getRank() == 1) {
     assert(rhsType.getRank() == 1 && "corrupt contraction");
-    Value zero = rewriter.create<ConstantOp>(loc, lhsType,
-                                             rewriter.getZeroAttr(lhsType));
-    Value fma = rewriter.create<vector::FMAOp>(loc, op.lhs(), op.rhs(), zero);
+    Value m = rewriter.create<MulFOp>(loc, op.lhs(), op.rhs());
     StringAttr kind = rewriter.getStringAttr("add");
-    return rewriter.create<vector::ReductionOp>(loc, resType, kind, fma,
+    return rewriter.create<vector::ReductionOp>(loc, resType, kind, m,
                                                 op.acc());
   }
   // Construct new iterator types and affine map array attribute.
   SmallVector<AffineMap, 4> lowIndexingMaps;
   lowIndexingMaps.push_back(adjustMap(iMap[0], iterIndex, rewriter));
   lowIndexingMaps.push_back(adjustMap(iMap[1], iterIndex, rewriter));
   lowIndexingMaps.push_back(adjustMap(iMap[2], iterIndex, rewriter));
   auto lowAffine = rewriter.getAffineMapArrayAttr(lowIndexingMaps);

mlir/test/Dialect/Vector/vector-contract-transforms.mlir

#dotp_trait = {
   indexing_maps = #dotp_accesses,
   iterator_types = ["reduction"]
 }
 
 // CHECK-LABEL: func @extract_contract1
 // CHECK-SAME: %[[A:.*0]]: vector<4xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<4xf32>,
 // CHECK-SAME: %[[C:.*2]]: f32
-// CHECK:      %[[Z:.*]] = constant dense<0.000000e+00> : vector<4xf32>
-// CHECK:      %[[F:.*]] = vector.fma %[[A]], %[[B]], %[[Z]] : vector<4xf32>
+// CHECK:      %[[F:.*]] = mulf %[[A]], %[[B]] : vector<4xf32>
 // CHECK:      %[[R:.*]] = vector.reduction "add", %[[F]], %[[C]] : vector<4xf32> into f32
 // CHECK:      return %[[R]] : f32
 
 func @extract_contract1(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: f32) -> f32 {
   %0 = vector.contract #dotp_trait %arg0, %arg1, %arg2
     : vector<4xf32>, vector<4xf32> into f32
   return %0 : f32
 }
 
 #matvec_accesses = [
   affine_map<(i, j) -> (i, j)>,
   affine_map<(i, j) -> (j)>,
   affine_map<(i, j) -> (i)>
 ]
 #matvec_trait = {
   indexing_maps = #matvec_accesses,
   iterator_types = ["parallel", "reduction"]
 }
 
 // CHECK-LABEL: func @extract_contract2
 // CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<3xf32>,
 // CHECK-SAME: %[[C:.*2]]: vector<2xf32>
 // CHECK:      %[[R:.*]] = constant dense<0.000000e+00> : vector<2xf32>
-// CHECK:      %[[Z:.*]] = constant dense<0.000000e+00> : vector<3xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[C]][0] : vector<2xf32>
-// CHECK:      %[[T2:.*]] = vector.fma %[[T0]], %[[B]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T2:.*]] = mulf %[[T0]], %[[B]] : vector<3xf32>
 // CHECK:      %[[T3:.*]] = vector.reduction "add", %[[T2]], %[[T1]] : vector<3xf32> into f32
 // CHECK:      %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : f32 into vector<2xf32>
 // CHECK:      %[[T5:.*]] = vector.extract %[[A]][1] : vector<2x3xf32>
 // CHECK:      %[[T6:.*]] = vector.extract %[[C]][1] : vector<2xf32>
-// CHECK:      %[[T7:.*]] = vector.fma %[[T5]], %[[B]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T7:.*]] = mulf %[[T5]], %[[B]] : vector<3xf32>
 // CHECK:      %[[T8:.*]] = vector.reduction "add", %[[T7]], %[[T6]] : vector<3xf32> into f32
 // CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : f32 into vector<2xf32>
 // CHECK:      return %[[T9]] : vector<2xf32>
 
 func @extract_contract2(%arg0: vector<2x3xf32>,
                         %arg1: vector<3xf32>,
       %arg2: vector<2xf32>) -> vector<2xf32> {
   %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2
#vecmat_trait = {
   iterator_types = ["parallel", "reduction"]
 }
 
 // CHECK-LABEL: func @extract_contract3
 // CHECK-SAME: %[[A:.*0]]: vector<3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,
 // CHECK-SAME: %[[C:.*2]]: vector<2xf32>
 // CHECK:      %[[R:.*]] = constant dense<0.000000e+00> : vector<2xf32>
-// CHECK:      %[[Z:.*]] = constant dense<0.000000e+00> : vector<3xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[B]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[C]][0] : vector<2xf32>
-// CHECK:      %[[T2:.*]] = vector.fma %[[A]], %[[T0]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T2:.*]] = mulf %[[A]], %[[T0]] : vector<3xf32>
 // CHECK:      %[[T3:.*]] = vector.reduction "add", %[[T2]], %[[T1]] : vector<3xf32> into f32
 // CHECK:      %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : f32 into vector<2xf32>
 // CHECK:      %[[T5:.*]] = vector.extract %[[B]][1] : vector<2x3xf32>
 // CHECK:      %[[T6:.*]] = vector.extract %[[C]][1] : vector<2xf32>
-// CHECK:      %[[T7:.*]] = vector.fma %[[A]], %[[T5]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T7:.*]] = mulf %[[A]], %[[T5]] : vector<3xf32>
 // CHECK:      %[[T8:.*]] = vector.reduction "add", %[[T7]], %[[T6]] : vector<3xf32> into f32
 // CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : f32 into vector<2xf32>
 // CHECK:      return %[[T9]] : vector<2xf32>
 
 func @extract_contract3(%arg0: vector<3xf32>,
                         %arg1: vector<2x3xf32>,
                         %arg2: vector<2xf32>) -> vector<2xf32> {
   %0 = vector.contract #vecmat_trait %arg0, %arg1, %arg2

 // CHECK:    %[[T1:.*]] = vector.extract %[[C]][0] : vector<2x2xf32>
 // CHECK:    %[[T2:.*]] = vector.extract %[[B]][0] : vector<2x2xf32>
 // CHECK:    %[[T3:.*]] = vector.extract %[[T2]][0] : vector<2xf32>
 // CHECK:    %[[T4:.*]] = vector.insert %[[T3]], %[[Z]] [0] : f32 into vector<2xf32>
 // CHECK:    %[[T5:.*]] = vector.extract %[[B]][1] : vector<2x2xf32>
 // CHECK:    %[[T6:.*]] = vector.extract %[[T5]][0] : vector<2xf32>
 // CHECK:    %[[T7:.*]] = vector.insert %[[T6]], %[[T4]] [1] : f32 into vector<2xf32>
 // CHECK:    %[[T8:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
-// CHECK:    %[[T9:.*]] = vector.fma %[[T0]], %[[T7]], %[[Z]] : vector<2xf32>
+// CHECK:    %[[T9:.*]] = mulf %[[T0]], %[[T7]] : vector<2xf32>
 // CHECK:    %[[T10:.*]] = vector.reduction "add", %[[T9]], %[[T8]] : vector<2xf32> into f32
 // CHECK:    %[[T11:.*]] = vector.insert %[[T10]], %[[Z]] [0] : f32 into vector<2xf32>
 // CHECK:    %[[T12:.*]] = vector.extract %[[B]][0] : vector<2x2xf32>
 // CHECK:    %[[T13:.*]] = vector.extract %[[T12]][1] : vector<2xf32>
 // CHECK:    %[[T14:.*]] = vector.insert %[[T13]], %[[Z]] [0] : f32 into vector<2xf32>
 // CHECK:    %[[T15:.*]] = vector.extract %[[B]][1] : vector<2x2xf32>
 // CHECK:    %[[T16:.*]] = vector.extract %[[T15]][1] : vector<2xf32>
 // CHECK:    %[[T17:.*]] = vector.insert %[[T16]], %[[T14]] [1] : f32 into vector<2xf32>
 // CHECK:    %[[T18:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
-// CHECK:    %[[T19:.*]] = vector.fma %[[T0]], %[[T17]], %[[Z]] : vector<2xf32>
+// CHECK:    %[[T19:.*]] = mulf %[[T0]], %[[T17]] : vector<2xf32>
 // CHECK:    %[[T20:.*]] = vector.reduction "add", %[[T19]], %[[T18]] : vector<2xf32> into f32
 // CHECK:    %[[T21:.*]] = vector.insert %[[T20]], %[[T11]] [1] : f32 into vector<2xf32>
 // CHECK:    %[[T22:.*]] = vector.insert %[[T21]], %[[R]] [0] : vector<2xf32> into vector<2x2xf32>
 // CHECK:    %[[T23:.*]] = vector.extract %[[A]][1] : vector<2x2xf32>
 // CHECK:    %[[T24:.*]] = vector.extract %[[C]][1] : vector<2x2xf32>
 // CHECK:    %[[T25:.*]] = vector.extract %[[B]][0] : vector<2x2xf32>
 // CHECK:    %[[T26:.*]] = vector.extract %[[T25]][0] : vector<2xf32>
 // CHECK:    %[[T27:.*]] = vector.insert %[[T26]], %[[Z]] [0] : f32 into vector<2xf32>
 // CHECK:    %[[T28:.*]] = vector.extract %[[B]][1] : vector<2x2xf32>
 // CHECK:    %[[T29:.*]] = vector.extract %[[T28]][0] : vector<2xf32>
 // CHECK:    %[[T30:.*]] = vector.insert %[[T29]], %[[T27]] [1] : f32 into vector<2xf32>
 // CHECK:    %[[T31:.*]] = vector.extract %[[T24]][0] : vector<2xf32>
-// CHECK:    %[[T32:.*]] = vector.fma %[[T23]], %[[T30]], %[[Z]] : vector<2xf32>
+// CHECK:    %[[T32:.*]] = mulf %[[T23]], %[[T30]] : vector<2xf32>
 // CHECK:    %[[T33:.*]] = vector.reduction "add", %[[T32]], %[[T31]] : vector<2xf32> into f32
 // CHECK:    %[[T34:.*]] = vector.insert %[[T33]], %[[Z]] [0] : f32 into vector<2xf32>
 // CHECK:    %[[T35:.*]] = vector.extract %[[B]][0] : vector<2x2xf32>
 // CHECK:    %[[T36:.*]] = vector.extract %[[T35]][1] : vector<2xf32>
 // CHECK:    %[[T37:.*]] = vector.insert %[[T36]], %[[Z]] [0] : f32 into vector<2xf32>
 // CHECK:    %[[T38:.*]] = vector.extract %[[B]][1] : vector<2x2xf32>
 // CHECK:    %[[T39:.*]] = vector.extract %[[T38]][1] : vector<2xf32>
 // CHECK:    %[[T40:.*]] = vector.insert %[[T39]], %[[T37]] [1] : f32 into vector<2xf32>
 // CHECK:    %[[T41:.*]] = vector.extract %[[T24]][1] : vector<2xf32>
-// CHECK:    %[[T42:.*]] = vector.fma %[[T23]], %[[T40]], %[[Z]] : vector<2xf32>
+// CHECK:    %[[T42:.*]] = mulf %[[T23]], %[[T40]] : vector<2xf32>
 // CHECK:    %[[T43:.*]] = vector.reduction "add", %[[T42]], %[[T41]] : vector<2xf32> into f32
 // CHECK:    %[[T44:.*]] = vector.insert %[[T43]], %[[T34]] [1] : f32 into vector<2xf32>
 // CHECK:    %[[T45:.*]] = vector.insert %[[T44]], %[[T22]] [1] : vector<2xf32> into vector<2x2xf32>
 // CHECK:    return %[[T45]] : vector<2x2xf32>
 
 func @extract_contract4(%arg0: vector<2x2xf32>,
                         %arg1: vector<2x2xf32>,
                         %arg2: vector<2x2xf32>) -> vector<2x2xf32> {
#contraction2d_trait = {
   indexing_maps = #contraction2d_accesses,
   iterator_types = ["reduction", "reduction"]
 }
 
 // CHECK-LABEL: func @full_contract1
 // CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,
 // CHECK-SAME: %[[C:.*2]]: f32
-// CHECK:      %[[Z:.*]] = constant dense<0.000000e+00> : vector<3xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[B]][0] : vector<2x3xf32>
-// CHECK:      %[[T2:.*]] = vector.fma %[[T0]], %[[T1]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T2:.*]] = mulf %[[T0]], %[[T1]] : vector<3xf32>
 // CHECK:      %[[T3:.*]] = vector.reduction "add", %[[T2]], %[[C]] : vector<3xf32> into f32
 // CHECK:      %[[T4:.*]] = vector.extract %[[A]][1] : vector<2x3xf32>
 // CHECK:      %[[T5:.*]] = vector.extract %[[B]][1] : vector<2x3xf32>
-// CHECK:      %[[T6:.*]] = vector.fma %[[T4]], %[[T5]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T6:.*]] = mulf %[[T4]], %[[T5]] : vector<3xf32>
 // CHECK:      %[[T7:.*]] = vector.reduction "add", %[[T6]], %[[T3]] : vector<3xf32> into f32
 // CHECK:      return %[[T7]] : f32
 
 func @full_contract1(%arg0: vector<2x3xf32>,
                      %arg1: vector<2x3xf32>,
          %arg2: f32) -> f32 {
   %0 = vector.contract #contraction2d_trait %arg0, %arg1, %arg2
     : vector<2x3xf32>, vector<2x3xf32> into f32

 // CHECK:      %[[T2:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
 // CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[Z]] [0] : f32 into vector<3xf32>
 // CHECK:      %[[T4:.*]] = vector.extract %[[B]][1] : vector<3x2xf32>
 // CHECK:      %[[T5:.*]] = vector.extract %[[T4]][0] : vector<2xf32>
 // CHECK:      %[[T6:.*]] = vector.insert %[[T5]], %[[T3]] [1] : f32 into vector<3xf32>
 // CHECK:      %[[T7:.*]] = vector.extract %[[B]][2] : vector<3x2xf32>
 // CHECK:      %[[T8:.*]] = vector.extract %[[T7]][0] : vector<2xf32>
 // CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T6]] [2] : f32 into vector<3xf32>
-// CHECK:      %[[T10:.*]] = vector.fma %[[T0]], %[[T9]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T10:.*]] = mulf %[[T0]], %[[T9]] : vector<3xf32>
 // CHECK:      %[[T11:.*]] = vector.reduction "add", %[[T10]], %[[C]] : vector<3xf32> into f32
 // CHECK:      %[[T12:.*]] = vector.extract %[[A]][1] : vector<2x3xf32>
 // CHECK:      %[[T13:.*]] = vector.extract %[[B]][0] : vector<3x2xf32>
 // CHECK:      %[[T14:.*]] = vector.extract %[[T13]][1] : vector<2xf32>
 // CHECK:      %[[T15:.*]] = vector.insert %[[T14]], %[[Z]] [0] : f32 into vector<3xf32>
 // CHECK:      %[[T16:.*]] = vector.extract %[[B]][1] : vector<3x2xf32>
 // CHECK:      %[[T17:.*]] = vector.extract %[[T16]][1] : vector<2xf32>
 // CHECK:      %[[T18:.*]] = vector.insert %[[T17]], %[[T15]] [1] : f32 into vector<3xf32>
 // CHECK:      %[[T19:.*]] = vector.extract %[[B]][2] : vector<3x2xf32>
 // CHECK:      %[[T20:.*]] = vector.extract %[[T19]][1] : vector<2xf32>
 // CHECK:      %[[T21:.*]] = vector.insert %[[T20]], %[[T18]] [2] : f32 into vector<3xf32>
-// CHECK:      %[[T22:.*]] = vector.fma %[[T12]], %[[T21]], %[[Z]] : vector<3xf32>
+// CHECK:      %[[T22:.*]] = mulf %[[T12]], %[[T21]] : vector<3xf32>
 // CHECK:      %[[T23:.*]] = vector.reduction "add", %[[T22]], %[[T11]] : vector<3xf32> into f32
 // CHECK:      return %[[T23]] : f32
 
 func @full_contract2(%arg0: vector<2x3xf32>,
                      %arg1: vector<3x2xf32>,
          %arg2: f32) -> f32 {
   %0 = vector.contract #contraction2d_trans_trait %arg0, %arg1, %arg2
     : vector<2x3xf32>, vector<3x2xf32> into f32