[LoopInfo] Introduce getUniqueNonLatchExitBlocks utility function

[skatkov]

Summary

Extract the code from LoopUnrollRuntime into utility function to
re-use it in D63923.

[reames]

LGTM

[fhahn]

IMO this seems like quite a specific function to add (there are other blocks that might be interesting to filter ). I think we could achieve the same thing with a filter_range on the caller side without much extra code and with it being slightly more explicit. But that’s quite subjective, so feel free to disregard.

[fhahn]

Regarding tests, I think we have unit tests for most functions in LoopInfo, it would be good to add a test there for the new function too, if you decide to add it

[skatkov]

Make getUniqueExitBlocksHelper more general with using filter predicate.
Unit test is added.

[skatkov]

I decided to avoid using public API with predicate due to it will require introducing template function in template class and effectively moving implementation to LoopInfo.h

The explicit function for explicit need seems to be fine to me. If someone needs to introduce new function or move this to public API all building blocks are in place.

llvm/include/llvm/Analysis/LoopInfo.h

#endif
   bool hasDedicatedExits() const;
 
   /// Return all unique successor blocks of this loop.
   /// These are the blocks _outside of the current loop_ which are branched to.
   /// This assumes that loop exits are in canonical form, i.e. all exits are
   /// dedicated exits.
   void getUniqueExitBlocks(SmallVectorImpl<BlockT *> &ExitBlocks) const;
 
+  /// Return all unique successor blocks of this loop except successors from
+  /// Latch block are not considered. If the exit comes from Latch has also
+  /// non Latch predecessor in a loop it will be added to ExitBlocks.
+  /// These are the blocks _outside of the current loop_ which are branched to.
+  /// This assumes that loop exits are in canonical form, i.e. all exits are
+  /// dedicated exits.
+  void getUniqueNonLatchExitBlocks(SmallVectorImpl<BlockT *> &ExitBlocks) const;
+
   /// If getUniqueExitBlocks would return exactly one block, return that block.
   /// Otherwise return null.
   BlockT *getUniqueExitBlock() const;
 
   /// Edge type.
   typedef std::pair<BlockT *, BlockT *> Edge;
 
   /// Return all pairs of (_inside_block_,_outside_block_).

llvm/include/llvm/Analysis/LoopInfoImpl.h

bool LoopBase<BlockT, LoopT>::hasDedicatedExits() const {
   for (BlockT *EB : ExitBlocks)
     for (BlockT *Predecessor : children<Inverse<BlockT *>>(EB))
       if (!contains(Predecessor))
         return false;
   // All the requirements are met.
   return true;
 }
 
