[mlir][vector] Add pattern to drop lead unit dim for Contraction Op

If the result operand has a unit leading dim it is removed from all operands.

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D119206

Author: nirvedhmeshram

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

@@ -200,7 +200,10 @@ def Vector_ContractionOp :
       "ArrayAttr":$indexingMaps, "ArrayAttr":$iteratorTypes)>,
     OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
       "ArrayRef<ArrayRef<AffineExpr>>":$indexingExprs,
-      "ArrayRef<StringRef>":$iteratorTypes)>
+      "ArrayRef<StringRef>":$iteratorTypes)>,
+    OpBuilder<(ins "Value":$lhs, "Value":$rhs, "Value":$acc,
+      "ArrayAttr":$indexingMaps, "ArrayAttr":$iteratorTypes,
+      "CombiningKind":$kind)>
   ];
   let extraClassDeclaration = [{
     VectorType getLhsType() {

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -502,13 +502,20 @@ void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
                                   Value lhs, Value rhs, Value acc,
                                   ArrayAttr indexingMaps,
                                   ArrayAttr iteratorTypes) {
+  build(builder, result, lhs, rhs, acc, indexingMaps, iteratorTypes,
+        ContractionOp::getDefaultKind());
+}
+
+void vector::ContractionOp::build(OpBuilder &builder, OperationState &result,
+                                  Value lhs, Value rhs, Value acc,
+                                  ArrayAttr indexingMaps,
+                                  ArrayAttr iteratorTypes, CombiningKind kind) {
   result.addOperands({lhs, rhs, acc});
   result.addTypes(acc.getType());
   result.addAttribute(getIndexingMapsAttrName(), indexingMaps);
   result.addAttribute(getIteratorTypesAttrName(), iteratorTypes);
   result.addAttribute(ContractionOp::getKindAttrName(),
-                      CombiningKindAttr::get(ContractionOp::getDefaultKind(),
-                                             builder.getContext()));
+                      CombiningKindAttr::get(kind, builder.getContext()));
 }
 
 ParseResult ContractionOp::parse(OpAsmParser &parser, OperationState &result) {

mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp

@@ -6,6 +6,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
 #include "mlir/Dialect/Vector/Utils/VectorUtils.h"
 #include "mlir/IR/Builders.h"
@@ -220,6 +221,128 @@ struct CastAwayTransferWriteLeadingOneDim
   }
 };
 
