[mlir] [VectorOps] Merge OUTER/AXPY vector.contract lowering into single case

[aartbik]

Summary

We temporarily had separate OUTER lowering (for matmat flavors) and
AXPY lowering (for matvec flavors). With the new generalized
"vector.outerproduct" semantics, these cases can be merged into
a single lowering method. This refactoring will simplify future
decisions on cost models and lowering heuristics.

[nicolasvasilache]

Much nicer now, thanks for cleaning this up!

mlir/include/mlir/Dialect/Vector/VectorOps.h

 /// Enum to control the lowering of `vector.contract` operations.
 enum class VectorContractLowering {
   /// Progressively lower to finer grained `vector.contract` and dot-products.
   Dot = 0,
   /// Lower to `vector.matrix_multiply`, maps 1-1 to LLVM matrix intrinsics.
   Matmul = 1,
   /// Lower to `vector.outerproduct`.
   OuterProduct = 2,
-  /// Lower to series of AXPY chained through FMA.
-  AXPY = 3,
 };
 /// Enum to control the lowering of `vector.transpose` operations.
 enum class VectorTransposeLowering {
   // Lower transpose into element-wise extract and inserts.
   EltWise = 0,
   /// Lower 2-D transpose to `vector.flat_transpose`, maps 1-1 to LLVM matrix
   /// intrinsics.
   Flat = 1,

mlir/include/mlir/Dialect/Vector/VectorTransforms.h

public:
   void rewrite(vector::ContractionOp op,
                PatternRewriter &rewriter) const override;
 
 private:
   /// Options to control the vector patterns.
   vector::VectorTransformsOptions vectorTransformsOptions;
 };
 
-/// Progressive lowering of a `vector.contract %a, %b, %c` with
-/// matvec semantics to series of AXPY operations that are chained
-/// through FMA operations.
-///
-/// This only kicks in when VectorTransformsOptions is set to AXPY.
-//
-// TODO: this is very similar, but not quite the same as the outerproduct
-// lowering above; merge the two?
-class ContractionOpToAXPYLowering
-    : public OpRewritePattern<vector::ContractionOp> {
-public:
-  using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;
-  ContractionOpToAXPYLowering(
-      vector::VectorTransformsOptions vectorTransformsOptions,
-      MLIRContext *context)
-      : OpRewritePattern<vector::ContractionOp>(context),
-        vectorTransformsOptions(vectorTransformsOptions) {}
-
-  LogicalResult match(vector::ContractionOp op) const override;
-  void rewrite(vector::ContractionOp op,
-               PatternRewriter &rewriter) const override;
-
-private:
-  /// Options to control the vector patterns.
-  vector::VectorTransformsOptions vectorTransformsOptions;
-};
-
 /// Progressive lowering of ContractionOp.
 ///
 /// 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
 ///   ..

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 ///    %bRowK = vector.extract %b[K]
 ///    %atRowK = vector.extract %at[K]
 ///    %cK = vector.outerproduct %atRowK, %bRowK, %cK-1
 /// ```
 ///
 /// This only kicks in when VectorTransformsOptions is set to OuterProduct but
 /// otherwise supports any layout permutation of the matrix-multiply.
 LogicalResult
