[mlir][vector] Add missing support for contract of integer lowering.

[ThomasRaoux]

Summary

Some of the lowering of vector.contract didn't support integer case. Since reduction of integer cannot accumulate we always break up the reduction op, it should be merged by a separate canonicalization if possible.

[aartbik]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2032

Minor, but ignorable suggestion: since we repeat this, how about making a helper somewhere that is called createAdd and createMul and pass in the isInt?

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

21–23

As you state in the description, we always have this breakup now, which is a bit unfortunate (the merged version tends to generate better code). Do you feel special cases codegen is too much of a maintenance burden? Can we at least get a commitment to get the follow up canonicalization in place

[nicolasvasilache]

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

21–23

For cases that reduce to FMA, I have one that I need to flush into core somewhere.
For cases like integer etc it'll prob be HW_specific.

[bondhugula]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2014

Nit: You can move this to further below - to just outside the for loop.

[ThomasRaoux]

Address feedback.

[ThomasRaoux]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

2014

Done.

2032

I was also thinking about it. I followed your suggestion.

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

21–23

@nicolasvasilache pushed this canonicalization so this shouldn't be a concern!

mlir/lib/Dialect/Vector/VectorTransforms.cpp

if (++idx[r] == tp.getDimSize(r)) {
       idx[r] = 0;
       incIdx(idx, tp, r - 1);
     }
   }
 };
 
 } // namespace
 
+/// Creates an AddIOp if `isInt` is true otherwise create an AddFOp using
+/// operands `x` and `y`.
+static Value createAdd(Location loc, Value x, Value y, bool isInt,
+                       PatternRewriter &rewriter) {
+  if (isInt)
+    return rewriter.create<AddIOp>(loc, x, y);
+  return rewriter.create<AddFOp>(loc, x, y);
+}
+
+/// Creates a MulIOp if `isInt` is true otherwise create an MulFOp using
+/// operands `x and `y`.
+static Value createMul(Location loc, Value x, Value y, bool isInt,
+                       PatternRewriter &rewriter) {
+  if (isInt)
+    return rewriter.create<MulIOp>(loc, x, y);
+  return rewriter.create<MulFOp>(loc, x, y);
+}
+
 namespace mlir {
 
 /// Progressively lower a `vector.contract %a, %b, %c` with row-major matmul
 /// semantics to:
 /// ```
 ///    %flattened_a = vector.shape_cast %a
 ///    %flattened_b = vector.shape_cast %b
 ///    %flattened_d = vector.matmul %flattened_a, %flattened_b
if (isParallelIterator(iteratorTypes[0]) &&
     } else {
       return failure();
     }
   } else {
     return failure();
   }
 
   VectorType dstType = op.getResultType().cast<VectorType>();
   assert(dstType.getRank() >= 1 && dstType.getRank() <= 2 &&
          "Expected dst type of rank 1 or 2");
 
   unsigned rank = dstType.getRank();
   unsigned dstRows = dstType.getShape()[0];
   unsigned dstColumns = rank == 1 ? 1 : dstType.getShape()[1];
 
   // ExtractOp does not allow dynamic indexing, we must unroll explicitly.
   Value res =
       rewriter.create<ConstantOp>(loc, dstType, rewriter.getZeroAttr(dstType));
+  bool isInt = dstType.getElementType().isa<IntegerType>();
   for (unsigned r = 0; r < dstRows; ++r) {
     Value a = rewriter.create<vector::ExtractOp>(op.getLoc(), lhs, r);
     for (unsigned c = 0; c < dstColumns; ++c) {
       Value b = rank == 1
                     ? rhs
                     : rewriter.create<vector::ExtractOp>(op.getLoc(), rhs, c);
-      Value m = rewriter.create<MulFOp>(op.getLoc(), a, b);
+      Value m = createMul(op.getLoc(), a, b, isInt, rewriter);
       Value reduced = rewriter.create<vector::ReductionOp>(
           op.getLoc(), dstType.getElementType(), rewriter.getStringAttr("add"),
           m, ValueRange{});
 
       SmallVector<int64_t, 2> pos = rank == 1 ? SmallVector<int64_t, 2>{r}
                                               : SmallVector<int64_t, 2>{r, c};
       res = rewriter.create<vector::InsertOp>(op.getLoc(), reduced, res, pos);
     }
   }
   if (auto acc = op.acc())
