[LoopDeletion] Remove dead loops with no exit blocks

[atmnpatel]

Summary

Currently, LoopDeletion refuses to remove dead loops with no exit blocks
because it cannot statically determine the control flow after it removes
the block. This leads to miscompiles if the loop is an infinite loop and
should've been removed.

[atmnpatel]

Add parent revision for buildkite.

[lebedev.ri]

llvm/include/llvm/Analysis/LoopInfoImpl.h

146

I'd put this directly after getExitBlocks(),
and potentially use getExitBlocks(), not getUniqueExitBlocks(),
since uniqification won't really affect the existence of said blocks.

[jdoerfert]

Logic is sound to me. Some nits below.

Please add a test in which an endless loop w/o exits and w/o side-effects can be replaced with unreachable but the loop has more instructions. So, a phi and some control flow that computes the new phi value.

llvm/include/llvm/Analysis/LoopInfoImpl.h

150

Nit: size==0 => empty

llvm/lib/Transforms/Utils/LoopUtils.cpp

587

Can't you move ExitBlock into the else branch. Or branch on ExitBlock and assert hasNoExitBlocks if it is null?

Also, do we need to do all the stuff below in the endless case? Could we just return from the then branch?

Nit: Not sure why you need to set the insert point again.

[atmnpatel]

• Adds test.
• Uses getExitBlocks() instead of getUniqueExitBlocks() and moved definition.
• Changed definition of hasNoExitBlocks() to use empty() instead of conditional
• Style changes in deleteDeadLoop for readibility.

[atmnpatel]

Update comments.

[jdoerfert]

LGTM, a few nits to be addressed prior to the merge.

llvm/include/llvm/Analysis/LoopInfoImpl.h

75

-Unique

llvm/lib/Transforms/Scalar/LoopDeletion.cpp

200

Nit: aka, at most one exit block ;)

llvm/lib/Transforms/Utils/LoopUtils.cpp

685

Assuming the only use of the DeadDebugInst vector is the code blow that we don't execute (no point to "move" the stuff to), this entire loop is useless for the no-exit case.

Line 668 is trivially true for the no-exit case and 677-684 not needed as described above.

Thus, move the if(ExitBlock) to 648.

[atmnpatel]

• Addressed nits
• Rebase
• Final Pre-commit testing

[atmnpatel]

using buildkite recommended fix

[atmnpatel]

removed non-parent revision.

[atmnpatel]

• Fixed failing test

llvm/include/llvm/Analysis/LoopInfo.h

#endif
   /// 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.
   void getUniqueNonLatchExitBlocks(SmallVectorImpl<BlockT *> &ExitBlocks) const;
 
   /// If getUniqueExitBlocks would return exactly one block, return that block.
   /// Otherwise return null.
   BlockT *getUniqueExitBlock() const;
 
+  /// Return true if this loop does not have any exit blocks.
+  bool hasNoExitBlocks() const;
+
   /// Edge type.
   typedef std::pair<BlockT *, BlockT *> Edge;
 
   /// Return all pairs of (_inside_block_,_outside_block_).
   void getExitEdges(SmallVectorImpl<Edge> &ExitEdges) const;
 
   /// If there is a preheader for this loop, return it. A loop has a preheader
   /// if there is only one edge to the header of the loop from outside of the

llvm/include/llvm/Analysis/LoopInfoImpl.h

void LoopBase<BlockT, LoopT>::getExitBlocks(
   assert(!isInvalid() && "Loop not in a valid state!");
   for (const auto BB : blocks())
     for (auto *Succ : children<BlockT *>(BB))
       if (!contains(Succ))
         // Not in current loop? It must be an exit block.
         ExitBlocks.push_back(Succ);
 }
 
