[mlir] [VectorOps] generalized vector.contract semantics

[aartbik]

Summary

Previously, vector.contract did not allow an empty set of
free or batch dimensions (K = 0) which defines a basic
reduction into a scalar (like a dot product). This CL
relaxes that restriction. Also adds constraints on
element type of operands and results. With tests.

[merge_guards_bot]

Unit tests: pass. 62465 tests passed, 0 failed and 845 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[rriddle]

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

151–152

nit: Can you format these to one line now?

mlir/lib/Dialect/VectorOps/VectorOps.cpp

206

nit: Can you just inline this function? It only has one use case and isn't very large.

227

Is this supposed to be except for?

[aartbik]

address comments

[aartbik]

mlir/lib/Dialect/VectorOps/VectorOps.cpp

206

Yes, of course. I had bigger plans for this, but then only used it once :-)

227

No, I meant "accept" as in continue without doing the tests below. But rephrased this into something better.

[merge_guards_bot]

Unit tests: pass. 62465 tests passed, 0 failed and 845 were skipped.

clang-tidy: pass.

clang-format: pass.

Build artifacts: diff.json, clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt, test-results.xml

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[andydavis1]

Looks good from my perspective.

[aartbik]

rebased with head

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: unknown.

clang-format: unknown.

Build artifacts: console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

[merge_guards_bot]

Unit tests: unknown.

clang-tidy: pass.

clang-format: fail. Please format your changes with clang-format by running git-clang-format HEAD^ or applying this patch.

Build artifacts: clang-tidy.txt, clang-format.patch, CMakeCache.txt, console-log.txt

Pre-merge checks is in beta. Report issue. Please join beta or enable it for your project.

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

class Vector_Op<string mnemonic, list<OpTrait> traits = []> :
   let printer = [{ return ::print(p, *this); }];
   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]>,
-    Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyVector:$acc,
+  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<"third operand acc and result have same element type",
+                 TCresVTEtIsSameAsOpBase<0, 1>>]>,
+    Arguments<(ins AnyVector:$lhs, AnyVector:$rhs, AnyType:$acc,
                Variadic<VectorOf<[I1]>>:$masks,
                AffineMapArrayAttr:$indexing_maps, ArrayAttr:$iterator_types)>,
-    Results<(outs AnyVector)> {
+    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
     intermediate result to the accumulator argument of rank K, and returns a
     vector result of rank K (where K = num_lhs_free_dims + num_rhs_free_dims +
-    num_batch_dims (see dimension type descriptions below)).
+    num_batch_dims (see dimension type descriptions below)). For K = 0 (no
+    free or batch dimensions), the accumulator and output are a scalar.
 
     Optional vector mask arguments (produced by CreateMaskOp or ConstantMaskOp)
     specify the dynamic dimension sizes of valid data within the lhs/rhs vector
     arguments.
 
     An iterator type attribute list must be specified, where each element of
     the list represents an iterator with one of the following types:
 
     *) "reduction": reduction dimensions are present in the lhs and rhs
-                    arguments but not in the output (or optional accumulator
+                    arguments but not in the output (and accumulator
                     argument). These are the dimensions along which the vector
                     contraction op computes the sum of products, and
                     contracting dimension pair dimension sizes must match
                     between lhs/rhs.
     *) "parallel": Batch dimensions are iterator type "parallel", and
                    are non-contracting dimensions present in the lhs, rhs and
                    output. The lhs/rhs co-iterate along the batch dimensions,
                    which should be expressed in their indexing maps.
 
                    Free dimensions are iterator type "parallel", and are
                    non-contraction, non-batch dimensions accessed by either the
                    lhs or rhs (but not both). The lhs and rhs free dimensions
                    are unrelated to each other and do not co-iterate, which
                    should be expressed in their indexing maps.
 
     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.
 
     Examples:
 
