[mlir][linalg] Vectorize linalg.pad_tensor consumed by transfer_write

[springerm]

Summary

Vectorize linalg.pad_tensor without generating a linalg.init_tensor when consumed by a transfer_write.

Depends On D103780

[springerm]

rebase

[springerm]

rebase

[springerm]

minor change

[nicolasvasilache]

Please add some tests.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

825

Missing {in_bounds = [true, true]} in the doc.

858

Note: there is probably a similar mlir::makeComposedAffineApply to be written for AffineMinOps which could reduce reliance on CSE in many cases, by construction.

863

Why not just minOp1->getOperands() == minOp2->getOperands() (or .equals, I forget) ?

[springerm]

• add tests
• handle more cases (dyn/static dims) in pad_tensor + transfer_write pattern

[springerm]

rebase

[springerm]

split into 3 smaller patterns with benefit

[springerm]

no change

[springerm]

use benefit as discussed

[nicolasvasilache]

This CL feels like it's doing too much at once, can we please split the generic pattern from the 3 inittensor-less patterns ?

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

889

More comments please.

/// These patterns are meant to apply in complementary fashion.
/// We use benefits and pattern ordering to encode the complementary case disjunction in a simpler way.
/// To handle the ordering without scattered magic constants, these patterns must be added with the following API.

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

691

Can all this logic be replaced by:

