[LoopInfo] Fix getUniqueNonLatchExitBlocks

[skatkov]

Summary

It is possible that exit block has two predecessors and one of them is a latch
block while another is not.

Current algorithm is based on the assumption that all exits are dedicated
and therefore we can check only first predecessor of loop exit to find all unique
exits.

However if we do not consider latch block and it is first predecessor of some
exit then this exit will be found.

Regression test is added.

As a side effect of algorithm re-writing, the restriction that all exits are dedicated
is eliminated

[efriedma]

LGTM

llvm/include/llvm/Analysis/LoopInfo.h

#endif
   BlockT *getExitBlock() const;
 
   /// Return true if no exit block for the loop has a predecessor that is
   /// outside the loop.
   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;

llvm/include/llvm/Analysis/LoopInfoImpl.h

 }
 
 // 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(L->hasDedicatedExits() &&
-         "getUniqueExitBlocks assumes the loop has canonical form exits!");
-
-  SmallVector<BlockT *, 32> SwitchExitBlocks;
+  assert(!L->isInvalid() && "Loop not in a valid state!");
+  SmallPtrSet<BlockT *, 32> Visited;
   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 (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)
-        continue;
-
-      // If a terminator has more then two successors, for example SwitchInst,
-      // then it is possible that there are multiple edges from current block to
-      // one exit block.
-      if (std::distance(BlockTraits::child_begin(Block),
-                        BlockTraits::child_end(Block)) <= 2) {
-        ExitBlocks.push_back(Successor);
-        continue;
-      }
-
-      // In case of multiple edges from current block to exit block, collect
-      // only one edge in ExitBlocks. Use switchExitBlocks to keep track of
-      // duplicate edges.
-      if (!is_contained(SwitchExitBlocks, Successor)) {
-        SwitchExitBlocks.push_back(Successor);
+  for (BlockT *BB : Filtered)
+    for (BlockT *Successor : children<BlockT *>(BB))
+      if (!L->contains(Successor))
+        if (Visited.insert(Successor).second)
           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; });
 }
 

llvm/unittests/Analysis/LoopInfoTest.cpp

runWithLoopInfo(*M, "foo", [&](Function &F, LoopInfo &LI) {
     L->getUniqueExitBlocks(Exits);
     EXPECT_TRUE(Exits.size() == 2);
     // And one unique non latch exit.
     Exits.clear();
     L->getUniqueNonLatchExitBlocks(Exits);
     EXPECT_TRUE(Exits.size() == 1);
   });
 }
+
+// Regression test for  getUniqueNonLatchExitBlocks functions.
+// It should detect the exit if it comes from both latch and non-latch blocks.
+TEST(LoopInfoTest, LoopNonLatchUniqueExitBlocks) {
+  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.end\n"
+      "for.inc:\n"
+      "  %inc = add nsw i32 %i.0, 1\n"
+      "  br i1 %cmp, label %for.cond, label %for.end, !llvm.loop !0\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 1 unique exit.
+    L->getUniqueExitBlocks(Exits);
+    EXPECT_TRUE(Exits.size() == 1);
+    // And one unique non latch exit.
+    Exits.clear();
+    L->getUniqueNonLatchExitBlocks(Exits);
+    EXPECT_TRUE(Exits.size() == 1);
+  });
+}