-      // 2D vector contraction with one contracting dimension (matmul).
+      // 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"]
+      }
+      %3 = vector.contract #contraction_trait %0, %1, %2
+        : vector<10xf32>, vector<10xf32> into f32
+
+      // 2D vector contraction with one contracting dimension (matmul, K = 2).
       #contraction_accesses = [
         affine_map<(i, j, k) -> (i, k)>,
         affine_map<(i, j, k) -> (k, j)>,
         affine_map<(i, j, k) -> (i, j)>
       ]
       #contraction_trait = {
         indexing_maps = #contraction_accesses,
         iterator_types = ["parallel", "parallel", "reduction"]
       }
 
       %3 = vector.contract #contraction_trait %0, %1, %2
         : vector<4x3xf32>, vector<3x7xf32> into vector<4x7xf32>
 
       // 4D to 3D vector contraction with two contracting dimensions and
-      // one batch dimension.
+      // one batch dimension (K = 3).
       #contraction_accesses = [
         affine_map<(b0, f0, f1, c0, c1) -> (c0, b0, c1, f0)>,
         affine_map<(b0, f0, f1, c0, c1) -> (b0, c1, c0, f1)>,
         affine_map<(b0, f0, f1, c0, c1) -> (b0, f0, f1)>
       ]
       #contraction_trait = {
         indexing_maps = #contraction_accesses,
         iterator_types = ["parallel", "parallel", "parallel",
-                    "reduction", "reduction"]
+                          "reduction", "reduction"]
       }
 
       %4 = vector.contract #contraction_trait %0, %1, %2
           : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
 
       // 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>
   }];
   let builders = [OpBuilder<
     "Builder *builder, OperationState &result, Value lhs, Value rhs, "
     "Value acc, ArrayAttr indexingMaps, ArrayAttr iteratorTypes">];
   let extraClassDeclaration = [{
     VectorType getLhsType() {
       return lhs().getType().cast<VectorType>();
     }
     VectorType getRhsType() {
       return rhs().getType().cast<VectorType>();
     }
-    VectorType getAccType() {
-      return acc().getType().cast<VectorType>();
-    }
+    Type getAccType() { return acc().getType(); }
     VectorType getLHSVectorMaskType() {
       if (llvm::size(masks()) != 2) return VectorType();
       return getOperand(3).getType().cast<VectorType>();
     }
     VectorType getRHSVectorMaskType() {
       if (llvm::size(masks()) != 2) return VectorType();
       return getOperand(4).getType().cast<VectorType>();
     }
-    VectorType getResultType() {
-      return getResult().getType().cast<VectorType>();
-    }
+    Type getResultType() { return getResult().getType(); }
     ArrayRef<StringRef> getTraitAttrNames();
     SmallVector<AffineMap, 4> getIndexingMaps();
     static unsigned getAccOperandIndex() { return 2; }
 
     // Returns the bounds of each dimension in the iteration space spanned
     // by the iterator types of this operation.
     void getIterationBounds(SmallVectorImpl<int64_t> &iterationBounds);
 

mlir/lib/Dialect/VectorOps/VectorOps.cpp

 
 static ParseResult parseContractionOp(OpAsmParser &parser,
                                       OperationState &result) {
   OpAsmParser::OperandType lhsInfo;
   OpAsmParser::OperandType rhsInfo;
   OpAsmParser::OperandType accInfo;
   SmallVector<OpAsmParser::OperandType, 2> masksInfo;
   SmallVector<Type, 2> types;
-  Type resultVectorType;
+  Type resultType;
   auto loc = parser.getCurrentLocation();
   DictionaryAttr dictAttr;
   // TODO(andydavis, ntv) Unify linalg op attribute parsing.
   if (parser.parseAttribute(dictAttr, "_", result.attributes) ||
       parser.parseOperand(lhsInfo) || parser.parseComma() ||
       parser.parseOperand(rhsInfo) || parser.parseComma() ||
       parser.parseOperand(accInfo) ||
       parser.parseTrailingOperandList(masksInfo) ||
       parser.parseOptionalAttrDict(result.attributes) ||
       parser.parseColonTypeList(types) ||
-      parser.parseKeywordType("into", resultVectorType) ||
+      parser.parseKeywordType("into", resultType) ||
       parser.resolveOperand(lhsInfo, types[0], result.operands) ||
       parser.resolveOperand(rhsInfo, types[1], result.operands) ||
-      parser.resolveOperand(accInfo, resultVectorType, result.operands) ||
-      parser.addTypeToList(resultVectorType, result.types))
+      parser.resolveOperand(accInfo, resultType, result.operands) ||
+      parser.addTypeToList(resultType, result.types))
     return failure();
   result.attributes.assign(dictAttr.getValue().begin(),
                            dictAttr.getValue().end());
   if (masksInfo.empty())
     return success();
   if (masksInfo.size() != 2)
     return parser.emitError(parser.getNameLoc(),
                             "expected zero or exactly 2 vector mask operands");
if (dimPair.first < 0 || dimPair.first >= lhsType.getRank() ||
         dimPair.second < 0 || dimPair.second >= rhsType.getRank() ||
         lhsType.getDimSize(dimPair.first) != rhsType.getDimSize(dimPair.second))
       return false;
   }
   return true;
 }
 
 static bool verifyOutputShape(
-    VectorType lhsType, VectorType rhsType, VectorType accType,
-    VectorType resType,
+    VectorType lhsType, VectorType rhsType, Type accType, Type resType,
     const std::vector<std::pair<int64_t, int64_t>> &contractingDimMap,
     const std::vector<std::pair<int64_t, int64_t>> &batchDimMap) {
   DenseSet<int64_t> lhsContractingDimSet;
   DenseSet<int64_t> rhsContractingDimSet;
   for (auto &dimPair : contractingDimMap) {
     lhsContractingDimSet.insert(dimPair.first);
     rhsContractingDimSet.insert(dimPair.second);
   }
static bool verifyOutputShape(
 
   // Add free dimensions from 'rhsType' to 'expectedResultDims'.
   for (int64_t i = 0, e = rhsType.getRank(); i < e; ++i) {
     if (rhsContractingDimSet.count(i) > 0 || rhsBatchDimSet.count(i) > 0)
       continue;
     expectedResultDims.push_back(rhsType.getDimSize(i));
   }
 
+  // Verify 'expectedResultDims'.
+  if (expectedResultDims.size() == 0) {
+    // No batch or free dimension implies a scalar result.
+    if (resType.isa<VectorType>() || accType.isa<VectorType>())
+      return false;
+
+  } else {
+    // At least one batch or free dimension implies a vector result.
+    auto resVectorType = resType.dyn_cast<VectorType>();
+    auto accVectorType = accType.dyn_cast<VectorType>();
+    if (!resVectorType || !accVectorType)
+      return false;
+
     // Verify dimension from 'resType' against 'expectedResultDims'.
-  if (resType.getShape().size() != expectedResultDims.size() ||
-      accType.getShape().size() != expectedResultDims.size())
+    if (resVectorType.getShape().size() != expectedResultDims.size() ||
+        accVectorType.getShape().size() != expectedResultDims.size())
       return false;
-  for (int64_t i = 0, e = resType.getRank(); i < e; ++i) {
-    if (resType.getDimSize(i) != expectedResultDims[i] ||
-        accType.getDimSize(i) != expectedResultDims[i])
+    for (int64_t i = 0, e = resVectorType.getRank(); i < e; ++i) {
+      if (resVectorType.getDimSize(i) != expectedResultDims[i] ||
+          accVectorType.getDimSize(i) != expectedResultDims[i])
         return false;
     }
+  }
   return true;
 }
 
 static LogicalResult verify(ContractionOp op) {
   auto lhsType = op.getLhsType();
   auto rhsType = op.getRhsType();
   auto accType = op.getAccType();
   auto resType = op.getResultType();
 
   // Verify that an indexing map was specified for each vector operand.
   if (op.indexing_maps().size() != 3)
     return op.emitOpError("expected an indexing map for each vector operand");
 
   // Verify that each index map has 'numIterators' inputs, no symbols, and
   // that the number of map outputs equals the rank of its associated
   // vector operand.
   unsigned numIterators = op.iterator_types().getValue().size();
   for (auto it : llvm::enumerate(op.indexing_maps())) {
     auto index = it.index();
     auto map = it.value().cast<AffineMapAttr>().getValue();
     if (map.getNumSymbols() != 0)
       return op.emitOpError("expected indexing map ")
              << index << " to have no symbols";
+    auto vectorType = op.getOperand(index).getType().dyn_cast<VectorType>();
+    unsigned rank = vectorType ? vectorType.getShape().size() : 0;
+    // Since (...) -> () is parsed into an empty map, we need to add
+    // a special case for this situation: continue the verification
+    // of an empty map if the resulting rank is indeed zero, i.e. this
+    // is a reduction into a scalar.
+    if (map.getNumDims() == 0 && map.getNumResults() == 0 && rank == 0)
+      continue;
+    // Verify that the map has the right number of inputs, outputs, and indices.
     if (map.getNumDims() != numIterators)
       return op.emitOpError("expected indexing map ")
              << index << " to have " << numIterators << " number of inputs";
-    auto operandType = op.getOperand(index).getType().cast<VectorType>();
-    unsigned rank = operandType.getShape().size();
     if (map.getNumResults() != rank)
       return op.emitOpError("expected indexing map ")
              << index << " to have " << rank << " number of outputs";
     if (!map.isProjectedPermutation())
       return op.emitOpError("expected indexing map ")
              << index << " to be a projected permutation of its inputs";
   }
 

 static int64_t getResultIndex(AffineMap map, AffineExpr targetExpr) {
   for (int64_t i = 0, e = map.getNumResults(); i < e; ++i)
     if (targetExpr == map.getResult(i))
       return i;
   return -1;
 }
 
 static std::vector<std::pair<int64_t, int64_t>>
 getDimMap(ArrayRef<AffineMap> indexingMaps, ArrayAttr iteratorTypes,
           StringRef targetIteratorTypeName, MLIRContext *context) {
   std::vector<std::pair<int64_t, int64_t>> dimMap;
   for (auto it : llvm::enumerate(iteratorTypes)) {
     auto iteratorTypeName = it.value().cast<StringAttr>().getValue();
     if (iteratorTypeName != targetIteratorTypeName)
       continue;
     // Search lhs/rhs map results for 'targetExpr'.
     auto targetExpr = getAffineDimExpr(it.index(), context);
     int64_t lhsDim = getResultIndex(indexingMaps[0], targetExpr);
     int64_t rhsDim = getResultIndex(indexingMaps[1], targetExpr);
     if (lhsDim >= 0 && rhsDim >= 0)
       dimMap.push_back({lhsDim, rhsDim});
   }
   return dimMap;
 }
 
 void ContractionOp::getIterationBounds(
     SmallVectorImpl<int64_t> &iterationBounds) {
   auto lhsShape = getLhsType().getShape();
-  auto resShape = getResultType().getShape();
+  auto resVectorType = getResultType().dyn_cast<VectorType>();
   SmallVector<AffineMap, 4> indexingMaps(getIndexingMaps());
   SmallVector<int64_t, 2> iterationShape;
   for (auto it : llvm::enumerate(iterator_types())) {
     // Search lhs/rhs map results for 'targetExpr'.
     auto targetExpr = getAffineDimExpr(it.index(), getContext());
     auto iteratorTypeName = it.value().cast<StringAttr>().getValue();
     if (iteratorTypeName == getReductionIteratorTypeName()) {
       // Get reduction dim size from lhs shape (same size in rhsShape).
       int64_t lhsDimIndex = getResultIndex(indexingMaps[0], targetExpr);
       assert(lhsDimIndex >= 0);
       iterationBounds.push_back(lhsShape[lhsDimIndex]);
       continue;
     }
     // Get parallel dimension size from result shape.
     int64_t resDimIndex = getResultIndex(indexingMaps[2], targetExpr);
     assert(resDimIndex >= 0);
-    iterationBounds.push_back(resShape[resDimIndex]);
+    assert(resVectorType != nullptr);
+    iterationBounds.push_back(resVectorType.getShape()[resDimIndex]);
   }
 }
 
 void ContractionOp::getIterationIndexMap(
     std::vector<DenseMap<int64_t, int64_t>> &iterationIndexMap) {
   unsigned numMaps = indexing_maps().getValue().size();
   iterationIndexMap.resize(numMaps);
   for (auto it : llvm::enumerate(indexing_maps())) {

mlir/test/Dialect/VectorOps/invalid.mlir

func @contraction(%arg0: vector<7x8x16x15xf32>, %arg1: vector<8x16x7x5xf32>,
   %lhs_mask = vector.constant_mask [7, 8, 16, 15] : vector<7x8x16x15xi1>
   %rhs_mask = vector.constant_mask [8, 16, 7, 5] : vector<8x16x7x5xi1>
   // expected-error@+1 {{expected zero or exactly 2 vector mask operands}}
   %0 = vector.contract #contraction_trait %arg0, %arg1, %arg2, %lhs_mask
       : vector<7x8x16x15xf32>, vector<8x16x7x5xf32> into vector<8x15x5xf32>
 }
 
 // -----
 
+#contraction_accesses = [
+        affine_map<(i, j, k) -> (i, k)>,
+        affine_map<(i, j, k) -> (k, j)>,
+        affine_map<(i, j, k) -> (i, j)>
+      ]
+#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}}
+  %0 = vector.contract #contraction_trait %arg0, %arg1, %arg2
+    : vector<4x3xi32>, vector<3x7xf32> into vector<4x7xf32>
+}
+
+// -----
+
 func @create_mask() {
   %c2 = constant 2 : index
   %c3 = constant 3 : index
   // expected-error@+1 {{must specify an operand for each result vector dimension}}
   %0 = vector.create_mask %c3, %c2 : vector<4x3x7xi1>
 }
 
 

mlir/test/Dialect/VectorOps/ops.mlir

 }
 
 // CHECK-LABEL: @strided_slice
 func @strided_slice(%arg0: vector<4x8x16xf32>) -> vector<2x2x16xf32> {
   // CHECK: vector.strided_slice %{{.*}} {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x8x16xf32>
   %1 = vector.strided_slice %arg0 {offsets = [2, 2], sizes = [2, 2], strides = [1, 1]} : vector<4x8x16xf32> to vector<2x2x16xf32>
   return %1: vector<2x2x16xf32>
 }
 
+#contraction_to_scalar_accesses = [
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> (i)>,
+  affine_map<(i) -> ()>
+]
+#contraction_to_scalar_trait = {
+  indexing_maps = #contraction_to_scalar_accesses,
+  iterator_types = ["reduction"]
+}
+// CHECK-LABEL: contraction_to_scalar
+func @contraction_to_scalar(%arg0: vector<10xf32>, %arg1: vector<10xf32>) -> f32 {
+  // CHECK:      %[[C0:.*]] = constant 0.000000e+00 : f32
+  %f0 = constant 0.0: f32
+  // CHECK:      %[[X:.*]] = vector.contract {indexing_maps = [#{{.*}}, #{{.*}}, #{{.*}}], iterator_types = ["reduction"]} %{{.*}}, %{{.*}}, %[[C0]] : vector<10xf32>, vector<10xf32> into f32
+  %0 = vector.contract #contraction_to_scalar_trait %arg0, %arg1, %f0
+    : vector<10xf32>, vector<10xf32> into f32
+  // CHECK:      return %[[X]] : f32
+  return %0 : f32
+}
+
 #contraction_accesses0 = [
   affine_map<(b0, f0, f1, c0, c1) -> (c0, b0, c1, f0)>,
   affine_map<(b0, f0, f1, c0, c1) -> (b0, c1, c0, f1)>,
   affine_map<(b0, f0, f1, c0, c1) -> (b0, f0, f1)>
 ]
 #contraction_trait0 = {
   indexing_maps = #contraction_accesses0,
   iterator_types = ["parallel", "parallel", "parallel", "reduction", "reduction"]