[mlir][scf] Add callback to annotate ops during pipelining

[ThomasRaoux]

Summary

This allow user to register a callback that can annotate operations
during software pipelining. This allows user potential annotate op to
know what part of the pipeline they correspond to.

This is analogue to what is done in loop unrolling: https://reviews.llvm.org/D107789

[antiagainst]

LGTM! Just a few nits.

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

130

I think for enum class we just need to use style like Prologue, Kernel, etc.? All caps are for enum and there is no need to have the prefix anymore given enum class would need namespace anyway.

137

Could you add documentation for the parameters to this function, esp. the last unsigned? Right now it's unclear what it means without looking at how it is used.

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

143

Don't need default for enum classes. New cases will be checked by the compilers.

[ThomasRaoux]

Address review comments

[ThomasRaoux]

Thanks @antiagainst

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

void populateSCFStructuralTypeConversionsAndLegality(
     ConversionTarget &target);
 
 /// Options to dictate how loops should be pipelined.
 struct PipeliningOption {
   /// Lambda returning all the operation in the forOp, with their stage, in the
   /// order picked for the pipelined loop.
   using GetScheduleFnType = std::function<void(
       scf::ForOp, std::vector<std::pair<Operation *, unsigned>> &)>;
-  GetScheduleFnType getScheduleFn;
+  GetScheduleFnType getScheduleFn = nullptr;
+  enum class PipelinerPart {
+    PIPELINER_PROLOGUE,
+    PIPELINER_KERNEL,
+    PIPELINER_EPILOGUE,
+  };
+  /// Lambda called by the pipeliner to allow the user to annotate the IR while
+  /// it is generated.
+  using AnnotationlFnType =
+      std::function<void(Operation *, PipelinerPart, unsigned)>;
+  AnnotationlFnType annotateFn = nullptr;
   // TODO: add option to decide if the prologue/epilogue 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

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;
 
   // 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(
   // source of all the scf.yield operands needs to be defined by operations in
   // the loop.
   if (llvm::any_of(forOp.getBody()->getTerminator()->getOperands(),
                    [this](Value operand) {
                      Operation *def = operand.getDefiningOp();
                      return !def || stages.find(def) == stages.end();
                    }))
     return false;
+  annotateFn = options.annotateFn;
   return true;
 }
 
 void LoopPipelinerInternal::emitPrologue(PatternRewriter &rewriter) {
   // Initialize the iteration argument to the loop initiale values.
   for (BlockArgument &arg : forOp.getRegionIterArgs()) {
     OpOperand &operand = forOp.getOpOperandForRegionIterArg(arg);
     setValueMapping(arg, operand.get(), 0);
   }
   auto yield = cast<scf::YieldOp>(forOp.getBody()->getTerminator());
   for (int64_t i = 0; i < maxStage; i++) {
     // special handling for induction variable as the increment is implicit.
     Value iv =
         rewriter.create<arith::ConstantIndexOp>(forOp.getLoc(), lb + i * step);
     setValueMapping(forOp.getInductionVar(), iv, i);
     for (Operation *op : opOrder) {
       if (stages[op] > i)
         continue;
       Operation *newOp = rewriter.clone(*op);
       for (unsigned opIdx = 0; opIdx < op->getNumOperands(); opIdx++) {
         auto it = valueMapping.find(op->getOperand(opIdx));
         if (it != valueMapping.end())
           newOp->setOperand(opIdx, it->second[i - stages[op]]);
       }
+      if (annotateFn)
+        annotateFn(newOp, PipeliningOption::PipelinerPart::PIPELINER_PROLOGUE,
+                   i);
       for (unsigned destId : llvm::seq(unsigned(0), op->getNumResults())) {
         setValueMapping(op->getResult(destId), newOp->getResult(destId),
                         i - stages[op]);
         // If the value is a loop carried dependency update the loop argument
         // mapping.
         for (OpOperand &operand : yield->getOpOperands()) {
           if (operand.get() != op->getResult(destId))
             continue;
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 (annotateFn)
+      annotateFn(newOp, PipeliningOption::PipelinerPart::PIPELINER_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
   // returned values that will be needed by the epilogue.
   llvm::SmallVector<Value> yieldOperands;
for (Operation *op : opOrder) {
       for (unsigned opIdx = 0; opIdx < op->getNumOperands(); opIdx++) {
         auto it = valueMapping.find(op->getOperand(opIdx));
         if (it != valueMapping.end()) {
           Value v = it->second[maxStage - stages[op] + i];
           assert(v);
           newOp->setOperand(opIdx, v);
         }
       }
+      if (annotateFn)
+        annotateFn(newOp, PipeliningOption::PipelinerPart::PIPELINER_EPILOGUE,
+                   i - 1);
       for (unsigned destId : llvm::seq(unsigned(0), op->getNumResults())) {
         setValueMapping(op->getResult(destId), newOp->getResult(destId),
                         maxStage - stages[op] + i);
         // If the value is a loop carried dependency update the loop argument
         // mapping and keep track of the last version to replace the original
         // forOp uses.
         for (OpOperand &operand :
              forOp.getBody()->getTerminator()->getOpOperands()) {

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
 
 // 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>

 //  CHECK-NEXT:     %[[IV2:.*]] = arith.addi %[[IV]], %[[C2]] : index
 //  CHECK-NEXT:     %[[L3:.*]] = memref.load %[[A]][%[[IV2]]] : memref<?xf32>
 //  CHECK-NEXT:     scf.yield %[[ADD1]], %[[L3]] : f32, f32
 //  CHECK-NEXT:   }
 // Epilogue:
 //  CHECK-NEXT:   memref.store %[[LR]]#0, %[[R]][%[[C2]]] : memref<?xf32>
 //  CHECK-NEXT:   %[[ADD2:.*]] = arith.addf %[[LR]]#1, %{{.*}} : f32
 //  CHECK-NEXT:   memref.store %[[ADD2]], %[[R]][%[[C3]]] : memref<?xf32>
+
+// Prologue:
+//  ANNOTATE:   memref.load {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "prologue"}
+//  ANNOTATE:   memref.load {{.*}} {__test_pipelining_iteration = 1 : i32, __test_pipelining_part = "prologue"}
+// Kernel:
+//  ANNOTATE:   scf.for
+//  ANNOTATE:     memref.store {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
+//  ANNOTATE:     arith.addf {{.*}} {__test_pipelining_iteration = 0 : i32, __test_pipelining_part = "kernel"}
+//  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"}
+
 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

 static const StringLiteral kTestPipeliningLoopMarker =
     "__test_pipelining_loop__";
 static const StringLiteral kTestPipeliningStageMarker =
     "__test_pipelining_stage__";
 /// Marker to express the order in which operations should be after pipelining.
 static const StringLiteral kTestPipeliningOpOrderMarker =
     "__test_pipelining_op_order__";
 
+static const StringLiteral kTestPipeliningAnnotationPart =
+    "__test_pipelining_part";
+static const StringLiteral kTestPipeliningAnnotationIteration =
+    "__test_pipelining_iteration";
+
 class TestSCFPipeliningPass
     : public PassWrapper<TestSCFPipeliningPass, OperationPass<FuncOp>> {
 public:
+  TestSCFPipeliningPass() = default;
+  TestSCFPipeliningPass(const TestSCFPipeliningPass &) {}
   StringRef getArgument() const final { return "test-scf-pipelining"; }
   StringRef getDescription() const final { return "test scf.forOp pipelining"; }
-  explicit TestSCFPipeliningPass() = default;
+
+  Option<bool> annotatePipeline{
+      *this, "annotate",
+      llvm::cl::desc("Annote operations during loop pipelining transformation"),
+      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()));
       }
     });
   }
 
+  static void annotate(Operation *op,
+                       mlir::scf::PipeliningOption::PipelinerPart part,
+                       unsigned iteration) {
+    OpBuilder b(op);
+    switch (part) {
+    case mlir::scf::PipeliningOption::PipelinerPart::PIPELINER_PROLOGUE:
+      op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("prologue"));
+      break;
+    case mlir::scf::PipeliningOption::PipelinerPart::PIPELINER_KERNEL:
+      op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("kernel"));
+      break;
+    case mlir::scf::PipeliningOption::PipelinerPart::PIPELINER_EPILOGUE:
+      op->setAttr(kTestPipeliningAnnotationPart, b.getStringAttr("epilogue"));
+      break;
+    default:
+      llvm_unreachable("unknown part");
+    }
+    op->setAttr(kTestPipeliningAnnotationIteration,
+                b.getI32IntegerAttr(iteration));
+  }
+
   void getDependentDialects(DialectRegistry &registry) const override {
     registry.insert<arith::ArithmeticDialect, StandardOpsDialect>();
   }
 
   void runOnOperation() override {
     RewritePatternSet patterns(&getContext());
     mlir::scf::PipeliningOption options;
     options.getScheduleFn = getSchedule;
-
+    if (annotatePipeline)
+      options.annotateFn = annotate;
     scf::populateSCFLoopPipeliningPatterns(patterns, options);
     (void)applyPatternsAndFoldGreedily(getOperation(), std::move(patterns));
     getOperation().walk([](Operation *op) {
       // Clean up the markers.
       op->removeAttr(kTestPipeliningStageMarker);
       op->removeAttr(kTestPipeliningOpOrderMarker);
     });
   }