+template <class BlockT, class LoopT>
+bool LoopBase<BlockT, LoopT>::hasNoExitBlocks() const {
+  SmallVector<BlockT *, 8> UniqueExitBlocks;
+  getExitBlocks(UniqueExitBlocks);
+  return UniqueExitBlocks.empty();
+}
+
 /// getExitBlock - If getExitBlocks would return exactly one block,
 /// return that block. Otherwise return null.
 template <class BlockT, class LoopT>
 BlockT *LoopBase<BlockT, LoopT>::getExitBlock() const {
   assert(!isInvalid() && "Loop not in a valid state!");
   SmallVector<BlockT *, 8> ExitBlocks;
   getExitBlocks(ExitBlocks);
   if (ExitBlocks.size() == 1)
BlockT *LoopBase<BlockT, LoopT>::getUniqueExitBlock() const {
   SmallVector<BlockT *, 8> UniqueExitBlocks;
   getUniqueExitBlocks(UniqueExitBlocks);
   if (UniqueExitBlocks.size() == 1)
     return UniqueExitBlocks[0];
   return nullptr;
 }
 
 /// getExitEdges - Return all pairs of (_inside_block_,_outside_block_).
 template <class BlockT, class LoopT>
 void LoopBase<BlockT, LoopT>::getExitEdges(
     SmallVectorImpl<Edge> &ExitEdges) const {
   assert(!isInvalid() && "Loop not in a valid state!");
   for (const auto BB : blocks())
     for (auto *Succ : children<BlockT *>(BB))
       if (!contains(Succ))
         // Not in current loop? It must be an exit block.
         ExitEdges.emplace_back(BB, Succ);
 }
 
 /// getLoopPreheader - If there is a preheader for this loop, return it.  A
 /// loop has a preheader if there is only one edge to the header of the loop

llvm/lib/Transforms/Scalar/LoopDeletion.cpp

static bool isLoopDead(Loop *L, ScalarEvolution &SE,
                        BasicBlock *Preheader) {
   // Make sure that all PHI entries coming from the loop are loop invariant.
   // Because the code is in LCSSA form, any values used outside of the loop
   // must pass through a PHI in the exit block, meaning that this check is
   // sufficient to guarantee that no loop-variant values are used outside
   // of the loop.
   bool AllEntriesInvariant = true;
   bool AllOutgoingValuesSame = true;
+  if (!L->hasNoExitBlocks()) {
     for (PHINode &P : ExitBlock->phis()) {
       Value *incoming = P.getIncomingValueForBlock(ExitingBlocks[0]);
 
       // Make sure all exiting blocks produce the same incoming value for the
       // block. If there are different incoming values for different exiting
-    // blocks, then it is impossible to statically determine which value should
-    // be used.
+      // blocks, then it is impossible to statically determine which value
+      // should be used.
       AllOutgoingValuesSame =
           all_of(makeArrayRef(ExitingBlocks).slice(1), [&](BasicBlock *BB) {
             return incoming == P.getIncomingValueForBlock(BB);
           });
 
       if (!AllOutgoingValuesSame)
         break;
 
       if (Instruction *I = dyn_cast<Instruction>(incoming))
         if (!L->makeLoopInvariant(I, Changed, Preheader->getTerminator())) {
           AllEntriesInvariant = false;
           break;
         }
     }
+  }
 
   if (Changed)
     SE.forgetLoopDispositions(L);
 
   if (!AllEntriesInvariant || !AllOutgoingValuesSame)
     return false;
 
   // Make sure that no instructions in the block have potential side-effects.
if (ExitBlock && isLoopNeverExecuted(L)) {
     return LoopDeletionResult::Deleted;
   }
 
   // The remaining checks below are for a loop being dead because all statements
   // in the loop are invariant.
   SmallVector<BasicBlock *, 4> ExitingBlocks;
   L->getExitingBlocks(ExitingBlocks);
 
-  // We require that the loop only have a single exit block.  Otherwise, we'd
-  // be in the situation of needing to be able to solve statically which exit
-  // block will be branched to, or trying to preserve the branching logic in
-  // a loop invariant manner.
-  if (!ExitBlock) {
-    LLVM_DEBUG(dbgs() << "Deletion requires single exit block\n");
+  // We require that the loop only have a single exit block or no exit blocks.
+  // Otherwise, we'd be in the situation of needing to be able to solve
+  // statically which exit block will be branched to, or trying to preserve the
+  // branching logic in a loop invariant manner.
+  if (!ExitBlock && !L->hasNoExitBlocks()) {
+    LLVM_DEBUG(
+        dbgs() << "Deletion requires single exit block or no exit blocks\n");
     return LoopDeletionResult::Unmodified;
   }
   // Finally, we have to check that the loop really is dead.
   bool Changed = false;
   if (!isLoopDead(L, SE, ExitingBlocks, ExitBlock, Changed, Preheader)) {
     LLVM_DEBUG(dbgs() << "Loop is not invariant, cannot delete.\n");
     return Changed ? LoopDeletionResult::Modified
                    : LoopDeletionResult::Unmodified;

llvm/lib/Transforms/Utils/LoopUtils.cpp

void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT, ScalarEvolution *SE,
   // we remove things is very important to avoid invalidation issues.
 
   // Tell ScalarEvolution that the loop is deleted. Do this before
   // deleting the loop so that ScalarEvolution can look at the loop
   // to determine what it needs to clean up.
   if (SE)
     SE->forgetLoop(L);
 
-  auto *ExitBlock = L->getUniqueExitBlock();
-  assert(ExitBlock && "Should have a unique exit block!");
-  assert(L->hasDedicatedExits() && "Loop should have dedicated exits!");
-
   auto *OldBr = dyn_cast<BranchInst>(Preheader->getTerminator());
   assert(OldBr && "Preheader must end with a branch");
   assert(OldBr->isUnconditional() && "Preheader must have a single successor");
   // Connect the preheader to the exit block. Keep the old edge to the header
   // around to perform the dominator tree update in two separate steps
   // -- #1 insertion of the edge preheader -> exit and #2 deletion of the edge
   // preheader -> header.
   //
void llvm::deleteDeadLoop(Loop *L, DominatorTree *DT, ScalarEvolution *SE,
   //
   // Even when the loop is never executed, we cannot remove the edge from the
   // source block to the exit block. Consider the case where the unexecuted loop
   // branches back to an outer loop. If we deleted the loop and removed the edge
   // coming to this inner loop, this will break the outer loop structure (by
   // deleting the backedge of the outer loop). If the outer loop is indeed a
   // non-loop, it will be deleted in a future iteration of loop deletion pass.
   IRBuilder<> Builder(OldBr);
+
+  auto *ExitBlock = L->getUniqueExitBlock();
+  if (ExitBlock) {
+    assert(ExitBlock && "Should have a unique exit block!");
+    assert(L->hasDedicatedExits() && "Loop should have dedicated exits!");
+
     Builder.CreateCondBr(Builder.getFalse(), L->getHeader(), ExitBlock);
     // Remove the old branch. The conditional branch becomes a new terminator.
     OldBr->eraseFromParent();
 
     // Rewrite phis in the exit block to get their inputs from the Preheader
     // instead of the exiting block.
     for (PHINode &P : ExitBlock->phis()) {
       // Set the zero'th element of Phi to be from the preheader and remove all
       // other incoming values. Given the loop has dedicated exits, all other
       // incoming values must be from the exiting blocks.
       int PredIndex = 0;
       P.setIncomingBlock(PredIndex, Preheader);
       // Removes all incoming values from all other exiting blocks (including
       // duplicate values from an exiting block).
       // Nuke all entries except the zero'th entry which is the preheader entry.
       // NOTE! We need to remove Incoming Values in the reverse order as done
       // below, to keep the indices valid for deletion (removeIncomingValues
-    // updates getNumIncomingValues and shifts all values down into the operand
-    // being deleted).
+      // updates getNumIncomingValues and shifts all values down into the
+      // operand being deleted).
       for (unsigned i = 0, e = P.getNumIncomingValues() - 1; i != e; ++i)
         P.removeIncomingValue(e - i, false);
 
       assert((P.getNumIncomingValues() == 1 &&
               P.getIncomingBlock(PredIndex) == Preheader) &&
              "Should have exactly one value and that's from the preheader!");
     }
 
     DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Eager);
     if (DT) {
       DTU.applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock}});
       if (MSSA) {
-      MSSAU->applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock}}, *DT);
+        MSSAU->applyUpdates({{DominatorTree::Insert, Preheader, ExitBlock}},
+                            *DT);
         if (VerifyMemorySSA)
           MSSA->verifyMemorySSA();
       }
     }
 
     // Disconnect the loop body by branching directly to its exit.
     Builder.SetInsertPoint(Preheader->getTerminator());
     Builder.CreateBr(ExitBlock);
     // Remove the old branch.
     Preheader->getTerminator()->eraseFromParent();
 
     if (DT) {
       DTU.applyUpdates({{DominatorTree::Delete, Preheader, L->getHeader()}});
       if (MSSA) {
-      MSSAU->applyUpdates({{DominatorTree::Delete, Preheader, L->getHeader()}},
-                          *DT);
+        MSSAU->applyUpdates(
+            {{DominatorTree::Delete, Preheader, L->getHeader()}}, *DT);
         SmallSetVector<BasicBlock *, 8> DeadBlockSet(L->block_begin(),
                                                      L->block_end());
         MSSAU->removeBlocks(DeadBlockSet);
         if (VerifyMemorySSA)
           MSSA->verifyMemorySSA();
       }
     }