-    res = rewriter.create<AddFOp>(op.getLoc(), res, acc);
+    res = createAdd(op.getLoc(), res, acc, isInt, rewriter);
   rewriter.replaceOp(op, res);
   return success();
 }
 
 /// Progressive lowering of ContractionOp.
 /// One:
 ///   %x = vector.contract with at least one free/batch dimension
 /// is replaced by:

 // Lower one reduction dimension.
 Value ContractionOpLowering::lowerReduction(vector::ContractionOp op,
                                             PatternRewriter &rewriter) const {
   auto loc = op.getLoc();
   VectorType lhsType = op.getLhsType();
   VectorType rhsType = op.getRhsType();
   Type resType = op.getResultType();
   assert(!resType.isa<VectorType>());
+  bool isInt = resType.isa<IntegerType>();
   // Use iterator index 0.
   int64_t iterIndex = 0;
   SmallVector<AffineMap, 4> iMap = op.getIndexingMaps();
   Optional<int64_t> lookupLhs = getResultIndex(iMap[0], iterIndex);
   Optional<int64_t> lookupRhs = getResultIndex(iMap[1], iterIndex);
   assert(lookupLhs.hasValue() && "missing LHS parallel index");
   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 m = rewriter.create<MulFOp>(loc, op.lhs(), op.rhs());
+    Value m = createMul(loc, op.lhs(), op.rhs(), isInt, rewriter);
     StringAttr kind = rewriter.getStringAttr("add");
-    return rewriter.create<vector::ReductionOp>(loc, resType, kind, m,
-                                                op.acc());
+    Value res = rewriter.create<vector::ReductionOp>(loc, resType, kind, m,
+                                                     ValueRange{});
+    if (auto acc = op.acc())
+      res = createAdd(op.getLoc(), res, acc, isInt, rewriter);
+    return res;
   }
   // Construct new iterator types and affine map array attribute.
   std::array<AffineMap, 3> lowIndexingMaps = {
       adjustMap(iMap[0], iterIndex, rewriter),
       adjustMap(iMap[1], iterIndex, rewriter),
       adjustMap(iMap[2], iterIndex, rewriter)};
   auto lowAffine = rewriter.getAffineMapArrayAttr(lowIndexingMaps);
   auto lowIter =

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

#dotp_trait = {
   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:      %[[F:.*]] = mulf %[[A]], %[[B]] : vector<4xf32>
-// CHECK:      %[[R:.*]] = vector.reduction "add", %[[F]], %[[C]] : vector<4xf32> into f32
-// CHECK:      return %[[R]] : f32
+// CHECK:      %[[R:.*]] = vector.reduction "add", %[[F]] : vector<4xf32> into f32
+// CHECK:      %[[ACC:.*]] = addf %[[R]], %[[C]] : f32
+// CHECK:      return %[[ACC]] : 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
 }
 
+// CHECK-LABEL: func @extract_contract1_int
+// CHECK-SAME: %[[A:.*0]]: vector<4xi32>,
+// CHECK-SAME: %[[B:.*1]]: vector<4xi32>,
+// CHECK-SAME: %[[C:.*2]]: i32
+// CHECK:      %[[F:.*]] = muli %[[A]], %[[B]] : vector<4xi32>
+// CHECK:      %[[R:.*]] = vector.reduction "add", %[[F]] : vector<4xi32> into i32
+// CHECK:      %[[ACC:.*]] = addi %[[R]], %[[C]] : i32
+// CHECK:      return %[[ACC]] : i32
+
+func @extract_contract1_int(%arg0: vector<4xi32>, %arg1: vector<4xi32>, %arg2: i32) -> i32 {
+  %0 = vector.contract #dotp_trait %arg0, %arg1, %arg2
+    : vector<4xi32>, vector<4xi32> into i32
+  return %0 : i32
+}
+
 #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"]

 func @extract_contract2(%arg0: vector<2x3xf32>,
                         %arg1: vector<3xf32>,
       %arg2: vector<2xf32>) -> vector<2xf32> {
   %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2
     : vector<2x3xf32>, vector<3xf32> into vector<2xf32>
   return %0 : vector<2xf32>
 }
 
+// CHECK-LABEL: func @extract_contract2_int
+// CHECK-SAME: %[[A:.*0]]: vector<2x3xi32>,
+// CHECK-SAME: %[[B:.*1]]: vector<3xi32>,
+// CHECK-SAME: %[[C:.*2]]: vector<2xi32>
+// CHECK:      %[[R:.*]] = constant dense<0> : vector<2xi32>
+// CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xi32>
+// CHECK:      %[[T2:.*]] = muli %[[T0]], %[[B]] : vector<3xi32>
+// CHECK:      %[[T3:.*]] = vector.reduction "add", %[[T2]] : vector<3xi32> into i32
+// CHECK:      %[[T4:.*]] = vector.insert %[[T3]], %[[R]] [0] : i32 into vector<2xi32>
+// CHECK:      %[[T5:.*]] = vector.extract %[[A]][1] : vector<2x3xi32>
+// CHECK:      %[[T7:.*]] = muli %[[T5]], %[[B]] : vector<3xi32>
+// CHECK:      %[[T8:.*]] = vector.reduction "add", %[[T7]] : vector<3xi32> into i32
+// CHECK:      %[[T9:.*]] = vector.insert %[[T8]], %[[T4]] [1] : i32 into vector<2xi32>
+// CHECK:      %[[T10:.*]] = addi %[[T9]], %[[C]] : vector<2xi32>
+// CHECK:      return %[[T10]] : vector<2xi32>
+func @extract_contract2_int(%arg0: vector<2x3xi32>,
+                        %arg1: vector<3xi32>,
+      %arg2: vector<2xi32>) -> vector<2xi32> {
+  %0 = vector.contract #matvec_trait %arg0, %arg1, %arg2
+    : vector<2x3xi32>, vector<3xi32> into vector<2xi32>
+  return %0 : vector<2xi32>
+}
+
 #vecmat_accesses = [
   affine_map<(i, j) -> (j)>,
   affine_map<(i, j) -> (i, j)>,
   affine_map<(i, j) -> (i)>
 ]
 #vecmat_trait = {
   indexing_maps = #vecmat_accesses,
   iterator_types = ["parallel", "reduction"]

 
 // CHECK-LABEL: func @full_contract1
 // CHECK-SAME: %[[A:.*0]]: vector<2x3xf32>,
 // CHECK-SAME: %[[B:.*1]]: vector<2x3xf32>,
 // CHECK-SAME: %[[C:.*2]]: f32
 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[B]][0] : vector<2x3xf32>
 // 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:.*]] = mulf %[[T4]], %[[T5]] : vector<3xf32>
