[mlir][scf] Add option to loop pipelining to not peel the epilogue

[ThomasRaoux]

Summary

Add an option to predicate the epilogue within the kernel instead of
peeling the epilogue. This is a useful option to prevent generating
large amount of code for deep pipeline. This currently require a user
lambda to implement operation predication.

[christopherbate]

Thanks, just one small question below, LGTM otherwise.

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

151

What do you mean by this? Why would it be better?

[bondhugula]

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

147

Nit: If epilogue -> If the epilogus
Comma after 'predicated'.
user -> the user
predicated version -> the predicated version

155

You don't need the PatternRewriter. The `OpBuilder is enough? That allows wider usage outside of pattern rewrites.

158

Please remove reference to the undef op -- it's completely vague/missing information as to why.

require -> requires

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

236

This second sentence is broken. Rewrite?

266

Nit: Create

326

predicated

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

127

I think "skip-epilogue-peeling" isn't fully accurate. no-epilogue-peeling?

149

version of op.

152

Yes, this last line is vague. Instead, please mark a TODO if there is a future scope for improvement.

[ThomasRaoux]

Address review comments

[ThomasRaoux]

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

155

I want to allow user to potentially erase the operation and return a new one which requires a PatternRewriter.

158

Removed.

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

151

What I meant is that in general it is probably better to speculatively execute operations without side effect than adding control flow. This was just to point out how it is meant to be used but I think this is more confusing than anything else so I removed this comment.

152

Removed, this is not about future improvements just a comment on how it is expected to be used but looks like it is more confusing than helping.

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

 class MLIRContext;
 class Region;
 class RewriterBase;
 class TypeConverter;
 class RewritePatternSet;
 class Operation;
 class Value;
 class ValueRange;
