[mlir] [VectorOps] Allow AXPY to be expressed as special case of OUTERPRODUCT

[aartbik]

Summary

This specialization allows sharing more code where an AXPY follows naturally
in cases where an OUTERPRODUCT on a scalar would be generated.

[aartbik]

fixed typo in example text

[aartbik]

favor rhs scalar

[nicolasvasilache]

Thanks for adopting what we discussed previously and special casing the impl!

Overall this looks good, I was expecting to see much more red on the VectorTransforms side as well as folding the AXPY into OuterProduct option.
Is this left for a future revision?
It seems the refactoring would be more compelling if it had all the red it deserves :)

[aartbik]

In D83453#2143781, @nicolasvasilache wrote:

Is this left for a future revision?

Yes, the merging of the outerproduct/AXPY lowering is already done, but in a separate CL.

I like to keep CL's small to make the review easier, and extending the semantics of outer with axpy seemed a logically contained unit.
Refactoring to use the new operation seemed another logical unit.

[nicolasvasilache]

Refactoring to use the new operation seemed another logical unit.

SGTM, thanks!

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

let description = [{
 
     An indexing map attribute list must be specified with an entry for lhs, rhs
     and acc arguments. An indexing map attribute specifies a mapping from each
     iterator in the iterator type list, to each dimension of an N-D vector.
 
     Example:
 
     ```mlir
-    // Simple dot product (K = 0).
+    // Simple DOT product (K = 0).
     #contraction_accesses = [
      affine_map<(i) -> (i)>,
      affine_map<(i) -> (i)>,
      affine_map<(i) -> ()>
     ]
     #contraction_trait = {
       indexing_maps = #contraction_accesses,
       iterator_types = ["reduction"]
let extraClassDeclaration = [{
     }
     VectorType getDestVectorType() {
       return dest().getType().cast<VectorType>();
     }
   }];
 }
 
 def Vector_OuterProductOp :
-  Vector_Op<"outerproduct", [NoSideEffect, SameOperandsAndResultElementType]>,
-    Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, Variadic<AnyVector>:$acc)>,
+  Vector_Op<"outerproduct", [NoSideEffect,
+    PredOpTrait<"lhs operand and result have same element type",
+                TCresVTEtIsSameAsOpBase<0, 0>>,
+    PredOpTrait<"rhs operand and result have same element type",
+                TCresVTEtIsSameAsOpBase<0, 1>>]>,
+    Arguments<(ins AnyVector:$lhs, AnyType:$rhs, Variadic<AnyVector>:$acc)>,
     Results<(outs AnyVector)> {
   let summary = "vector outerproduct with optional fused add";
   let description = [{
-    Takes 2 1-D vectors and returns the 2-D vector containing the outer-product.
-
-    An optional extra 2-D vector argument may be specified in which case the
-    operation returns the sum of the outer-product and the extra vector. In this
-    multiply-accumulate scenario, the rounding mode is that obtained by
-    guaranteeing that a fused-multiply add operation is emitted. When lowered to
-    the LLVMIR dialect, this form emits `llvm.intr.fma`, which is guaranteed to
-    lower to actual `fma` instructions on x86.
+    Takes 2 1-D vectors and returns the 2-D vector containing the outer-product,
+    as illustrated below:
+    ```
+     outer |   [c, d]
+     ------+------------
+       [a, | [ [a*c, a*d],
+        b] |   [b*c, b*d] ]
+    ```
+    This operation also accepts a 1-D vector lhs and a scalar rhs. In this
+    case a simple AXPY operation is performed, which returns a 1-D vector.
+    ```
+        [a, b] * c = [a*c, b*c]
+    ```
+
+    An optional extra vector argument with the same shape as the output
+    vector may be specified in which case the operation returns the sum of
+    the outer-product and the extra vector. In this multiply-accumulate
+    scenario for floating-point arguments, the rounding mode is enforced
+    by guaranteeing that a fused-multiply add operation is emitted. When
+    lowered to the LLVMIR dialect, this form emits `llvm.intr.fma`, which
+    is guaranteed to lower to actual `fma` instructions on x86.
 
     Example:
 
     ```
     %2 = vector.outerproduct %0, %1: vector<4xf32>, vector<8xf32>
     return %2: vector<4x8xf32>
 
     %3 = vector.outerproduct %0, %1, %2:
       vector<4xf32>, vector<8xf32>, vector<4x8xf32>
     return %3: vector<4x8xf32>
+
+    %6 = vector.outerproduct %4, %5: vector<10xf32>, f32
+    return %6: vector<10xf32>
+
     ```
   }];
   let builders = [
     // Build an op without mask, use the type of `acc` as the return type.
     OpBuilder<
     "OpBuilder &builder, OperationState &result, Value lhs, Value rhs, "
     "Value acc">];
   let extraClassDeclaration = [{
     VectorType getOperandVectorTypeLHS() {
       return lhs().getType().cast<VectorType>();
     }
-    VectorType getOperandVectorTypeRHS() {
-      return rhs().getType().cast<VectorType>();
+    Type getOperandTypeRHS() {
+      return rhs().getType();
     }
     VectorType getOperandVectorTypeACC() {
-      return (llvm::size(acc()) == 0) ? VectorType() :
-        (*acc().begin()).getType().cast<VectorType>();
+      return (llvm::size(acc()) == 0)
+        ? VectorType()
+        : (*acc().begin()).getType().cast<VectorType>();
     }
     VectorType getVectorType() {
       return getResult().getType().cast<VectorType>();
     }
   }];
 }
 
 // TODO: Add transformation which decomposes ReshapeOp into an optimized

mlir/integration_test/Dialect/Vector/CPU/test-outerproduct-f32.mlir

 // RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 !vector_type_A = type vector<8xf32>
 !vector_type_B = type vector<8xf32>
 !vector_type_C = type vector<8x8xf32>
 
 !vector_type_X = type vector<2xf32>
 !vector_type_Y = type vector<3xf32>
 !vector_type_Z = type vector<2x3xf32>
 
+!vector_type_R = type vector<7xf32>
+
 func @vector_outerproduct_splat_8x8(%fa: f32, %fb: f32, %fc: f32) -> !vector_type_C {
   %a = splat %fa: !vector_type_A
   %b = splat %fb: !vector_type_B
   %c = splat %fc: !vector_type_C
   %d = vector.outerproduct %a, %b, %c : !vector_type_A, !vector_type_B
   return %d: !vector_type_C
 }
 
 func @vector_outerproduct_vec_2x3(%x : !vector_type_X,
                                   %y : !vector_type_Y) -> !vector_type_Z {
   %o = vector.outerproduct %x, %y : !vector_type_X, !vector_type_Y
   return %o: !vector_type_Z
 }
 
 func @vector_outerproduct_vec_2x3_acc(%x : !vector_type_X,
                                       %y : !vector_type_Y,
                                       %z : !vector_type_Z) -> !vector_type_Z {
   %o = vector.outerproduct %x, %y, %z : !vector_type_X, !vector_type_Y
   return %o: !vector_type_Z
 }
 
 func @entry() {
+  %f0 = constant 0.0: f32
   %f1 = constant 1.0: f32
   %f2 = constant 2.0: f32
   %f3 = constant 3.0: f32
   %f4 = constant 4.0: f32
   %f5 = constant 5.0: f32
   %f10 = constant 10.0: f32
 
   // Simple case, splat scalars into vectors, then take outer product.
func @entry() {
   %q = call @vector_outerproduct_vec_2x3_acc(%x, %y, %p)
       : (!vector_type_X, !vector_type_Y, !vector_type_Z) -> (!vector_type_Z)
   vector.print %q : !vector_type_Z
   //
   // outer product 2x3:
   //
   // CHECK: ( ( 6, 8, 10 ), ( 12, 16, 20 ) )
 
+  %3 = vector.broadcast %f0 : f32 to !vector_type_R
+  %4 = vector.insert %f1,  %3[1] : f32 into !vector_type_R
+  %5 = vector.insert %f2,  %4[2] : f32 into !vector_type_R
+  %6 = vector.insert %f3,  %5[3] : f32 into !vector_type_R
+  %7 = vector.insert %f4,  %6[4] : f32 into !vector_type_R
+  %8 = vector.insert %f5,  %7[5] : f32 into !vector_type_R
+  %9 = vector.insert %f10, %8[6] : f32 into !vector_type_R
+
+  %o = vector.broadcast %f1 : f32 to !vector_type_R
+
+  %axpy1 = vector.outerproduct %9, %f2     : !vector_type_R, f32
+  %axpy2 = vector.outerproduct %9, %f2, %o : !vector_type_R, f32
+
+  vector.print %axpy1 : !vector_type_R
+  vector.print %axpy2 : !vector_type_R
+  //
+  // axpy operations:
+  //
+  // CHECK: ( 0, 2, 4, 6, 8, 10, 20 )
+  // CHECK: ( 1, 3, 5, 7, 9, 11, 21 )
+
   return
 }

mlir/integration_test/Dialect/Vector/CPU/test-outerproduct-i64.mlir

 // RUN: mlir-opt %s -convert-scf-to-std -convert-vector-to-llvm -convert-std-to-llvm | \
 // RUN: mlir-cpu-runner -e entry -entry-point-result=void  \
 // RUN:   -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
 // RUN: FileCheck %s
 
 !vector_type_A = type vector<8xi64>
 !vector_type_B = type vector<8xi64>
 !vector_type_C = type vector<8x8xi64>
 
 !vector_type_X = type vector<2xi64>
 !vector_type_Y = type vector<3xi64>
 !vector_type_Z = type vector<2x3xi64>
 
+!vector_type_R = type vector<7xi64>
+
 func @vector_outerproduct_splat_8x8(%ia: i64, %ib: i64, %ic: i64) -> !vector_type_C {
   %a = splat %ia: !vector_type_A
   %b = splat %ib: !vector_type_B
   %c = splat %ic: !vector_type_C
   %d = vector.outerproduct %a, %b, %c : !vector_type_A, !vector_type_B
   return %d: !vector_type_C
 }
 
 func @vector_outerproduct_vec_2x3(%x : !vector_type_X,
                                   %y : !vector_type_Y) -> !vector_type_Z {
   %o = vector.outerproduct %x, %y : !vector_type_X, !vector_type_Y
   return %o: !vector_type_Z
 }
 
 func @vector_outerproduct_vec_2x3_acc(%x : !vector_type_X,
                                       %y : !vector_type_Y,
                                       %z : !vector_type_Z) -> !vector_type_Z {
   %o = vector.outerproduct %x, %y, %z : !vector_type_X, !vector_type_Y
   return %o: !vector_type_Z
 }
 
 func @entry() {
+  %i0 = constant 0: i64
   %i1 = constant 1: i64
   %i2 = constant 2: i64
   %i3 = constant 3: i64
   %i4 = constant 4: i64
   %i5 = constant 5: i64
   %i10 = constant 10: i64
 
   // Simple case, splat scalars into vectors, then take outer product.
func @entry() {
   %q = call @vector_outerproduct_vec_2x3_acc(%x, %y, %p)
       : (!vector_type_X, !vector_type_Y, !vector_type_Z) -> (!vector_type_Z)
   vector.print %q : !vector_type_Z
   //
   // outer product 2x3:
   //
   // CHECK: ( ( 6, 8, 10 ), ( 12, 16, 20 ) )
 
+  %3 = vector.broadcast %i0 : i64 to !vector_type_R
+  %4 = vector.insert %i1,  %3[1] : i64 into !vector_type_R
+  %5 = vector.insert %i2,  %4[2] : i64 into !vector_type_R
+  %6 = vector.insert %i3,  %5[3] : i64 into !vector_type_R
+  %7 = vector.insert %i4,  %6[4] : i64 into !vector_type_R
+  %8 = vector.insert %i5,  %7[5] : i64 into !vector_type_R
+  %9 = vector.insert %i10, %8[6] : i64 into !vector_type_R
+
+  %o = vector.broadcast %i1 : i64 to !vector_type_R
+
+  %axpy1 = vector.outerproduct %9, %i2     : !vector_type_R, i64
+  %axpy2 = vector.outerproduct %9, %i2, %o : !vector_type_R, i64
+
+  vector.print %axpy1 : !vector_type_R
+  vector.print %axpy2 : !vector_type_R
+  //
+  // axpy operations:
+  //
+  // CHECK: ( 0, 2, 4, 6, 8, 10, 20 )
+  // CHECK: ( 1, 3, 5, 7, 9, 11, 21 )
+
   return
 }

mlir/lib/Dialect/Vector/VectorOps.cpp

static ParseResult parseOuterProductOp(OpAsmParser &parser,
   if (parser.parseOperandList(operandsInfo) || parser.parseColonType(tLHS) ||
       parser.parseComma() || parser.parseType(tRHS))
     return failure();
   if (operandsInfo.size() < 2)
     return parser.emitError(parser.getNameLoc(),
                             "expected at least 2 operands");
   VectorType vLHS = tLHS.dyn_cast<VectorType>();
   VectorType vRHS = tRHS.dyn_cast<VectorType>();
-  if (!vLHS || !vRHS)
-    return parser.emitError(parser.getNameLoc(), "expected 2 vector types");
-  VectorType resType = VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
-                                       vLHS.getElementType());
+  if (!vLHS)
+    return parser.emitError(parser.getNameLoc(),
+                            "expected vector type for operand #1");
+  VectorType resType =
+      vRHS ? VectorType::get({vLHS.getDimSize(0), vRHS.getDimSize(0)},
+                             vLHS.getElementType())
+           : VectorType::get({vLHS.getDimSize(0)}, vLHS.getElementType());
   return failure(
       parser.resolveOperand(operandsInfo[0], tLHS, result.operands) ||
       parser.resolveOperand(operandsInfo[1], tRHS, result.operands) ||
       (operandsInfo.size() > 2 &&
        parser.resolveOperand(operandsInfo[2], resType, result.operands)) ||
       parser.addTypeToList(resType, result.types));
 }
 
 static LogicalResult verify(OuterProductOp op) {
+  Type tRHS = op.getOperandTypeRHS();
   VectorType vLHS = op.getOperandVectorTypeLHS(),
-             vRHS = op.getOperandVectorTypeRHS(),
+             vRHS = tRHS.dyn_cast<VectorType>(),
              vACC = op.getOperandVectorTypeACC(), vRES = op.getVectorType();
+
   if (vLHS.getRank() != 1)
     return op.emitOpError("expected 1-d vector for operand #1");
+
+  if (vRHS) {
+    // Proper OUTER operation.
     if (vRHS.getRank() != 1)
       return op.emitOpError("expected 1-d vector for operand #2");
     if (vRES.getRank() != 2)
       return op.emitOpError("expected 2-d vector result");
     if (vLHS.getDimSize(0) != vRES.getDimSize(0))
       return op.emitOpError("expected #1 operand dim to match result dim #1");
     if (vRHS.getDimSize(0) != vRES.getDimSize(1))
       return op.emitOpError("expected #2 operand dim to match result dim #2");
+  } else {
+    // An AXPY operation.
+    if (vRES.getRank() != 1)
+      return op.emitOpError("expected 1-d vector result");
+    if (vLHS.getDimSize(0) != vRES.getDimSize(0))
+      return op.emitOpError("expected #1 operand dim to match result dim #1");
+  }
+
   if (vACC && vACC != vRES)
     return op.emitOpError("expected operand #3 of same type as result type");
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ReshapeOp
 //===----------------------------------------------------------------------===//

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 /// Progressive lowering of OuterProductOp.
 /// One:
 ///   %x = vector.outerproduct %lhs, %rhs, %acc
 /// is replaced by:
 ///   %z = zero-result
 ///   %0 = vector.extract %lhs[0]
 ///   %1 = vector.broadcast %0
 ///   %2 = vector.extract %acc[0]
-///   %3 = vector.fma %1, %arg1, %2
+///   %3 = vector.fma %1, %rhs, %2
 ///   %4 = vector.insert %3, %z[0]
 ///   ..
 ///   %x = vector.insert %.., %..[N-1]
 ///
 class OuterProductOpLowering : public OpRewritePattern<vector::OuterProductOp> {
 public:
   using OpRewritePattern<vector::OuterProductOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::OuterProductOp op,
                                 PatternRewriter &rewriter) const override {
     auto loc = op.getLoc();
 
-    VectorType rhsType = op.getOperandVectorTypeRHS();
+    VectorType lhsType = op.getOperandVectorTypeLHS();
+    VectorType rhsType = op.getOperandTypeRHS().dyn_cast<VectorType>();
     VectorType resType = op.getVectorType();
     Type eltType = resType.getElementType();
+    bool isInt = eltType.isa<IntegerType>();
     Value acc = (op.acc().empty()) ? nullptr : op.acc()[0];
 
+    if (!rhsType) {
+      // Special case: AXPY operation.
+      Value b = rewriter.create<vector::BroadcastOp>(loc, lhsType, op.rhs());
+      rewriter.replaceOp(op, genMult(loc, op.lhs(), b, acc, isInt, rewriter));
+      return success();
+    }
+
     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 e = rewriter.create<vector::ExtractOp>(loc, rhsType, acc, pos);
-        if (eltType.isa<IntegerType>())
-          m = rewriter.create<AddIOp>(
-              loc, rewriter.create<MulIOp>(loc, b, op.rhs()), e);
-        else
-          m = rewriter.create<vector::FMAOp>(loc, b, op.rhs(), e);
-      } else {
-        if (eltType.isa<IntegerType>())
-          m = rewriter.create<MulIOp>(loc, b, op.rhs());
-        else
-          m = rewriter.create<MulFOp>(loc, b, op.rhs());
-      }
+      Value a = rewriter.create<vector::BroadcastOp>(loc, rhsType, x);
+      Value r = nullptr;
+      if (acc)
+        r = rewriter.create<vector::ExtractOp>(loc, rhsType, acc, pos);
+      Value m = genMult(loc, a, op.rhs(), r, isInt, rewriter);
       result = rewriter.create<vector::InsertOp>(loc, resType, m, result, pos);
     }
     rewriter.replaceOp(op, result);
     return success();
   }
+
+private:
+  static Value genMult(Location loc, Value x, Value y, Value acc, bool isInt,
+                       PatternRewriter &rewriter) {
+    if (acc) {
+      if (isInt)
+        return rewriter.create<AddIOp>(loc, rewriter.create<MulIOp>(loc, x, y),
+                                       acc);
+      return rewriter.create<vector::FMAOp>(loc, x, y, acc);
+    }
+    if (isInt)
+      return rewriter.create<MulIOp>(loc, x, y);
+    return rewriter.create<MulFOp>(loc, x, y);
+  }
 };
 
 /// Progressive lowering of ConstantMaskOp.
 /// One:
 ///   %x = vector.constant_mask [a,b]
 /// is replaced by:
 ///   %z = zero-result
 ///   %l = vector.constant_mask [b]

mlir/test/Dialect/Vector/invalid.mlir

 
 func @outerproduct_num_operands(%arg0: f32) {
   // expected-error@+1 {{expected at least 2 operands}}
   %1 = vector.outerproduct %arg0 : f32, f32
 }
 // -----
 
 func @outerproduct_non_vector_operand(%arg0: f32) {
-  // expected-error@+1 {{expected 2 vector types}}
+  // expected-error@+1 {{expected vector type for operand #1}}
   %1 = vector.outerproduct %arg0, %arg0 : f32, f32
 }
 
 // -----
 
 func @outerproduct_operand_1(%arg0: vector<4xf32>, %arg1: vector<4x8xf32>) {
   // expected-error@+1 {{expected 1-d vector for operand #1}}
   %1 = vector.outerproduct %arg1, %arg1 : vector<4x8xf32>, vector<4x8xf32>

 
 func @outerproduct_operand_2_dim_generic(%arg0: vector<4xf32>, %arg1: vector<8xf32>) {
   // expected-error@+1 {{expected #2 operand dim to match result dim #2}}
   %1 = "vector.outerproduct" (%arg0, %arg1) : (vector<4xf32>, vector<8xf32>) -> (vector<4x16xf32>)
 }
 
 // -----
 
+func @outerproduct_axpy_operand(%arg0: vector<4x8xf32>, %arg1: f32) {
+  // expected-error@+1 {{expected 1-d vector for operand #1}}
+  %1 = vector.outerproduct %arg0, %arg1 : vector<4x8xf32>, f32
+}
+
+// -----
+
+func @outerproduct_axpy_result_generic(%arg0: vector<4xf32>, %arg1: f32) {
+  // expected-error@+1 {{expected 1-d vector result}}
+  %1 = "vector.outerproduct" (%arg0, %arg1) : (vector<4xf32>, f32) -> (vector<4x8xf32>)
+}
+
+// -----
+
+func @outerproduct_axpy_operand_dim_generic(%arg0: vector<8xf32>, %arg1: f32) {
+  // expected-error@+1 {{expected #1 operand dim to match result dim #1}}
+  %1 = "vector.outerproduct" (%arg0, %arg1) : (vector<8xf32>, f32) -> (vector<16xf32>)
+}
+
+// -----
+
 func @outerproduct_operand_3_result_type_generic(%arg0: vector<4xf32>, %arg1: vector<8xf32>, %arg2: vector<4x16xf32>) {
   // expected-error@+1 {{expected operand #3 of same type as result type}}
   %1 = "vector.outerproduct" (%arg0, %arg1, %arg2) : (vector<4xf32>, vector<8xf32>, vector<4x16xf32>) -> (vector<4x8xf32>)
 }
 
 // -----
 
 func @test_vector.transfer_read(%arg0: memref<?x?xf32>) {

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

 // CHECK:      return %[[T11]] : vector<2x3xi32>
 func @outerproduct_acc_int(%arg0: vector<2xi32>,
                            %arg1: vector<3xi32>,
                            %arg2: vector<2x3xi32>) -> vector<2x3xi32> {
   %0 = vector.outerproduct %arg0, %arg1, %arg2 : vector<2xi32>, vector<3xi32>
   return %0: vector<2x3xi32>
 }
 
+// CHECK-LABEL: func @axpy_fp(
+// CHECK-SAME: %[[A:.*0]]: vector<16xf32>,
+// CHECK-SAME: %[[B:.*1]]: f32)
+// CHECK: %[[T0:.*]] = splat %[[B]] : vector<16xf32>
+// CHECK: %[[T1:.*]] = mulf %[[A]], %[[T0]] : vector<16xf32>
+// CHECK: return %[[T1]] : vector<16xf32>
+func @axpy_fp(%arg0: vector<16xf32>, %arg1: f32) -> vector<16xf32> {
+   %0 = vector.outerproduct %arg0, %arg1: vector<16xf32>, f32
+   return %0: vector<16xf32>
+}
+
+// CHECK-LABEL: func @axpy_fp_add(
+// CHECK-SAME: %[[A:.*0]]: vector<16xf32>,
+// CHECK-SAME: %[[B:.*1]]: f32,
+// CHECK-SAME: %[[C:.*2]]: vector<16xf32>)
+// CHECK: %[[T0:.*]] = splat %[[B]] : vector<16xf32>
+// CHECK: %[[T1:.*]] = vector.fma %[[A]], %[[T0]], %[[C]] : vector<16xf32>
+// CHECK: return %[[T1]] : vector<16xf32>
+func @axpy_fp_add(%arg0: vector<16xf32>, %arg1: f32, %arg2 : vector<16xf32>) -> vector<16xf32> {
+   %0 = vector.outerproduct %arg0, %arg1, %arg2: vector<16xf32>, f32
+   return %0: vector<16xf32>
+}
+
+// CHECK-LABEL: func @axpy_int(
+// CHECK-SAME: %[[A:.*0]]: vector<16xi32>,
+// CHECK-SAME: %[[B:.*1]]: i32)
+// CHECK: %[[T0:.*]] = splat %[[B]] : vector<16xi32>
+// CHECK: %[[T1:.*]] = muli %[[A]], %[[T0]] : vector<16xi32>
+// CHECK: return %[[T1]] : vector<16xi32>
+func @axpy_int(%arg0: vector<16xi32>, %arg1: i32) -> vector<16xi32> {
+   %0 = vector.outerproduct %arg0, %arg1: vector<16xi32>, i32
+   return %0: vector<16xi32>
+}
+
+// CHECK-LABEL: func @axpy_int_add(
+// CHECK-SAME: %[[A:.*0]]: vector<16xi32>,
+// CHECK-SAME: %[[B:.*1]]: i32,
+// CHECK-SAME: %[[C:.*2]]: vector<16xi32>)
+// CHECK: %[[T0:.*]] = splat %[[B]] : vector<16xi32>
+// CHECK: %[[T1:.*]] = muli %[[A]], %[[T0]] : vector<16xi32>
+// CHECK: %[[T2:.*]] = addi %[[T1]], %[[C]] : vector<16xi32>
+// CHECK: return %[[T2]] : vector<16xi32>
+func @axpy_int_add(%arg0: vector<16xi32>, %arg1: i32, %arg2: vector<16xi32>) -> vector<16xi32> {
+   %0 = vector.outerproduct %arg0, %arg1, %arg2: vector<16xi32>, i32
+   return %0: vector<16xi32>
+}
+
 // CHECK-LABEL: func @transpose23
 // CHECK-SAME: %[[A:.*]]: vector<2x3xf32>
 // CHECK:      %[[Z:.*]] = constant dense<0.000000e+00> : vector<3x2xf32>
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0, 0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.insert %[[T0]], %[[Z]] [0, 0] : f32 into vector<3x2xf32>
 // CHECK:      %[[T2:.*]] = vector.extract %[[A]][1, 0] : vector<2x3xf32>
 // CHECK:      %[[T3:.*]] = vector.insert %[[T2]], %[[T1]] [0, 1] : f32 into vector<3x2xf32>
 // CHECK:      %[[T4:.*]] = vector.extract %[[A]][0, 1] : vector<2x3xf32>