[mlir][linalg] Tile PadTensorOp

[springerm]

Summary

Tiling can be enabled with linalg-tile-pad-tensor-ops. Only scf::ForOp can be generated at the moment.

Depends On D105602

[springerm]

small update

[springerm]

more static types

[nicolasvasilache]

mlir/include/mlir/Dialect/Linalg/Passes.td

174

Can this be folded into the existing tiling pass ?

These are pure mechanics and don't have any policy / heuristics logic anyway.
They are not meant to drive anything other than tests.
Moving them to purely test patterns would also be appropriate.

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

47

s/asOpFoldResult/getAsOpFoldResult ?

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

158

s/writeTileToTensor/insertSliceIntoTensor?

371

The part that performs the tiling rewrite logic should be separated in a standalone transformation.

The pattern should call that transformation, similarly to other similar pattterns.
The part about using filters should be completely separated out and put in the same places as the other patterns that use filters.

The main reason is that these filters are meant to go away in the future when we have something better to control pattern composition.

398

This is a good observation that we have also discussed in other generalization contexts.

Viewing tiling as: given a subset of data, give me the subset of the op that produces the data, generalizes better as rewrite extract_slice(op(x, y, z)) into op(f(x), g(y), h(z)).
Bonus points if in a followup revision you are able to update linalg.generic tiling to operate similarly :) .

525

Can these just be folded in the pass above?

[springerm]

rebase

[springerm]

rebase

[springerm]

no change

[springerm]

address comments

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

371

I removed the filter and instead set an attribute manually. This attribute is needed, so that the pattern does not apply infinitely. Basically, this pattern just generates the loop structure. So once the structure is generated, it should not be generated a second time.

What do you mean by "tiling rewrite logic"? Do you mean how to rewrite the PadTensorOp inside the loop? That part is already as minimal as possible. Basically, I just move the PadTensorOp inside the loop nest (by cloning it and replacing the old one with the result of the loop nest). The actual rewrite logic specific to the PadTensorOp is handled by ExtractSliceOfPadTensorSwapPattern.

[springerm]

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

398

Probably requires a larger refactoring, but then it should be possible to use the exact same code path for generating the for-loop for each operation (PadTensorOp and LinalgOp). Looking into this afterwards.

[springerm]

address comments

[springerm]

Moved the tiling functionality into a separate function, as discussed.

[nicolasvasilache]

Cool thanks for this!

[springerm]

fix test

mlir/include/mlir/Dialect/Linalg/Passes.h

 
 std::unique_ptr<Pass> createLinalgFusionOfTensorOpsPass();
 std::unique_ptr<Pass> createFoldReshapeOpsByLinearizationPass();
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
+createLinalgPadTensorTilingPass(ArrayRef<int64_t> tileSizes = {});
+
+std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes = {},
                                   ArrayRef<StringRef> distributionTypes = {});
 
 std::unique_ptr<OperationPass<FuncOp>>
 createLinalgPromotionPass(bool dynamicBuffers, bool useAlloca);

mlir/include/mlir/Dialect/Linalg/Passes.td

let dependentDialects = [
     "scf::SCFDialect"
   ];
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
 }
 
+def LinalgPadTensorOpTiling : FunctionPass<"linalg-tile-pad-tensor-ops"> {
+  let summary = "Tile linalg.pad_tensor operations";
+  let constructor = "mlir::createLinalgPadTensorTilingPass()";
+  let dependentDialects = [
+    "AffineDialect",
+    "linalg::LinalgDialect",
+    "memref::MemRefDialect",
+    "scf::SCFDialect"
+  ];
+  let options = [
+    ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
+               "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
+  ];
+}
+
 def LinalgTilingToParallelLoops
     : FunctionPass<"linalg-tile-to-parallel-loops"> {
   let summary = "Tile operations in the linalg dialect to parallel loops";
   let constructor = "mlir::createLinalgTilingToParallelLoopsPass()";
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t", "Tile sizes",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 void applyPermutationToVector(SmallVector<T, N> &inVec,
                               ArrayRef<unsigned> permutation) {
   SmallVector<T, N> auxVec(inVec.size());
   for (unsigned i = 0; i < permutation.size(); ++i)
     auxVec[i] = inVec[permutation[i]];
   inVec = auxVec;
 }
 
