Loop peeling: check that latch is conditional branch

[JosephTremoulet]

Summary

Loop peeling assumes that the loop's latch is a conditional branch. Add
a check to canPeel that explicitly checks for this, and a testcase that
otherwise fails an assertion when trying to peel a loop whose back-edge
is the non-unwind edge of an invoke.

[JosephTremoulet]

This is meant to address a regression from rG7f69860243e8, which moved a check for whether dynamic cast to BranchInst succeeds in UnrollLoop from before to after the call to peelLoop. Adding a check to canPeel seemed like the most logical option to me.

[JosephTremoulet]

llvm/lib/Transforms/Utils/LoopPeel.cpp

664

Sorry, I jumped the gun in posting this -- this change makes the assertion here fail in several of the loop peeling lit tests. I'm taking a look...

[JosephTremoulet]

• Cast the terminator, not the block

[JosephTremoulet]

llvm/lib/Transforms/Utils/LoopPeel.cpp

664

Sorry, I jumped the gun in posting this -- this change makes the assertion here fail in several of the loop peeling lit tests. I'm taking a look...

Fixed.

[fhahn]

llvm/lib/Transforms/Utils/LoopPeel.cpp

123

this does only check for branch instructions, not conditional branches. Can you update the comment?

[fhahn]

Looks like this issue has also been reported in https://bugs.llvm.org/show_bug.cgi?id=48812

[JosephTremoulet]

• Match comment to code

[JosephTremoulet]

llvm/lib/Transforms/Utils/LoopPeel.cpp

123

this does only check for branch instructions, not conditional branches. Can you update the comment?

Updated.

[fhahn]

LGTM, thanks! Could you also add the switch test case from PR48812?

[JosephTremoulet]

• Add testcase from PR48812

[JosephTremoulet]

Could you also add the switch test case from PR48812?

Done.

[fhahn]

In D94995#2509609, @JosephTremoulet wrote:

Could you also add the switch test case from PR48812?

Done.

Thanks. looks good

llvm/lib/Transforms/Utils/LoopPeel.cpp

bool llvm::canPeel(Loop *L) {
   // Only peel loops that contain a single exit
   if (!L->getExitingBlock() || !L->getUniqueExitBlock())
     return false;
 
   // Don't try to peel loops where the latch is not the exiting block.
   // This can be an indication of two different things:
   // 1) The loop is not rotated.
   // 2) The loop contains irreducible control flow that involves the latch.
-  if (L->getLoopLatch() != L->getExitingBlock())
+  const BasicBlock *Latch = L->getLoopLatch();
+  if (Latch != L->getExitingBlock())
+    return false;
+
+  // Peeling is only supported if the latch is a branch.
+  if (!isa<BranchInst>(Latch->getTerminator()))
     return false;
 
   return true;
 }
 
 // This function calculates the number of iterations after which the given Phi
 // becomes an invariant. The pre-calculated values are memorized in the map. The
 // function (shortcut is I) is calculated according to the following definition:

 /// For loops that dynamically execute \p PeelCount iterations or less
 /// this provides a benefit, since the peeled off iterations, which account
 /// for the bulk of dynamic execution, can be further simplified by scalar
 /// optimizations.
 bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
                     ScalarEvolution *SE, DominatorTree *DT, AssumptionCache *AC,
                     bool PreserveLCSSA) {
   assert(PeelCount > 0 && "Attempt to peel out zero iterations?");
   assert(canPeel(L) && "Attempt to peel a loop which is not peelable?");
 
   LoopBlocksDFS LoopBlocks(L);
   LoopBlocks.perform(LI);
 
   BasicBlock *Header = L->getHeader();
   BasicBlock *PreHeader = L->getLoopPreheader();
   BasicBlock *Latch = L->getLoopLatch();
   SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> ExitEdges;

llvm/test/Transforms/LoopUnroll/peel-loop-conditions.ll

for.inc:
   %inc = add nsw i32 %i.05, 1
   %cmp = icmp slt i32 %inc, %k
   br i1 %cmp, label %for.body, label %for.end
 
 for.end:
   ret void
 }
 
+; Invoke is not a conditional branch that we can optimize,
+; so this shouldn't be peeled at all.  This is a reproducer
+; for a bug where evaluating the loop would fail an assertion.
+define void @test17() personality i8* undef{
+; CHECK-LABEL: @test17(
+; CHECK-NEXT:  body:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[CONST:%.*]] = phi i64 [ -33, [[LOOP]] ], [ -20, [[BODY:%.*]] ]
+; CHECK-NEXT:    invoke void @f1()
+; CHECK-NEXT:    to label [[LOOP]] unwind label [[EH_UNW_LOOPEXIT:%.*]]
+; CHECK:       eh.Unw.loopexit:
+; CHECK-NEXT:    [[LPAD_LOOPEXIT:%.*]] = landingpad { i8*, i32 }
+; CHECK-NEXT:    catch i8* null
+; CHECK-NEXT:    ret void
+;
+body:
+  br label %loop
+
+loop:
+  %const = phi i64 [ -33, %loop ], [ -20, %body ]
+  invoke void @f1()
+  to label %loop unwind label %eh.Unw.loopexit
+
+eh.Unw.loopexit:
+  %lpad.loopexit = landingpad { i8*, i32 }
+  catch i8* null
+  ret void
+}
+
+; Testcase reduced from PR48812.  We expect no peeling
+; because the latch terminator is a switch.
+define void @test18(i32* %p) {
+; CHECK-LABEL: @test18(
+; CHECK-NEXT:  init:
+; CHECK-NEXT:    br label [[LOOP:%.*]]
+; CHECK:       loop:
+; CHECK-NEXT:    [[CONST:%.*]] = phi i32 [ 40, [[INIT:%.*]] ], [ 0, [[LATCH:%.*]] ]
+; CHECK-NEXT:    br label [[LATCH]]
+; CHECK:       latch:
+; CHECK-NEXT:    [[CONTROL:%.*]] = load volatile i32, i32* [[P:%.*]], align 4
+; CHECK-NEXT:    switch i32 [[CONTROL]], label [[EXIT:%.*]] [
+; CHECK-NEXT:    i32 2, label [[LOOP]]
+; CHECK-NEXT:    ]
+; CHECK:       exit:
+; CHECK-NEXT:    ret void
+;
+init:
+  br label %loop
+
+loop:
+  %const = phi i32 [ 40, %init ], [ 0, %latch ]
+  br label %latch
+
+latch:
+  %control = load volatile i32, i32* %p
+  switch i32 %control, label %exit [
+  i32 2, label %loop
+  ]
+
+exit:
+  ret void
+}
+
 declare void @init()
 declare void @sink()