[MLIR] Make isPerfectlyNested check more efficient

[bondhugula]

Summary

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

[rriddle]

mlir/lib/Transforms/Utils/LoopUtils.cpp

715

nit: You could use llvm::hasNItems(block->begin(), block->end(), 2). We could also add a range based version to make it more usable.

[bondhugula]

mlir/lib/Transforms/Utils/LoopUtils.cpp

715

This would lead to one extra conditional check since it wouldn't factor in the additional information that this block is known to have at least one element. But that method is useful in general - thanks.

mlir/lib/Transforms/Utils/LoopUtils.cpp

bool mlir::isValidLoopInterchangePermutation(ArrayRef<AffineForOp> loops,
   // rooted at 'loops[0]', at loop depths in range [1, maxLoopDepth].
   assert(loopPermMap.size() == loops.size());
   unsigned maxLoopDepth = loops.size();
   std::vector<SmallVector<DependenceComponent, 2>> depCompsVec;
   getDependenceComponents(loops[0], maxLoopDepth, &depCompsVec);
   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;
 }
 
 // input[i] should move from position i -> permMap[i]. Returns the position in
 // `input` that becomes the new outermost loop.
 unsigned mlir::permuteLoops(MutableArrayRef<AffineForOp> input,
                             ArrayRef<unsigned> permMap) {