[mlir][SCF] Do not peel loops inside partial iterations

[springerm]

Summary

Do not apply loop peeling to loops that are contained in the partial iteration of an already peeled loop. This is to avoid code explosion when dealing with large loop nests. Can be controlled with a new pass option skip-partial.

[herhut]

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

452–456

nit: you could also do this with a single walk.

[springerm]

addressed comments

mlir/include/mlir/Dialect/SCF/Passes.td

def SCFBufferize : FunctionPass<"scf-bufferize"> {
   let constructor = "mlir::createSCFBufferizePass()";
   let dependentDialects = ["memref::MemRefDialect"];
 }
 
 def SCFForLoopPeeling
     : FunctionPass<"for-loop-peeling"> {
   let summary = "Peel `for` loops at their upper bounds.";
   let constructor = "mlir::createForLoopPeelingPass()";
+  let options = [
+    Option<"skipPartial", "skip-partial", "bool",
+           /*default=*/"true",
+           "Do not peel loops inside of the last, partial iteration of another "
+           "already peeled loop.">
+  ];
   let dependentDialects = ["AffineDialect"];
 }
 
 def SCFForLoopSpecialization
     : FunctionPass<"for-loop-specialization"> {
   let summary = "Specialize `for` loops for vectorization";
   let constructor = "mlir::createForLoopSpecializationPass()";
 }

mlir/include/mlir/Dialect/SCF/Transforms.h

 
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
 /// into one scf.parallel operations. Uses a naive aliasing and dependency
 /// analysis.
 void naivelyFuseParallelOps(Region &region);
 
 /// Rewrite a for loop with bounds/step that potentially do not divide evenly
 /// into a for loop where the step divides the iteration space evenly, followed
-/// by an scf.if for the last (partial) iteration (if any). This transformation
-/// is called "loop peeling".
+/// by an scf.if for the last (partial) iteration (if any; returned via `ifOp`).
+/// This transformation is called "loop peeling".
 ///
 /// This transformation is beneficial for a wide range of transformations such
 /// as vectorization or loop tiling: It enables additional canonicalizations
 /// inside the peeled loop body such as rewriting masked loads into unmaked
 /// loads.
 ///
 /// E.g., assuming a lower bound of 0 (for illustration purposes):
 /// ```

 /// not rewritten if:
 /// * Loop step size is 1 or
 /// * Loop bounds and step size are static, and step already divides the
 ///   iteration space evenly.
 ///
 /// Note: This function rewrites the given scf.for loop in-place and creates a
 /// new scf.if operation for the last iteration. It replaces all uses of the
 /// unpeeled loop with the results of the newly generated scf.if.
-LogicalResult peelAndCanonicalizeForLoop(RewriterBase &rewriter, ForOp forOp);
+LogicalResult peelAndCanonicalizeForLoop(RewriterBase &rewriter, ForOp forOp,
+                                         scf::IfOp &ifOp);
 
 /// Tile a parallel loop of the form
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4, %arg5)
 ///
 /// into
 ///   scf.parallel (%i0, %i1) = (%arg0, %arg1) to (%arg2, %arg3)
 ///                                             step (%arg4*tileSize[0],

mlir/lib/Dialect/SCF/Transforms/LoopSpecialization.cpp

if (insideLoop) {
     // loop would run for at least one more iteration.)
     constraints.addInequality({1, -1, 1, -1});
   }
 
   return canonicalizeAffineMinOp(rewriter, minOp, constraints);
 }
 
 LogicalResult mlir::scf::peelAndCanonicalizeForLoop(RewriterBase &rewriter,
-                                                    ForOp forOp) {
+                                                    ForOp forOp,
+                                                    scf::IfOp &ifOp) {
   Value ub = forOp.upperBound();
-  scf::IfOp ifOp;
   Value splitBound;
   if (failed(peelForLoop(rewriter, forOp, ifOp, splitBound)))
     return failure();
 
   // Rewrite affine.min ops.
   forOp.walk([&](AffineMinOp minOp) {
     (void)rewritePeeledAffineOp(rewriter, minOp, forOp.getInductionVar(), ub,
                                 forOp.step(), /*insideLoop=*/true);
   });
   ifOp.walk([&](AffineMinOp minOp) {
     (void)rewritePeeledAffineOp(rewriter, minOp, splitBound, ub, forOp.step(),
                                 /*insideLoop=*/false);
   });
 
   return success();
 }
 
 static constexpr char kPeeledLoopLabel[] = "__peeled_loop__";