-ContractionOpToOuterProductOpLowering ::match(vector::ContractionOp op) const {
+ContractionOpToOuterProductOpLowering::match(vector::ContractionOp op) const {
   // TODO: implement masks
   if (llvm::size(op.masks()) != 0)
     return failure();
 
   if (vectorTransformsOptions.vectorContractLowering !=
       vector::VectorContractLowering::OuterProduct)
     return failure();
 
-  // Transpose arguments to make them ready for lowering to OuterProduct. The
-  // constraint to match is that we must load full rows at a time with
-  // vector::ExtractOp.
+  // Determine if the parallel/reduction structure matches something
+  // that can be expressed a reduction_size unrolled sequence.
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
   AffineExpr m, n, k;
   bindDims(op.getContext(), m, n, k);
   auto iteratorTypes = op.iterator_types().getValue();
-  if (!isParallelIterator(iteratorTypes[0]) ||
-      !isParallelIterator(iteratorTypes[1]) ||
-      !isReductionIterator(iteratorTypes[2]))
-    return failure();
   SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
+  if (isParallelIterator(iteratorTypes[0]) &&
+      isParallelIterator(iteratorTypes[1]) &&
+      isReductionIterator(iteratorTypes[2])) {
+    //
+    // Two outer parallel, one inner reduction (matmat flavor).
     // When lowering to outerproduct we can support all permutations.
+    //
     if (maps != infer({{m, k}, {k, n}, {m, n}}) &&
         maps != infer({{m, k}, {n, k}, {m, n}}) &&
         maps != infer({{k, m}, {k, n}, {m, n}}) &&
         maps != infer({{k, m}, {n, k}, {m, n}}) &&
         maps != infer({{m, k}, {k, n}, {n, m}}) &&
         maps != infer({{m, k}, {n, k}, {n, m}}) &&
         maps != infer({{k, m}, {k, n}, {n, m}}) &&
         maps != infer({{k, m}, {n, k}, {n, m}}))
       return failure();
     return success();
+  } else if (isParallelIterator(iteratorTypes[0]) &&
+             isReductionIterator(iteratorTypes[1])) {
+    //
+    // One outer parallel, one inner reduction (matvec flavor)
+    // See if a series of AXPY operations chained through FMA operations
+    // could replace the default DOT implementation.
+    //
+    if (maps != infer({{m, n}, {n}, {m}}) && // mat-vec
+        maps != infer({{n, m}, {n}, {m}}) && // mat-trans-vec
+        maps != infer({{n}, {m, n}, {m}}) && // vec-mat
+        maps != infer({{n}, {n, m}, {m}}))   // vec-mat-trans
+      return failure();
+    return success();
+  }
+  return failure();
 }
 
 void ContractionOpToOuterProductOpLowering::rewrite(
     vector::ContractionOp op, PatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   int64_t reductionSize = 0;
   VectorType lhsType = op.getLhsType();
   Value lhs = op.lhs(), rhs = op.rhs(), res = op.acc();
 
-  // Transpose arguments to make them ready for lowering to OuterProduct. The
-  // constraint to match is that we must load full rows at a time with
-  // vector::ExtractOp.
+  // Set up the parallel/reduction structure in right form.
   using MapList = ArrayRef<ArrayRef<AffineExpr>>;
   auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
   AffineExpr m, n, k;
   bindDims(rewriter.getContext(), m, n, k);
   SmallVector<int64_t, 2> perm{1, 0};
+  auto iteratorTypes = op.iterator_types().getValue();
   SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
-  // First batch of cases, no need to output permute.
+  if (isParallelIterator(iteratorTypes[0]) &&
+      isParallelIterator(iteratorTypes[1]) &&
+      isReductionIterator(iteratorTypes[2])) {
+    //
+    // Two outer parallel, one inner reduction (matmat flavor).
+    //
     if (maps == infer({{m, k}, {k, n}, {m, n}})) {
       // This is the classical row-major matmul. Just permute the lhs.
       reductionSize = lhsType.getDimSize(1);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{m, k}, {n, k}, {m, n}})) {
       // TODO: may be better to fail and use some vector<k> -> scalar reduction.
       reductionSize = lhsType.getDimSize(1);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
     } else if (maps == infer({{k, m}, {k, n}, {m, n}})) {
       // No need to permute anything.
       reductionSize = lhsType.getDimSize(0);
     } else if (maps == infer({{k, m}, {n, k}, {m, n}})) {
       // Just permute the rhs.
       reductionSize = lhsType.getDimSize(0);
       rhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
-  }
-  // Second batch of cases, reshuffle to avoid output permute.
-  else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
+    } else if (maps == infer({{m, k}, {k, n}, {n, m}})) {
       // This is the classical row-major matmul. Just permute the lhs.
       reductionSize = lhsType.getDimSize(1);
       Value tmp = rhs;
       rhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
       lhs = tmp;
     } else if (maps == infer({{m, k}, {n, k}, {n, m}})) {
       // TODO: may be better to fail and use some vector<k> -> scalar reduction.
       reductionSize = lhsType.getDimSize(1);
       Value tmp = rhs;
       rhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
       lhs = rewriter.create<vector::TransposeOp>(loc, tmp, perm);
     } else if (maps == infer({{k, m}, {k, n}, {n, m}})) {
       // No need to permute anything, but still swap lhs and rhs.
       reductionSize = lhsType.getDimSize(0);
       std::swap(lhs, rhs);
     } else if (maps == infer({{k, m}, {n, k}, {n, m}})) {
       // Just permute the rhs.
       reductionSize = lhsType.getDimSize(0);
       Value tmp = lhs;
       lhs = rewriter.create<vector::TransposeOp>(loc, rhs, perm);
       rhs = tmp;
     }
