[mlir][spirv][vector] Enable vector3 when converting to dot product

[antiagainst]

Summary

It's common to see such cases for contraction from convolution with
input channel as 3. Although we aren't utilizing all 4 lanes for
dot product, it should still be better than performing the multiply
and reduction separately.

[kuhar]

LGTM but I wonder if we should generalize this to support any vector width > 2. For widths > 4, we could unroll it into a chain of 4-element dot products.

[antiagainst]

In D148642#4278321, @kuhar wrote:

LGTM but I wonder if we should generalize this to support any vector width > 2. For widths > 4, we could unroll it into a chain of 4-element dot products.

For vectors with a size larger than 4, we could already relying on unrolling at vector level, right?

[kuhar]

In D148642#4278533, @antiagainst wrote:

In D148642#4278321, @kuhar wrote:

LGTM but I wonder if we should generalize this to support any vector width > 2. For widths > 4, we could unroll it into a chain of 4-element dot products.

For vectors with a size larger than 4, we could already relying on unrolling at vector level, right?

That's independent of the unrolling IMO -- whatever the unrolling scheme used is, we can efficiently lower any reduction of vectors of size >= 3.

[antiagainst]

In D148642#4278562, @kuhar wrote:

In D148642#4278533, @antiagainst wrote:

In D148642#4278321, @kuhar wrote:

LGTM but I wonder if we should generalize this to support any vector width > 2. For widths > 4, we could unroll it into a chain of 4-element dot products.

For vectors with a size larger than 4, we could already relying on unrolling at vector level, right?

That's independent of the unrolling IMO -- whatever the unrolling scheme used is, we can efficiently lower any reduction of vectors of size >= 3.

Yeah, agreed. We can do that when a use case arise later.

mlir/lib/Conversion/VectorToSPIRV/VectorToSPIRV.cpp

