[LoopUnroll] Don't peel loops where the latch isn't the exiting block

[mkuper]

Summary

Peeling assumed this doesn't happen, but didn't check it.

This fixes PR32178.

[mkuper]

ping

[mzolotukhin]

Looks good to me, but, as you noticed in the PR - do we want to work with such loops at all? Did you check what happens if we just bail out on such loops?

Michael

[mkuper]

In D30757#702345, @mzolotukhin wrote:

Looks good to me, but, as you noticed in the PR - do we want to work with such loops at all? Did you check what happens if we just bail out on such loops?

Michael

No, but I'm tempted to, because I think we have more bugs related to irreducible control flow that this patch doesn't fix.

Do you know if we already have a utility somewhere that detects this? There's no way to get this information out of LoopInfo that I'm aware of.
LV rolls its own ( hasCyclesInLoopBody() ), which would be fine for my original use case, but r296898 started handling non-innermost loops.

[mzolotukhin]

Do you know if we already have a utility somewhere that detects this?

Can we require the loop to be in a simplified form and its latch to be exiting? Will it suffice?

[mkuper]

In D30757#702406, @mzolotukhin wrote:

Do you know if we already have a utility somewhere that detects this?

Can we require the loop to be in a simplified form and its latch to be exiting? Will it suffice?

Maybe?

It will certainly squash this bug, not sure about the others.
We already sort of try to do this - UnrollLoop() bails if the latch is not conditional. I assume the expectation is that one of the edges is the backedge, and the other is the exiting edge, but I don't think this is actually enforced, and it looks like unrolling (as opposed to peeling) manages to deal with it.

[mzolotukhin]

This version looks good to me too.

MichaelZ

PS: Can we have a similar problem in the regular loop unrolling?

test/Transforms/LoopUnroll/peel-loop-irreducible.ll

37–38

I'd prefer to have LCSSA-ed version of the loop in the test. And, by the way, do we need %sum, %plus and %incsum at all? Maybe remove them to make the test even smaller?

[mkuper]

Thanks!

Re the PS - I don't know.
As I wrote above - UnrollLoop() verifies the latch is conditional, but doesn't verify the two edges are actually going to the header and the exit block. But we haven't actually *seen* any issues with this, as far as I know.

test/Transforms/LoopUnroll/peel-loop-irreducible.ll

37–38

Sure, I'll see if it minimizes more.

lib/Transforms/Utils/LoopUnrollPeel.cpp

static void cloneLoopBlocks(Loop *L, unsigned IterNumber, BasicBlock *InsertTop,
 
   // Hook-up the control flow for the newly inserted blocks.
   // The new header is hooked up directly to the "top", which is either
   // the original loop preheader (for the first iteration) or the previous
   // iteration's exiting block (for every other iteration)
   InsertTop->getTerminator()->setSuccessor(0, cast<BasicBlock>(VMap[Header]));
 
   // Similarly, for the latch:
-  // The original exiting edge is still hooked up to the loop exit.
   // The backedge now goes to the "bottom", which is either the loop's real
   // header (for the last peeled iteration) or the copied header of the next
   // iteration (for every other iteration)
   BasicBlock *NewLatch = cast<BasicBlock>(VMap[Latch]);
   BranchInst *LatchBR = cast<BranchInst>(NewLatch->getTerminator());
   unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
   LatchBR->setSuccessor(HeaderIdx, InsertBot);
-  LatchBR->setSuccessor(1 - HeaderIdx, Exit);
   if (DT)
     DT->changeImmediateDominator(InsertBot, NewLatch);
 
   // The new copy of the loop body starts with a bunch of PHI nodes
   // that pick an incoming value from either the preheader, or the previous
   // loop iteration. Since this copy is no longer part of the loop, we
   // resolve this statically:
   // For the first iteration, we use the value from the preheader directly.
if (IterNumber == 0) {
       else
         VMap[&*I] = LatchVal;
     }
     cast<BasicBlock>(VMap[Header])->getInstList().erase(NewPHI);
   }
 
   // Fix up the outgoing values - we need to add a value for the iteration
   // we've just created. Note that this must happen *after* the incoming
-  // values are adjusted, since the value going out of the latch may also be
+  // values are adjusted, since the value going out of the exit may also be
   // a value coming into the header.