-  assert(reductionSize > 0);
-
-  // ExtractOp does not allow dynamic indexing, we must unroll explicitly.
-  for (unsigned k = 0; k < reductionSize; ++k) {
-    Value a = rewriter.create<vector::ExtractOp>(op.getLoc(), lhs, k);
-    Value b = rewriter.create<vector::ExtractOp>(op.getLoc(), rhs, k);
-    res = rewriter.create<vector::OuterProductOp>(op.getLoc(), a, b, res);
-  }
-  rewriter.replaceOp(op, res);
-}
-
-/// Progressive lowering of a `vector.contract %a, %b, %c` with
-/// matvec semantics to series of AXPY operations that are chained
-/// through FMA operations.
-///
-/// This only kicks in when VectorTransformsOptions is set to AXPY.
+  } else {
     //
-// TODO: this is very similar, but not quite the same as the outerproduct
-// lowering above; merge the two?
-LogicalResult
-ContractionOpToAXPYLowering::match(vector::ContractionOp op) const {
-  // TODO: implement masks
-  if (llvm::size(op.masks()) != 0)
-    return failure();
-
-  if (vectorTransformsOptions.vectorContractLowering !=
-      vector::VectorContractLowering::AXPY)
-    return failure();
-
-  auto iteratorTypes = op.iterator_types().getValue();
-  if (!isParallelIterator(iteratorTypes[0]) ||
-      !isReductionIterator(iteratorTypes[1]))
-    return failure();
-
-  // See if a series of AXPY operations chained through FMA operations
-  // could replace the default DOT implementation.
-  using MapList = ArrayRef<ArrayRef<AffineExpr>>;
-  auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
-  AffineExpr m, n;
-  bindDims(op.getContext(), m, n);
-  SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
-  if (maps != infer({{m, n}, {n}, {m}}) && // mat-vec
-      maps != infer({{n, m}, {n}, {m}}) && // mat-trans-vec
-      maps != infer({{n}, {m, n}, {m}}) && // vec-mat
-      maps != infer({{n}, {n, m}, {m}}))   // vec-mat-trans
-    return failure();
-  return success();
-}
-
-void ContractionOpToAXPYLowering::rewrite(vector::ContractionOp op,
-                                          PatternRewriter &rewriter) const {
-  Location loc = op.getLoc();
-  VectorType lhsType = op.getLhsType();
-  Value lhs = op.lhs(), rhs = op.rhs(), res = op.acc();
-
-  using MapList = ArrayRef<ArrayRef<AffineExpr>>;
-  auto infer = [](MapList m) { return AffineMap::inferFromExprList(m); };
-  AffineExpr m, n;
-  bindDims(op.getContext(), m, n);
-  SmallVector<int64_t, 2> perm{1, 0};
+    // One outer parallel, one inner reduction (matvec flavor)
     //
-  SmallVector<AffineMap, 4> maps = op.getIndexingMaps();
-  int64_t reductionSize = 0;
+    assert(isParallelIterator(iteratorTypes[0]) &&
+           isReductionIterator(iteratorTypes[1]));
     if (maps == infer({{m, n}, {n}, {m}})) {
       // Case mat-vec: transpose.
       reductionSize = lhsType.getDimSize(1);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{n, m}, {n}, {m}})) {
       // Case mat-trans-vec: ready to go.
       reductionSize = lhsType.getDimSize(0);
     } else if (maps == infer({{n}, {m, n}, {m}})) {
       // Case vec-mat: swap and transpose.
       reductionSize = lhsType.getDimSize(0);
       std::swap(lhs, rhs);
       lhs = rewriter.create<vector::TransposeOp>(loc, lhs, perm);
     } else if (maps == infer({{n}, {n, m}, {m}})) {
       // Case vec-mat-trans: swap and ready to go.
       reductionSize = lhsType.getDimSize(0);
       std::swap(lhs, rhs);
     }
