[PATCH] D78428: [MLIR] Make isPerfectlyNested check more efficient

[Uday Bondhugula via Phabricator via llvm-commits]

bondhugula created this revision.
bondhugula added reviewers: mehdi_amini, andydavis1.
Herald added subscribers: llvm-commits, frgossen, grosul1, Joonsoo, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, burmako, jpienaar, rriddle.
Herald added a project: LLVM.
bondhugula added a child revision: D78423: [MLIR] NFC affine for op tiling cleanup / utility rename.

Make mlir::isPerfectlyNested more efficient; use O(1) check instead of
O(N) size() method.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D78428

Files:
  mlir/lib/Transforms/Utils/LoopUtils.cpp


Index: mlir/lib/Transforms/Utils/LoopUtils.cpp
===================================================================
--- mlir/lib/Transforms/Utils/LoopUtils.cpp
+++ mlir/lib/Transforms/Utils/LoopUtils.cpp
@@ -705,17 +705,25 @@
   return checkLoopInterchangeDependences(depCompsVec, loops, loopPermMap);
 }
 
-/// Return true if `loops` is a perfect nest.
+/// Returns true if `loops` is a perfectly nested loop nest, where loops appear
+/// in it from outermost to innermost.
 static bool LLVM_ATTRIBUTE_UNUSED
 isPerfectlyNested(ArrayRef<AffineForOp> loops) {
-  auto outerLoop = loops.front();
+  assert(!loops.empty() && "no loops provided");
+
+  // We already know that the block can't be empty.
+  auto hasTwoElements = [](Block *block) {
+    auto secondOpIt = std::next(block->begin());
+    return secondOpIt != block->end() && &*secondOpIt == &block->back();
+  };
+
+  auto enclosingLoop = loops.front();
   for (auto loop : loops.drop_front()) {
     auto parentForOp = dyn_cast<AffineForOp>(loop.getParentOp());
     // parentForOp's body should be just this loop and the terminator.
-    if (parentForOp != outerLoop ||
-        parentForOp.getBody()->getOperations().size() != 2)
+    if (parentForOp != enclosingLoop || !hasTwoElements(parentForOp.getBody()))
       return false;
-    outerLoop = loop;
+    enclosingLoop = loop;
   }
   return true;
 }


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D78428.258521.patch
Type: text/x-patch
Size: 1369 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200418/67fb46be/attachment.bin>