[LoopDeletion] Update exits correctly when multiple duplicate edges from an exiting block

[anna]

Summary

Currently, we incorrectly update exit blocks of loops when there are multiple
edges from a single exiting block to the exit block. This can happen when we
have switches as the terminator of the exiting blocks.
The fix here is to correctly update the phi nodes in the exit block, and remove
all incoming values *except* for one which is from the preheader.

Note: Currently, this error can manifest only while deleting non-executed loops. However, it
is possible to trigger this error in invariant loops, once we enhance the logic
around the exit conditions for the loop check.

[efriedma]

lib/Transforms/Scalar/LoopDeletion.cpp

250 ↗

(On Diff #103586)

Do you need to call setIncomingValue() in the case where LoopIsNeverExecuted is false?

253 ↗

(On Diff #103586)

There is no variable named ExitingBlock.

[anna]

lib/Transforms/Scalar/LoopDeletion.cpp

250 ↗

(On Diff #103586)

We don't need to. When LoopIsNeverExecuted is false, all the exiting blocks have the same 'value', which gets moved to the preheader. This is also the reason why in the original implementation (before this patch), we keep just one exiting block and remove all others.

As a follow-on NFC, I'm going to move this undef value setting into a prepass before calling this function deleteDeadLoop.

253 ↗

(On Diff #103586)

Oops, stale comment. I'll remove the mention of ExitingBlock.

[efriedma]

LGTM with the comment fixed.

lib/Transforms/Scalar/LoopDeletion.cpp

250 ↗

(On Diff #103586)

Okay, that makes sense.

llvm/trunk/lib/Transforms/Scalar/LoopDeletion.cpp

static void deleteDeadLoop(Loop *L, DominatorTree &DT, ScalarEvolution &SE,
   // 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.
   SE.forgetLoop(L);
 
   auto *ExitBlock = L->getUniqueExitBlock();
   assert(ExitBlock && "Should have a unique exit block!");
 
+  assert(L->hasDedicatedExits() && "Loop should have dedicated exits!");
+
   // Connect the preheader directly to the exit block.
   // 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.
   Preheader->getTerminator()->replaceUsesOfWith(L->getHeader(), ExitBlock);
 
-  SmallVector<BasicBlock *, 4> ExitingBlocks;
-  L->getExitingBlocks(ExitingBlocks);
   // Rewrite phis in the exit block to get their inputs from the Preheader
   // instead of the exiting block.
-  BasicBlock *ExitingBlock = ExitingBlocks[0];
   BasicBlock::iterator BI = ExitBlock->begin();
   while (PHINode *P = dyn_cast<PHINode>(BI)) {
-    int j = P->getBasicBlockIndex(ExitingBlock);
-    assert(j >= 0 && "Can't find exiting block in exit block's phi node!");
+    // 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;
     if (LoopIsNeverExecuted)
-      P->setIncomingValue(j, UndefValue::get(P->getType()));
-    P->setIncomingBlock(j, Preheader);
-    for (unsigned i = 1; i < ExitingBlocks.size(); ++i)
-      P->removeIncomingValue(ExitingBlocks[i]);
+      P->setIncomingValue(PredIndex, UndefValue::get(P->getType()));
+    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).
+    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!");
     ++BI;
   }
 
   // Update the dominator tree and remove the instructions and blocks that will
   // be deleted from the reference counting scheme.
   SmallVector<DomTreeNode*, 8> ChildNodes;
   for (Loop::block_iterator LI = L->block_begin(), LE = L->block_end();
        LI != LE; ++LI) {

llvm/trunk/test/Transforms/LoopDeletion/unreachable-loops.ll

L3:
   %y.L3.next = add i64 %y.L3, 1
   %cond2 = icmp slt i64 %y.L3, %n
   br i1 %cond2, label %L3, label %L2
 
 exit:
  ret void
 }
 
+
+; 2 edges from a single exiting block to the exit block.
+define i64 @test12(i64 %n){
+;CHECK-LABEL: @test12
+; CHECK-NOT: L1:
+; CHECK-NOT: L1Latch:
+; CHECK-LABEL: L1.preheader:
+; CHECK-NEXT:    br label %exit
+; CHECK-LABEL: exit:
+; CHECK-NEXT:    %y.phi = phi i64 [ undef, %L1.preheader ]
+; CHECK-NEXT:    ret i64 %y.phi
+
+entry:
+  br i1 true, label %exit1, label %L1
+
+exit1:
+  ret i64 42
+
+L1:                                               ; preds = %L1Latch, %entry
+  %y.next = phi i64 [ 0, %entry ], [ %y.add, %L1Latch ]
+  br i1 true, label %L1Latch, label %exit
+
+L1Latch:                                          ; preds = %L1
+  %y = phi i64 [ %y.next, %L1 ]
+  %y.add = add i64 %y, %n
+  %cond2 = icmp eq i64 %y.add, 42
+  switch i64 %n, label %L1 [
+    i64 10, label %exit
+    i64 20, label %exit
+  ]
+
+exit:                                             ; preds = %L1Latch, %L1Latch
+  %y.phi = phi i64 [ 10, %L1Latch ], [ 10, %L1Latch ], [ %y.next, %L1]
+  ret i64 %y.phi
+}
+
+; multiple edges to exit block from the same exiting blocks
+define i64 @test13(i64 %n) {
+; CHECK-LABEL: @test13
+; CHECK-NOT: L1:
+; CHECK-NOT: L1Latch:
+; CHECK-LABEL: L1.preheader:
+; CHECK-NEXT:    br label %exit
+; CHECK-LABEL: exit:
+; CHECK-NEXT:    %y.phi = phi i64 [ undef, %L1.preheader ]
+; CHECK-NEXT:    ret i64 %y.phi
+
+entry:
+  br i1 true, label %exit1, label %L1
+
+exit1:
+  ret i64 42
+
+L1:                                               ; preds = %L1Latch, %entry
+  %y.next = phi i64 [ 0, %entry ], [ %y.add, %L1Latch ]
+  br i1 true, label %L1Block, label %exit
+
+L1Block:                                          ; preds = %L1
+  %y = phi i64 [ %y.next, %L1 ]
+  %y.add = add i64 %y, %n
+  %cond2 = icmp eq i64 %y.add, 42
+  switch i64 %n, label %L1Latch [
+    i64 10, label %exit
+    i64 20, label %exit
+  ]
+
+L1Latch:
+  switch i64 %n, label %L1 [
+    i64 30, label %exit
+    i64 40, label %exit
+  ]
+
+exit:                                             ; preds = %L1Block, %L1, %L1Latch
+  %y.phi = phi i64 [ 10, %L1Block ], [ 10, %L1Block ], [ %y.next, %L1 ], [ 30, %L1Latch ], [ 30, %L1Latch ]
+  ret i64 %y.phi
+}