[mlir][linalg] Add scf.for loop peeling to codegen strategy

[springerm]

Summary

Only scf.for loops are supported at the moment. linalg.tiled_loop support will be added in a subsequent commit.

Only static tensor sizes are fully supported. Loops for dynamic tensor sizes can be peeled, but the generated code is not optimal due to a missing canonicalization pattern.

[springerm]

rebase

[springerm]

no change

[springerm]

Do not use markers anymore. Instead, peel loops right after generating them.

[nicolasvasilache]

I'll let this one simmer until you had a chance to decide whether you update scf.for peeling to be consistent with linalg.tiled_loop peeling.
Ping me if you want me to review.

[springerm]

rebase

[springerm]

This is ready for review now.

[nicolasvasilache]

LGTM, thanks!

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

struct LinalgTilingOptions {
   /// shaped and can thus be vectorized.
   PaddingValueComputationFunction paddingValueComputationFunction = nullptr;
 
   LinalgTilingOptions &
   setPaddingValueComputationFunction(PaddingValueComputationFunction fun) {
     paddingValueComputationFunction = std::move(fun);
     return *this;
   }
+
+  /// Add label attributes to tiled loops, so that they can be easily found
+  /// and processed by other patterns.
+  bool markLoops;
+
+  LinalgTilingOptions &setMarkLoops(bool b) {
+    markLoops = b;
+    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.
 RewritePatternSet getLinalgTilingCanonicalizationPatterns(MLIRContext *ctx);
 void populateLinalgTilingCanonicalizationPatterns(RewritePatternSet &patterns);
 
struct LinalgBaseTilingPattern : public RewritePattern {
   // Entry point to match a specific Linalg op.
   LinalgBaseTilingPattern(
       StringRef opName, MLIRContext *context, LinalgTilingOptions options,
       LinalgTransformationFilter filter = LinalgTransformationFilter(),
       PatternBenefit benefit = 1);
   LogicalResult matchAndRewriteBase(Operation *op, PatternRewriter &rewriter,
                                     TiledLinalgOp &result) const;
 
+  // Walk `op` and remove all markers from tiled loops.
+  static void removeMarkers(Operation *op);
+
 private:
   /// LinalgTransformMarker handles special attribute manipulations.
   LinalgTransformationFilter filter;
   /// Options to control tiling;
   LinalgTilingOptions options;
 };
 
 template <typename OpTy>
LogicalResult matchAndRewrite(Operation *op,
     if (tiledLinalgOp.tensorResults.empty())
       rewriter.eraseOp(op);
     else
       rewriter.replaceOp(op, tiledLinalgOp.tensorResults);
     return success();
   }
 };
 
+template <typename OpTy>
+struct PeelTiledLoopsPattern : public OpRewritePattern<OpTy> {
+  PeelTiledLoopsPattern(MLIRContext *ctx, int64_t loopId)
+      : OpRewritePattern<OpTy>(ctx), loopId(loopId) {}
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override;
+
+  const int64_t loopId;
+};
+
 struct LinalgFusionOptions {
   /// List of operands indices to use for fusion.
   llvm::SmallSet<unsigned, 1> indicesToFuse = {};
   LinalgFusionOptions &setIndicesToFuse(ArrayRef<int64_t> operands) {
     indicesToFuse.insert(operands.begin(), operands.end());
     return *this;
   }
 };

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

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Affine/Utils.h"
 #include "mlir/Dialect/Linalg/Analysis/DependenceAnalysis.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/SCF/Transforms.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.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"
: RewritePattern(opName, benefit, context), filter(filter),
       options(options) {}
 
