Create a LoopUtil function to return perfectly nested loop set

--

PiperOrigin-RevId: 242019230

3 files changed, 18 insertions, 15 deletions

diff --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index 2ed159ca647..39ed5a100a1 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -1062,18 +1062,9 @@ computeLoopInterchangePermutation(ArrayRef<Operation *> ops,
 // pushing loop carried dependence to a greater depth in the loop nest.
 static void sinkSequentialLoops(MemRefDependenceGraph::Node *node) {
   assert(node->op->isa<AffineForOp>());
-  // Get perfectly nested sequence of loops starting at root of loop nest
-  // (the first op being another AffineFor, and the second op - a terminator).
-  // TODO(andydavis,bondhugula) Share this with similar code in loop tiling.
   SmallVector<AffineForOp, 4> loops;
   AffineForOp curr = node->op->cast<AffineForOp>();
-  loops.push_back(curr);
-  auto *currBody = curr.getBody();
-  while (currBody->begin() == std::prev(currBody->end(), 2) &&
-         (curr = curr.getBody()->front().dyn_cast<AffineForOp>())) {
-    loops.push_back(curr);
-    currBody = curr.getBody();
-  }
+  getPerfectlyNestedLoops(loops, curr);
   if (loops.size() < 2)
     return;
 
diff --git a/mlir/lib/Transforms/LoopTiling.cpp b/mlir/lib/Transforms/LoopTiling.cpp
index 956d50ec26f..c215fa34172 100644
--- a/mlir/lib/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Transforms/LoopTiling.cpp
@@ -270,11 +270,7 @@ static void getTileableBands(Function &f,
   // (inclusive).
   auto getMaximalPerfectLoopNest = [&](AffineForOp root) {
     SmallVector<AffineForOp, 6> band;
-    AffineForOp currInst = root;
-    do {
-      band.push_back(currInst);
-    } while (currInst.getBody()->getOperations().size() == 2 &&
-             (currInst = currInst.getBody()->front().dyn_cast<AffineForOp>()));
+    getPerfectlyNestedLoops(band, root);
     bands->push_back(band);
   };
 
diff --git a/mlir/lib/Transforms/Utils/LoopUtils.cpp b/mlir/lib/Transforms/Utils/LoopUtils.cpp
index 2b17f4b84fa..1e9697a9ad5 100644
--- a/mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ b/mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -353,6 +353,22 @@ LogicalResult mlir::instBodySkew(AffineForOp forOp, ArrayRef<uint64_t> shifts,
   return success();
 }
 
+/// Get perfectly nested sequence of loops starting at root of loop nest
+/// (the first op being another AffineFor, and the second op - a terminator).
+/// A loop is perfectly nested iff: the first op in the loop's body is another
+/// AffineForOp, and the second op is a terminator).
+void mlir::getPerfectlyNestedLoops(SmallVectorImpl<AffineForOp> &nestedLoops,
+                                   AffineForOp root) {
+  AffineForOp curr = root;
+  nestedLoops.push_back(curr);
+  auto *currBody = curr.getBody();
+  while (currBody->begin() == std::prev(currBody->end(), 2) &&
+         (curr = curr.getBody()->front().dyn_cast<AffineForOp>())) {
+    nestedLoops.push_back(curr);
+    currBody = curr.getBody();
+  }
+}
+
 /// Unrolls this loop completely.
 LogicalResult mlir::loopUnrollFull(AffineForOp forOp) {
   Optional<uint64_t> mayBeConstantTripCount = getConstantTripCount(forOp);