+  }
   assert(reductionSize > 0);
 
-  // A direct series of AXPY operations, chained through FMA.
-  Type resType = op.getResultType();
+  // Unroll outer-products along reduction.
   for (int64_t k = 0; k < reductionSize; ++k) {
-    Value a = rewriter.create<vector::ExtractOp>(loc, lhs, k);
-    Value s = rewriter.create<vector::ExtractOp>(loc, rhs, k);
-    Value b = rewriter.create<vector::BroadcastOp>(loc, resType, s);
-    res = rewriter.create<vector::FMAOp>(loc, a, b, res);
+    Value a = rewriter.create<vector::ExtractOp>(op.getLoc(), lhs, k);
+    Value b = rewriter.create<vector::ExtractOp>(op.getLoc(), rhs, k);
+    res = rewriter.create<vector::OuterProductOp>(op.getLoc(), a, b, res);
   }
   rewriter.replaceOp(op, res);
 }
 
 /// Progressive lowering of ContractionOp.
 /// One:
 ///   %x = vector.contract with at least one free/batch dimension
 /// is replaced by:
ContractionOpLowering::matchAndRewrite(vector::ContractionOp op,
   // TODO: implement benefits, cost models.
   MLIRContext *ctx = op.getContext();
   ContractionOpToMatmulOpLowering pat1(vectorTransformsOptions, ctx);
   if (succeeded(pat1.match(op)))
     return failure();
   ContractionOpToOuterProductOpLowering pat2(vectorTransformsOptions, ctx);
   if (succeeded(pat2.match(op)))
     return failure();
-  ContractionOpToAXPYLowering pat3(vectorTransformsOptions, ctx);
-  if (succeeded(pat3.match(op)))
-    return failure();
 
   // Find first batch dimension in LHS/RHS, and lower when found.
   std::vector<std::pair<int64_t, int64_t>> batchDimMap = op.getBatchDimMap();
   if (!batchDimMap.empty()) {
     int64_t lhsIndex = batchDimMap[0].first;
     int64_t rhsIndex = batchDimMap[0].second;
     rewriter.replaceOp(op, lowerParallel(op, lhsIndex, rhsIndex, rewriter));
     return success();
patterns.insert<BroadcastOpLowering,
                   ConstantMaskOpLowering,
                   OuterProductOpLowering,
                   ShapeCastOp2DDownCastRewritePattern,
                   ShapeCastOp2DUpCastRewritePattern,
                   ShapeCastOpRewritePattern>(context);
   patterns.insert<TransposeOpLowering,
                   ContractionOpLowering,
                   ContractionOpToMatmulOpLowering,
-                  ContractionOpToOuterProductOpLowering,
-                  ContractionOpToAXPYLowering>(parameters, context);
+                  ContractionOpToOuterProductOpLowering>(parameters, context);
   // clang-format on
 }

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

-// RUN: mlir-opt %s -test-vector-contraction-conversion=vector-axpy=1 | FileCheck %s
+// RUN: mlir-opt %s -test-vector-contraction-conversion=vector-outerproduct=1 | FileCheck %s
 
 #matvec_accesses = [
   affine_map<(i, j) -> (i, j)>,
   affine_map<(i, j) -> (j)>,
   affine_map<(i, j) -> (i)>
 ]
 #matvec_trait = {
   indexing_maps = #matvec_accesses,
#vecmattrans_trait = {
   indexing_maps = #vecmattrans_accesses,
   iterator_types = ["parallel", "reduction"]
 }
 
 // CHECK-LABEL: func @matvec2x2
 // CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
 // CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
 // CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
-// CHECK: %[[C0:.*]] = constant dense<0.000000e+00> : vector<2x2xf32>
 // CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
 // CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
 // CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