+class PatternRewriter;
 
 namespace scf {
 
 class IfOp;
 class ForOp;
 class ParallelOp;
 
 /// Fuses all adjacent scf.parallel operations with identical bounds and step
struct PipeliningOption {
   /// it is generated.
   /// The callback passes the operation created along with the part of the
   /// pipeline and the iteration index. The iteration index is always 0 for the
   /// kernel. For the prologue and epilogue, it corresponds to the iteration
   /// peeled out of the loop in the range [0, maxStage[.
   using AnnotationlFnType =
       std::function<void(Operation *, PipelinerPart, unsigned)>;
   AnnotationlFnType annotateFn = nullptr;
-  // TODO: add option to decide if the prologue/epilogue should be peeled.
+
+  /// Control whether the epilogue should be peeled out of the loop or
+  /// operations should be predicated to skip the early stages in the last loop
+  /// iterations. If the epilogue is predicated; the user needs to provide a
+  /// lambda to generate the predicated version of operations.
+  bool peelEpilogue = true;
+
+  // Lamdba to predicate operations when the prologue or epilogue are not
+  // peeled. This takes the original operation, an i1 predicate value and the
+  // pattern rewriter.
+  using PredicateOpFn =
+      std::function<Operation *(Operation *, Value, PatternRewriter &)>;
+  PredicateOpFn predicateFn = nullptr;
+
+  // TODO: add option to decide if the prologue should be peeled.
 };
 
 /// Populate patterns for SCF software pipelining transformation.
 /// This transformation generates the pipelined loop and doesn't do any
 /// assumptions on the schedule dictated by the option structure.
 /// Software pipelining is usually done in two part. The first part of
 /// pipelining is to schedule the loop and assign a stage and cycle to each
 /// operations. This is highly dependent on the target and is implemented as an

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

protected:
   ForOp forOp;
   unsigned maxStage = 0;
   DenseMap<Operation *, unsigned> stages;
   std::vector<Operation *> opOrder;
   int64_t ub;
   int64_t lb;
   int64_t step;
   PipeliningOption::AnnotationlFnType annotateFn = nullptr;
+  bool peelEpilogue;
+  PipeliningOption::PredicateOpFn predicateFn = nullptr;
 
   // When peeling the kernel we generate several version of each value for
   // different stage of the prologue. This map tracks the mapping between
   // original Values in the loop and the different versions
   // peeled from the loop.
   DenseMap<Value, llvm::SmallVector<Value>> valueMapping;
 
   /// Assign a value to `valueMapping`, this means `val` represents the version
bool LoopPipelinerInternal::initializeLoopInfo(
   auto lowerBoundCst =
       forOp.getLowerBound().getDefiningOp<arith::ConstantIndexOp>();
   auto stepCst = forOp.getStep().getDefiningOp<arith::ConstantIndexOp>();
   if (!upperBoundCst || !lowerBoundCst || !stepCst)
     return false;
   ub = upperBoundCst.value();
   lb = lowerBoundCst.value();
   step = stepCst.value();
+  peelEpilogue = options.peelEpilogue;
+  predicateFn = options.predicateFn;
+  if (!peelEpilogue && predicateFn == nullptr)
+    return false;
   int64_t numIteration = ceilDiv(ub - lb, step);
   std::vector<std::pair<Operation *, unsigned>> schedule;
   options.getScheduleFn(forOp, schedule);
   if (schedule.empty())
     return false;
 
   opOrder.reserve(schedule.size());
   for (auto &opSchedule : schedule) {
for (unsigned stageIdx = 0; stageIdx < info.lastUseStage - info.defStage;
           valueMapping[value][maxStage - info.lastUseStage + stageIdx];
       assert(valueVersion);
       newLoopArg.push_back(valueVersion);
       loopArgMap[std::make_pair(value, info.lastUseStage - info.defStage -
                                            stageIdx)] = newLoopArg.size() - 1;
     }
   }
 
-  // Create the new kernel loop. Since we need to peel `numStages - 1`
-  // iteration we change the upper bound to remove those iterations.
-  Value newUb = rewriter.create<arith::ConstantIndexOp>(forOp.getLoc(),
-                                                        ub - maxStage * step);
+  // Create the new kernel loop. When we peel the epilgue we need to peel
+  // `numStages - 1` iterations. Then we adjust the upper bound to remove those
+  // iterations.
+  Value newUb = forOp.getUpperBound();
+  if (peelEpilogue)
+    newUb = rewriter.create<arith::ConstantIndexOp>(forOp.getLoc(),
   auto newForOp =
       rewriter.create<scf::ForOp>(forOp.getLoc(), forOp.getLowerBound(), newUb,
                                   forOp.getStep(), newLoopArg);
   return newForOp;
 }
 
 void LoopPipelinerInternal::createKernel(
     scf::ForOp newForOp,
     const llvm::MapVector<Value, LoopPipelinerInternal::LiverangeInfo>
         &crossStageValues,
     const llvm::DenseMap<std::pair<Value, unsigned>, unsigned> &loopArgMap,
     PatternRewriter &rewriter) {
   valueMapping.clear();
 
   // Create the kernel, we clone instruction based on the order given by
   // user and remap operands coming from a previous stages.
   rewriter.setInsertionPoint(newForOp.getBody(), newForOp.getBody()->begin());
   BlockAndValueMapping mapping;
   mapping.map(forOp.getInductionVar(), newForOp.getInductionVar());
   for (const auto &arg : llvm::enumerate(forOp.getRegionIterArgs())) {
     mapping.map(arg.value(), newForOp.getRegionIterArgs()[arg.index()]);
   }
+  SmallVector<Value> predicates(maxStage + 1, nullptr);
+  if (!peelEpilogue) {
+    // Create a predicate for each stage except the last stage.
+    for (unsigned i = 0; i < maxStage; i++) {
+      Value c = rewriter.create<arith::ConstantIndexOp>(
+          newForOp.getLoc(), ub - (maxStage - i) * step);
+      Value pred = rewriter.create<arith::CmpIOp>(
+          newForOp.getLoc(), arith::CmpIPredicate::slt,
+          newForOp.getInductionVar(), c);
+      predicates[i] = pred;
+    }
+  }
   for (Operation *op : opOrder) {
     int64_t useStage = stages[op];
     auto *newOp = rewriter.clone(*op, mapping);
     for (OpOperand &operand : op->getOpOperands()) {
       // Special case for the induction variable uses. We replace it with a
       // version incremented based on the stage where it is used.
       if (operand.get() == forOp.getInductionVar()) {
         rewriter.setInsertionPoint(newOp);
for (OpOperand &operand : op->getOpOperands()) {
       if (stageDef == stages.end() || stageDef->second == useStage)
         continue;
       auto remap = loopArgMap.find(
           std::make_pair(operand.get(), useStage - stageDef->second));
       assert(remap != loopArgMap.end());
       newOp->setOperand(operand.getOperandNumber(),
                         newForOp.getRegionIterArgs()[remap->second]);
     }
+    if (predicates[useStage]) {
+      newOp = predicateFn(newOp, predicates[useStage], rewriter);
+      // Remap the results to the new predicated one.
+      for (auto values : llvm::zip(op->getResults(), newOp->getResults()))
+        mapping.map(std::get<0>(values), std::get<1>(values));
+    }
+    rewriter.setInsertionPointAfter(newOp);
     if (annotateFn)
       annotateFn(newOp, PipeliningOption::PipelinerPart::Kernel, 0);
   }
 
   // Collect the Values that need to be returned by the forOp. For each
   // value we need to have `LastUseStage - DefStage` number of versions
   // returned.
   // We create a mapping between original values and the associated loop
LogicalResult matchAndRewrite(ForOp forOp,
     llvm::DenseMap<std::pair<Value, unsigned>, unsigned> loopArgMap;
     // 3. Create the new kernel loop and return the block arguments mapping.
     ForOp newForOp =
         pipeliner.createKernelLoop(crossStageValues, rewriter, loopArgMap);
     // Create the kernel block, order ops based on user choice and remap
     // operands.
     pipeliner.createKernel(newForOp, crossStageValues, loopArgMap, rewriter);
 
+    llvm::SmallVector<Value> returnValues =
+        newForOp.getResults().take_front(forOp->getNumResults());
+    if (options.peelEpilogue) {
       // 4. Emit the epilogue after the new forOp.
       rewriter.setInsertionPointAfter(newForOp);
-    llvm::SmallVector<Value> returnValues = pipeliner.emitEpilogue(rewriter);
-
+      returnValues = pipeliner.emitEpilogue(rewriter);
+    }
     // 5. Erase the original loop and replace the uses with the epilogue output.
     if (forOp->getNumResults() > 0)
       rewriter.replaceOp(forOp, returnValues);
     else
       rewriter.eraseOp(forOp);
 
     return success();
   }

mlir/test/Dialect/SCF/loop-pipelining.mlir

 // RUN: mlir-opt %s -test-scf-pipelining -split-input-file | FileCheck %s
 // RUN: mlir-opt %s -test-scf-pipelining=annotate -split-input-file | FileCheck %s --check-prefix ANNOTATE
+// RUN: mlir-opt %s -test-scf-pipelining=no-epilogue-peeling -split-input-file | FileCheck %s --check-prefix NOEPILOGUE
 
 // CHECK-LABEL: simple_pipeline(
 //  CHECK-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
 //   CHECK-DAG:   %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG:   %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG:   %[[C3:.*]] = arith.constant 3 : index
 // Prologue:
 //       CHECK:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>

 //  ANNOTATE:     memref.load {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
 //  ANNOTATE:     scf.yield
 //  ANNOTATE:   }
 // Epilogue:
 //  ANNOTATE:   memref.store {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "epilogue"}
 //  ANNOTATE:   arith.addf {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "epilogue"}
 //  ANNOTATE:   memref.store {{.*}} {__test_pipelining_iteration = 1 : i32, __test_pipelining_part = "epilogue"}
 
+// NOEPILOGUE-LABEL: three_stage(
+//  NOEPILOGUE-SAME:   %[[A:.*]]: memref<?xf32>, %[[R:.*]]: memref<?xf32>) {
+//   NOEPILOGUE-DAG:   %[[C0:.*]] = arith.constant 0 : index
+//   NOEPILOGUE-DAG:   %[[C1:.*]] = arith.constant 1 : index
+//   NOEPILOGUE-DAG:   %[[C2:.*]] = arith.constant 2 : index
+//   NOEPILOGUE-DAG:   %[[C3:.*]] = arith.constant 3 : index
+//   NOEPILOGUE-DAG:   %[[C4:.*]] = arith.constant 4 : index
+//   NOEPILOGUE-DAG:   %[[CF:.*]] = arith.constant 0.000000e+00 : f32
+// Prologue:
+//       NOEPILOGUE:   %[[L0:.*]] = memref.load %[[A]][%[[C0]]] : memref<?xf32>
+//  NOEPILOGUE-NEXT:   %[[ADD0:.*]] = arith.addf %[[L0]], %{{.*}} : f32
+//  NOEPILOGUE-NEXT:   %[[L1:.*]] = memref.load %[[A]][%[[C1]]] : memref<?xf32>
+// Kernel:
+//  NOEPILOGUE-NEXT:   %[[LR:.*]]:2 = scf.for %[[IV:.*]] = %[[C0]] to %[[C4]]
+//  NOEPILOGUE-SAME:     step %[[C1]] iter_args(%[[ADDARG:.*]] = %[[ADD0]],
+//  NOEPILOGUE-SAME:     %[[LARG:.*]] = %[[L1]]) -> (f32, f32) {
+//   NOEPILOGUE-DAG:     %[[S0:.*]] = arith.cmpi slt, %[[IV]], %[[C2]] : index
+//   NOEPILOGUE-DAG:     %[[S1:.*]] = arith.cmpi slt, %[[IV]], %[[C3]] : index
+//  NOEPILOGUE-NEXT:     memref.store %[[ADDARG]], %[[R]][%[[IV]]] : memref<?xf32>
+//  NOEPILOGUE-NEXT:     %[[ADD1:.*]] = scf.if %[[S1]] -> (f32) {
+//  NOEPILOGUE-NEXT:       %[[PADD:.*]] = arith.addf %[[LARG]], %{{.*}} : f32
+//  NOEPILOGUE-NEXT:       scf.yield %[[PADD]] : f32
+//  NOEPILOGUE-NEXT:     } else {
+//  NOEPILOGUE-NEXT:       scf.yield %[[CF]] : f32
+//  NOEPILOGUE-NEXT:     }
+//  NOEPILOGUE-NEXT:     %[[IV2:.*]] = arith.addi %[[IV]], %[[C2]] : index
+//  NOEPILOGUE-NEXT:     %[[L3:.*]] = scf.if %[[S0]] -> (f32) {
+//  NOEPILOGUE-NEXT:       %[[PL:.*]] = memref.load %[[A]][%[[IV2]]] : memref<?xf32>
+//  NOEPILOGUE-NEXT:       scf.yield %[[PL]] : f32
+//  NOEPILOGUE-NEXT:     } else {
+//  NOEPILOGUE-NEXT:       scf.yield %[[CF]] : f32
+//  NOEPILOGUE-NEXT:     }
+//  NOEPILOGUE-NEXT:     scf.yield %[[ADD1]], %[[L3]] : f32, f32
+//  NOEPILOGUE-NEXT:   }
+// No epilogue should be generated.
+//   NOEPILOGUE-NOT:   memref.store
+//       NOEPILOGUE:   return
+
 func.func @three_stage(%A: memref<?xf32>, %result: memref<?xf32>) {
   %c0 = arith.constant 0 : index
   %c1 = arith.constant 1 : index
   %c4 = arith.constant 4 : index
   %cf = arith.constant 1.0 : f32
   scf.for %i0 = %c0 to %c4 step %c1 {
     %A_elem = memref.load %A[%i0] { __test_pipelining_stage__ = 0, __test_pipelining_op_order__ = 2 } : memref<?xf32>
     %A1_elem = arith.addf %A_elem, %cf { __test_pipelining_stage__ = 1, __test_pipelining_op_order__ = 1 } : f32

mlir/test/lib/Dialect/SCF/TestSCFUtils.cpp

struct TestSCFPipeliningPass
   StringRef getArgument() const final { return "test-scf-pipelining"; }
   StringRef getDescription() const final { return "test scf.forOp pipelining"; }
 
   Option<bool> annotatePipeline{
       *this, "annotate",
       llvm::cl::desc("Annote operations during loop pipelining transformation"),
       llvm::cl::init(false)};
 
+  Option<bool> noEpiloguePeeling{
+      *this, "no-epilogue-peeling",
+      llvm::cl::desc("Use predicates instead of peeling the epilogue."),
+      llvm::cl::init(false)};
+
   static void
   getSchedule(scf::ForOp forOp,
               std::vector<std::pair<Operation *, unsigned>> &schedule) {
     if (!forOp->hasAttr(kTestPipeliningLoopMarker))
       return;
     schedule.resize(forOp.getBody()->getOperations().size() - 1);
     forOp.walk([&schedule](Operation *op) {
       auto attrStage =
           op->getAttrOfType<IntegerAttr>(kTestPipeliningStageMarker);
       auto attrCycle =
           op->getAttrOfType<IntegerAttr>(kTestPipeliningOpOrderMarker);
       if (attrCycle && attrStage) {
         schedule[attrCycle.getInt()] =
             std::make_pair(op, unsigned(attrStage.getInt()));
       }
     });
   }
 
+  /// Helper to generate "predicated" version of `op`. For simplicity we just
+  /// wrap the operation in a scf.ifOp operation.
+  static Operation *predicateOp(Operation *op, Value pred,
+                                PatternRewriter &rewriter) {
+    Location loc = op->getLoc();
+    auto ifOp =
+        rewriter.create<scf::IfOp>(loc, op->getResultTypes(), pred, true);
+    // True branch.
+    op->moveBefore(&ifOp.getThenRegion().front(),
+                   ifOp.getThenRegion().front().end());
+    rewriter.setInsertionPointAfter(op);
+    rewriter.create<scf::YieldOp>(loc, op->getResults());
+    // False branch.
+    rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
+    SmallVector<Value> zeros;
+    for (Type type : op->getResultTypes()) {
+      zeros.push_back(
+          rewriter.create<arith::ConstantOp>(loc, rewriter.getZeroAttr(type)));
+    }
+    rewriter.create<scf::YieldOp>(loc, zeros);
+    return ifOp.getOperation();
+  }
+
   static void annotate(Operation *op,
                        mlir::scf::PipeliningOption::PipelinerPart part,
                        unsigned iteration) {
     OpBuilder b(op);
     switch (part) {
     case mlir::scf::PipeliningOption::PipelinerPart::Prologue:
       op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("prologue"));
       break;
struct TestSCFPipeliningPass
   }
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     mlir::scf::PipeliningOption options;
     options.getScheduleFn = getSchedule;
     if (annotatePipeline)
       options.annotateFn = annotate;
+    if (noEpiloguePeeling) {
+      options.peelEpilogue = false;
+      options.predicateFn = predicateOp;
+    }
     scf::populateSCFLoopPipeliningPatterns(patterns, options);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
     getOperation().walk([](Operation *op) {
       // Clean up the markers.
       op->removeAttr(kTestPipeliningStageMarker);
       op->removeAttr(kTestPipeliningOpOrderMarker);
     });
   }