1 files changed, 1 insertions, 5 deletions
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);
   };