-// CHECK: %[[T3:.*]] = vector.extract %[[T0]][0, 0] : vector<2x2xf32>
-// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[C0]] [0, 0] : f32 into vector<2x2xf32>
-// CHECK: %[[T5:.*]] = vector.extract %[[T0]][1, 0] : vector<2x2xf32>
-// CHECK: %[[T6:.*]] = vector.insert %[[T5]], %[[T4]] [0, 1] : f32 into vector<2x2xf32>
-// CHECK: %[[T7:.*]] = vector.extract %[[T0]][0, 1] : vector<2x2xf32>
-// CHECK: %[[T8:.*]] = vector.insert %[[T7]], %[[T6]] [1, 0] : f32 into vector<2x2xf32>
-// CHECK: %[[T9:.*]] = vector.extract %[[T0]][1, 1] : vector<2x2xf32>
-// CHECK: %[[T10:.*]] = vector.insert %[[T9]], %[[T8]] [1, 1] : f32 into vector<2x2xf32>
-// CHECK: %[[T11:.*]] = vector.extract %[[T10]][0] : vector<2x2xf32>
-// CHECK: %[[T12:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
-// CHECK: %[[T13:.*]] = splat %[[T12]] : vector<2xf32>
-// CHECK: %[[T14:.*]] = vector.fma %[[T11]], %[[T13]], %[[T2]] : vector<2xf32>
-// CHECK: %[[T15:.*]] = vector.extract %[[T10]][1] : vector<2x2xf32>
-// CHECK: %[[T16:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
-// CHECK: %[[T17:.*]] = splat %[[T16]] : vector<2xf32>
-// CHECK: %[[T18:.*]] = vector.fma %[[T15]], %[[T17]], %[[T14]] : vector<2xf32>
-// CHECK: store %[[T18]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: %[[T3:.*]] = vector.transpose %[[T0]], [1, 0] : vector<2x2xf32> to vector<2x2xf32>
+// CHECK: %[[T4:.*]] = vector.extract %[[T3]][0] : vector<2x2xf32>
+// CHECK: %[[T5:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
+// CHECK: %[[T6:.*]] = vector.outerproduct %[[T4]], %[[T5]], %[[T2]] : vector<2xf32>, f32
+// CHECK: %[[T7:.*]] = vector.extract %[[T3]][1] : vector<2x2xf32>
+// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
+// CHECK: %[[T9:.*]] = vector.outerproduct %[[T7]], %[[T8]], %[[T6]] : vector<2xf32>, f32
+// CHECK: store %[[T9]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: return
 func @matvec2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
                                                 %arg2: memref<vector<2xf32>>) {
   %A = load %arg0[] : memref<vector<2x2xf32>>
   %x = load %arg1[] : memref<vector<2xf32>>
   %b = load %arg2[] : memref<vector<2xf32>>
   %0 = vector.contract #matvec_trait %A, %x, %b : vector<2x2xf32>, vector<2xf32> into vector<2xf32>
   store %0, %arg2[] : memref<vector<2xf32>>
   return
 }
 
 // CHECK-LABEL: func @mattransvec2x2
 // CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
 // CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
 // CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
 // CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
 // CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
 // CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
 // CHECK: %[[T3:.*]] = vector.extract %[[T0]][0] : vector<2x2xf32>
 // CHECK: %[[T4:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
-// CHECK: %[[T5:.*]] = splat %[[T4]] : vector<2xf32>
-// CHECK: %[[T6:.*]] = vector.fma %[[T3]], %[[T5]], %[[T2]] : vector<2xf32>
-// CHECK: %[[T7:.*]] = vector.extract %[[T0]][1] : vector<2x2xf32>
-// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
-// CHECK: %[[T9:.*]] = splat %[[T8]] : vector<2xf32>
-// CHECK: %[[T10:.*]] = vector.fma %[[T7]], %[[T9]], %[[T6]] : vector<2xf32>
-// CHECK: store %[[T10]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: %[[T5:.*]] = vector.outerproduct %[[T3]], %[[T4]], %[[T2]] : vector<2xf32>, f32
+// CHECK: %[[T6:.*]] = vector.extract %[[T0]][1] : vector<2x2xf32>
+// CHECK: %[[T7:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
+// CHECK: %[[T8:.*]] = vector.outerproduct %[[T6]], %[[T7]], %[[T5]] : vector<2xf32>, f32
+// CHECK: store %[[T8]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: return
 func @mattransvec2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
                                                      %arg2: memref<vector<2xf32>>) {
   %A = load %arg0[] : memref<vector<2x2xf32>>
   %x = load %arg1[] : memref<vector<2xf32>>
   %b = load %arg2[] : memref<vector<2xf32>>
   %0 = vector.contract #mattransvec_trait %A, %x, %b : vector<2x2xf32>, vector<2xf32> into vector<2xf32>
   store %0, %arg2[] : memref<vector<2xf32>>
   return
 }
 
 // CHECK-LABEL: func @vecmat2x2
 // CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
 // CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
 // CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