+static constexpr char kPartialIterationLabel[] = "__partial_iteration__";
 
 namespace {
 struct ForLoopPeelingPattern : public OpRewritePattern<ForOp> {
-  using OpRewritePattern<ForOp>::OpRewritePattern;
+  ForLoopPeelingPattern(MLIRContext *ctx, bool skipPartial)
+      : OpRewritePattern<ForOp>(ctx), skipPartial(skipPartial) {}
 
   LogicalResult matchAndRewrite(ForOp forOp,
                                 PatternRewriter &rewriter) const override {
+    // Do not peel already peeled loops.
     if (forOp->hasAttr(kPeeledLoopLabel))
       return failure();
-    if (failed(peelAndCanonicalizeForLoop(rewriter, forOp)))
+    if (skipPartial) {
+      // No peeling of loops inside the partial iteration (scf.if) of another
+      // peeled loop.
+      Operation *op = forOp.getOperation();
+      while ((op = op->getParentOfType<scf::IfOp>())) {
+        if (op->hasAttr(kPartialIterationLabel))
+          return failure();
+      }
+    }
+    // Apply loop peeling.
+    scf::IfOp ifOp;
+    if (failed(peelAndCanonicalizeForLoop(rewriter, forOp, ifOp)))
       return failure();
     // Apply label, so that the same loop is not rewritten a second time.
     rewriter.updateRootInPlace(forOp, [&]() {
       forOp->setAttr(kPeeledLoopLabel, rewriter.getUnitAttr());
     });
+    ifOp->setAttr(kPartialIterationLabel, rewriter.getUnitAttr());
     return success();
   }
+
+  /// If set to true, loops inside partial iterations of another peeled loop
+  /// are not peeled. This reduces the size of the generated code. Partial
+  /// iterations are not usually performance critical.
+  /// Note: Takes into account the entire chain of parent operations, not just
+  /// the direct parent.
+  bool skipPartial;
 };
 } // namespace
 
 namespace {
 struct ParallelLoopSpecialization
     : public SCFParallelLoopSpecializationBase<ParallelLoopSpecialization> {
   void runOnFunction() override {
     getFunction().walk(
         [](ParallelOp op) { specializeParallelLoopForUnrolling(op); });
   }
 };
 
 struct ForLoopSpecialization
     : public SCFForLoopSpecializationBase<ForLoopSpecialization> {
   void runOnFunction() override {
     getFunction().walk([](ForOp op) { specializeForLoopForUnrolling(op); });
   }
 };
 
 struct ForLoopPeeling : public SCFForLoopPeelingBase<ForLoopPeeling> {
   void runOnFunction() override {
     FuncOp funcOp = getFunction();
     MLIRContext *ctx = funcOp.getContext();
     RewritePatternSet patterns(ctx);
-    patterns.add<ForLoopPeelingPattern>(ctx);
+    patterns.add<ForLoopPeelingPattern>(ctx, skipPartial);
     (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 
-    // Drop the marker.
-    funcOp.walk([](ForOp op) { op->removeAttr(kPeeledLoopLabel); });
+    // Drop the markers.
+    funcOp.walk([](Operation *op) {
+      op->removeAttr(kPeeledLoopLabel);
+      op->removeAttr(kPartialIterationLabel);
+    });
   }
 };
 } // namespace
 
 std::unique_ptr<Pass> mlir::createParallelLoopSpecializationPass() {
   return std::make_unique<ParallelLoopSpecialization>();
 }
 
 std::unique_ptr<Pass> mlir::createForLoopSpecializationPass() {
   return std::make_unique<ForLoopSpecialization>();
 }
 
 std::unique_ptr<Pass> mlir::createForLoopPeelingPass() {
   return std::make_unique<ForLoopPeeling>();
 }

mlir/test/Dialect/SCF/for-loop-peeling.mlir

 // RUN: mlir-opt %s -for-loop-peeling -canonicalize -split-input-file | FileCheck %s
+// RUN: mlir-opt %s -for-loop-peeling=skip-partial=false -canonicalize -split-input-file | FileCheck %s -check-prefix=CHECK-NO-SKIP
 
 //  CHECK-DAG: #[[MAP0:.*]] = affine_map<()[s0, s1, s2] -> (s1 - (s1 - s0) mod s2)>
 //  CHECK-DAG: #[[MAP1:.*]] = affine_map<()[s0, s1, s2] -> (-(s0 - (s0 - s1) mod s2) + s0)>
 //      CHECK: func @fully_dynamic_bounds(
 // CHECK-SAME:     %[[LB:.*]]: index, %[[UB:.*]]: index, %[[STEP:.*]]: index
 //      CHECK:   %[[C0_I32:.*]] = constant 0 : i32
 //      CHECK:   %[[NEW_UB:.*]] = affine.apply #[[MAP0]]()[%[[LB]], %[[UB]], %[[STEP]]]
 //      CHECK:   %[[LOOP:.*]] = scf.for %[[IV:.*]] = %[[LB]] to %[[NEW_UB]]
scf.for %iv = %lb to %ub step %step {
 
     // min(%ub - %iv, %step, %some_val) cannot be simplified because the range
     // of %some_val is unknown.
     %m4 = affine.min #map4(%ub, %iv, %some_val)[%step]
     memref.store %m4, %d[%c4] : memref<?xindex>
   }
   return
 }