LogicalResult matchAndRewrite(vector::ReductionOp op,
     if (!resultType)
       return rewriter.notifyMatchFailure(op, "result is not an integer");
 
     int64_t resultBitwidth = resultType.getIntOrFloatBitWidth();
     if (!llvm::is_contained({32, 64}, resultBitwidth))
       return rewriter.notifyMatchFailure(op, "unsupported integer bitwidth");
 
     VectorType inVecTy = op.getSourceVectorType();
-    if (inVecTy.getNumElements() != 4 || inVecTy.getShape().size() != 1 ||
-        inVecTy.isScalable())
+    if (!llvm::is_contained({4, 3}, inVecTy.getNumElements()) ||
+        inVecTy.getShape().size() != 1 || inVecTy.isScalable())
       return rewriter.notifyMatchFailure(op, "unsupported vector shape");
 
     auto mul = op.getVector().getDefiningOp<arith::MulIOp>();
     if (!mul)
       return rewriter.notifyMatchFailure(
           op, "reduction operand is not 'arith.muli'");
 
     if (succeeded(handleCase<arith::ExtSIOp, arith::ExtSIOp, spirv::SDotOp,
struct VectorReductionToDotProd final : OpRewritePattern<vector::ReductionOp> {
   }
 
 private:
   template <typename LhsExtensionOp, typename RhsExtensionOp, typename DotOp,
             typename DotAccOp, bool SwapOperands>
   static LogicalResult handleCase(vector::ReductionOp op, arith::MulIOp mul,
                                   PatternRewriter &rewriter) {
     auto lhs = mul.getLhs().getDefiningOp<LhsExtensionOp>();
-    if (!lhs || !getElementTypeOrSelf(lhs.getIn().getType()).isInteger(8))
+    if (!lhs)
+      return failure();
+    Value lhsIn = lhs.getIn();
+    auto lhsInType = cast<VectorType>(lhsIn.getType());
+    if (!lhsInType.getElementType().isInteger(8))
       return failure();
 
     auto rhs = mul.getRhs().getDefiningOp<RhsExtensionOp>();
-    if (!rhs || !getElementTypeOrSelf(rhs.getIn().getType()).isInteger(8))
+    if (!rhs)
       return failure();
-
-    Value lhsIn = lhs.getIn();
     Value rhsIn = rhs.getIn();
+    auto rhsInType = cast<VectorType>(rhsIn.getType());
+    if (!rhsInType.getElementType().isInteger(8))
+      return failure();
+
+    if (op.getSourceVectorType().getNumElements() == 3) {
+      IntegerType i8Type = rewriter.getI8Type();
+      auto v4i8Type = VectorType::get({4}, i8Type);
+      Location loc = op.getLoc();
+      Value zero = spirv::ConstantOp::getZero(i8Type, loc, rewriter);
+      lhsIn = rewriter.create<spirv::CompositeConstructOp>(
+          loc, v4i8Type, ValueRange{lhsIn, zero});
+      rhsIn = rewriter.create<spirv::CompositeConstructOp>(
+          loc, v4i8Type, ValueRange{rhsIn, zero});
+    }
 
     // There's no variant of dot prod ops for unsigned LHS and signed RHS, so
     // we have to swap operands instead in that case.
     if (SwapOperands)
       std::swap(lhsIn, rhsIn);
 
     if (Value acc = op.getAcc()) {
       rewriter.replaceOpWithNewOp<DotAccOp>(op, op.getType(), lhsIn, rhsIn, acc,

mlir/test/Conversion/VectorToSPIRV/vector-reduction-to-spirv-dot-prod.mlir

 func.func @to_unsigned_signed_dot_acc(%arg0: vector<4xi8>, %arg1: vector<4xi8>, %acc: i32) -> i32 {
   %lhs = arith.extui %arg0 : vector<4xi8> to vector<4xi32>
   %rhs = arith.extsi %arg1 : vector<4xi8> to vector<4xi32>
   %mul = arith.muli %lhs, %rhs : vector<4xi32>
   %red = vector.reduction <add>, %mul, %acc : vector<4xi32> into i32
   return %red : i32
 }
 
+// CHECK-LABEL: func.func @to_sdot_vector3
+//  CHECK-SAME: (%[[ARG0:.+]]: vector<3xi8>, %[[ARG1:.+]]: vector<3xi8>)
+//       CHECK:   %[[ZERO:.+]] = spirv.Constant 0 : i8
+//       CHECK:   %[[LHS:.+]] = spirv.CompositeConstruct %[[ARG0]], %[[ZERO]] : (vector<3xi8>, i8) -> vector<4xi8>
+//       CHECK:   %[[RHS:.+]] = spirv.CompositeConstruct %[[ARG1]], %[[ZERO]] : (vector<3xi8>, i8) -> vector<4xi8>
+//       CHECK:   %[[SDOT:.+]] = spirv.SDot %[[LHS]], %[[RHS]] : (vector<4xi8>, vector<4xi8>) -> i32
+//       CHECK:   return %[[SDOT]]
+func.func @to_sdot_vector3(%arg0: vector<3xi8>, %arg1: vector<3xi8>) -> i32 {
+  %lhs = arith.extsi %arg0 : vector<3xi8> to vector<3xi32>
+  %rhs = arith.extsi %arg1 : vector<3xi8> to vector<3xi32>
+  %mul = arith.muli %lhs, %rhs : vector<3xi32>
+  %red = vector.reduction <add>, %mul : vector<3xi32> into i32
+  return %red : i32
+}
+
 // -----
+
 // Negative tests.
 
 // CHECK-LABEL: func.func @too_short
-//  CHECK-SAME:   ([[ARG0:%.+]]: vector<3xi8>, [[ARG1:%.+]]: vector<3xi8>)
+//  CHECK-SAME:   ([[ARG0:%.+]]: vector<2xi8>, [[ARG1:%.+]]: vector<2xi8>)
 //  CHECK:        [[RED:%.+]] = vector.reduction
 //  CHECK-NEXT:   return [[RED]] : i32
-func.func @too_short(%arg0: vector<3xi8>, %arg1: vector<3xi8>) -> i32 {
-  %lhs = arith.extsi %arg0 : vector<3xi8> to vector<3xi32>
-  %rhs = arith.extsi %arg1 : vector<3xi8> to vector<3xi32>
-  %mul = arith.muli %lhs, %rhs : vector<3xi32>
-  %red = vector.reduction <add>, %mul : vector<3xi32> into i32
+func.func @too_short(%arg0: vector<2xi8>, %arg1: vector<2xi8>) -> i32 {
+  %lhs = arith.extsi %arg0 : vector<2xi8> to vector<2xi32>
+  %rhs = arith.extsi %arg1 : vector<2xi8> to vector<2xi32>
+  %mul = arith.muli %lhs, %rhs : vector<2xi32>
+  %red = vector.reduction <add>, %mul : vector<2xi32> into i32
   return %red : i32
 }
 
 // CHECK-LABEL: func.func @too_long
 //  CHECK-SAME:   ([[ARG0:%.+]]: vector<6xi8>, [[ARG1:%.+]]: vector<6xi8>)
 //  CHECK:        [[RED:%.+]] = vector.reduction
 //  CHECK-NEXT:   return [[RED]] : i32
 func.func @too_long(%arg0: vector<6xi8>, %arg1: vector<6xi8>) -> i32 {