[mlir][linalg] TiledLoopOp peeling: Do not peel partial iterations

[springerm]

Summary

Extend the unit test with an option for skipping partial iterations during loop peeling.

[springerm]

Addressing comment: https://reviews.llvm.org/D108270#inline-1035954

[huajsj]

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

608

Check if the loop has already peeled.

613

If the loop nests is partial iterations, then don't peel.

623–624

one more line after 623?

639–640

If set to true, TiledLoopOps that are the partial iteration of another peeled TiledLoopOp will not be peeled.

[herhut]

Thanks, just a nit.

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

653–657

These two walks can be combined.

[springerm]

address comments

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

613

Reformulated: The loop nest represented by the TiledLoopOp should not be peeled if one of its loops is a partial iteration.

mlir/test/Dialect/Linalg/tiled-loop-peeling.mlir

 // RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-tiled-loop-peeling=2 -split-input-file | FileCheck %s -check-prefix=CHECK-TILE-2
 // RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-tiled-loop-peeling=0,1,2 -split-input-file | FileCheck %s -check-prefix=CHECK-TILE-012
+// RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns="test-tiled-loop-peeling=0,1,2 skip-partial" -split-input-file | FileCheck %s -check-prefix=CHECK-TILE-012-SKIP-PARTIAL
 
 // CHECK-TILE-2-LABEL: func @tiled_loop_3d_tensor(
 //  CHECK-TILE-2-SAME:     %[[input:.*]]: tensor<?x?x?xf32>, %[[s0:.*]]: index, %[[s1:.*]]: index, %[[s2:.*]]: index
 //   CHECK-TILE-2-DAG:   %[[c0:.*]] = constant 0 : index
 //   CHECK-TILE-2-DAG:   %[[c1:.*]] = constant 1 : index
 //   CHECK-TILE-2-DAG:   %[[c2:.*]] = constant 2 : index
 //       CHECK-TILE-2:   %[[dim0:.*]] = tensor.dim %[[input]], %[[c0]]
 //       CHECK-TILE-2:   %[[dim1:.*]] = tensor.dim %[[input]], %[[c1]]

 //       CHECK-TILE-012:   linalg.tiled_loop {{.*}} {
 //       CHECK-TILE-012:     linalg.yield
 //       CHECK-TILE-012:   }
 //       CHECK-TILE-012:   linalg.tiled_loop {{.*}} {
 //       CHECK-TILE-012:     linalg.yield
 //       CHECK-TILE-012:   }
 //   CHECK-TILE-012-NOT: linalg.tiled_loop
 