if (isa<ConstantOp>(block.getTerminator()->getOperand(0))
  return the attribute?
return null;

?

I'd expect can then decide on the caller side if you need to clone the constant?

[springerm]

rebase

[springerm]

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

691

"Constant" means "constant within the block" here, i.e., defined outside of the block. Extended the comment to make it more clear.

[springerm]

simplify

[springerm]

rebase

[springerm]

small change

[springerm]

small change

[springerm]

rebase

[springerm]

rebase

[springerm]

no change

[springerm]

small improvements

[springerm]

no change

[springerm]

rebase

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 
 /// PadTensorOp does not implement the LinalgStructuredOpInterface `LinalgOp`,
 /// it needs a specific pattern to vectorize.
 struct PadTensorOpVectorizationPattern : public OpRewritePattern<PadTensorOp> {
   using OpRewritePattern<PadTensorOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(PadTensorOp padOp,
                                 PatternRewriter &rewriter) const override;
+
+private:
+  LogicalResult tryVectorizeUser(PatternRewriter &rewriter,
+                                 linalg::PadTensorOp padOp,
+                                 SubTensorInsertOp insertOp,
+                                 Value padValue) const;
+
+  LogicalResult tryVectorizeUser(PatternRewriter &rewriter,
+                                 linalg::PadTensorOp padOp,
+                                 vector::TransferReadOp xferOp,
+                                 Value padValue) const;
+
+  LogicalResult tryVectorizeUser(PatternRewriter &rewriter,
+                                 linalg::PadTensorOp padOp,
+                                 vector::TransferWriteOp xferOp) const;
 };
 
 /// Match and rewrite for the pattern:
 /// ```
 ///    %alloc = ...
 ///    [optional] %view = memref.view %alloc ...
 ///    %subView = subview %allocOrView ...
 ///    [optional] linalg.fill(%allocOrView, %cst) ...
 ///    ...

mlir/lib/Dialect/Linalg/Transforms/Vectorization.cpp

LogicalResult PadTensorOpVectorizationPattern::matchAndRewrite(
 
   // Only support padding with a constant for now, i.e. either:
   //   1. A BBarg from a different block.
   //   2. A value defined outside of the current block.
   Block &block = padOp.region().front();
   auto yieldOp = cast<YieldOp>(block.getTerminator());
   assert(yieldOp.getNumOperands() == 1 && "expected single operand yield");
   Value padValue = yieldOp.values().front();
   Operation *definingOp = padValue.getDefiningOp();
   if (definingOp && definingOp->getBlock() == &block)
     return failure();
   if (!definingOp && padValue.cast<BlockArgument>().getOwner() == &block)
     return failure();
 
+  // Shortcut: Try to avoid creating an InitTensorOp.
+  bool changed = false;
+  // Insert users in vector, because some users may be replaced/removed.
+  for (auto *user : llvm::to_vector<4>(padOp->getUsers())) {
+    if (auto insertOp = dyn_cast<SubTensorInsertOp>(user)) {
+      if (insertOp.source() == padOp.result())
+        changed |=
+            succeeded(tryVectorizeUser(rewriter, padOp, insertOp, padValue));
+    } else if (auto readOp = dyn_cast<vector::TransferReadOp>(user)) {
+      changed |= succeeded(tryVectorizeUser(rewriter, padOp, readOp, padValue));
+    } else if (auto writeOp = dyn_cast<vector::TransferWriteOp>(user)) {
+      changed |= succeeded(tryVectorizeUser(rewriter, padOp, writeOp));
+    }
+  }
+  if (changed)
+    return success();
+  // End of shortcut: Must create an InitTensorOp...
+
   // TODO: if any pad_high is not a static 0, needs a mask. For now, just bail.
   if (llvm::any_of(padOp.getMixedHighPad(),
                    [&](OpFoldResult ofr) { return isNotZeroIndex(ofr); }))
     return failure();
 
   // Now we can rewrite as InitTensorOp + TransferReadOp@[0..0] +
   // TransferWriteOp@[0..0].
   SmallVector<Value> indices(
       resultShapedType.getRank(),
       rewriter.create<ConstantIndexOp>(padOp.getLoc(), 0));
   Value read = rewriter.create<vector::TransferReadOp>(
       padOp.getLoc(), vectorType, padOp.source(), indices, padValue);
   Value init =
       rewriter.create<InitTensorOp>(padOp.getLoc(), resultShapedType.getShape(),
                                     resultShapedType.getElementType());
   rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(padOp, read, init,
                                                        indices);
 
   return success();
 }
 
+/// Helper function that retrieves the value of an IntegerAttr.
+static int64_t getIntFromAttr(Attribute attr) {
+  return attr.cast<IntegerAttr>().getInt();
+}
+
+/// Rewrite use of PadTensorOp result in SubtensorInsertOp. E.g.:
+/// ```
+/// %0 = linalg.pad_tensor %src ... : tensor<?x?xf32> to tensor<17x5xf32>
+/// %r = subtensor_insert %0 into %dest[%a, %b, 0, 0] [1, 1, 17, 5] [1, 1, 1, 1]
+///     : tensor<17x5xf32> into tensor<?x?x17x5xf32>
+/// ```
+/// is rewritten to:
+/// ```
+/// %0 = vector.transfer_read %src[%c0, %c0], %padding
+///     : tensor<?x?xf32>, vector<17x5xf32>
+/// %r = vector.transfer_write %0, %dest[%a, %b, %c0, %c0]
+///     {in_bounds = [true, true]} : vector<17x5xf32>, tensor<?x?x17x5xf32>
+/// ```
+LogicalResult PadTensorOpVectorizationPattern::tryVectorizeUser(
+    PatternRewriter &rewriter, linalg::PadTensorOp padOp,
+    SubTensorInsertOp insertOp, Value padValue) const {
+  auto vecType = VectorType::get(padOp.getType().getShape(),
+                                 padOp.getType().getElementType());
+  unsigned vecRank = vecType.getRank();
+  unsigned tensorRank = insertOp.getType().getRank();
+
+  // Only unit stride supported.
+  if (llvm::any_of(insertOp.getMixedStrides(), [](auto s) {
+        return s.template is<Value>() ||
+               getIntFromAttr(s.template get<Attribute>()) != 1;
+      }))
+    return failure();
+
+  // Check if sizes match: Insert the entire tensor into most minor dims.
+  auto sizes = insertOp.getMixedSizes();
+  for (unsigned i = 0; i < tensorRank - vecRank; ++i) {
+    if (sizes[i].is<Value>() || getIntFromAttr(sizes[i].get<Attribute>()) != 1)
+      return failure();
+  }
+  for (unsigned i = tensorRank - vecRank; i < tensorRank; ++i) {
+    if (sizes[i].is<Value>() ||
+        getIntFromAttr(sizes[i].get<Attribute>()) !=
+            vecType.getDimSize(i + vecRank - tensorRank))
+      return failure();
+  }
+
+  // Read is out-of-bounds and will be padded.
+  SmallVector<bool> outOfBounds(vecRank, false);
+  auto readMap = AffineMapAttr::get(rewriter.getMultiDimIdentityMap(vecRank));
+  // Assuming that low indices of PadTensorOp are all %c0. Must use a different
+  // starting point + masking for the vector read when the pattern is extended.
+  auto read = rewriter.create<vector::TransferReadOp>(
+      padOp.getLoc(), vecType, padOp.source(), padOp.low(), readMap, padValue,
+      /*mask=*/Value(), rewriter.getBoolArrayAttr(outOfBounds));
+
+  // Compute indices of TransferWriteOp.
+  SmallVector<Value> indices;
+  llvm::for_each(insertOp.getMixedOffsets(), [&](auto o) {
+    if (o.template is<Value>()) {
+      indices.push_back(o.template get<Value>());
+    } else {
+      // Convert int64 attr to index attr.
+      auto intAttr =
+          rewriter.getIndexAttr(getIntFromAttr(o.template get<Attribute>()));
+      indices.push_back(rewriter.create<ConstantOp>(
+          padOp.getLoc(), rewriter.getIndexType(), intAttr));
+    }
+  });
+
+  // Write is fully in-bounds.
+  SmallVector<bool> inBounds(vecRank, true);
+  // Write to the most minor dimensions of the tensor.
+  auto writeMap = AffineMapAttr::get(AffineMap::getMinorIdentityMap(
+      tensorRank, vecRank, rewriter.getContext()));
+  rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
+      insertOp, insertOp.getType(), read.getResult(), insertOp.dest(), indices,
+      writeMap, /*mask=*/Value(), rewriter.getBoolArrayAttr(inBounds));
+
+  return success();
+}
+
+/// Rewrite use of PadTensorOp result in TransferReadOp. E.g.:
+/// ```
+/// %0 = linalg.pad_tensor %src ... : tensor<?x?xf32> to tensor<17x5xf32>
+/// %r = vector.transfer_read %0[%c0, %c0], %cst
+///     {in_bounds = [true, true]} : tensor<17x5xf32>, vector<17x5xf32>
+/// ```
+/// is rewritten to:
+/// ```
+/// %r = vector.transfer_read %src[%c0, %c0], %padding
+///     : tensor<?x?xf32>, vector<17x5xf32>
+/// ```
+/// Note: By restricting this pattern to in-bounds TransferReadOps, we can be
+/// sure that the original padding value %cst was never used.
+LogicalResult PadTensorOpVectorizationPattern::tryVectorizeUser(
+    PatternRewriter &rewriter, linalg::PadTensorOp padOp,
+    vector::TransferReadOp xferOp, Value padValue) const {
+  if (xferOp.hasOutOfBoundsDim())
+    return failure();
+  if (xferOp.mask())
+    return failure();
+
+  rewriter.updateRootInPlace(xferOp, [&]() {
+    SmallVector<bool> inBounds(xferOp.getVectorType().getRank(), false);
+    xferOp->setAttr(xferOp.getInBoundsAttrName(),
+                    rewriter.getBoolArrayAttr(inBounds));
+    xferOp.sourceMutable().assign(padOp.source());
+    xferOp.paddingMutable().assign(padValue);
+  });
+
+  return success();
+}
+
+/// Rewrite use of PadTensorOp result in TransferWriteOp.
+/// This pattern rewrites TransferWriteOps that write to a padded tensor value,
+/// where the same amount of padding is immediately removed again after the
+/// write. In such cases, the TransferWriteOp can write to the non-padded tensor
+/// value and apply out-of-bounds masking. E.g.:
+/// ```
+/// %0 = subtensor ...[...] [%s0, %s1] [1, 1] : tensor<...> to tensor<?x?xf32>
+/// %1 = linalg.pad_tensor %0 ... : tensor<?x?xf32> to tensor<17x5xf32>
+/// %2 = vector.transfer_write %vec, %1[...]
+///     : vector<17x5xf32>, tensor<17x5xf32>
+/// %r = subtensor %2[0, 0] [%s0, %s1] [1, 1]
+///     : tensor<17x5xf32> to tensor<?x?xf32>
+/// ```
+/// is rewritten to:
+/// ```
+/// %0 = subtensor ...[...] [%s0, %s1] [1, 1] : tensor<...> to tensor<?x?xf32>
+/// %r = vector.transfer_write %vec, %0[...] : vector<17x5xf32>, tensor<?x?xf32>
+/// ```
+/// Note: It is important that the SubTensorOp %r resizes the result of the
+/// TransferWriteOp to the same size as the input of the TensorPadOp (or an even
+/// smaller size). Otherwise, %r's new (dynamic) dimensions would differ from
+/// %r's old dimensions.
+LogicalResult PadTensorOpVectorizationPattern::tryVectorizeUser(
+    PatternRewriter &rewriter, linalg::PadTensorOp padOp,
+    vector::TransferWriteOp xferOp) const {
+  OpBuilder::InsertionGuard guard(rewriter);
+
+  // TransferWriteOp result must be directly consumed by a SubTensorOp.
+  if (!xferOp->hasOneUse())
+    return failure();
+  auto trimPadding = dyn_cast<SubTensorOp>(*xferOp->user_begin());
+  if (!trimPadding)
+    return failure();
+  // Get size of input of PadTensorOp. (Assuming that it is a SubTensorOp.)
+  auto *padInput = padOp.source().getDefiningOp();
+  if (!padInput)
+    return failure();
+  auto subtensorOp = dyn_cast<SubTensorOp>(padInput);
+  if (!subtensorOp)
+    return failure();
+  // Sizes of both SubTensorOps must be identical.
+  if (subtensorOp.sizes().size() != trimPadding.sizes().size())
+    return failure();
+  if (llvm::any_of(llvm::zip(subtensorOp.sizes(), trimPadding.sizes()),
+                   [](auto d) { return std::get<0>(d) != std::get<1>(d); }))
+    return failure();
+
+  rewriter.setInsertionPoint(xferOp);
+  SmallVector<bool> inBounds(xferOp.getVectorType().getRank(), false);
+  auto newXferOp = rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
+      xferOp, subtensorOp.getType(), xferOp.vector(), subtensorOp,
+      xferOp.indices(), xferOp.permutation_mapAttr(), xferOp.mask(),
+      rewriter.getBoolArrayAttr(inBounds));
+  rewriter.replaceOp(trimPadding, newXferOp->getResult(0));
+
+  return success();
+}
+
 // TODO: cleanup all the convolution vectorization patterns.
 template <class ConvOp, int N>
 LogicalResult ConvOpVectorization<ConvOp, N>::matchAndRewrite(
     ConvOp op, PatternRewriter &rewriter) const {
   Location loc = op.getLoc();
   MLIRContext *context = op.getContext();
 
   ShapedType inShapeType = op.getInputShapedType(0);