[mlir][Linalg] Add a padding option to Linalg tiling

[nicolasvasilache]

Summary

This revision allows the base Linalg tiling pattern to optionally require padding to
a constant bounding shape.
When requested, a simple analysis is performed, similar to buffer promotion.
A temporary linalg.simple_pad op is added to model padding for the purpose of
connecting the dots. This will be replaced by a more fleshed out linalg.pad_tensor
op when it is available.
In the meantime, this temporary op serves the purpose of exhibiting the necessary
properties required from a more fleshed out pad op, to compose with transformations
properly.

[rriddle]

mlir/include/mlir/Interfaces/ViewLikeInterface.h

25

This typedef pollutes the entire mlir namespace.

[nicolasvasilache]

mlir/include/mlir/Interfaces/ViewLikeInterface.h

25

I can just use OpFoldResult.
OTOH, I'd welcome renaming OpFoldResult to ValueOrAttr which is much more descriptive IMO.
Thoughts?

[nicolasvasilache]

Drop ValueOrAttr.

[nicolasvasilache]

Attempt to appease win.

[ftynse]

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

384–388

It is still helpful to document the requirements and conventions of this op, instead of it being just a placeholder with unknown semantics.

393–395

Just drop the verifier field.

mlir/include/mlir/Interfaces/ViewLikeInterface.h

17–22

Leftover unnecessary changes.

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

116

We usually pass builder by-reference. Do you need a copy for some reason?

126

Nit: SmallVector can now infer the default number of stack elements, so if there is no reason to require 4 specifically, it can be just omitted.

133

size.get<Attribute> (there's no need to dyn_cast after a check).

138–139

I'm not a fan of emitting warnings _inside_ transformations. IMO, errors and warnings are there for the validity state of the op in general, not for its suitability to a particular transformation. Remarks may be fine (the user will have to ask for them explicitly to see them). rewriter.notifyMatchFailure() is probably preferred.

150

Write this comment? :)

[hanchung]

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

122

use auto since getDefiningOp already states the type.

151–153

Please add documentation because this is not a trivial function, e.g., the result will be stored to res if success.

[mravishankar]

This feels like it is promotion on tensors with the linalg.simple_pad (for now and linalg.pad later) being used instead of the alloc(...) ; copy(...) method. So the tensor "promotion" and buffer promotion are living in entirely different places. Ill go deeper into this code, but my read of the tests informs this comment.

[nicolasvasilache]

Address comments.

[nicolasvasilache]

In D95149#2516098, @mravishankar wrote:

This feels like it is promotion on tensors with the linalg.simple_pad (for now and linalg.pad later) being used instead of the alloc(...) ; copy(...) method. So the tensor "promotion" and buffer promotion are living in entirely different places. Ill go deeper into this code, but my read of the tests informs this comment.

There are similarities and differences.
Promotion is guaranteed to alloc and copy but is not guaranteed to create static shapes.

This in contrast has no guarantees of alloc/copy and is guaranteed to create static shapes.
This is meant to be used mostly as an enabler for vectorization as well as packing.
Packing is more closely related to promotion with the big difference that we keep SSA use-def chains and hoisting is significantly simpler to implement.

In any case I expect we'll need to experiment before we can come to a conclusion, this gives us the ability to experiment.

[mravishankar]

Promotion is guaranteed to alloc and copy but is not guaranteed to create static shapes.
This in contrast has no guarantees of alloc/copy and is guaranteed to create static shapes.

That seems to be a small difference. I agree that promotion itself doesn't gaurantee static shape but it does do a best effort (and actually in IREE it always ends up being static shaped if the tile size is static). It should be easy to extend promotion to abort if it doesnt figure out the static sizes when "promotion" is applied to ops with tensor semantics.

[nicolasvasilache]

It should be easy to extend promotion to abort if it doesnt figure out the static sizes when "promotion" is applied to ops with tensor semantics.

The fact that promotion tries to create static buffers is very much due to history and needs to go away.

Here is what happens concretely today to inject static sizes in the memory pipeline:

1. Promotion creates alloc + fill + copy after some best effort local analysis to try and create a static buffer.
2. Fingers are crossed so that vectorization succeeds.
3. After vectorization, LinalgCopyVTRForwardingPattern and LinalgCopyVTWForwardingPattern are applied with a best effort dependence analysis to undo 1. and replace by vector.transfer.
4. Later masked/unmasked vector.transfer splitting is performed to isolate parts.

In contrast, thanks to the progress on transformations on tensors we can:

1. Tile on tensors with padding at the level(s) we are interested in.
2. Hoist, as we wish to create and amortize packing.
3. Decide which packed tensors materialize in memory, which ones turn into simple subview + control flow and which ones turn into vector transfers.

The implications of padding at the tiling on tensor level are far reaching as you have already seen in IREE, phases are being reordered and Promotion is a bystanding casualty.
See the followup https://reviews.llvm.org/D95243 to get a better idea how things look like after padding + packing.
I would not be surprised that Promotion just gets deleted / becomes part of some bufferization in a not too distant future: the ease with which we can perform transformations thanks to SSA use-def chains is just a no-brainer.

[nicolasvasilache]

Attempt to appeas win, again..

[nicolasvasilache]

Attempt to appease win, 3rd take..

[mravishankar]

In D95149#2518130, @nicolasvasilache wrote:

It should be easy to extend promotion to abort if it doesnt figure out the static sizes when "promotion" is applied to ops with tensor semantics.

The fact that promotion tries to create static buffers is very much due to history and needs to go away.

Here is what happens concretely today to inject static sizes in the memory pipeline:

1. Promotion creates alloc + fill + copy after some best effort local analysis to try and create a static buffer.
2. Fingers are crossed so that vectorization succeeds.
3. After vectorization, LinalgCopyVTRForwardingPattern and LinalgCopyVTWForwardingPattern are applied with a best effort dependence analysis to undo 1. and replace by vector.transfer.
4. Later masked/unmasked vector.transfer splitting is performed to isolate parts.

In contrast, thanks to the progress on transformations on tensors we can:

1. Tile on tensors with padding at the level(s) we are interested in.
2. Hoist, as we wish to create and amortize packing.
3. Decide which packed tensors materialize in memory, which ones turn into simple subview + control flow and which ones turn into vector transfers.

The implications of padding at the tiling on tensor level are far reaching as you have already seen in IREE, phases are being reordered and Promotion is a bystanding casualty.
See the followup https://reviews.llvm.org/D95243 to get a better idea how things look like after padding + packing.
I would not be surprised that Promotion just gets deleted / becomes part of some bufferization in a not too distant future: the ease with which we can perform transformations thanks to SSA use-def chains is just a no-brainer.

I am not disputing that this is useful, but at least within IREE, I am not sure we can actually set a static tile size at the time of tiling. In IREE with the linalg on tensors path the tile size has to by dynamic during tile and distribute since it is done at the Flow dialect level where no backend specific information is available. The static tile size picked while tiling might not be one that is best for the backend that is targeted (information that is available only at the HAL level). So making this padding as part of the tiling transformation might not work out for IREE. My understanding then is that if we separate out the padding parts from the tiling transformation itself then it has the same constraints/restrictions as promotion, but with the advantage of SSA use-def chains.

[nicolasvasilache]

I am not disputing that this is useful, but at least within IREE, I am not sure we can actually set a static tile size at the time of tiling. In IREE with the linalg on tensors path the tile size has to by dynamic during tile and distribute since it is done at the Flow dialect level where no backend specific information is available. The static tile size picked while tiling might not be one that is best for the backend that is targeted (information that is available only at the HAL level). So making this padding as part of the tiling transformation might not work out for IREE. My understanding then is that if we separate out the padding parts from the tiling transformation itself then it has the same constraints/restrictions as promotion, but with the advantage of SSA use-def chains.

In the particular context of IREE, I see 2 things we can do easily:

1. Pad at the level of tile and fuse and vectorize on tensors at the HAL level to enable vectorization (i.e. at the k^-th level of padding)
2. If we have a need for it, this is not hard to generalize to "pad to the next multiple of X" where X can either be a constant or an SSA value.

The fact that the option exists to pad does not mean you have to do it at the same time as "tile-and-distribute."