-// CHECK:      %[[T7:.*]] = vector.reduction "add", %[[T6]], %[[T3]] : vector<3xf32> into f32
-// CHECK:      return %[[T7]] : f32
+// CHECK:      %[[T3:.*]] = vector.reduction "add", %[[T2]] : vector<3xf32> into f32
+// CHECK:      %[[T4:.*]] = addf %[[T3]], %[[C]] : f32
+// CHECK:      %[[T5:.*]] = vector.extract %[[A]][1] : vector<2x3xf32>
+// CHECK:      %[[T6:.*]] = vector.extract %[[B]][1] : vector<2x3xf32>
+// CHECK:      %[[T7:.*]] = mulf %[[T5]], %[[T6]] : vector<3xf32>
+// CHECK:      %[[T8:.*]] = vector.reduction "add", %[[T7]] : vector<3xf32> into f32
+// CHECK:      %[[T9:.*]] = addf %[[T8]], %[[T4]] : f32
+// CHECK:      return %[[T9]] : 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
   return %0 : f32
 }

 // CHECK:      %[[T0:.*]] = vector.extract %[[A]][0] : vector<2x3xf32>
 // CHECK:      %[[T1:.*]] = vector.extract %[[B]][0, 0] : vector<3x2xf32>
 // CHECK:      %[[T3:.*]] = vector.insert %[[T1]], %[[Z]] [0] : f32 into vector<3xf32>
 // CHECK:      %[[T4:.*]] = vector.extract %[[B]][1, 0] : vector<3x2xf32>
 // CHECK:      %[[T6:.*]] = vector.insert %[[T4]], %[[T3]] [1] : f32 into vector<3xf32>
 // CHECK:      %[[T7:.*]] = vector.extract %[[B]][2, 0] : vector<3x2xf32>
 // CHECK:      %[[T9:.*]] = vector.insert %[[T7]], %[[T6]] [2] : f32 into vector<3xf32>
 // CHECK:      %[[T10:.*]] = mulf %[[T0]], %[[T9]] : vector<3xf32>
-// CHECK:      %[[T11:.*]] = vector.reduction "add", %[[T10]], %[[C]] : vector<3xf32> into f32
+// CHECK:      %[[T11:.*]] = vector.reduction "add", %[[T10]] : vector<3xf32> into f32
+// CHECK:      %[[ACC0:.*]] = addf %[[T11]], %[[C]] : f32
 //
 // CHECK:      %[[T12:.*]] = vector.extract %[[A]][1] : vector<2x3xf32>
 // CHECK:      %[[T13:.*]] = vector.extract %[[B]][0, 1] : vector<3x2xf
 // CHECK:      %[[T15:.*]] = vector.insert %[[T13]], %[[Z]] [0] : f32 into vector<3xf32>
 // CHECK:      %[[T16:.*]] = vector.extract %[[B]][1, 1] : vector<3x2xf32>
 // CHECK:      %[[T18:.*]] = vector.insert %[[T16]], %[[T15]] [1] : f32 into vector<3xf32>
 // CHECK:      %[[T19:.*]] = vector.extract %[[B]][2, 1] : vector<3x2xf32>
 // CHECK:      %[[T21:.*]] = vector.insert %[[T19]], %[[T18]] [2] : f32 into vector<3xf32>
 // CHECK:      %[[T22:.*]] = mulf %[[T12]], %[[T21]] : vector<3xf32>
-// CHECK:      %[[T23:.*]] = vector.reduction "add", %[[T22]], %[[T11]] : vector<3xf32> into f32
-// CHECK:      return %[[T23]] : f32
+// CHECK:      %[[T23:.*]] = vector.reduction "add", %[[T22]] : vector<3xf32> into f32
+// CHECK:      %[[ACC1:.*]] = addf %[[T23]], %[[ACC0]] : f32
+// CHECK:      return %[[ACC1]] : 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
   return %0 : f32
 }