+  } else {
+    assert(L->hasNoExitBlocks() &&
+           "Loop should have either zero or one exit blocks.");
+    Builder.SetInsertPoint(OldBr);
+    Builder.CreateUnreachable();
+    Preheader->getTerminator()->eraseFromParent();
+  }
 
   // Use a map to unique and a vector to guarantee deterministic ordering.
   llvm::SmallDenseSet<std::pair<DIVariable *, DIExpression *>, 4> DeadDebugSet;
   llvm::SmallVector<DbgVariableIntrinsic *, 4> DeadDebugInst;
 
   // Given LCSSA form is satisfied, we should not have users of instructions
   // within the dead loop outside of the loop. However, LCSSA doesn't take
   // unreachable uses into account. We handle them here.
for (Instruction &I : *Block) {
       auto *DVI = dyn_cast<DbgVariableIntrinsic>(&I);
       if (!DVI)
         continue;
       auto Key = DeadDebugSet.find({DVI->getVariable(), DVI->getExpression()});
       if (Key != DeadDebugSet.end())
         continue;
       DeadDebugSet.insert({DVI->getVariable(), DVI->getExpression()});
       DeadDebugInst.push_back(DVI);
     }
 
   // After the loop has been deleted all the values defined and modified
   // inside the loop are going to be unavailable.
   // Since debug values in the loop have been deleted, inserting an undef
   // dbg.value truncates the range of any dbg.value before the loop where the
   // loop used to be. This is particularly important for constant values.