-// CHECK: %[[C0:.*]] = constant dense<0.000000e+00> : vector<2x2xf32>
 // CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
 // CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
 // CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
-// CHECK: %[[T3:.*]] = vector.extract %[[T0]][0, 0] : vector<2x2xf32>
-// CHECK: %[[T4:.*]] = vector.insert %[[T3]], %[[C0]] [0, 0] : f32 into vector<2x2xf32>
-// CHECK: %[[T5:.*]] = vector.extract %[[T0]][1, 0] : vector<2x2xf32>
-// CHECK: %[[T6:.*]] = vector.insert %[[T5]], %[[T4]] [0, 1] : f32 into vector<2x2xf32>
-// CHECK: %[[T7:.*]] = vector.extract %[[T0]][0, 1] : vector<2x2xf32>
-// CHECK: %[[T8:.*]] = vector.insert %[[T7]], %[[T6]] [1, 0] : f32 into vector<2x2xf32>
-// CHECK: %[[T9:.*]] = vector.extract %[[T0]][1, 1] : vector<2x2xf32>
-// CHECK: %[[T10:.*]] = vector.insert %[[T9]], %[[T8]] [1, 1] : f32 into vector<2x2xf32>
-// CHECK: %[[T11:.*]] = vector.extract %[[T10]][0] : vector<2x2xf32>
-// CHECK: %[[T12:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
-// CHECK: %[[T13:.*]] = splat %[[T12]] : vector<2xf32>
-// CHECK: %[[T14:.*]] = vector.fma %[[T11]], %[[T13]], %[[T2]] : vector<2xf32>
-// CHECK: %[[T15:.*]] = vector.extract %[[T10]][1] : vector<2x2xf32>
-// CHECK: %[[T16:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
-// CHECK: %[[T17:.*]] = splat %[[T16]] : vector<2xf32>
-// CHECK: %[[T18:.*]] = vector.fma %[[T15]], %[[T17]], %[[T14]] : vector<2xf32>
-// CHECK: store %[[T18]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: %[[T3:.*]] = vector.transpose %[[T0]], [1, 0] : vector<2x2xf32> to vector<2x2xf32>
+// CHECK: %[[T4:.*]] = vector.extract %[[T3]][0] : vector<2x2xf32>
+// CHECK: %[[T5:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
+// CHECK: %[[T6:.*]] = vector.outerproduct %[[T4]], %[[T5]], %[[T2]] : vector<2xf32>, f32
+// CHECK: %[[T7:.*]] = vector.extract %[[T3]][1] : vector<2x2xf32>
+// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
+// CHECK: %[[T9:.*]] = vector.outerproduct %[[T7]], %[[T8]], %[[T6]] : vector<2xf32>, f32
+// CHECK: store %[[T9]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: return
 func @vecmat2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
                                                 %arg2: memref<vector<2xf32>>) {
   %A = load %arg0[] : memref<vector<2x2xf32>>
   %x = load %arg1[] : memref<vector<2xf32>>
   %b = load %arg2[] : memref<vector<2xf32>>
   %0 = vector.contract #vecmat_trait %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
   store %0, %arg2[] : memref<vector<2xf32>>
   return
 }
 
 // CHECK-LABEL: func @vecmattrans2x2
 // CHECK-SAME: %[[A:.*0]]: memref<vector<2x2xf32>>
 // CHECK-SAME: %[[B:.*1]]: memref<vector<2xf32>>
 // CHECK-SAME: %[[C:.*2]]: memref<vector<2xf32>>
 // CHECK: %[[T0:.*]] = load %[[A]][] : memref<vector<2x2xf32>>
 // CHECK: %[[T1:.*]] = load %[[B]][] : memref<vector<2xf32>>
 // CHECK: %[[T2:.*]] = load %[[C]][] : memref<vector<2xf32>>
 // CHECK: %[[T3:.*]] = vector.extract %[[T0]][0] : vector<2x2xf32>
 // CHECK: %[[T4:.*]] = vector.extract %[[T1]][0] : vector<2xf32>