-template <class BlockT, class LoopT>
-void LoopBase<BlockT, LoopT>::getUniqueExitBlocks(
-    SmallVectorImpl<BlockT *> &ExitBlocks) const {
+// Helper function to get unique loop exits. Pred is a predicate pointing to
+// BasicBlocks in a loop which should be considered to find loop exits.
+template <class BlockT, class LoopT, typename PredicateT>
+void getUniqueExitBlocksHelper(const LoopT *L,
+                               SmallVectorImpl<BlockT *> &ExitBlocks,
+                               PredicateT Pred) {
   typedef GraphTraits<BlockT *> BlockTraits;
   typedef GraphTraits<Inverse<BlockT *>> InvBlockTraits;
 
-  assert(hasDedicatedExits() &&
+  assert(L->hasDedicatedExits() &&
          "getUniqueExitBlocks assumes the loop has canonical form exits!");
 
   SmallVector<BlockT *, 32> SwitchExitBlocks;
-  for (BlockT *Block : this->blocks()) {
+  auto Filtered = make_filter_range(L->blocks(), Pred);
+  for (BlockT *Block : Filtered) {
     SwitchExitBlocks.clear();
     for (BlockT *Successor : children<BlockT *>(Block)) {
       // If block is inside the loop then it is not an exit block.
-      if (contains(Successor))
+      if (L->contains(Successor))
         continue;
 
       BlockT *FirstPred = *InvBlockTraits::child_begin(Successor);
 
       // If current basic block is this exit block's first predecessor then only
       // insert exit block in to the output ExitBlocks vector. This ensures that
       // same exit block is not inserted twice into ExitBlocks vector.
       if (Block != FirstPred)
for (BlockT *Successor : children<BlockT *>(Block)) {
         SwitchExitBlocks.push_back(Successor);
         ExitBlocks.push_back(Successor);
       }
     }
   }
 }
 
 template <class BlockT, class LoopT>
+void LoopBase<BlockT, LoopT>::getUniqueExitBlocks(
+    SmallVectorImpl<BlockT *> &ExitBlocks) const {
+  getUniqueExitBlocksHelper(this, ExitBlocks,
+                            [](const BlockT *BB) { return true; });
+}
+
+template <class BlockT, class LoopT>
+void LoopBase<BlockT, LoopT>::getUniqueNonLatchExitBlocks(
+    SmallVectorImpl<BlockT *> &ExitBlocks) const {
+  const BlockT *Latch = getLoopLatch();
+  assert(Latch && "Latch block must exists");
+  getUniqueExitBlocksHelper(this, ExitBlocks,
+                            [Latch](const BlockT *BB) { return BB != Latch; });
+}
+
+template <class BlockT, class LoopT>
 BlockT *LoopBase<BlockT, LoopT>::getUniqueExitBlock() const {
   SmallVector<BlockT *, 8> UniqueExitBlocks;
   getUniqueExitBlocks(UniqueExitBlocks);
   if (UniqueExitBlocks.size() == 1)
     return UniqueExitBlocks[0];
   return nullptr;
 }
 

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

if (CreateRemainderLoop) {
     return NewLoop;
   }
   else
     return nullptr;
 }
 
 /// Returns true if we can safely unroll a multi-exit/exiting loop. OtherExits
 /// is populated with all the loop exit blocks other than the LatchExit block.
-static bool
-canSafelyUnrollMultiExitLoop(Loop *L, SmallVectorImpl<BasicBlock *> &OtherExits,
-                             BasicBlock *LatchExit, bool PreserveLCSSA,
+static bool canSafelyUnrollMultiExitLoop(Loop *L, BasicBlock *LatchExit,
+                                         bool PreserveLCSSA,
                                          bool UseEpilogRemainder) {
 
   // We currently have some correctness constrains in unrolling a multi-exit
   // loop. Check for these below.
 
   // We rely on LCSSA form being preserved when the exit blocks are transformed.
   if (!PreserveLCSSA)
     return false;
-  SmallVector<BasicBlock *, 4> Exits;
-  L->getUniqueExitBlocks(Exits);
-  for (auto *BB : Exits)
-    if (BB != LatchExit)
-      OtherExits.push_back(BB);
 
   // TODO: Support multiple exiting blocks jumping to the `LatchExit` when
   // UnrollRuntimeMultiExit is true. This will need updating the logic in
   // connectEpilog/connectProlog.
   if (!LatchExit->getSinglePredecessor()) {
     LLVM_DEBUG(
         dbgs() << "Bailout for multi-exit handling when latch exit has >1 "
                   "predecessor.\n");

 
 /// Returns true if we can profitably unroll the multi-exit loop L. Currently,
 /// we return true only if UnrollRuntimeMultiExit is set to true.
 static bool canProfitablyUnrollMultiExitLoop(
     Loop *L, SmallVectorImpl<BasicBlock *> &OtherExits, BasicBlock *LatchExit,
     bool PreserveLCSSA, bool UseEpilogRemainder) {
 
 #if !defined(NDEBUG)
-  SmallVector<BasicBlock *, 8> OtherExitsDummyCheck;
-  assert(canSafelyUnrollMultiExitLoop(L, OtherExitsDummyCheck, LatchExit,
-                                      PreserveLCSSA, UseEpilogRemainder) &&
+  assert(canSafelyUnrollMultiExitLoop(L, LatchExit, PreserveLCSSA,
+                                      UseEpilogRemainder) &&
          "Should be safe to unroll before checking profitability!");
 #endif
 
   // Priority goes to UnrollRuntimeMultiExit if it's supplied.
   if (UnrollRuntimeMultiExit.getNumOccurrences())
     return UnrollRuntimeMultiExit;
 
   // The main pain point with multi-exit loop unrolling is that once unrolled,
if (L->contains(LatchExit)) {
     LLVM_DEBUG(
         dbgs()
         << "One of the loop latch successors must be the exit block.\n");
     return false;
   }
 
   // These are exit blocks other than the target of the latch exiting block.
   SmallVector<BasicBlock *, 4> OtherExits;
+  L->getUniqueNonLatchExitBlocks(OtherExits);
   bool isMultiExitUnrollingEnabled =
-      canSafelyUnrollMultiExitLoop(L, OtherExits, LatchExit, PreserveLCSSA,
+      canSafelyUnrollMultiExitLoop(L, LatchExit, PreserveLCSSA,
                                    UseEpilogRemainder) &&
       canProfitablyUnrollMultiExitLoop(L, OtherExits, LatchExit, PreserveLCSSA,
                                        UseEpilogRemainder);
   // Support only single exit and exiting block unless multi-exit loop unrolling is enabled.
   if (!isMultiExitUnrollingEnabled &&
       (!L->getExitingBlock() || OtherExits.size())) {
     LLVM_DEBUG(
         dbgs()

llvm/unittests/Analysis/LoopInfoTest.cpp

runWithLoopInfoPlus(
         PHINode &Instruction_usedoutside = cast<PHINode>(*(++II));
         EXPECT_FALSE(
             L->isAuxiliaryInductionVariable(Instruction_usedoutside, SE));
         PHINode &Instruction_mulopcode = cast<PHINode>(*(++II));
         EXPECT_FALSE(
             L->isAuxiliaryInductionVariable(Instruction_mulopcode, SE));
       });
 }
+
+// Examine getUniqueExitBlocks/getUniqueNonLatchExitBlocks functions.
+TEST(LoopInfoTest, LoopUniqueExitBlocks) {
+  const char *ModuleStr =
+      "target datalayout = \"e-m:o-i64:64-f80:128-n8:16:32:64-S128\"\n"
+      "define void @foo(i32 %n, i1 %cond) {\n"
+      "entry:\n"
+      "  br label %for.cond\n"
+      "for.cond:\n"
+      "  %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.inc ]\n"
+      "  %cmp = icmp slt i32 %i.0, %n\n"
+      "  br i1 %cond, label %for.inc, label %for.end1\n"
+      "for.inc:\n"
+      "  %inc = add nsw i32 %i.0, 1\n"
+      "  br i1 %cmp, label %for.cond, label %for.end2, !llvm.loop !0\n"
+      "for.end1:\n"
+      "  br label %for.end\n"
+      "for.end2:\n"
+      "  br label %for.end\n"
+      "for.end:\n"
+      "  ret void\n"
+      "}\n"
+      "!0 = distinct !{!0, !1}\n"
+      "!1 = !{!\"llvm.loop.distribute.enable\", i1 true}\n";
+
+  // Parse the module.
+  LLVMContext Context;
+  std::unique_ptr<Module> M = makeLLVMModule(Context, ModuleStr);
+
+  runWithLoopInfo(*M, "foo", [&](Function &F, LoopInfo &LI) {
+    Function::iterator FI = F.begin();
+    // First basic block is entry - skip it.
+    BasicBlock *Header = &*(++FI);
+    assert(Header->getName() == "for.cond");
+    Loop *L = LI.getLoopFor(Header);
+
+    SmallVector<BasicBlock *, 2> Exits;
+    // This loop has 2 unique exits.
+    L->getUniqueExitBlocks(Exits);
+    EXPECT_TRUE(Exits.size() == 2);
+    // And one unique non latch exit.
+    Exits.clear();
+    L->getUniqueNonLatchExitBlocks(Exits);
+    EXPECT_TRUE(Exits.size() == 1);
+  });
+}