[LoopPipelining] factor out applyLoopPipelining() method (#20)

include/Transforms/Passes.h

@@ -27,6 +27,10 @@ bool applyAffineLoopPerfection(AffineForOp loop, OpBuilder &builder);
 /// Apply remove variable bound to all inner loops of the input loop.
 bool applyRemoveVariableBound(AffineForOp loop, OpBuilder &builder);
 
+/// Apply loop pipelining to the input loop, all inner loops are automatically
+/// fully unrolled.
+bool applyLoopPipelining(AffineForOp loop, OpBuilder &builder);
+
 //===----------------------------------------------------------------------===//
 // Optimization Pass Entries
 //===----------------------------------------------------------------------===//

include/Transforms/Passes.td

@@ -32,7 +32,7 @@ def LoopPipelining : Pass<"loop-pipelining", "FuncOp"> {
   let constructor = "mlir::scalehls::createLoopPipeliningPass()";
 
   let options = [
-    Option<"pipelineLevel", "pipeline-level", "unsigned", /*default=*/"1", 
+    Option<"pipelineLevel", "pipeline-level", "unsigned", /*default=*/"0", 
            "Positive number: loop level to be pipelined (from innermost)">
   ];
 }

lib/Transforms/AffineLoopPerfection.cpp

@@ -46,7 +46,7 @@ bool scalehls::applyAffineLoopPerfection(AffineForOp innermostLoop,
   while (true) {
     // Get the parent loop of the child loop.
     auto childLoop = loops.back();
-    auto loop = dyn_cast<AffineForOp>(childLoop.getParentOp());
+    auto loop = childLoop.getParentOfType<AffineForOp>();
 
     // Break the procedure if the parent operation is no longer a loop.
     if (!loop)

lib/Transforms/AffineStoreForward.cpp

@@ -16,6 +16,9 @@
 using namespace mlir;
 using namespace scalehls;
 
+// The difference between this pass and built-in memref-dataflow-opt is this
+// pass support to forward the StoreOps that are conditionally executed.
+
 namespace {
 // The store to load forwarding relies on three conditions:
 //

lib/Transforms/ArrayPartition.cpp

@@ -137,6 +137,9 @@ void ArrayPartition::runOnOperation() {
       // TODO: how to decide which to pick?
       applyArrayPartition<mlir::AffineLoadOp>(loadMap, builder);
       applyArrayPartition<mlir::AffineStoreOp>(storeMap, builder);
+
+      // TODO: how to handle the case when different sub-functions have
+      // different array partition strategy selected?
     }
   }
 

lib/Transforms/LoopPipelining.cpp

@@ -2,6 +2,7 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "Analysis/Utils.h"
 #include "Transforms/Passes.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
@@ -13,60 +14,80 @@ using namespace scalehls;
 
 namespace {
 struct LoopPipelining : public LoopPipeliningBase<LoopPipelining> {
-  void runOnOperation() override;
-};
-} // namespace
-
-void LoopPipelining::runOnOperation() {
+  void runOnOperation() override {
     auto func = getOperation();
     auto builder = OpBuilder(func);
 
-  // Walk through loops in the function.
-  for (auto forOp : func.getOps<mlir::AffineForOp>()) {
-    // TODO: support more fine-grained pipeline insertion configuration.
-    SmallVector<mlir::AffineForOp, 4> nestedLoops;
-    forOp.walk([&](mlir::AffineForOp loop) { nestedLoops.push_back(loop); });
+    // Walk through all loops.
+    for (auto forOp : func.getOps<AffineForOp>()) {
+      // Collect all innermost loops.
+      SmallVector<AffineForOp, 4> innermostLoops;
+      forOp.walk([&](AffineForOp loop) {
+        if (getChildLoopNum(loop) == 0)
+          innermostLoops.push_back(loop);
+      });
 
-    auto targetLoop = nestedLoops.back();
-    if (nestedLoops.size() > pipelineLevel)
-      targetLoop = *std::next(nestedLoops.begin(), pipelineLevel);
+      // Apply loop pipelining to coresponding level of each innermost loop.
+      for (auto loop : innermostLoops) {
+        auto currentLoop = loop;
+        unsigned loopLevel = 0;
+        while (true) {
+          auto parentLoop = currentLoop.getParentOfType<AffineForOp>();
 
+          // If meet the outermost loop, pipeline the current loop.
+          if (!parentLoop || pipelineLevel == loopLevel) {
+            applyLoopPipelining(currentLoop, builder);
+            break;
+          }
+
+          // Move to the next loop level.
+          currentLoop = parentLoop;
+          ++loopLevel;
+        }
+      }
+    }
+
+    // Canonicalize the IR after loop pipelining.
+    OwningRewritePatternList patterns;
+    for (auto *op : builder.getContext()->getRegisteredOperations())
+      op->getCanonicalizationPatterns(patterns, builder.getContext());
+
+    applyPatternsAndFoldGreedily(func.getRegion(), std::move(patterns));
+  }
+};
+} // namespace
+
+/// Apply loop pipelining to the input loop, all inner loops are automatically
+/// fully unrolled.
+bool scalehls::applyLoopPipelining(AffineForOp targetLoop, OpBuilder &builder) {
   targetLoop.setAttr("pipeline", builder.getBoolAttr(true));
 
   // All inner loops of the pipelined loop are automatically unrolled.
-    targetLoop.walk([&](mlir::AffineForOp loop) {
+  targetLoop.walk([&](AffineForOp loop) {
     if (loop != targetLoop)
       loopUnrollFull(loop);
   });
 
   // All outer loops that perfect nest the pipelined loop can be flattened.
-    SmallVector<mlir::AffineForOp, 4> flattenedLoops;
+  SmallVector<AffineForOp, 4> flattenedLoops;
   flattenedLoops.push_back(targetLoop);
   while (true) {
     auto currentLoop = flattenedLoops.back();
-      if (auto outerLoop = currentLoop.getParentOfType<mlir::AffineForOp>()) {
-        // Only if the current loop is the only child loop of the outer loop,
-        // the outer loop can be flattened into the current loop.
+    if (auto outerLoop = currentLoop.getParentOfType<AffineForOp>()) {
+      // Only if the current loop is the only child loop of the outer loop, the
+      // outer loop can be flattened into the current loop.
       auto &body = outerLoop.getLoopBody().front();
       if (&body.front() == currentLoop && body.getOperations().size() == 2) {
         flattenedLoops.push_back(outerLoop);
-          outerLoop.setAttr("flatten", builder.getBoolAttr("true"));
+        outerLoop.setAttr("flatten", builder.getBoolAttr(true));
       } else
         break;
     } else
       break;
   }
-  }
 
-  // Canonicalize the IR after loop unrolling.
-  OwningRewritePatternList patterns;
-
-  auto *context = &getContext();
-  for (auto *op : context->getRegisteredOperations())
-    op->getCanonicalizationPatterns(patterns, context);
-
-  applyPatternsAndFoldGreedily(func.getOperation()->getRegions(),
-                               std::move(patterns));
+  // For now, this method will always success.
+  return true;
 }
 
 std::unique_ptr<mlir::Pass> scalehls::createLoopPipeliningPass() {