-// CHECK: %[[T5:.*]] = splat %[[T4]] : vector<2xf32>
-// CHECK: %[[T6:.*]] = vector.fma %[[T3]], %[[T5]], %[[T2]] : vector<2xf32>
-// CHECK: %[[T7:.*]] = vector.extract %[[T0]][1] : vector<2x2xf32>
-// CHECK: %[[T8:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
-// CHECK: %[[T9:.*]] = splat %[[T8]] : vector<2xf32>
-// CHECK: %[[T10:.*]] = vector.fma %[[T7]], %[[T9]], %[[T6]] : vector<2xf32>
-// CHECK: store %[[T10]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: %[[T5:.*]] = vector.outerproduct %[[T3]], %[[T4]], %[[T2]] : vector<2xf32>, f32
+// CHECK: %[[T6:.*]] = vector.extract %[[T0]][1] : vector<2x2xf32>
+// CHECK: %[[T7:.*]] = vector.extract %[[T1]][1] : vector<2xf32>
+// CHECK: %[[T8:.*]] = vector.outerproduct %[[T6]], %[[T7]], %[[T5]] : vector<2xf32>, f32
+// CHECK: store %[[T8]], %[[C]][] : memref<vector<2xf32>>
+// CHECK: return
 func @vecmattrans2x2(%arg0: memref<vector<2x2xf32>>, %arg1: memref<vector<2xf32>>,
                                                      %arg2: memref<vector<2xf32>>) {
   %A = load %arg0[] : memref<vector<2x2xf32>>
   %x = load %arg1[] : memref<vector<2xf32>>
   %b = load %arg2[] : memref<vector<2xf32>>
   %0 = vector.contract #vecmattrans_trait %x, %A, %b : vector<2xf32>, vector<2x2xf32> into vector<2xf32>
   store %0, %arg2[] : memref<vector<2xf32>>
   return
 }

mlir/test/lib/Transforms/TestVectorTransforms.cpp

struct TestVectorContractionConversion
   Option<bool> lowerToFlatTranspose{
       *this, "vector-flat-transpose",
       llvm::cl::desc("Lower 2-D vector.transpose to vector.flat_transpose"),
       llvm::cl::init(false)};
   Option<bool> lowerToOuterProduct{
       *this, "vector-outerproduct",
       llvm::cl::desc("Lower vector.contract to vector.outerproduct"),
       llvm::cl::init(false)};
-  Option<bool> lowerToAXPY{*this, "vector-axpy",
-                           llvm::cl::desc("Lower vector.contract to AXPY"),
-                           llvm::cl::init(false)};
 
   void runOnFunction() override {
     OwningRewritePatternList patterns;
 
     // Test on one pattern in isolation.
     if (lowerToOuterProduct) {
       VectorContractLowering lowering = VectorContractLowering::OuterProduct;
       VectorTransformsOptions options{lowering};
       patterns.insert<ContractionOpToOuterProductOpLowering>(options,
                                                              &getContext());
       applyPatternsAndFoldGreedily(getFunction(), patterns);
       return;
     }
 
     // Test on all contract lowering patterns.
     VectorContractLowering contractLowering = VectorContractLowering::Dot;
     if (lowerToFlatMatrix)
       contractLowering = VectorContractLowering::Matmul;
-    else if (lowerToAXPY)
-      contractLowering = VectorContractLowering::AXPY;
     VectorTransposeLowering transposeLowering =
         VectorTransposeLowering::EltWise;
     if (lowerToFlatTranspose)
       transposeLowering = VectorTransposeLowering::Flat;
     VectorTransformsOptions options{contractLowering, transposeLowering};
     populateVectorContractLoweringPatterns(patterns, &getContext(), options);
     applyPatternsAndFoldGreedily(getFunction(), patterns);
   }