+  if (ExitBlock) {
     DIBuilder DIB(*ExitBlock->getModule());
     Instruction *InsertDbgValueBefore = ExitBlock->getFirstNonPHI();
     assert(InsertDbgValueBefore &&
            "There should be a non-PHI instruction in exit block, else these "
            "instructions will have no parent.");
     for (auto *DVI : DeadDebugInst)
       DIB.insertDbgValueIntrinsic(UndefValue::get(Builder.getInt32Ty()),
                                   DVI->getVariable(), DVI->getExpression(),
                                   DVI->getDebugLoc(), InsertDbgValueBefore);
+  }
 
   // Remove the block from the reference counting scheme, so that we can
   // delete it freely later.
   for (auto *Block : L->blocks())
     Block->dropAllReferences();
 
   if (MSSA && VerifyMemorySSA)
     MSSA->verifyMemorySSA();

llvm/test/Transforms/LoopDeletion/mustprogress.ll

 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes
 ; RUN: opt < %s -loop-deletion -S | FileCheck %s
 
 ;; Original C Code:
 ;;  void unknown_tripcount_mustprogress_attr_mustprogress_loopmd(int a, int b) {
 ;;    for (; a < b;) ;
 ;;    for (;;) ;
 ;;  }
 
 define void @unknown_tripcount_mustprogress_attr_mustprogress_loopmd(i32 %a, i32 %b) #0 {
 ; CHECK: Function Attrs: mustprogress
 ; CHECK-LABEL: define {{[^@]+}}@unknown_tripcount_mustprogress_attr_mustprogress_loopmd
 ; CHECK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]]) [[ATTR0:#.*]] {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_END:%.*]]
 ; CHECK:       for.end:
-; CHECK-NEXT:    br label [[FOR_COND1:%.*]]
-; CHECK:       for.cond1:
-; CHECK-NEXT:    br label [[FOR_COND1]]
+; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %for.cond
 for.cond:
   %cmp = icmp slt i32 %a, %b
   br i1 %cmp, label %for.body, label %for.end
 for.body:
   br label %for.cond, !llvm.loop !2

 
 define void @unknown_tripcount_mustprogress_attr_no_mustprogess_loopmd(i32 %a, i32 %b) #0 {
 ; CHECK: Function Attrs: mustprogress
 ; CHECK-LABEL: define {{[^@]+}}@unknown_tripcount_mustprogress_attr_no_mustprogess_loopmd
 ; CHECK-SAME: (i32 [[A:%.*]], i32 [[B:%.*]]) [[ATTR0]] {
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    br label [[FOR_END:%.*]]
 ; CHECK:       for.end:
-; CHECK-NEXT:    br label [[FOR_COND1:%.*]]
-; CHECK:       for.cond1:
-; CHECK-NEXT:    br label [[FOR_COND1]]
+; CHECK-NEXT:    unreachable
 ;
 entry:
   br label %for.cond
 for.cond:
   %cmp = icmp slt i32 %a, %b
   br i1 %cmp, label %for.body, label %for.end
 for.body:
   br label %for.cond

llvm/test/Transforms/LoopDeletion/no-exit-blocks.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --function-signature --check-attributes
+; RUN: opt < %s -loop-deletion -S | FileCheck %s
+
+define void @f() #0 {
+; CHECK: Function Attrs: mustprogress
+; CHECK-LABEL: define {{[^@]+}}@f
+; CHECK-SAME: () [[ATTR0:#.*]] {
+; CHECK-NEXT:    unreachable
+;
+  br label %1
+
+  %.01 = phi i32 [ 1, %0 ], [ %3, %2 ]
+  %.0 = phi i32 [ 1, %0 ], [ %3, %2 ]
+  br label %2
+
+  %3 = add nsw i32 %.01, %.0
+  br label %1
+}
+
+attributes #0 = { mustprogress }