+/// Given a value, try to extract a constant integer as an Attribute.
+/// If this fails, return the original value.
+OpFoldResult asOpFoldResult(OpBuilder &builder, Value val);
+
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
 

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

for (IndexOp indexOp : block.getOps<linalg::IndexOp>()) {
     bindDims(b.getContext(), index, iv);
     AffineApplyOp applyOp = b.create<AffineApplyOp>(
         indexOp.getLoc(), index + iv,
         ValueRange{indexOp.getResult(), ivs[rangeIndex->second]});
     indexOp.getResult().replaceAllUsesExcept(applyOp, applyOp);
   }
 }
 
+// Insert a tile `source` into the destination tensor `dest`. The position at
+// which the tile is inserted (as well as size of tile) is taken from a given
+// ExtractSliceOp `sliceOp`.
+static Value writeTileToTensor(OpBuilder &b, Location loc,
+                               tensor::ExtractSliceOp sliceOp, Value source,
+                               Value dest) {
+  return b.create<tensor::InsertSliceOp>(
+      loc, sliceOp.source().getType(), source, dest, sliceOp.offsets(),
+      sliceOp.sizes(), sliceOp.strides(), sliceOp.static_offsets(),
+      sliceOp.static_sizes(), sliceOp.static_strides());
+}
+
 template <typename LoopTy>
 static Optional<TiledLinalgOp>
 tileLinalgOpImpl(OpBuilder &b, LinalgOp op, ValueRange tileSizes,
                  const LinalgTilingOptions &options) {
   auto nLoops = op.getNumLoops();
   // Initial tile sizes may be too big, only take the first nLoops.
   tileSizes = tileSizes.take_front(nLoops);
 
auto tiledLoopBodyBuilder = [&](OpBuilder &b, Location loc,
 
     // Insert a insert_slice for each output tensor.
     unsigned resultIdx = 0;
     for (OpOperand *opOperand : op.getOutputTensorOperands()) {
       // TODO: use an interface/adaptor to avoid leaking position in
       // `tiledOperands`.
       Value outputTensor = tiledOperands[opOperand->getOperandNumber()];
       if (auto sliceOp = outputTensor.getDefiningOp<tensor::ExtractSliceOp>()) {
-        tensorResults.push_back(b.create<tensor::InsertSliceOp>(
-            loc, sliceOp.source().getType(), res->getResult(resultIdx),
-            sliceOp.source(), sliceOp.offsets(), sliceOp.sizes(),
-            sliceOp.strides(), sliceOp.static_offsets(), sliceOp.static_sizes(),
-            sliceOp.static_strides()));
+        tensorResults.push_back(writeTileToTensor(
+            b, loc, sliceOp, res->getResult(resultIdx), sliceOp.source()));
       } else {
         tensorResults.push_back(res->getResult(resultIdx));
       }
       ++resultIdx;
     }
     return scf::ValueVector(tensorResults.begin(), tensorResults.end());
   };
   GenerateLoopNest<LoopTy>::doit(b, op.getLoc(), loopRanges, op, iteratorTypes,
mlir::linalg::tileLinalgOp(OpBuilder &b, LinalgOp op,
   case LinalgTilingLoopType::TiledLoops:
     return tileLinalgOpImpl<linalg::TiledLoopOp>(b, op, options);
   default:;
   }
   return llvm::None;
 }
 
 namespace {
+struct PadTensorOpTilingPattern : public OpRewritePattern<PadTensorOp> {
+  PadTensorOpTilingPattern(MLIRContext *ctx, LinalgTilingOptions opt)
+      : OpRewritePattern<PadTensorOp>(ctx), options(opt),
+        filter(LinalgTransformationFilter(ArrayRef<Identifier>{},
+                                          Identifier::get("tiled", ctx))) {}
+
+  LogicalResult matchAndRewrite(PadTensorOp op,
+                                PatternRewriter &rewriter) const override {
+    // Can tile only PadTensorOp that have an output operand.
+    if (!op.output())
+      return failure();
+    if (failed(filter.checkAndNotify(rewriter, op)))
+      return failure();
+    Location loc = op.getLoc();
+    OpBuilder::InsertionGuard g(rewriter);
+    rewriter.setInsertionPoint(op);
+
+    // Clone PadTensorOp so that the existing op can be replaced more easily.
+    Operation *newPadOp = rewriter.clone(*op.getOperation());
+    // Get rank and tile sizes.
+    int64_t rank = op.getResultType().getRank();
+    SmallVector<Value> tileSizes =
+        options.tileSizeComputationFunction(rewriter, op);
+    assert(tileSizes.size() == rank);
+    // Compute lower and upper bounds of the loop nest.
+    SmallVector<Value> lbs, dims, steps;
+    for (int64_t i = 0; i < rank; ++i) {
+      if (!isZero(tileSizes[i])) {
+        lbs.push_back(rewriter.create<ConstantIndexOp>(loc, 0));
+        dims.push_back(rewriter.create<tensor::DimOp>(loc, op.output(), i));
+        steps.push_back(tileSizes[i]);
+      }
+    }
+    // Generate loop nest: One loop per dimension.
+    LoopNest loopNest = mlir::scf::buildLoopNest(
+        rewriter, loc, lbs, /*ubs=*/dims, steps, ValueRange(op.output()),
+        [&](OpBuilder &b, Location loc, ValueRange localIvs,
+            ValueRange iterArgs) -> scf::ValueVector {
+          // Compute offsets and sizes of ExtractSliceOp.
+          SmallVector<Value> offsets =
+              computeTileOffsets(b, loc, localIvs, tileSizes);
+          SmallVector<Value> sizes =
+              computeTileSizes(b, loc, localIvs, tileSizes, dims);
+          // Create ExtractSliceOp: Extract a tile from the PadTensorOp.
+          // Note: The PadTensorOp is located outside of the loop nest. It is
+          // later moved inside by ExtractSliceOfPadTensorSwapPattern.
+          auto map = AffineMap::getMultiDimIdentityMap(rank, b.getContext());
+          Value tiledOutput = makeTiledShape(b, loc, newPadOp->getResult(0),
+                                             tileSizes, map, offsets, sizes);
+          auto sliceOp = tiledOutput.getDefiningOp<tensor::ExtractSliceOp>();
+          assert(sliceOp && "expected ExtractSliceOp");
+          // Insert the tile into the output tensor.
+          Value yieldValue =
+              writeTileToTensor(b, loc, sliceOp, sliceOp, iterArgs[0]);
+          return scf::ValueVector({yieldValue});
+        });
+    // Replace all uses of the original PadTensorOp.
+    rewriter.replaceOp(op, loopNest.getResults()[0]);
+    filter.replaceLinalgTransformationFilter(rewriter, newPadOp);
+    return success();
+  }
+
+  LinalgTilingOptions options;
+  LinalgTransformationFilter filter;
+};
+} // namespace
+
+namespace {
 /// Helper classes for type list expansion.
 template <typename... OpTypes>
 class CanonicalizationPatternList;
 
 template <>
 class CanonicalizationPatternList<> {
 public:
   static void insert(RewritePatternSet &patterns) {}
void mlir::linalg::populateLinalgTilingCanonicalizationPatterns(
   AffineMaxOp::getCanonicalizationPatterns(patterns, ctx);
   scf::ForOp::getCanonicalizationPatterns(patterns, ctx);
   scf::ParallelOp::getCanonicalizationPatterns(patterns, ctx);
   ConstantIndexOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::ExtractSliceOp::getCanonicalizationPatterns(patterns, ctx);
   memref::SubViewOp::getCanonicalizationPatterns(patterns, ctx);
   tensor::CastOp::getCanonicalizationPatterns(patterns, ctx);
   memref::ViewOp::getCanonicalizationPatterns(patterns, ctx);
+  PadTensorOp::getCanonicalizationPatterns(patterns, ctx);
   ctx->getLoadedDialect<LinalgDialect>()->getCanonicalizationPatterns(patterns);
   CanonicalizationPatternList<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
       >::insert(patterns);
 }
 
 /// Populate the given list with patterns that apply Linalg tiling.
applyTilingToLoopPatterns(LinalgTilingLoopType loopType, FuncOp funcOp,
   (void)applyPatternsAndFoldGreedily(
       funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
   // Drop the marker.
   funcOp.walk([](LinalgOp op) {
     op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
   });
 }
 
+static void applyPadTensorOpTilingToLoopPatterns(
+    LinalgTilingLoopType loopType, FuncOp funcOp, ArrayRef<int64_t> tileSizes,
+    ArrayRef<StringRef> distributionTypes = {}) {
+  auto options = LinalgTilingOptions()
+                     .setTileSizes(tileSizes)
+                     .setLoopType(loopType)
+                     .setDistributionTypes(distributionTypes);
+  MLIRContext *ctx = funcOp.getContext();
+  RewritePatternSet patterns(ctx);
+  patterns.add<PadTensorOpTilingPattern>(patterns.getContext(), options);
+  patterns.add<ExtractSliceOfPadTensorSwapPattern>(patterns.getContext());
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+  (void)applyPatternsAndFoldGreedily(
+      funcOp, getLinalgTilingCanonicalizationPatterns(ctx));
+  // Drop the marker.
+  funcOp.walk([](PadTensorOp op) {
+    op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
+  });
+}
+
 namespace {
 struct LinalgTilingPass : public LinalgTilingBase<LinalgTilingPass> {
   LinalgTilingPass() = default;
   LinalgTilingPass(ArrayRef<int64_t> sizes) { tileSizes = sizes; }
 
   void runOnFunction() override {
     applyTilingToLoopPatterns(LinalgTilingLoopType::Loops, getFunction(),
                               tileSizes);
   }
 };
 
+struct LinalgPadTensorOpTilingPass
+    : public LinalgPadTensorOpTilingBase<LinalgPadTensorOpTilingPass> {
+  LinalgPadTensorOpTilingPass() = default;
+  LinalgPadTensorOpTilingPass(ArrayRef<int64_t> sizes) { tileSizes = sizes; }
+
+  void runOnFunction() override {
+    applyPadTensorOpTilingToLoopPatterns(LinalgTilingLoopType::Loops,
+                                         getFunction(), tileSizes);
+  }
+};
+
 struct LinalgTilingToParallelLoopsPass
     : public LinalgTilingToParallelLoopsBase<LinalgTilingToParallelLoopsPass> {
   LinalgTilingToParallelLoopsPass() = default;
   LinalgTilingToParallelLoopsPass(ArrayRef<int64_t> sizes) {
     tileSizes = sizes;
   }
 
   void runOnFunction() override {

 } // namespace
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTilingPass(ArrayRef<int64_t> tileSizes) {
   return std::make_unique<LinalgTilingPass>(tileSizes);
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
+mlir::createLinalgPadTensorTilingPass(ArrayRef<int64_t> tileSizes) {
+  return std::make_unique<LinalgPadTensorOpTilingPass>(tileSizes);
+}
+
+std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTilingToParallelLoopsPass(ArrayRef<int64_t> tileSizes) {
   return std::make_unique<LinalgTilingToParallelLoopsPass>(tileSizes);
 }
 
 std::unique_ptr<OperationPass<FuncOp>>
 mlir::createLinalgTilingToTiledLoopPass(ArrayRef<int64_t> tileSizes,
                                         ArrayRef<StringRef> distributionTypes) {
   return std::make_unique<LinalgTilingToTiledLoopsPass>(tileSizes,
                                                         distributionTypes);
 }

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

 static Value asValue(OpBuilder &builder, Location loc, OpFoldResult ofr) {
   if (auto val = ofr.dyn_cast<Value>())
     return val;
   auto intVal = getConstantIntValue(ofr);
   assert(intVal && "expected Value or IntegerAttr");
   return builder.create<ConstantIndexOp>(loc, *intVal);
 }
 
-/// Given a value, try to extract a constant index-type integer as an Attribute.
-/// If this fails, return the original value.
-static OpFoldResult asOpFoldResult(OpBuilder &builder, Value val) {
-  if (auto constInt = getConstantIntValue(val))
-    return builder.getIndexAttr(*constInt);
-  return val;
-}
-
 LogicalResult ExtractSliceOfPadTensorSwapPattern::matchAndRewrite(
     tensor::ExtractSliceOp sliceOp, PatternRewriter &rewriter) const {
   auto padOp = sliceOp.source().getDefiningOp<PadTensorOp>();
   if (!padOp)
     return failure();
   // Only unit stride supported.
   if (!sliceOp.hasUnitStride())
     return failure();

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

for (Range range : ranges) {
     ubs.emplace_back(range.size);
     steps.emplace_back(range.stride);
   }
 }
 
 namespace mlir {
 namespace linalg {
 
+/// Given a value, try to extract a constant integer as an Attribute.
+/// If this fails, return the original value.
+OpFoldResult asOpFoldResult(OpBuilder &builder, Value val) {
+  if (auto constInt = getConstantIntValue(val))
+    return builder.getIndexAttr(*constInt);
+  return val;
+}
+
 /// Helper function that creates a memref::DimOp or tensor::DimOp depending on
 /// the type of `source`.
 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim) {
   if (source.getType().isa<UnrankedMemRefType, MemRefType>())
     return b.createOrFold<memref::DimOp>(loc, source, dim);
   if (source.getType().isa<UnrankedTensorType, RankedTensorType>())
     return b.createOrFold<tensor::DimOp>(loc, source, dim);
   llvm_unreachable("Expected MemRefType or TensorType");
Value makeTiledShape(OpBuilder &builder, Location loc, Value valueToTile,
   offsets.reserve(rank);
   sizes.reserve(rank);
   strides.reserve(rank);
   for (unsigned r = 0; r < rank; ++r) {
     LLVM_DEBUG(llvm::dbgs() << "makeTiledShape: for dim#" << r);
     if (!isTiled(map.getSubMap({r}), tileSizes)) {
       offsets.push_back(builder.getIndexAttr(0));
       Value dim = createOrFoldDimOp(builder, loc, valueToTile, r);
-      sizes.push_back(dim);
+      sizes.push_back(asOpFoldResult(builder, dim));
       strides.push_back(builder.getIndexAttr(1));
       LLVM_DEBUG(llvm::dbgs() << ": not tiled: use size: " << dim << "\n");
       continue;
     }
     LLVM_DEBUG(llvm::dbgs() << ": tiled: figure out subsize...\n");
 
     // Tiling creates a new slice at the proper index, the slice step is 1
     // (i.e. the op does not subsample, stepping occurs in the loop).

mlir/test/Dialect/Linalg/tile-pad-tensor-op.mlir

+// RUN: mlir-opt %s -linalg-tile-pad-tensor-ops="linalg-tile-sizes=2,3" -cse -split-input-file | \
+// RUN: FileCheck %s -check-prefix=TILE2
+// RUN: mlir-opt %s -linalg-tile-pad-tensor-ops="linalg-tile-sizes=0,3" -cse -split-input-file | \
+// RUN: FileCheck %s -check-prefix=TILE1
+
+// TILE2-LABEL: func @dynamic_pad_tensor(
+//  TILE2-SAME:     %[[IN:.*]]: tensor<?x?xf32>, %[[OUT:.*]]: tensor<?x?xf32>
+//   TILE2-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE2-DAG:   %[[C1:.*]] = constant 1 : index
+//   TILE2-DAG:   %[[C2:.*]] = constant 2 : index
+//   TILE2-DAG:   %[[C3:.*]] = constant 3 : index
+//       TILE2:   %[[DIM0:.*]] = tensor.dim %[[OUT]], %[[C0]]
+//       TILE2:   %[[DIM1:.*]] = tensor.dim %[[OUT]], %[[C1]]
+//       TILE2:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM0]] step %[[C2]]
+//       TILE2:     scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE2:       %[[SWAP_RESULT:.*]] = scf.if
+//       TILE2:         tensor.generate
+//       TILE2:       else
+//       TILE2:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:         %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]]
+//       TILE2:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:   return %[[RESULT]]
+
+// TILE1-LABEL: func @dynamic_pad_tensor(
+//  TILE1-SAME:     %[[IN:.*]]: tensor<?x?xf32>, %[[OUT:.*]]: tensor<?x?xf32>
+//   TILE1-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE1-DAG:   %[[C1:.*]] = constant 1 : index
+//   TILE1-DAG:   %[[C3:.*]] = constant 3 : index
+//       TILE1:   %[[DIM1:.*]] = tensor.dim %[[OUT]], %[[C1]]
+//       TILE1:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[DIM1]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE1:     %[[DIM0:.*]] = tensor.dim %[[OUT]], %[[C0]]
+//       TILE1:     %[[SWAP_RESULT:.*]] = scf.if
+//       TILE1:       tensor.generate
+//       TILE1:     else
+//       TILE1:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE1:       %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]] low[3, %{{.*}}] high[{{.*}}, {{.*}}]
+//       TILE1:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [%[[DIM0]], {{.*}}] [1, 1]
+//       TILE1:   return %[[RESULT]]
+
+func @dynamic_pad_tensor(%input_tensor: tensor<?x?xf32>,
+                         %output_tensor: tensor<?x?xf32>,
+                         %pad_value: f32) -> tensor<?x?xf32> {
+  %0 = linalg.pad_tensor %input_tensor, %output_tensor
+    low[3, 4] high[5, 3] {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad_value : f32
+    } : tensor<?x?xf32> to tensor<?x?xf32>
+  return %0 : tensor<?x?xf32>
+}
+
+// -----
+
+// TILE2-LABEL: func @static_pad_tensor(
+//  TILE2-SAME:     %[[IN:.*]]: tensor<7x9xf32>, %[[OUT:.*]]: tensor<15x16xf32>
+//   TILE2-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE2-DAG:   %[[C2:.*]] = constant 2 : index
+//   TILE2-DAG:   %[[C3:.*]] = constant 3 : index
+//   TILE2-DAG:   %[[C15:.*]] = constant 15 : index
+//   TILE2-DAG:   %[[C16:.*]] = constant 16 : index
+//       TILE2:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C15]] step %[[C2]]
+//       TILE2:     scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE2:       %[[SWAP_RESULT:.*]] = scf.if
+//       TILE2:         tensor.generate
+//       TILE2:       else
+//       TILE2:         %[[SLICE:.*]] = tensor.extract_slice %[[IN]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:         %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]]
+//       TILE2:       tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][{{.*}}, {{.*}}] [{{.*}}, {{.*}}] [1, 1]
+//       TILE2:   return %[[RESULT]]
+
+
+// TILE1-LABEL: func @static_pad_tensor(
+//  TILE1-SAME:     %[[IN:.*]]: tensor<7x9xf32>, %[[OUT:.*]]: tensor<15x16xf32>
+//   TILE1-DAG:   %[[C0:.*]] = constant 0 : index
+//   TILE1-DAG:   %[[C3:.*]] = constant 3 : index
+//   TILE1-DAG:   %[[C16:.*]] = constant 16 : index
+//       TILE1:   %[[RESULT:.*]] = scf.for {{.*}} = %[[C0]] to %[[C16]] step %[[C3]] iter_args(%[[INNER_OUT:.*]] =
+//       TILE1:     %[[SWAP_RESULT:.*]] = scf.if
+//       TILE1:       tensor.generate
+//       TILE1:     else
+//       TILE1:       %[[SLICE:.*]] = tensor.extract_slice %[[IN]][0, {{.*}}] [7, {{.*}}] [1, 1]
+//       TILE1:       %[[PAD:.*]] = linalg.pad_tensor %[[SLICE]] low[3, %{{.*}}] high[5, {{.*}}]
+//       TILE1:     tensor.insert_slice %[[SWAP_RESULT]] into %[[INNER_OUT]][0, {{.*}}] [15, {{.*}}] [1, 1]
+//       TILE1:   return %[[RESULT]]
+
+func @static_pad_tensor(%input_tensor: tensor<7x9xf32>,
+                        %output_tensor: tensor<15x16xf32>,
+                        %pad_value: f32) -> tensor<15x16xf32> {
+  %0 = linalg.pad_tensor %input_tensor, %output_tensor
+    low[3, 4] high[5, 3] {
+    ^bb0(%arg1: index, %arg2: index):
+      linalg.yield %pad_value : f32
+    } : tensor<7x9xf32> to tensor<15x16xf32>
+  return %0 : tensor<15x16xf32>
+}