+//      CHECK-TILE-012-SKIP-PARTIAL: func @tiled_loop_3d_tensor(
+// CHECK-TILE-012-SKIP-PARTIAL-SAME:     %[[input:.*]]: tensor<?x?x?xf32>
+//  CHECK-TILE-012-SKIP-PARTIAL-DAG:   %[[c0:.*]] = constant 0 : index
+//  CHECK-TILE-012-SKIP-PARTIAL-DAG:   %[[c1:.*]] = constant 1 : index
+//  CHECK-TILE-012-SKIP-PARTIAL-DAG:   %[[c2:.*]] = constant 2 : index
+//  CHECK-TILE-012-SKIP-PARTIAL-DAG:   %[[dim0:.*]] = tensor.dim %[[input]], %[[c0]]
+//  CHECK-TILE-012-SKIP-PARTIAL-DAG:   %[[dim1:.*]] = tensor.dim %[[input]], %[[c1]]
+//  CHECK-TILE-012-SKIP-PARTIAL-DAG:   %[[dim2:.*]] = tensor.dim %[[input]], %[[c2]]
+//      CHECK-TILE-012-SKIP-PARTIAL:   %[[p0:.*]] = affine.apply #{{.*}}()[%[[dim0]]
+//      CHECK-TILE-012-SKIP-PARTIAL:   %[[p1:.*]] = affine.apply #{{.*}}()[%[[dim1]]
+//      CHECK-TILE-012-SKIP-PARTIAL:   %[[p2:.*]] = affine.apply #{{.*}}()[%[[dim2]]
+//      CHECK-TILE-012-SKIP-PARTIAL:   linalg.tiled_loop {{.*}} = (%[[c0]], %[[c0]], %[[c0]]) to (%[[p0]], %[[p1]], %[[p2]])
+//      CHECK-TILE-012-SKIP-PARTIAL:   linalg.tiled_loop {{.*}} = (%[[c0]], %[[c0]], %[[p2]]) to (%[[p0]], %[[p1]], %[[dim2]])
+//      CHECK-TILE-012-SKIP-PARTIAL:   linalg.tiled_loop {{.*}} = (%[[c0]], %[[p1]], %[[c0]]) to (%[[p0]], %[[dim1]], %[[dim2]])
+//      CHECK-TILE-012-SKIP-PARTIAL:   linalg.tiled_loop {{.*}} = (%[[p0]], %[[c0]], %[[c0]]) to (%[[dim0]], %[[dim1]], %[[dim2]])
 func @tiled_loop_3d_tensor(%arg0: tensor<?x?x?xf32>, %s0: index, %s1: index,
                            %s2: index) -> tensor<?x?x?xf32> {
   %cst = constant 0.000000e+00 : f32
   %c0 = constant 0 : index
   %c1 = constant 1 : index
   %c2 = constant 2 : index
   %c8 = constant 8 : index
   %dim0 = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>

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

ListOption<int64_t> tileSizesForPadding{
       llvm::cl::MiscFlags::CommaSeparated};
   ListOption<unsigned> testInterchangePattern{
       *this, "test-interchange-pattern", llvm::cl::MiscFlags::CommaSeparated,
       llvm::cl::desc("Test the interchange pattern.")};
   ListOption<unsigned> testTiledLoopPeeling{
       *this, "test-tiled-loop-peeling",
       llvm::cl::desc("Test peeling of linalg.tiled_loop ops"),
       llvm::cl::OneOrMore, llvm::cl::MiscFlags::CommaSeparated};
+  Option<bool> skipPartial{
+      *this, "skip-partial",
+      llvm::cl::desc("Skip loops inside partial iterations during peeling"),
+      llvm::cl::init(false)};
 };
 } // end anonymous namespace
 
 static void applyPatterns(FuncOp funcOp) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
 
   //===--------------------------------------------------------------------===//
static void applyInterchangePattern(FuncOp funcOp,
   interchangePattern.add<GenericOpInterchangePattern>(
       context, interchangeVector,
       LinalgTransformationFilter(ArrayRef<Identifier>{},
                                  Identifier::get("interchange", context)));
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(interchangePattern));
 }
 
 static constexpr char kPeeledLoopsLabel[] = "__peeled_loops__";
