[mlir] Allow vector.contract to have mixed types operands

[ThomasRaoux]

Summary

Allow lhs and rhs to have different type than accumulator/destination. Some hardware like GPUs support natively operations like uint8xuint8xuint32.

[aartbik]

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

44

should lhs/rhs still have same element type, at least?

145

perhaps add a mixed element type example?

mlir/test/Dialect/Vector/invalid.mlir

763

please don't remove the whole test, perhaps now test that acc/result have same type

mlir/test/Dialect/Vector/ops.mlir

179

can't we just leave this out, it seemed unused in the original method already?

[ThomasRaoux]

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

44

Yes that makes sense. I added back this constraint.

mlir/test/Dialect/Vector/invalid.mlir

763

I had tried to test acc/result have the same type but it would fail in the parser in a first place as it thinks the acc is used with a different type than declared. I ended up changing this to test that rhs and lhs have the same element type.

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

class Vector_Op<string mnemonic, list<OpTrait> traits = []> :
   let verifier = [{ return ::verify(*this); }];
   let parser = [{ return ::parse$cppClass(parser, result); }];
 }
 
 // TODO(andydavis, ntv) Add an attribute to specify a different algebra
 // with operators other than the current set: {*, +}.
 def Vector_ContractionOp :
   Vector_Op<"contract", [NoSideEffect,
-     PredOpTrait<"first operand lhs and result have same element type",
-                 TCresVTEtIsSameAsOpBase<0, 0>>,
-     PredOpTrait<"second operand rhs and result have same element type",
-                 TCresVTEtIsSameAsOpBase<0, 1>>,
+     PredOpTrait<"lhs and rhs have same element type", TCopVTEtIsSameAs<0, 1>>,
      PredOpTrait<"third operand acc and result have same element type",
-                 TCresVTEtIsSameAsOpBase<0, 1>>]>,
+                 TCresVTEtIsSameAsOpBase<0, 2>>]>,
     Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyType:$acc,
                Variadic<VectorOf<[I1]>>:$masks,
                AffineMapArrayAttr:$indexing_maps, ArrayAttr:$iterator_types)>,
     Results<(outs AnyType)> {
   let summary = "vector contraction operation";
   let description = [{
     Computes the sum of products of vector elements along contracting
     dimension pairs from 2 vectors of rank M and N respectively, adds this
let description = [{
 
     // 4D vector contraction with two contracting dimensions and optional
     // vector mask arguments.
     %lhs_mask = vector.constant_mask [7, 8, 16, 15] : vector<7x8x16x15xi1>
     %rhs_mask = vector.constant_mask [8, 16, 7, 5] : vector<8x16x7x5xi1>
 
     %5 = vector.contract #contraction_trait %0, %1, %2, %lhs_mask, %rhs_mask
        : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x8x5xf32>
+
+    // Vector contraction with mixed typed. lhs/rhs have different element
+    // types than accumulator/result.
+    %6 = vector.contract #contraction_trait %0, %1, %2
+      : vector<10xf16>, vector<10xf16> into f32
     ```
   }];
   let builders = [OpBuilder<
     "OpBuilder &builder, OperationState &result, Value lhs, Value rhs, "
     "Value acc, ArrayAttr indexingMaps, ArrayAttr iteratorTypes">,
     OpBuilder<
       "OpBuilder &builder, OperationState &result, Value lhs, Value rhs, "
       "Value acc, ArrayRef<ArrayRef<AffineExpr>> indexingExprs, "
       "ArrayRef<StringRef> iteratorTypes">];

mlir/lib/Dialect/Vector/VectorTransforms.cpp

 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/IR/PatternMatch.h"
+#include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
 
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "vector-to-vector"
 

 /// TODO(ajcbik): investigate lowering order impact on performance
 LogicalResult
 ContractionOpLowering::matchAndRewrite(vector::ContractionOp op,
                                        PatternRewriter &rewriter) const {
 
   // TODO(ajcbik): implement masks.
   if (llvm::size(op.masks()) != 0)
     return failure();
+  // TODO(thomasraoux): support mixed mode contract lowering.
+  if (op.getLhsType().getElementType() !=
+          getElementTypeOrSelf(op.getAccType()) ||
+      op.getRhsType().getElementType() != getElementTypeOrSelf(op.getAccType()))
+    return failure();
 
   // TODO(ntv, ajcbik): 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)))

mlir/test/Dialect/Vector/invalid.mlir

#contraction_accesses = [
       ]
 #contraction_trait = {
         indexing_maps = #contraction_accesses,
         iterator_types = ["parallel", "parallel", "reduction"]
       }
 func @contraction(%arg0: vector<4x3xi32>,
                   %arg1: vector<3x7xf32>,
                   %arg2: vector<4x7xf32>) -> vector<4x7xf32> {
-  // expected-error@+1 {{'vector.contract' op failed to verify that first operand lhs and result have same element type}}
+  // expected-error@+1 {{'vector.contract' op failed to verify that lhs and rhs have same element type}}
   %0 = vector.contract #contraction_trait %arg0, %arg1, %arg2
     : vector<4x3xi32>, vector<3x7xf32> into vector<4x7xf32>
 }
 
 // -----
 
 #contraction_accesses = [
   affine_map<(m, n, k) -> (m, k)>,

mlir/test/Dialect/Vector/ops.mlir

 #contraction_trait1 = {
   indexing_maps = #contraction_accesses1,
   iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction",
                     "reduction"]
 }
 // CHECK-LABEL: @contraction
 func @contraction(%arg0 : vector<7x8x16x15xf32>, %arg1 : vector<8x16x7x5xf32>,
                   %arg2 : vector<8x15x5xf32>, %arg3 : vector<8x8x15x5xf32>,
-                  %arg4 : index) {
+                  %arg4 : vector<7x8x16x15xf16>, %arg5 : vector<8x16x7x5xf16>) {
   // Test contraction with batch and contracting dims.
   // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
   %0 = vector.contract #contraction_trait0 %arg0, %arg1, %arg2
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
   // Test contraction with only contracting dims. In this case the lhs/rhs
   // dimension of size 8 will be considered a parallel dim for lhs/rhs and will
   // appear twice in the output.
   // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   %1 = vector.contract #contraction_trait1 %arg0, %arg1, %arg3
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   // Test contraction with optional vector mask arguments.
   %lhs_mask = vector.constant_mask [7, 8, 16, 15] : vector<7x8x16x15xi1>
   %rhs_mask = vector.constant_mask [8, 16, 7, 5] : vector<8x16x7x5xi1>
   // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
   %2 = vector.contract #contraction_trait1 %arg0, %arg1, %arg3, %lhs_mask,
                                            %rhs_mask
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x8x15x5xf32>
+  // Test contraction with mixed type.
+  // CHECK: vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction"]} {{.*}}, {{.*}}, {{.*}} : vector<7x8x16x15xf16>, vector<8x16x7x5xf16> into vector<8x8x15x5xf32>
+  %3 = vector.contract #contraction_trait1 %arg4, %arg5, %arg3
+      : vector<7x8x16x15xf16>, vector<8x16x7x5xf16> into vector<8x8x15x5xf32>
   return
 }
 
 // CHECK-LABEL: @create_vector_mask
 func @create_vector_mask() {
   // CHECK:      %[[C2:.*]] = constant 2 : index
   %c2 = constant 2 : index
   // CHECK-NEXT: %[[C3:.*]] = constant 3 : index