+  BasicBlock *Exiting = L->getExitingBlock();
+  assert(Exiting && "Can't peel loops with more than one exiting block.");
   for (BasicBlock::iterator I = Exit->begin(); isa<PHINode>(I); ++I) {
     PHINode *PHI = cast<PHINode>(I);
-    Value *LatchVal = PHI->getIncomingValueForBlock(Latch);
-    Instruction *LatchInst = dyn_cast<Instruction>(LatchVal);
-    if (LatchInst && L->contains(LatchInst))
-      LatchVal = VMap[LatchVal];
-    PHI->addIncoming(LatchVal, cast<BasicBlock>(VMap[Latch]));
+    Value *ExitVal = PHI->getIncomingValueForBlock(Exiting);
+    Instruction *ExitInst = dyn_cast<Instruction>(ExitVal);
+    if (ExitInst && L->contains(ExitInst))
+      ExitVal = VMap[ExitVal];
+    PHI->addIncoming(ExitVal, cast<BasicBlock>(VMap[Exiting]));
   }
 
   // LastValueMap is updated with the values for the current loop
   // which are used the next time this function is called.
   for (const auto &KV : VMap)
     LVMap[KV.first] = KV.second;
 }
 
for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     if (ExitWeight < CurHeaderWeight)
       CurHeaderWeight -= ExitWeight;
     else
       CurHeaderWeight = 1;
 
     cloneLoopBlocks(L, Iter, InsertTop, InsertBot, Exit,
                     NewBlocks, LoopBlocks, VMap, LVMap, DT, LI);
     if (DT) {
-      // Latches of the cloned loops dominate over the loop exit, so idom of the
-      // latter is the first cloned loop body, as original PreHeader dominates
-      // the original loop body.
+      // The original exit block is reachable either directly from the first
+      // iteration's exiting block (if the first iteration exists), or through
+      // the first iteration's latch (if it doesn't). So it's dominated by
+      // the latch's and the exiting block's common dominator.
       if (Iter == 0)
-        DT->changeImmediateDominator(Exit, cast<BasicBlock>(LVMap[Latch]));
+        DT->changeImmediateDominator(
+            Exit, DT->findNearestCommonDominator(
+                      cast<BasicBlock>(LVMap[Latch]),
+                      cast<BasicBlock>(LVMap[L->getExitingBlock()])));
 #ifndef NDEBUG
       if (VerifyDomInfo)
         DT->verifyDomTree();
 #endif
     }
 
     updateBranchWeights(InsertBot, cast<BranchInst>(VMap[LatchBR]), Iter,
                         PeelCount, ExitWeight);

test/Transforms/LoopUnroll/peel-loop-irreducible.ll

+; RUN: opt < %s -S -loop-unroll | FileCheck %s
+
+; Check we can properly handle loops where the latch and the exiting block
+; are different
+; CHECK-LABEL: @invariant_backedge_irreducible
+; CHECK: header.peel:
+; CHECK: br i1 {{.*}} label %latch.peel, label %exiting.peel
+; CHECK: latch.peel:
+; CHECK: br i1 {{.*}} label %header.peel.next, label %exiting.peel
+; CHECK: exiting.peel:
+; CHECK: br i1 {{.*}} label %latch.peel, label %exit
+
+define i32 @invariant_backedge_irreducible(i32 %a, i32 %b) {
+entry:
+  br label %header
+
+header:
+  %i = phi i32 [ 0, %entry ], [ %inc, %latch ]
+  %sum = phi i32 [ 0, %entry ], [ %incsum, %latch ]
+  %plus = phi i32 [ %a, %entry ], [ %b, %latch ]
+  %cmp.phi = phi i1 [ false, %entry ], [ %cmp, %latch ]
+  br i1 %cmp.phi, label %latch, label %exiting
+
+latch:
+  %incsum = add i32 %sum, %plus
+  %inc = add i32 %i, 1
+  %cmp = icmp slt i32 %i, 1000    
+  br i1 %cmp, label %header, label %exiting
+
+exiting:
+  %cmp.exiting = phi i1 [ %cmp.phi, %header ], [ %cmp, %latch]
+  br i1 %cmp.exiting, label %latch, label %exit
+
+exit:
+  ret i32 %sum
+}

test/Transforms/LoopUnroll/peel-loop-not-forced.ll

 ; RUN: opt < %s -S -loop-unroll -unroll-threshold=4 | FileCheck %s
 
 define i32 @invariant_backedge_1(i32 %a, i32 %b) {
 ; CHECK-LABEL: @invariant_backedge_1
 ; CHECK-NOT:     %plus = phi
 ; CHECK:       loop.peel:
+; CHECK:         br i1 {{.*}}, label %loop.peel.next, label %exit
 ; CHECK:       loop:
 ; CHECK:         %i = phi
 ; CHECK:         %sum = phi
 entry:
   br label %loop
 
 loop:
   %i = phi i32 [ 0, %entry ], [ %inc, %loop ]