+/// Turns vector.contract on vector with leading 1 dimensions into
+/// vector.extract followed by vector.contract on vector without leading
+/// 1 dimensions. Also performs tranpose of lhs and rhs operands if required
+/// prior to extract.
+struct CastAwayContractionLeadingOneDim
+    : public OpRewritePattern<vector::ContractionOp> {
+  using OpRewritePattern::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
+                                PatternRewriter &rewriter) const override {
+    VectorType oldAccType = contractOp.getAccType().dyn_cast<VectorType>();
+    if (oldAccType == nullptr)
+      return failure();
+    if (oldAccType.getRank() < 2)
+      return failure();
+    // TODO: implement masks.
+    if (llvm::size(contractOp.masks()) != 0)
+      return failure();
+    if (oldAccType.getShape()[0] != 1)
+      return failure();
+    // currently we support only dropping one dim but the pattern can be applied
+    // greedily to drop more.
+    int64_t dropDim = 1;
+
+    auto oldIndexingMaps = contractOp.getIndexingMaps();
+    SmallVector<AffineMap> newIndexingMaps;
+
+    auto oldIteratorTypes = contractOp.iterator_types();
+    SmallVector<Attribute> newIteratorTypes;
+
+    int64_t dimToDrop = oldIndexingMaps[2].getDimPosition(0);
+
+    if (!isParallelIterator(oldIteratorTypes[dimToDrop]))
+      // only parallel type iterators can be dropped.
+      return failure();
+
+    for (const auto &it : llvm::enumerate(oldIteratorTypes)) {
+      int64_t currDim = it.index();
+      if (currDim == dimToDrop)
+        continue;
+      newIteratorTypes.push_back(it.value());
+    }
+
+    SmallVector<Value> operands = {contractOp.lhs(), contractOp.rhs(),
+                                   contractOp.acc()};
+    SmallVector<Value> newOperands;
+
+    for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
+      // Check if the dim to be dropped exists as a leading dim in the operand
+      // if it does then we use vector.extract to drop it.
+      bool validExtract = false;
+      SmallVector<AffineExpr> results;
+      auto map = it.value();
+      int64_t orginalZeroDim = it.value().getDimPosition(0);
+      if (orginalZeroDim != dimToDrop) {
+        // There are two reasons to be in this path, 1. We need to
+        // tranpose the operand to make the dim to be dropped
+        // leading. 2. The dim to be dropped does not exist and in
+        // that case we dont want to add a unit tranpose but we must
+        // check all the indices to make sure this is the case.
+        bool tranposeNeeded = false;
+        SmallVector<int64_t> perm;
+        SmallVector<AffineExpr> transposeResults;
+
+        for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
+          int64_t currDim = map.getDimPosition(i);
+          if (currDim == dimToDrop) {
+            tranposeNeeded = true;
+            perm.insert(perm.begin(), i);
+            auto targetExpr = rewriter.getAffineDimExpr(currDim);
+            transposeResults.insert(transposeResults.begin(), targetExpr);
+          } else {
+            perm.push_back(i);
+            auto targetExpr = rewriter.getAffineDimExpr(currDim);
+            transposeResults.push_back(targetExpr);
+          }
+        }
+        // Do the tranpose now if needed so that we can drop the
+        // correct dim using extract later.
+        if (tranposeNeeded) {
+          map = AffineMap::get(map.getNumDims(), 0, transposeResults,
+                               contractOp.getContext());
+          operands[it.index()] = rewriter.create<vector::TransposeOp>(
+              contractOp.getLoc(), operands[it.index()], perm);
+        }
+      }
+      // We have taken care to have the dim to be dropped be
+      // the leading dim. If its still not leading that means it
+      // does not exist in this operand and hence we do not need
+      // an extract.
+      if (map.getDimPosition(0) == dimToDrop)
+        validExtract = true;
+
+      for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
+        int64_t currDim = map.getDimPosition(i);
+        if (currDim == dimToDrop)
+          // This is the dim we are dropping.
+          continue;
+        auto targetExpr = rewriter.getAffineDimExpr(
+            currDim < dimToDrop ? currDim : currDim - 1);
+        results.push_back(targetExpr);
+      }
+      newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
+                                               contractOp.getContext()));
+      // Extract if its a valid extraction, otherwise use the operand
+      // without extraction.
+      newOperands.push_back(validExtract
+                                ? rewriter.create<vector::ExtractOp>(
+                                      contractOp.getLoc(), operands[it.index()],
+                                      splatZero(dropDim))
+                                : operands[it.index()]);
+    }
+    auto newContractOp = rewriter.create<vector::ContractionOp>(
+        contractOp.getLoc(), newOperands[0], newOperands[1], newOperands[2],
+        rewriter.getAffineMapArrayAttr(newIndexingMaps),
+        rewriter.getArrayAttr(newIteratorTypes), contractOp.kind());
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
+        contractOp, contractOp->getResultTypes()[0], newContractOp);
+    return success();
+  }
+};
+
 class CastAwayElementwiseLeadingOneDim : public RewritePattern {
 public:
   CastAwayElementwiseLeadingOneDim(MLIRContext *context)
@@ -260,10 +383,11 @@ class CastAwayElementwiseLeadingOneDim : public RewritePattern {
 
 void mlir::vector::populateCastAwayVectorLeadingOneDimPatterns(
     RewritePatternSet &patterns) {
-  patterns.add<CastAwayExtractStridedSliceLeadingOneDim,
-               CastAwayInsertStridedSliceLeadingOneDim,
-               CastAwayTransferReadLeadingOneDim,
-               CastAwayTransferWriteLeadingOneDim,
-               CastAwayElementwiseLeadingOneDim>(patterns.getContext());
+  patterns
+      .add<CastAwayExtractStridedSliceLeadingOneDim,
+           CastAwayInsertStridedSliceLeadingOneDim,
+           CastAwayTransferReadLeadingOneDim,
+           CastAwayTransferWriteLeadingOneDim, CastAwayElementwiseLeadingOneDim,
+           CastAwayContractionLeadingOneDim>(patterns.getContext());
   populateShapeCastFoldingPatterns(patterns);
 }