+static constexpr char kPartialIterationLabel[] = "__partial_iteration__";
 
 namespace {
 /// Peel TiledLoopOps, i.e., split them into two loops: One loop where the
 /// `idx`-th loop contains only "full" iterations and a second loop for the
 /// remaining partial iteration (if any).
 struct TiledLoopPeelingPattern : public OpRewritePattern<TiledLoopOp> {
-  TiledLoopPeelingPattern(MLIRContext *ctx, int64_t idx)
-      : OpRewritePattern<TiledLoopOp>(ctx), idx(idx) {}
+  TiledLoopPeelingPattern(MLIRContext *ctx, int64_t idx, bool skipPartial)
+      : OpRewritePattern<TiledLoopOp>(ctx), idx(idx), skipPartial(skipPartial) {
+  }
 
   LogicalResult matchAndRewrite(TiledLoopOp loopOp,
                                 PatternRewriter &rewriter) const override {
     SmallVector<int64_t> peeledLoops;
     if (loopOp->hasAttr(kPeeledLoopsLabel)) {
       auto attr = loopOp->getAttr(kPeeledLoopsLabel).cast<ArrayAttr>();
       peeledLoops =
           llvm::to_vector<4>(llvm::map_range(attr, [](Attribute attr) {
             return attr.cast<IntegerAttr>().getInt();
           }));
       // Check if the loop was already peeled.
       if (llvm::find(peeledLoops, idx) != peeledLoops.end())
         return failure();
     }
+    if (skipPartial && loopOp->hasAttr(kPartialIterationLabel))
+      // No peeling of loop nests with a partial iteration.
+      return failure();
 
     if (static_cast<int64_t>(loopOp.iterator_types().size()) <= idx)
       return failure();
 
     // Peel loop and canonicalize.
     TiledLoopOp result;
     if (failed(linalg::peelAndCanonicalizeTiledLoop(rewriter, loopOp, idx,
                                                     result)))
       return failure();
-    peeledLoops.push_back(idx);
 
     // Apply label, so that the same loop is not rewritten a second time.
+    peeledLoops.push_back(idx);
     rewriter.updateRootInPlace(loopOp, [&]() {
       loopOp->setAttr(kPeeledLoopsLabel, rewriter.getI64ArrayAttr(peeledLoops));
     });
     result->setAttr(kPeeledLoopsLabel, rewriter.getI64ArrayAttr(peeledLoops));
+    result->setAttr(kPartialIterationLabel, rewriter.getUnitAttr());
+
     return success();
   }
 
   /// Index of loop to peel.
   int64_t idx;
+
+  /// If set to true, do not peel TiledLoopOps with a partial iteration.
+  bool skipPartial;
 };
 } // namespace
 
 static void applyTiledLoopPeelingPattern(FuncOp funcOp,
-                                         ArrayRef<unsigned> loops) {
+                                         ArrayRef<unsigned> loops,
+                                         bool skipPartial) {
   MLIRContext *ctx = funcOp.getContext();
   RewritePatternSet patterns(ctx);
   for (unsigned idx : loops)
-    patterns.add<TiledLoopPeelingPattern>(ctx, idx);
+    patterns.add<TiledLoopPeelingPattern>(ctx, idx, skipPartial);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 
-  // Drop the marker.
-  funcOp.walk([](TiledLoopOp op) { op->removeAttr(kPeeledLoopsLabel); });
+  // Drop the markers.
+  funcOp.walk([](TiledLoopOp op) {
+    op->removeAttr(kPeeledLoopsLabel);
+    op->removeAttr(kPartialIterationLabel);
+  });
 }
 
 /// Apply transformations specified as patterns.
 void TestLinalgTransforms::runOnFunction() {
   auto lambda = [&](void *) {
     getFunction().walk([](LinalgOp op) {
       op->removeAttr(LinalgTransforms::kLinalgTransformMarker);
     });
if (testGenericToVectorPattern)
     return applyLinalgToVectorPatterns(getFunction());
   if (testTransformPadTensor)
     return applyPadTensorToGenericPatterns(getFunction());
   if (testGeneralizePadTensor)
     return applyGeneralizePadTensorPatterns(getFunction());
   if (testSwapSubTensorPadTensor)
     return applyExtractSliceOfPadTensorSwapPattern(getFunction());
   if (testTiledLoopPeeling.hasValue())
-    return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling);
+    return applyTiledLoopPeelingPattern(getFunction(), testTiledLoopPeeling,
+                                        skipPartial);
   if (testTileAndPadPattern)
     return applyTileAndPadPattern(getFunction(), tileSizesForPadding);
   if (testHoistPadding) {
     getFunction().walk([&](linalg::PadTensorOp padTensorOp) {
       (void)linalg::hoistPaddingOnTensors(padTensorOp, testHoistPadding);
     });
   }
   if (testInterchangePattern.hasValue())