[mlir][vector] NFC: Expose castAwayContractionLeadingOneDim

[antiagainst]

Summary

This commit exposes the transformation behind the pattern.
It is useful for more targeted application on a specific op
for once.

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

 class MLIRContext;
 class RewritePatternSet;
 
 namespace arith {
 enum class AtomicRMWKind : uint64_t;
 } // namespace arith
 
 namespace vector {
+class ContractionOp;
 class TransferReadOp;
 class TransferWriteOp;
 class VectorDialect;
 
 namespace detail {
 struct BitmaskEnumStorage;
 } // namespace detail
 

 /// Collect a set of vector-to-vector canonicalization patterns.
 void populateVectorToVectorCanonicalizationPatterns(RewritePatternSet &patterns,
                                                     PatternBenefit benefit = 1);
 
 /// Collect a set of vector.shape_cast folding patterns.
 void populateShapeCastFoldingPatterns(RewritePatternSet &patterns,
                                       PatternBenefit benefit = 1);
 
+/// Cast away the leading unit dim, if exists, for the given contract op.
+/// Return success if the transformation applies; return failure otherwise.
+LogicalResult castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
+                                               RewriterBase &rewriter);
+
 /// Collect a set of leading one dimension removal patterns.
 ///
 /// These patterns insert vector.shape_cast to remove leading one dimensions
 /// to expose more canonical forms of read/write/insert/extract operations.
 /// With them, there are more chances that we can cancel out extract-insert
 /// pairs or forward write-read pairs.
 void populateCastAwayVectorLeadingOneDimPatterns(RewritePatternSet &patterns,
                                                  PatternBenefit benefit = 1);

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

LogicalResult matchAndRewrite(vector::TransferWriteOp write,
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         write, newVector, write.getSource(), write.getIndices(),
         AffineMapAttr::get(newMap), inBoundsAttr);
 
     return success();
   }
 };
 
-/// 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;
+} // namespace
 
-  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();
-    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;
+LogicalResult
+mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
+                                               RewriterBase &rewriter) {
 
-    auto oldIndexingMaps = contractOp.getIndexingMapsArray();
-    SmallVector<AffineMap> newIndexingMaps;
 
-    auto oldIteratorTypes = contractOp.getIteratorTypes();
-    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.getLhs(), contractOp.getRhs(),
-                                   contractOp.getAcc()};
-    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()],
+    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.getKind());
-    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
-        contractOp, contractOp->getResultTypes()[0], newContractOp);
-    return success();
-  }
+
+namespace {
+
+/// 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 {
+    return castAwayContractionLeadingOneDim(contractOp, rewriter);
+  }
 };
 
 class CastAwayElementwiseLeadingOneDim : public RewritePattern {
 public:
   CastAwayElementwiseLeadingOneDim(MLIRContext *context,
                                    PatternBenefit benefit = 1)
       : RewritePattern(MatchAnyOpTypeTag(), benefit, context) {}