I wouldn't pattern match this first implementation to how general the approach is: we also want to use it in the context of sparse, so I am really not scared of some dynamic values here and there.
After we connect this to linalg.pad_tensors and have gained some experience, we can decide where we want to push this as we will also need better types to support more advanced behaviors.

Still, I don't see promotion evolving to support tensors, which is the suggestion I seemed to get from your previous comment?

[nicolasvasilache]

Fix use after move.

[nicolasvasilache]

Undo bad fix and fix use-after-move properly.

[mravishankar]

THanks Nicolas for the explanation. Code looks fine to me. Still think that this is effectively adding promotion during tiling and doing this on tensors only is strange. Just trying to see if there is a more commonality between buffers and tensors so as to reuse things instead of completely new paths. Lets land and experiment. I understand how this would subsume promotion though.

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

383

Do we still need this. I think the linalg.pad operation landed already

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

183

Cant we just use shapedOutputOperands for this? Maybe just add a method for getting the shaped output types?

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

SmallVector<Value, 8> getRanges() {
           res.push_back(operand);
       return res;
     }
   }];
 
   let hasFolder = 1;
 }
 
+def Linalg_SimplePadOp : Linalg_Op<"simple_pad", [NoSideEffect]> {
+  let summary = "TODO: replace with pad_tensors when ready.";
+
+  let description = [{
+    `linalg.simple_pad` is a tmp placeholder for padding and packing on tensors.
+  }];
+
+  let arguments = (ins AnyRankedTensor:$tensor, AnyType:$padding);
+  let results = (outs AnyRankedTensor:$result);
+
+  // TODO: all static result, some dynamic input, static shapes match,
+  // element types match, ranks match etc. Use pad_tensors when ready.
+  let verifier = [{ return success(); }];
+
+  let extraClassDeclaration = [{
+    RankedTensorType getSourceType() {
+      return tensor().getType().cast<RankedTensorType>(); }
+    RankedTensorType getResultType() {
+      return getResult().getType().cast<RankedTensorType>(); }
+   }];
+
+  let assemblyFormat = [{
+    $tensor `pad` $padding attr-dict `:`
+      type($tensor) `to` type($result) `pad` type($padding)
+  }];
+}
+
 def Linalg_YieldOp : Linalg_Op<"yield", [NoSideEffect, ReturnLike, Terminator]>,
     Arguments<(ins Variadic<AnyType>:$values)> {
   let summary = "Linalg yield operation";
   let description = [{
     `linalg.yield` is a special terminator operation for blocks inside regions
     in `linalg` generic ops. It returns values to the immediately enclosing
     `linalg` generic op.
 

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

enum class LinalgTilingLoopType {
   Loops = 0,
   AffineLoops = 1,
   ParallelLoops = 2,
 };
 
 using TileSizeComputationFunction =
     std::function<SmallVector<Value, 4>(OpBuilder &, Operation *)>;
 
+using PaddingValueComputationFunction =
+    std::function<Value(OpBuilder &, Operation *)>;
+
 struct LinalgTilingOptions {
   /// Computation function that returns the tile sizes for each operation.
   /// Delayed construction of constant tile sizes should occur to interoperate
   /// with folding.
   TileSizeComputationFunction tileSizeComputationFunction = nullptr;
 
   LinalgTilingOptions &
   setTileSizeComputationFunction(TileSizeComputationFunction fun) {
struct LinalgTilingOptions {
   /// processors.
   Optional<LinalgLoopDistributionOptions> distribution = None;
 
   LinalgTilingOptions &
   setDistributionOptions(LinalgLoopDistributionOptions distributionOptions) {
     distribution = std::move(distributionOptions);
     return *this;
   }
+
+  /// Computation function that returns a padding value to use when padding to
+  /// force static sizes. When `paddingValueComputationFunction` is set, padding
+  /// operations are introduced, that guarantee the underlying op is statically
+  /// shaped and can thus be vectorized.
+  PaddingValueComputationFunction paddingValueComputationFunction = nullptr;
+
+  LinalgTilingOptions &
+  setPaddingValueComputationFunction(PaddingValueComputationFunction fun) {
+    paddingValueComputationFunction = std::move(fun);
+    return *this;
+  }
 };
 
 /// Canonicalization patterns relevant to apply after tiling patterns. These are
 /// applied automatically by the tiling pass but need to be applied manually
 /// when tiling is called programmatically.
 OwningRewritePatternList
 getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx);
 void populateLinalgTilingCanonicalizationPatterns(
     OwningRewritePatternList &patterns, MLIRContext *ctx);
 
+/// Base pattern that applied the tiling transformation specified by `options`.
+/// Abort and return failure in 2 cases:
+///   1. if the tiling specification is invalid and tiling fails to occur.
+///   2. if tiling occurs but `options.paddingValueComputationFunction` is set
+///      and some operand shape cannot be bounded statically.
 struct LinalgBaseTilingPattern : public RewritePattern {
   // Entry point to match any LinalgOp OpInterface.
   LinalgBaseTilingPattern(LinalgTilingOptions options,
                           LinalgMarker marker = LinalgMarker(),
                           PatternBenefit benefit = 1);
   // Entry point to match a specific Linalg op.
   LinalgBaseTilingPattern(StringRef opName, MLIRContext *context,
                           LinalgTilingOptions options,

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 //===- ViewLikeInterface.h - View-like operations interface ---------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the operation interface for view-like operations.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
 #define MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
 
 #include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpImplementation.h"
 
 namespace mlir {
+
+/// Represent a single Value or Attr to simplify the manipulation of mixed
+/// static and dynamic quantities.
+using ValueOrAttr = llvm::PointerUnion<Attribute, Value>;
+
 /// Auxiliary range data structure to unpack the offset, size and stride
 /// operands into a list of triples. Such a list can be more convenient to
 /// manipulate.
 struct Range {
   Value offset;
   Value size;
   Value stride;
 };

mlir/include/mlir/Interfaces/ViewLikeInterface.td

InterfaceMethod<
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.sizes();
       }]
     >,
     InterfaceMethod<
       /*desc=*/[{
+        Return a vector of all the static or dynamic sizes of the op.
+      }],
+      /*retTy=*/"SmallVector<ValueOrAttr, 4>",
+      /*methodName=*/"getMixedSizes",
+      /*args=*/(ins),
+      /*methodBody=*/"",
+      /*defaultImplementation=*/[{
+        SmallVector<ValueOrAttr, 4> res;
+        std::array<unsigned, 3> ranks = $_op.getArrayAttrRanks();
+        unsigned numDynamic = 0;
+        unsigned count = ranks[getOffsetOperandGroupPosition()];
+        for (unsigned idx = 0; idx < count; ++idx) {
+          if (isDynamicSize(idx))
+            res.push_back($_op.sizes()[numDynamic++]);
+          else
+            res.push_back($_op.static_sizes()[idx]);
+        }
+        return res;
+      }]
+    >,
+    InterfaceMethod<
+      /*desc=*/[{
         Return the dynamic stride operands.
       }],
       /*retTy=*/"OperandRange",
       /*methodName=*/"strides",
       /*args=*/(ins),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.strides();
InterfaceMethod<
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         return $_op.getOperand(getIndexOfDynamicStride(idx));
       }]
     >,
   ];
 
   let extraClassDeclaration = [{
+    static unsigned getOffsetOperandGroupPosition() { return 0; }
+    static unsigned getSizeOperandGroupPosition() { return 1; }
+    static unsigned getStrideOperandGroupPosition() { return 2; }
     static StringRef getStaticOffsetsAttrName() {
       return "static_offsets";
     }
     static StringRef getStaticSizesAttrName() {
       return "static_sizes";
     }
     static StringRef getStaticStridesAttrName() {
       return "static_strides";

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

 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/EDSC/Intrinsics.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/ScopeExit.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <type_traits>
 
 #define DEBUG_TYPE "linalg-transforms"
 
 using namespace mlir;
 using namespace mlir::edsc;
tileSizeComputationFunction = [tileSizes](OpBuilder &b, Operation *op) {
     return llvm::to_vector<4>(map_range(tileSizes, [&](int64_t s) {
       Value v = b.create<ConstantIndexOp>(op->getLoc(), s);
       return v;
     }));
   };
   return *this;
 }
 
+/// Try to compute a static bounding box for `operand`
+/// Return success if either:
+///   1. The operand is already statically shaped, `result` is left unchanged.
+///   2. The operand is (partially) dynamic, `result` is the result of a freshly
+///      created SimplePadOp.
+/// Return failure if the operand cannot be padded to a static shape.
+static LogicalResult padOperandToSmallestStaticBoundingBox(
+    OpBuilder b, linalg::LinalgOp opToPad, Value operand,
+    const LinalgTilingOptions &options, Value &result) {
+  auto tensorType = operand.getType().cast<RankedTensorType>();
+  // Already static shape, no need to pad.
+  if (tensorType.hasStaticShape())
+    return success();
+  SubTensorOp subtensor = operand.getDefiningOp<SubTensorOp>();
+  // Not a subtensor, cannot construct a static bounding box.
+  if (!subtensor)
+    return failure();
+  SmallVector<int64_t, 4> staticSizes;
+  staticSizes.reserve(tensorType.getRank());
+  auto shapedOp =
+      cast<OffsetSizeAndStrideOpInterface>(subtensor.getOperation());
+  for (auto size : shapedOp.getMixedSizes()) {
+    auto indexAttr =
+        size.is<Attribute>()
+            ? size.dyn_cast<Attribute>().dyn_cast<IntegerAttr>()
+            : linalg::getSmallestBoundingIndex(size.dyn_cast<Value>());
+    // SmallestBoundingIndex must exist for all sizes.
+    // For now return an error if we can't find it.
+    if (!indexAttr)
+      return opToPad->emitWarning(
+          "No constant bounding box can be found for padding");
+    staticSizes.push_back(indexAttr.getInt());
+  }
+  Value pad = options.paddingValueComputationFunction(b, opToPad);
+  auto staticTensorType =
+      RankedTensorType::get(staticSizes, tensorType.getElementType());
+  result = b.create<linalg::SimplePadOp>(opToPad->getLoc(), staticTensorType,
+                                         operand, pad);
+  return success();
+}
+
+//
+static LogicalResult rewriteAsPaddedOp(PatternRewriter &rewriter,
+                                       TiledLinalgOp &res,
+                                       const LinalgTilingOptions &options) {
+  LinalgOp opToPad = res.op;
+  Location loc = opToPad->getLoc();
+
+  // If the op is fully static, it does not need padding.
+  // TODO: there are cases where we may still want to pad to larger sizes.
+  if (llvm::all_of(opToPad.getShapedOperands(), [](Value v) {
+        return v.getType().cast<RankedTensorType>().hasStaticShape();
+      }))
+    return success();
+
+  OpBuilder::InsertionGuard g(rewriter);
+  // Set IP after op because we also take the dims of the original output.
+  rewriter.setInsertionPointAfter(opToPad);
+  // Make a copy of the shaped operands and update it.
+  SmallVector<Value, 4> operands = opToPad.getShapedOperands();
+  for (Value &v : operands) {
+    Value paddedOperand;
+    // If padding was requested but the shape cannot be bounded statically then
+    // the pattern fails to apply.
+    if (failed(padOperandToSmallestStaticBoundingBox(rewriter, opToPad, v,
+                                                     options, paddedOperand))) {
+      return failure();
+    }
+    // Update v if we indeed got a padded operand.
+    v = paddedOperand ? paddedOperand : v;
+  }
+
+  // Clone `opToPad` to operate on the statically padded shapes.
+  TypeRange resultTensorTypes =
+      ValueRange(operands).take_back(opToPad.getNumOutputs());
+  ValueRange otherOperands = opToPad.getAssumedNonShapedOperands();
+  operands.append(otherOperands.begin(), otherOperands.end());
+  linalg::LinalgOp paddedOp =
+      opToPad.clone(rewriter, loc, resultTensorTypes, operands);
+
+  // Recover the subtensor out of the new static results. This keeps the
+  // original linalg op around because it uses the dims of the original results.
+  // This later folds away.
+  SmallVector<Value, 4> paddedSubviewResults;
+  paddedSubviewResults.reserve(opToPad->getNumResults());
+  Value zero = rewriter.create<ConstantIndexOp>(loc, 0);
+  Value one = rewriter.create<ConstantIndexOp>(loc, 1);
+  llvm::SetVector<Operation *> newUsersOfOpToPad;
+  for (auto it : llvm::zip(opToPad->getResults(), paddedOp->getResults())) {
+    auto rank = std::get<0>(it).getType().cast<RankedTensorType>().getRank();
+    SmallVector<Value, 4> offsets(rank, zero);
+    auto sizes = llvm::to_vector<4>(
+        llvm::map_range(llvm::seq<unsigned>(0, rank), [&](unsigned d) -> Value {
+          auto dimOp = rewriter.create<DimOp>(loc, std::get<0>(it), d);
+          newUsersOfOpToPad.insert(dimOp);
+          return dimOp;
+        }));
+    SmallVector<Value, 4> strides(rank, one);
+    paddedSubviewResults.push_back(rewriter.create<SubTensorOp>(
+        loc, std::get<1>(it), offsets, sizes, strides));
+  }
+  // Replace the transient `opToPad` locally, except for uses that we just
+  // created for the purpose of extracting the dims.
+  rewriter.replaceOpWithIf(opToPad, paddedSubviewResults, [&](OpOperand &opOp) {
+    return !newUsersOfOpToPad.contains(opOp.getOwner());
+  });
+
+  res = TiledLinalgOp{paddedOp, res.loops, res.tensorResults};
+  return success();
+}
+
 /// Linalg base tiling pattern.
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     StringRef opName, MLIRContext *context, LinalgTilingOptions options,
     LinalgMarker marker, PatternBenefit benefit)
     : RewritePattern(opName, {}, benefit, context), marker(marker),
       options(options) {}
 
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
   if (failed(marker.checkAndNotify(rewriter, linalgOp)))
     return failure();
 
   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
 
   if (!res)
     return failure();
 
+  // Setup RAII guard to return properly.
+  bool succeeded = true;
+  LinalgOp tiledOp = res->op;
+  auto guard = llvm::make_scope_exit([&]() {
+    if (!succeeded)
+      return;
     // Return relevant information to derived pattern.
     result = *res;
+    // Replace marker on both tiledOp and tiledAndPaddedOp, if necessary.
+    marker.replaceLinalgMarker(rewriter, tiledOp);
+    if (tiledOp != res->op)
+      marker.replaceLinalgMarker(rewriter, res->op);
+  });
 
-  // New marker if specified.
-  marker.replaceLinalgMarker(rewriter, res->op.getOperation());
+  // Consider padding on the fly only if the op has tensor semantics.
+  if (!options.paddingValueComputationFunction ||
+      !linalgOp.hasTensorSemantics())
+    return success();
+
+  // Try to pad on the fly by rewriting res->op as a padded op.
+  if (failed(rewriteAsPaddedOp(rewriter, *res, options))) {
+    // Set so RAII guard does not propagate TiledLinalgOp to `result`.
+    succeeded = false;
+    return failure();
+  }
+
+  // Do not perform replacement of `linalgOp`, let the derived patterns
+  // do this as they see fit, from the resulting TiledLinalgOp.
   return success();
 }
 
 mlir::linalg::LinalgBaseTileAndFusePattern::LinalgBaseTileAndFusePattern(
     StringRef opName, MLIRContext *context,
     const LinalgDependenceGraph &dependenceGraph,
     LinalgTilingOptions tilingOptions, LinalgFusionOptions fusionOptions,
     LinalgMarker marker, LinalgMarker fusedOpMarker,

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

if (auto fromElements = dyn_cast_or_null<tensor::GenerateOp>(definingOp)) {
     auto dynExtents = fromElements.dynamicExtents().begin();
     for (auto dim : resultType.getShape().take_front(index.getInt()))
       if (dim == RankedTensorType::kDynamicSize)
         dynExtents++;
 
     return Value{*dynExtents};
   }
 
+  // The size at the given index is now known to be a dynamic size.
+  unsigned unsignedIndex = index.getValue().getZExtValue();
+
+  if (auto subtensor = dyn_cast_or_null<SubTensorOp>(definingOp)) {
+    assert(subtensor.isDynamicSize(unsignedIndex) &&
+           "Expected dynamic subtensor size");
+    return subtensor.getDynamicSize(unsignedIndex);
+  }
+
   // Fold dim to the size argument for an `AllocOp`, `ViewOp`, or `SubViewOp`.
   auto memrefType = argTy.dyn_cast<MemRefType>();
   if (!memrefType)
     return {};
 
-  // The size at the given index is now known to be a dynamic size of a memref.
-  unsigned unsignedIndex = index.getValue().getZExtValue();
   if (auto alloc = dyn_cast_or_null<AllocOp>(definingOp))
     return *(alloc.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 
   if (auto view = dyn_cast_or_null<ViewOp>(definingOp))
     return *(view.getDynamicSizes().begin() +
              memrefType.getDynamicDimIndex(unsignedIndex));
 

mlir/test/Dialect/Linalg/roundtrip.mlir

%0 = linalg.reshape %arg0
     memref<?x?x?x4x?xf32> into memref<?x?x?xf32>
   return %0 : memref<?x?x?xf32>
 }
 // CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0)>
 // CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1)>
 // CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3, d4)>
 //     CHECK: func @legal_collapsing_reshape_dynamic_memref
 //     CHECK:   linalg.reshape %{{.+}} [#[[MAP0]], #[[MAP1]], #[[MAP2]]]
+
+// -----
+
+// TODO: this op should disappear once pad_tensors is available.
+// CHECK-LABEL: func @simple_pad
+func @simple_pad(%0: tensor<?x4x?xf32>, %pad: f32) {
+//     CHECK:   linalg.simple_pad %{{.+}} pad %{{.+}}: tensor<?x4x?xf32> to tensor<8x4x8xf32>
+  %1 = linalg.simple_pad %0 pad %pad: tensor<?x4x?xf32> to tensor<8x4x8xf32> pad f32
+  return
+}

mlir/test/Dialect/Linalg/tile-and-pad-tensors.mlir

+// RUN: mlir-opt %s -test-linalg-transform-patterns=test-tile-and-pad-pattern -canonicalize | FileCheck %s
+
+// CHECK-LABEL: func @matmul_tensors(
+// CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<?x?xf32>
+// CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<?x?xf32>
+// CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<?x?xf32>) -> tensor<?x?xf32> {
+func @matmul_tensors(
+  %arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>, %arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32> {
+//      CHECK: %[[TD0:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC0:.*]] = %[[TC]]) -> (tensor<?x?xf32>) {
+//      CHECK:   %[[TD1:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC1:.*]] = %[[TC0]]) -> (tensor<?x?xf32>) {
+//      CHECK:     %[[TD2:.*]] = scf.for {{.*}} to {{.*}} step {{.*}} iter_args(%[[TC2:.*]] = %[[TC1]]) -> (tensor<?x?xf32>) {
+//      CHECK:       %[[sTA:.*]] = subtensor %[[TA]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
+//      CHECK:       %[[sTB:.*]] = subtensor %[[TB]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
+//      CHECK:       %[[sTC:.*]] = subtensor %[[TC2]][{{.*}}] : tensor<?x?xf32> to tensor<?x?xf32>
+
+// Dynamic op has been canonicalized away.
+//  CHECK-NOT:       linalg.matmul {{.*}} tensor<?x?xf32>
+
+// Padding injects static information.
+//      CHECK:       %[[pA:.*]] = linalg.simple_pad %[[sTA]] pad %{{.*}} : tensor<?x?xf32> to tensor<2x4xf32> pad f32
+//      CHECK:       %[[pB:.*]] = linalg.simple_pad %[[sTB]] pad %{{.*}} : tensor<?x?xf32> to tensor<4x3xf32> pad f32
+//      CHECK:       %[[pC:.*]] = linalg.simple_pad %[[sTC]] pad %{{.*}} : tensor<?x?xf32> to tensor<2x3xf32> pad f32
+//      CHECK:       %[[pD:.*]] = linalg.matmul ins(%[[pA]], %[[pB]] : tensor<2x4xf32>, tensor<4x3xf32>)
+// CHECK-SAME:                                  outs(%[[pC]] : tensor<2x3xf32>)  -> tensor<2x3xf32>
+//      CHECK:       %[[sTD:.*]] = subtensor %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xf32> to tensor<?x?xf32>
+//      CHECK:       %[[TD:.*]] = subtensor_insert %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<?x?xf32> into tensor<?x?xf32>
+//      CHECK:       scf.yield %[[TD]] : tensor<?x?xf32>
+//      CHECK:     scf.yield %[[TD2]] : tensor<?x?xf32>
+//      CHECK:   scf.yield %[[TD1]] : tensor<?x?xf32>
+  %0 = linalg.matmul {__internal_linalg_transform__ = "tile-and-pad"}
+      ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+     outs(%arg2: tensor<?x?xf32>)
+    -> tensor<?x?xf32>
+
+//      CHECK: return %[[TD0]] : tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

Option<bool> testGenericToVectorPattern{
       *this, "test-linalg-to-vector-patterns",
       llvm::cl::desc("Test a set of patterns that rewrite a linalg contraction "
                      "in vector.contract form"),
       llvm::cl::init(false)};
   Option<bool> testAffineMinSCFCanonicalizationPatterns{
       *this, "test-affine-min-scf-canonicalization-patterns",
       llvm::cl::desc("Test affine-min + scf canonicalization patterns."),
       llvm::cl::init(false)};
+  Option<bool> testTileAndPadPattern{
+      *this, "test-tile-and-pad-pattern",
+      llvm::cl::desc("Test tile and pad pattern"), llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   OwningRewritePatternList patterns;
 
   //===--------------------------------------------------------------------===//
static void applyAffineMinSCFCanonicalizationPatterns(FuncOp funcOp) {
   FrozenRewritePatternList frozenPatterns(std::move(foldPattern));
 
   // Explicitly walk and apply the pattern locally to avoid more general folding
   // on the rest of the IR.
   funcOp.walk([&frozenPatterns](AffineMinOp minOp) {
     applyOpPatternsAndFold(minOp, frozenPatterns);
   });
 }
+
+// For now, just assume it is the zero of type.
+// In the future, it should be the zero of type + op.
+static Value getNeutralOfLinalgOp(OpBuilder &b, Operation *op) {
+  auto t = op->getResult(0).getType().cast<ShapedType>().getElementType();
+  return b.create<ConstantOp>(op->getLoc(), t, b.getZeroAttr(t));
+}
+
+static void applyTileAndPadPattern(FuncOp funcOp) {
+  MLIRContext *context = funcOp.getContext();
+  OwningRewritePatternList tilingPattern;
+  auto linalgTilingOptions =
+      linalg::LinalgTilingOptions()
+          .setTileSizes({2, 3, 4})
+          .setPaddingValueComputationFunction(getNeutralOfLinalgOp);
+  tilingPattern.insert<linalg::LinalgTilingPattern<linalg::MatmulOp>>(
+      context, linalgTilingOptions,
+      linalg::LinalgMarker(Identifier::get("tile-and-pad", context)));
+  applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
+
+  // Explicitly walk and apply the pattern locally to avoid more general folding
+  // on the rest of the IR.
+  FrozenRewritePatternList frozenPatterns(std::move(tilingPattern));
+  funcOp.walk([&frozenPatterns](AffineMinOp minOp) {
+    applyOpPatternsAndFold(minOp, frozenPatterns);
+  });
+}
+
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {
   auto lambda = [&](void *) {
     getFunction().walk([](LinalgOp op) {
       op.removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
   };
   std::unique_ptr<void, decltype(lambda)> cleanupGuard{(void *)1, lambda};
return applyMatmulToVectorPatterns(getFunction(),
                                        testMatmulToVectorPatterns1dTiling,
                                        testMatmulToVectorPatterns2dTiling);
   if (testVectorTransferForwardingPatterns)
     return applyVectorTransferForwardingPatterns(getFunction());
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getFunction());
   if (testAffineMinSCFCanonicalizationPatterns)
     return applyAffineMinSCFCanonicalizationPatterns(getFunction());
+  if (testTileAndPadPattern)
+    return applyTileAndPadPattern(getFunction());
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms> testTransformPatternsPass(
       "test-linalg-transform-patterns",
       "Test Linalg transformation patterns by applying them greedily.");
 }
 } // namespace test
 } // namespace mlir