+
+// -----
+
+//     CHECK: func @nested_loops
+//     CHECK:   scf.for {{.*}} {
+//     CHECK:     scf.for {{.*}} {
+//     CHECK:     }
+//     CHECK:     scf.if {{.*}} {
+//     CHECK:     }
+//     CHECK:   }
+//     CHECK:   scf.if {{.*}} {
+//     CHECK:     scf.for {{.*}} {
+//     CHECK:     }
+// CHECK-NOT:     scf.if
+//     CHECK:   }
+
+//     CHECK-NO-SKIP: func @nested_loops
+//     CHECK-NO-SKIP:   scf.for {{.*}} {
+//     CHECK-NO-SKIP:     scf.for {{.*}} {
+//     CHECK-NO-SKIP:     }
+//     CHECK-NO-SKIP:     scf.if {{.*}} {
+//     CHECK-NO-SKIP:     }
+//     CHECK-NO-SKIP:   }
+//     CHECK-NO-SKIP:   scf.if {{.*}} {
+//     CHECK-NO-SKIP:     scf.for {{.*}} {
+//     CHECK-NO-SKIP:     }
+//     CHECK-NO-SKIP:     scf.if {{.*}} {
+//     CHECK-NO-SKIP:     }
+//     CHECK-NO-SKIP:   }
+#map = affine_map<(d0, d1)[s0] -> (s0, d0 - d1)>
+func @nested_loops(%lb0: index, %lb1 : index, %ub0: index, %ub1: index,
+                   %step: index) -> i32 {
+  %c0 = constant 0 : i32
+  %r0 = scf.for %iv0 = %lb0 to %ub0 step %step iter_args(%arg0 = %c0) -> i32 {
+    %r1 = scf.for %iv1 = %lb1 to %ub1 step %step iter_args(%arg1 = %arg0) -> i32 {
+      %s = affine.min #map(%ub1, %iv1)[%step]
+      %casted = index_cast %s : index to i32
+      %0 = addi %arg1, %casted : i32
+      scf.yield %0 : i32
+    }
+    %1 = addi %arg0, %r1 : i32
+    scf.yield %1 : i32
+  }
+  return %r0 : i32
+}