 mlir::linalg::LinalgBaseTilingPattern::LinalgBaseTilingPattern(
     MLIRContext *context, LinalgTilingOptions options,
     LinalgTransformationFilter filter, PatternBenefit benefit)
     : RewritePattern(MatchAnyOpTypeTag(), benefit, context), filter(filter),
       options(options) {}
 
+static const char kTiledLoopMarker[] = "__linalg_internal_tiled_loop__";
+
 LogicalResult mlir::linalg::LinalgBaseTilingPattern::matchAndRewriteBase(
     Operation *op, PatternRewriter &rewriter, TiledLinalgOp &result) const {
   LinalgOp linalgOp = dyn_cast<LinalgOp>(op);
   if (!linalgOp)
     return failure();
   if (failed(filter.checkAndNotify(rewriter, linalgOp)))
     return failure();
 
   Optional<TiledLinalgOp> res = tileLinalgOp(rewriter, linalgOp, options);
 
   if (!res)
     return failure();
 
+  // Mark tiled loops.
+  if (options.markLoops)
+    for (auto it : llvm::enumerate(res->loops))
+      it.value()->setAttr(kTiledLoopMarker,
+                          rewriter.getI64IntegerAttr(it.index()));
+
   // Setup RAII guard to return properly.
   LinalgOp tiledOp = res->op;
   auto guard = llvm::make_scope_exit([&]() {
     // Return relevant information to derived pattern.
     result = *res;
     // Replace filter on both tiledOp and tiledAndPaddedOp, if necessary.
     filter.replaceLinalgTransformationFilter(rewriter, tiledOp);
     if (tiledOp != res->op)
if (succeeded(rewriteAsPaddedOp(rewriter, res->op,
     // Do not perform replacement of `linalgOp`, let the derived patterns
     // do this as they see fit, from the resulting TiledLinalgOp.
     return success();
   }
   // Set so RAII guard does not propagate TiledLinalgOp to `result`.
   return failure();
 }
 
+void LinalgBaseTilingPattern::removeMarkers(Operation *op) {
+  op->walk([](Operation *c) { c->removeAttr(kTiledLoopMarker); });
+}
+
+template <>
+LogicalResult PeelTiledLoopsPattern<scf::ForOp>::matchAndRewrite(
+    scf::ForOp op, PatternRewriter &rewriter) const {
+  if (!op->hasAttr(kTiledLoopMarker))
+    return failure();
+  auto markerAttr = op->getAttr(kTiledLoopMarker).dyn_cast<IntegerAttr>();
+  assert(markerAttr && "expected integer attribute");
+  if (markerAttr.getInt() != loopId)
+    return failure();
+
+  scf::IfOp ifOp;
+  LogicalResult status = scf::peelAndCanonicalizeForLoop(rewriter, op, ifOp);
+  // Remove marker to avoid loop from getting peeled a second time.
+  rewriter.updateRootInPlace(op, [&]() { op->removeAttr(kTiledLoopMarker); });
+
+  if (succeeded(status)) {
+    // Do not tile loops inside partial iterations.
+    // TODO: This should be configurable.
+    LinalgBaseTilingPattern::removeMarkers(ifOp);
+    return success();
+  }
+  assert(!ifOp && "expected that loop was not peeled");
+  return failure();
+}
+
 static ValueRange getTiledOpResult(TiledLinalgOp tiledOp) {
   if (tiledOp.loops.empty())
     return tiledOp.op.getOperation()->getResults();
   return tiledOp.loops.front()->getResults();
 }
 
 static ValueRange
 getTiledAndFusedOpResult(TiledAndFusedLinalgOps tiledAndFusedOp) {

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

-// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-and-pad-pattern tile-sizes-for-padding=2,3,4" -canonicalize | FileCheck %s
-// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-and-pad-pattern tile-sizes-for-padding=2,3" -canonicalize | FileCheck %s -check-prefix=CHECK-1DIM-TILE
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern pad-tiles tile-sizes=2,3,4" -canonicalize | FileCheck %s
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern pad-tiles tile-sizes=2,3" -canonicalize | FileCheck %s -check-prefix=CHECK-1DIM-TILE
 
 // CHECK-LABEL: func @matmul_tensors(
 // CHECK-SAME:    %[[TA:[0-9a-z]+]]: tensor<?x?xi8>
 // CHECK-SAME:    %[[TB:[0-9a-z]+]]: tensor<?x?xi8>
 // CHECK-SAME:    %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
 func @matmul_tensors(
   %arg0: tensor<?x?xi8>, %arg1: tensor<?x?xi8>, %arg2: tensor<?x?xi32>)
     -> tensor<?x?xi32> {

 //      CHECK:         : tensor<?x?xi32> to tensor<2x3xi32>
 //      CHECK:       %[[pD:.*]] = linalg.matmul_i8_i8_i32 ins(%[[pA]], %[[pB]] : tensor<2x4xi8>, tensor<4x3xi8>)
 // CHECK-SAME:                                           outs(%[[pC]] : tensor<2x3xi32>)  -> tensor<2x3xi32>
 //      CHECK:       %[[sTD:.*]] = tensor.extract_slice %[[pD]][0, 0] [%{{.*}}, %{{.*}}] [1, 1] : tensor<2x3xi32> to tensor<?x?xi32>
 //      CHECK:       %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<?x?xi32> into tensor<?x?xi32>
 //      CHECK:       scf.yield %[[TD]] : tensor<?x?xi32>
 //      CHECK:     scf.yield %[[TD2]] : tensor<?x?xi32>
 //      CHECK:   scf.yield %[[TD1]] : tensor<?x?xi32>
-  %0 = linalg.matmul_i8_i8_i32 {__internal_linalg_transform__ = "tile-and-pad"}
+  %0 = linalg.matmul_i8_i8_i32 {__internal_linalg_transform__ = "tile"}
       ins(%arg0, %arg1: tensor<?x?xi8>, tensor<?x?xi8>)
      outs(%arg2: tensor<?x?xi32>)
     -> tensor<?x?xi32>
 
 //      CHECK: return %[[TD0]] : tensor<?x?xi32>
   return %0 : tensor<?x?xi32>
 }
 

 //      CHECK:       %[[TD:.*]] = tensor.insert_slice %[[sTD]] into %[[TC2]][{{.*}}]  : tensor<?x?x?xf32> into tensor<?x?x?xf32>
 //      CHECK:       scf.yield %[[TD]] : tensor<?x?x?xf32>
 //      CHECK:     scf.yield %[[TD2]] : tensor<?x?x?xf32>
 //      CHECK:   scf.yield %[[TD1]] : tensor<?x?x?xf32>
       %0 = linalg.generic {
       indexing_maps =  [ affine_map<(d0, d1, d2) -> ()>,
                         affine_map<(d0, d1, d2) -> (d0, d1, d2)> ],
       iterator_types = ["parallel", "parallel", "parallel"]}
-      {__internal_linalg_transform__ = "tile-and-pad"}
+      {__internal_linalg_transform__ = "tile"}
      ins(%arg1 : f32)
     outs(%arg0: tensor<?x?x?xf32>) {
       ^bb(%0: f32, %1: f32) :
         linalg.yield %0 : f32
     } -> tensor<?x?x?xf32>
   return %0 : tensor<?x?x?xf32>
 }
 
 // CHECK-1DIM-TILE: func @matmul_tensors(
 // CHECK-1DIM-TILE:    %[[TA:[0-9a-z]+]]: tensor<?x?xi8>
 // CHECK-1DIM-TILE:    %[[TB:[0-9a-z]+]]: tensor<?x?xi8>
 // CHECK-1DIM-TILE:    %[[TC:[0-9a-z]+]]: tensor<?x?xi32>) -> tensor<?x?xi32> {
 // CHECK-1DIM-TILE-NOT: scf.for
 // CHECK-1DIM-TILE: linalg.matmul_i8_i8_i32 ins(%[[TA]], %[[TB]] : tensor<?x?xi8>, tensor<?x?xi8>) outs(%[[TC]] : tensor<?x?xi32>) -> tensor<?x?xi32>
 
 func @matmul_partially_padded_tensors(
   %arg0: tensor<?x8xi8>, %arg1: tensor<8x?xi8>, %arg2: tensor<?x?xi32>)
     -> tensor<?x?xi32> {
-  %0 = linalg.matmul_i8_i8_i32 {__internal_linalg_transform__ = "tile-and-pad"}
+  %0 = linalg.matmul_i8_i8_i32 {__internal_linalg_transform__ = "tile"}
       ins(%arg0, %arg1: tensor<?x8xi8>, tensor<8x?xi8>)
      outs(%arg2: tensor<?x?xi32>)
     -> tensor<?x?xi32>
   return %0 : tensor<?x?xi32>
 }
 // CHECK-LABEL: func @matmul_partially_padded_tensors(
 // CHECK: linalg.matmul_i8_i8_i32 ins({{.*}}, {{.*}} : tensor<2x4xi8>, tensor<4x3xi8>) outs({{.*}} : tensor<2x3xi32>) -> tensor<2x3xi32>
 

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

+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern mark-tiled-loops tile-sizes=256,128,512" -canonicalize -for-loop-canonicalization -cse \
+// RUN:             -test-linalg-transform-patterns="test-peel-tiled-loops-pattern peeled-loops=0" -canonicalize -split-input-file |\
+// RUN:     FileCheck %s -check-prefix=CHECK-PEEL-0
+
+// RUN: mlir-opt %s -test-linalg-transform-patterns="test-tile-pattern mark-tiled-loops tile-sizes=256,128,512" -canonicalize -for-loop-canonicalization -cse \
+// RUN:             -test-linalg-transform-patterns="test-peel-tiled-loops-pattern peeled-loops=1,2" -canonicalize -split-input-file | \
+// RUN:     FileCheck %s -check-prefix=CHECK-PEEL-12
+
+//     CHECK-PEEL-0: func @matmul_static_tensor
+// CHECK-PEEL-0-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-PEEL-0-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-PEEL-0-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-PEEL-0-DAG:   %[[c512:.*]] = constant 512 : index
+// CHECK-PEEL-0-DAG:   %[[c1280:.*]] = constant 1280 : index
+// CHECK-PEEL-0-DAG:   %[[c1600:.*]] = constant 1600 : index
+// CHECK-PEEL-0-DAG:   %[[c1700:.*]] = constant 1700 : index
+//     CHECK-PEEL-0:   scf.for %{{.*}} = %[[c0]] to %[[c1280]] step %[[c256]] {{.*}} {
+//     CHECK-PEEL-0:     scf.for %{{.*}} = %[[c0]] to %[[c1700]] step %[[c128]] {{.*}} {
+//     CHECK-PEEL-0:       scf.for %{{.*}} = %[[c0]] to %[[c1600]] step %[[c512]] {{.*}} {
+//     CHECK-PEEL-0:         linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<256x?xf32>)
+//     CHECK-PEEL-0:       }
+//     CHECK-PEEL-0:     }
+//     CHECK-PEEL-0:   }
+//     CHECK-PEEL-0:   scf.for %{{.*}} = %[[c0]] to %[[c1700]] step %[[c128]] {{.*}} {
+//     CHECK-PEEL-0:     scf.for %{{.*}} = %[[c0]] to %[[c1600]] step %[[c512]] {{.*}} {
+//     CHECK-PEEL-0:       linalg.matmul ins({{.*}} : tensor<220x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<220x?xf32>)
+//     CHECK-PEEL-0:     }
+//     CHECK-PEEL-0:   }
+
+//     CHECK-PEEL-12: func @matmul_static_tensor
+// CHECK-PEEL-12-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-PEEL-12-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-PEEL-12-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-PEEL-12-DAG:   %[[c512:.*]] = constant 512 : index
+// CHECK-PEEL-12-DAG:   %[[c1500:.*]] = constant 1500 : index
+// CHECK-PEEL-12-DAG:   %[[c1536:.*]] = constant 1536 : index
+// CHECK-PEEL-12-DAG:   %[[c1600:.*]] = constant 1600 : index
+// CHECK-PEEL-12-DAG:   %[[c1664:.*]] = constant 1664 : index
+//     CHECK-PEEL-12:   scf.for %{{.*}} = %[[c0]] to %[[c1500]] step %[[c256]] {{.*}} {
+//     CHECK-PEEL-12:     scf.for %{{.*}} = %[[c0]] to %[[c1664]] step %[[c128]] {{.*}} {
+//     CHECK-PEEL-12:       scf.for %{{.*}} = %[[c0]] to %[[c1536]] step %[[c512]] {{.*}} {
+//     CHECK-PEEL-12:         linalg.matmul ins({{.*}} : tensor<?x512xf32>, tensor<512x128xf32>) outs({{.*}} : tensor<?x128xf32>)
+//     CHECK-PEEL-12:       }
+//     CHECK-PEEL-12:       linalg.matmul ins({{.*}} : tensor<?x64xf32>, tensor<64x128xf32>) outs({{.*}} : tensor<?x128xf32>)
+//     CHECK-PEEL-12:     }
+//     CHECK-PEEL-12:     scf.for %{{.*}} = %[[c0]] to %[[c1600]] step %[[c512]] {{.*}} {
+//     CHECK-PEEL-12:       linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x36xf32>) outs({{.*}} : tensor<?x36xf32>)
+//     CHECK-PEEL-12:     }
+//     CHECK-PEEL-12:   }
+func @matmul_static_tensor(%arg0: tensor<1500x1600xf32>, %arg1: tensor<1600x1700xf32>)
+    -> tensor<1500x1700xf32> {
+  %out = linalg.init_tensor [1500, 1700] : tensor<1500x1700xf32>
+  %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
+      ins(%arg0, %arg1: tensor<1500x1600xf32>, tensor<1600x1700xf32>)
+      outs(%out: tensor<1500x1700xf32>) -> tensor<1500x1700xf32>
+  return %r : tensor<1500x1700xf32>
+}
+
+// -----
+
+//     CHECK-PEEL-0: func @matmul_dynamic_tensor
+// CHECK-PEEL-0-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-PEEL-0-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-PEEL-0-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-PEEL-0-DAG:   %[[c512:.*]] = constant 512 : index
+//     CHECK-PEEL-0:   scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c256]] {{.*}} {
+//     CHECK-PEEL-0:     scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c128]] {{.*}} {
+//     CHECK-PEEL-0:       scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c512]] {{.*}} {
+//     CHECK-PEEL-0:         linalg.matmul ins({{.*}} : tensor<256x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-PEEL-0:       }
+//     CHECK-PEEL-0:     }
+//     CHECK-PEEL-0:   }
+//     CHECK-PEEL-0:   scf.if %{{.*}} {
+//     CHECK-PEEL-0:     scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c128]] {{.*}} {
+//     CHECK-PEEL-0:       scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c512]] {{.*}} {
+//     CHECK-PEEL-0:         linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-PEEL-0:       }
+//     CHECK-PEEL-0:     }
+//     CHECK-PEEL-0:   }
+
+//     CHECK-PEEL-12: func @matmul_dynamic_tensor
+// CHECK-PEEL-12-DAG:   %[[c0:.*]] = constant 0 : index
+// CHECK-PEEL-12-DAG:   %[[c128:.*]] = constant 128 : index
+// CHECK-PEEL-12-DAG:   %[[c256:.*]] = constant 256 : index
+// CHECK-PEEL-12-DAG:   %[[c512:.*]] = constant 512 : index
+//     CHECK-PEEL-12:   scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c256]] {{.*}} {
+//     CHECK-PEEL-12:     scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c128]] {{.*}} {
+//     CHECK-PEEL-12:       scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c512]] {{.*}} {
+//     CHECK-PEEL-12:         linalg.matmul ins({{.*}} : tensor<?x512xf32>, tensor<?x128xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-PEEL-12:       }
+//     CHECK-PEEL-12:       scf.if %{{.*}} {
+//     CHECK-PEEL-12:         linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x128xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-PEEL-12:       }
+//     CHECK-PEEL-12:     }
+//     CHECK-PEEL-12:     scf.if %{{.*}} {
+//     CHECK-PEEL-12:       scf.for %{{.*}} = %[[c0]] to %{{.*}} step %[[c512]] {{.*}} {
+//     CHECK-PEEL-12:         linalg.matmul ins({{.*}} : tensor<?x?xf32>, tensor<?x?xf32>) outs({{.*}} : tensor<?x?xf32>)
+//     CHECK-PEEL-12:       }
+//     CHECK-PEEL-12:     }
+//     CHECK-PEEL-12:   }
+func @matmul_dynamic_tensor(%arg0: tensor<?x?xf32>, %arg1: tensor<?x?xf32>)
+    -> tensor<?x?xf32> {
+  %c0 = constant 0 : index
+  %c1 = constant 1 : index
+  %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>
+  %d1 = tensor.dim %arg1, %c1 : tensor<?x?xf32>
+  %out = linalg.init_tensor [%d0, %d1] : tensor<?x?xf32>
+  %r = linalg.matmul {__internal_linalg_transform__ = "tile"}
+      ins(%arg0, %arg1: tensor<?x?xf32>, tensor<?x?xf32>)
+      outs(%out: tensor<?x?xf32>) -> tensor<?x?xf32>
+  return %r : tensor<?x?xf32>
+}

mlir/test/lib/Dialect/Linalg/TestLinalgTransforms.cpp

Option<bool> testVectorTransferForwardingPatterns{
       llvm::cl::desc(
           "Test a fused pass that forwards linalg.copy to vector.transfer"),
       llvm::cl::init(false)};
   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> testTileAndPadPattern{
-      *this, "test-tile-and-pad-pattern",
-      llvm::cl::desc("Test tile and pad pattern"), llvm::cl::init(false)};
+  Option<bool> testPeelTiledLoopsPattern{
+      *this, "test-peel-tiled-loops-pattern",
+      llvm::cl::desc("Test peeling of marked tiled loops"),
+      llvm::cl::init(false)};
+  Option<bool> testTilePattern{*this, "test-tile-pattern",
+                               llvm::cl::desc("Test tile pattern"),
+                               llvm::cl::init(false)};
   Option<int> testHoistPadding{*this, "test-hoist-padding",
                                llvm::cl::desc("Test hoist padding"),
                                llvm::cl::init(0)};
   Option<bool> testTransformPadTensor{
       *this, "test-transform-pad-tensor",
       llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
       llvm::cl::init(false)};
   Option<bool> testGeneralizePadTensor{
       *this, "test-generalize-pad-tensor",
       llvm::cl::desc("Test transform pad tensor by copying with generic ops"),
       llvm::cl::init(false)};
   Option<bool> testSwapSubTensorPadTensor{
       *this, "test-swap-subtensor-padtensor",
       llvm::cl::desc("Test rewrite of subtensor(pad_tensor) into "
                      "pad_tensor(subtensor)"),
       llvm::cl::init(false)};
-  ListOption<int64_t> tileSizesForPadding{
-      *this, "tile-sizes-for-padding",
-      llvm::cl::desc("Linalg tile sizes when tile+pad"), llvm::cl::ZeroOrMore,
-      llvm::cl::MiscFlags::CommaSeparated};
+  Option<bool> markTiledLoops{
+      *this, "mark-tiled-loops",
+      llvm::cl::desc("Mark tiled loops during tile+pad"),
+      llvm::cl::init(false)};
+  Option<bool> padTiles{*this, "pad-tiles",
+                        llvm::cl::desc("Pad tiles when test-tile-pattern"),
+                        llvm::cl::init(false)};
+  ListOption<int64_t> peeledLoops{
+      *this, "peeled-loops", llvm::cl::desc("Loops to be peeled"),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
+  ListOption<int64_t> tileSizes{
+      *this, "tile-sizes",
+      llvm::cl::desc("Linalg tile sizes for test-tile-pattern"),
+      llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated};
   ListOption<unsigned> testInterchangePattern{
       *this, "test-interchange-pattern", llvm::cl::MiscFlags::CommaSeparated,
       llvm::cl::desc("Test the interchange pattern.")};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();

 
 // 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, OpOperand &op) {
   auto t = getElementTypeOrSelf(op.get());
   return b.create<ConstantOp>(op.getOwner()->getLoc(), t, b.getZeroAttr(t));
 }
 
-static void applyTileAndPadPattern(FuncOp funcOp, ArrayRef<int64_t> tileSizes) {
+static void applyPeelTiledLoopsPattern(FuncOp funcOp,
+                                       ArrayRef<int64_t> peeledLoops) {
+  MLIRContext *context = funcOp.getContext();
+  RewritePatternSet peelingPattern(context);
+  for (int64_t loop : peeledLoops)
+    peelingPattern.add<linalg::PeelTiledLoopsPattern<scf::ForOp>>(context,
+                                                                  loop);
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(peelingPattern));
+  // Remove markers from tiled loops.
+  LinalgBaseTilingPattern::removeMarkers(funcOp);
+}
+
+static void applyTilePattern(FuncOp funcOp, ArrayRef<int64_t> tileSizes,
+                             bool markTiledLoops, bool padTiles) {
   MLIRContext *context = funcOp.getContext();
   RewritePatternSet tilingPattern(context);
   auto linalgTilingOptions =
-      linalg::LinalgTilingOptions()
-          .setTileSizes(tileSizes)
-          .setPaddingValueComputationFunction(getNeutralOfLinalgOp);
-  tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulI8I8I32Op>,
+      linalg::LinalgTilingOptions().setTileSizes(tileSizes).setMarkLoops(
+          markTiledLoops);
+  if (padTiles)
+    linalgTilingOptions.setPaddingValueComputationFunction(
+        getNeutralOfLinalgOp);
+
+  tilingPattern.add<linalg::LinalgTilingPattern<linalg::MatmulOp>,
+                    linalg::LinalgTilingPattern<linalg::MatmulI8I8I32Op>,
                     linalg::LinalgTilingPattern<linalg::GenericOp>>(
       context, linalgTilingOptions,
-      linalg::LinalgTransformationFilter(
-          Identifier::get("tile-and-pad", context)));
+      linalg::LinalgTransformationFilter(Identifier::get("tile", context)));
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(tilingPattern));
 }
 
 static void applyInterchangePattern(FuncOp funcOp,
                                     ArrayRef<unsigned> interchangeVector) {
   MLIRContext *context = funcOp.getContext();
   RewritePatternSet interchangePattern(context);
   interchangePattern.add<GenericOpInterchangePattern>(
void TestLinalgTransforms::runOnFunction() {
   if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getFunction());
   if (testTransformPadTensor)
     return applyPadTensorToGenericPatterns(getFunction());
   if (testGeneralizePadTensor)
     return applyGeneralizePadTensorPatterns(getFunction());
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getFunction());
-  if (testTileAndPadPattern)
-    return applyTileAndPadPattern(getFunction(), tileSizesForPadding);
+  if (testPeelTiledLoopsPattern)
+    return applyPeelTiledLoopsPattern(getFunction(), peeledLoops);
+  if (testTilePattern)
+    return applyTilePattern(getFunction(), tileSizes, markTiledLoops, padTiles);
   if (testHoistPadding) {
     getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
       (void)linalg::hoistPaddingOnTensors(padTensorOp, testHoistPadding);
     });
   }
   if (testInterchangePattern.hasValue())
     return applyInterchangePattern(getFunction(), testInterchangePattern);
 }
 
 namespace mlir {
 namespace test {
 void registerTestLinalgTransforms() {
   PassRegistration<TestLinalgTransforms>();
 }
 } // namespace test
 } // namespace mlir