[LoopUnroll] Fix a crash

[skatkov]

Summary

Loop peeling as a last step triggers loop simplification and this
can change the loop structure. As a result all cashed values like
latch branch becomes invalid.

Patch re-structure the code to take into account the possible
changes caused by peeling.

[Meinersbur]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

365–370

Could you add a comment about that these have to be done after peelLoop?

370–371

[nit] re-add newline like the where it was moved from

llvm/test/Transforms/LoopUnroll/unroll-after-peel.ll

8

Consider removing !prof !0

48

Is this call required for the test?

[skatkov]

Handled comments,

[skatkov]

llvm/test/Transforms/LoopUnroll/unroll-after-peel.ll

48

Peeling guard requires that all non-latch exits must end-up by call to deoptimize.
So in other words peeling will not happen if I remove this call.

[fhahn]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

404

Does this need updating? At this point, we could have peeled the loop and thus change it I think.

Also, by moving the early exit to after peeling, we now expose more cases to peeling. Can peeling handle the case of the early exit properly?

llvm/test/Transforms/LoopUnroll/unroll-after-peel.ll

6

the triple should not be needed?

40

Would it be possible to use a concrete condition here? (branch on undef is UB)

[skatkov]

Test is updated.

[skatkov]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

404

I'm sorry I do not follow what you mean here. Could you please elaborate on what should be changed?

The only earlier return we do not do now for peeling is
!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)

Loop peeling has its own guard canPeel() which is also triggered in computePeelCount.
Peeling guard requires that latch is a conditional exiting block.
So nothing new we are introducing.

[fhahn]

llvm/lib/Transforms/Utils/LoopUnroll.cpp

404

Loop peeling has its own guard canPeel() which is also triggered in computePeelCount.
Peeling guard requires that latch is a conditional exiting block.

If we have a similar check for peeling everything is fine, thanks for checking! Can you add an assertion here that we did not peel, so we can catch the case early when we need to update the returned status?

[skatkov]

requested assert is added

[Meinersbur]

LGTM, thank you.

llvm/lib/Transforms/Utils/LoopUnroll.cpp

365–366

[grammar] All these values should be taken only after peeling because they might have changed.

[fhahn]

Thanks for adding the assert! LGTM

llvm/lib/Transforms/Utils/LoopUnroll.cpp

 /// required and not fully unrolled).
 LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
                                   ScalarEvolution *SE, DominatorTree *DT,
                                   AssumptionCache *AC,
                                   const TargetTransformInfo *TTI,
                                   OptimizationRemarkEmitter *ORE,
                                   bool PreserveLCSSA, Loop **RemainderLoop) {
 
-  BasicBlock *Preheader = L->getLoopPreheader();
-  if (!Preheader) {
+  if (!L->getLoopPreheader()) {
     LLVM_DEBUG(dbgs() << "  Can't unroll; loop preheader-insertion failed.\n");
     return LoopUnrollResult::Unmodified;
   }
 
-  BasicBlock *LatchBlock = L->getLoopLatch();
-  if (!LatchBlock) {
+  if (!L->getLoopLatch()) {
     LLVM_DEBUG(dbgs() << "  Can't unroll; loop exit-block-insertion failed.\n");
     return LoopUnrollResult::Unmodified;
   }
 
   // Loops with indirectbr cannot be cloned.
   if (!L->isSafeToClone()) {
     LLVM_DEBUG(dbgs() << "  Can't unroll; Loop body cannot be cloned.\n");
     return LoopUnrollResult::Unmodified;
   }
 
-  // The current loop unroll pass can unroll loops that have
-  // (1) single latch; and
-  // (2a) latch is unconditional; or
-  // (2b) latch is conditional and is an exiting block
-  // FIXME: The implementation can be extended to work with more complicated
-  // cases, e.g. loops with multiple latches.
-  BasicBlock *Header = L->getHeader();
-  BranchInst *LatchBI = dyn_cast<BranchInst>(LatchBlock->getTerminator());
-
-  // A conditional branch which exits the loop, which can be optimized to an
-  // unconditional branch in the unrolled loop in some cases.
-  BranchInst *ExitingBI = nullptr;
-  bool LatchIsExiting = L->isLoopExiting(LatchBlock);
-  if (LatchIsExiting)
-    ExitingBI = LatchBI;
-  else if (BasicBlock *ExitingBlock = L->getExitingBlock())
-    ExitingBI = dyn_cast<BranchInst>(ExitingBlock->getTerminator());
-  if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
-    LLVM_DEBUG(
-        dbgs() << "Can't unroll; a conditional latch must exit the loop");
-    return LoopUnrollResult::Unmodified;
-  }
-  LLVM_DEBUG({
-    if (ExitingBI)
-      dbgs() << "  Exiting Block = " << ExitingBI->getParent()->getName()
-             << "\n";
-    else
-      dbgs() << "  No single exiting block\n";
-  });
-
-  if (Header->hasAddressTaken()) {
+  if (L->getHeader()->hasAddressTaken()) {
     // The loop-rotate pass can be helpful to avoid this in many cases.
     LLVM_DEBUG(
         dbgs() << "  Won't unroll loop: address of header block is taken.\n");
     return LoopUnrollResult::Unmodified;
   }
 
   if (ULO.TripCount != 0)
     LLVM_DEBUG(dbgs() << "  Trip Count = " << ULO.TripCount << "\n");
LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
   }
 
   assert(ULO.Count > 0);
   assert(ULO.TripMultiple > 0);
   assert(ULO.TripCount == 0 || ULO.TripCount % ULO.TripMultiple == 0);
 
   // Are we eliminating the loop control altogether?
   bool CompletelyUnroll = ULO.Count == ULO.TripCount;
-  SmallVector<BasicBlock *, 4> ExitBlocks;
-  L->getExitBlocks(ExitBlocks);
-  std::vector<BasicBlock*> OriginalLoopBlocks = L->getBlocks();
-
-  // Go through all exits of L and see if there are any phi-nodes there. We just
-  // conservatively assume that they're inserted to preserve LCSSA form, which
-  // means that complete unrolling might break this form. We need to either fix
-  // it in-place after the transformation, or entirely rebuild LCSSA. TODO: For
-  // now we just recompute LCSSA for the outer loop, but it should be possible
-  // to fix it in-place.
-  bool NeedToFixLCSSA = PreserveLCSSA && CompletelyUnroll &&
-                        any_of(ExitBlocks, [](const BasicBlock *BB) {
-                          return isa<PHINode>(BB->begin());
-                        });
 
   // We assume a run-time trip count if the compiler cannot
   // figure out the loop trip count and the unroll-runtime
   // flag is specified.
   bool RuntimeTripCount =
       (ULO.TripCount == 0 && ULO.Count > 0 && ULO.AllowRuntime);
 
   assert((!RuntimeTripCount || !ULO.PeelCount) &&
          "Did not expect runtime trip-count unrolling "
          "and peeling for the same loop");
 
   bool Peeled = false;
   if (ULO.PeelCount) {
     Peeled = peelLoop(L, ULO.PeelCount, LI, SE, DT, AC, PreserveLCSSA);
 
     // Successful peeling may result in a change in the loop preheader/trip
     // counts. If we later unroll the loop, we want these to be updated.
     if (Peeled) {
       // According to our guards and profitability checks the only
       // meaningful exit should be latch block. Other exits go to deopt,
       // so we do not worry about them.
       BasicBlock *ExitingBlock = L->getLoopLatch();
       assert(ExitingBlock && "Loop without exiting block?");
       assert(L->isLoopExiting(ExitingBlock) && "Latch is not exiting?");
-      Preheader = L->getLoopPreheader();
       ULO.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock);
       ULO.TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
     }
   }
 
+  // All these values should be taken only after peeling because they might have
+  // changed.
+  BasicBlock *Preheader = L->getLoopPreheader();
+  BasicBlock *Header = L->getHeader();
+  BasicBlock *LatchBlock = L->getLoopLatch();
+  SmallVector<BasicBlock *, 4> ExitBlocks;
+  L->getExitBlocks(ExitBlocks);
+  std::vector<BasicBlock *> OriginalLoopBlocks = L->getBlocks();
+
+  // Go through all exits of L and see if there are any phi-nodes there. We just
+  // conservatively assume that they're inserted to preserve LCSSA form, which
+  // means that complete unrolling might break this form. We need to either fix
+  // it in-place after the transformation, or entirely rebuild LCSSA. TODO: For
+  // now we just recompute LCSSA for the outer loop, but it should be possible
+  // to fix it in-place.
+  bool NeedToFixLCSSA =
+      PreserveLCSSA && CompletelyUnroll &&
+      any_of(ExitBlocks,
+             [](const BasicBlock *BB) { return isa<PHINode>(BB->begin()); });
+
+  // The current loop unroll pass can unroll loops that have
+  // (1) single latch; and
+  // (2a) latch is unconditional; or
+  // (2b) latch is conditional and is an exiting block
+  // FIXME: The implementation can be extended to work with more complicated
+  // cases, e.g. loops with multiple latches.
+  BranchInst *LatchBI = dyn_cast<BranchInst>(LatchBlock->getTerminator());
+
+  // A conditional branch which exits the loop, which can be optimized to an
+  // unconditional branch in the unrolled loop in some cases.
+  BranchInst *ExitingBI = nullptr;
+  bool LatchIsExiting = L->isLoopExiting(LatchBlock);
+  if (LatchIsExiting)
+    ExitingBI = LatchBI;
+  else if (BasicBlock *ExitingBlock = L->getExitingBlock())
+    ExitingBI = dyn_cast<BranchInst>(ExitingBlock->getTerminator());
+  if (!LatchBI || (LatchBI->isConditional() && !LatchIsExiting)) {
+    // If the peeling guard is changed this assert may be relaxed or even
+    // deleted.
+    assert(!Peeled && "Peeling guard changed!");
+    LLVM_DEBUG(
+        dbgs() << "Can't unroll; a conditional latch must exit the loop");
+    return LoopUnrollResult::Unmodified;
+  }
+  LLVM_DEBUG({
+    if (ExitingBI)
+      dbgs() << "  Exiting Block = " << ExitingBI->getParent()->getName()
+             << "\n";
+    else
+      dbgs() << "  No single exiting block\n";
+  });
+
   // Loops containing convergent instructions must have a count that divides
   // their TripMultiple.
   LLVM_DEBUG(
       {
         bool HasConvergent = false;
         for (auto &BB : L->blocks())
           for (auto &I : *BB)
             if (auto *CB = dyn_cast<CallBase>(&I))

llvm/test/Transforms/LoopUnroll/unroll-after-peel.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -passes=loop-unroll -S | FileCheck %s
+
+define i64 @hoge(i1 %c) {
+; CHECK-LABEL: @hoge(
+; CHECK-NEXT:  bb:
+; CHECK-NEXT:    br label [[BB1_PEEL_BEGIN:%.*]]
+; CHECK:       bb1.peel.begin:
+; CHECK-NEXT:    br label [[BB1_PEEL:%.*]]
+; CHECK:       bb1.peel:
+; CHECK-NEXT:    br i1 [[C:%.*]], label [[BB2_PEEL:%.*]], label [[BB4:%.*]]
+; CHECK:       bb2.peel:
+; CHECK-NEXT:    [[TMP3_PEEL:%.*]] = icmp slt i32 0, 9
+; CHECK-NEXT:    br i1 [[TMP3_PEEL]], label [[BB1_PEEL_NEXT:%.*]], label [[BB4]]
+; CHECK:       bb1.peel.next:
+; CHECK-NEXT:    br label [[BB1_PEEL_NEXT1:%.*]]
+; CHECK:       bb1.peel.next1:
+; CHECK-NEXT:    br label [[BB_PEEL_NEWPH:%.*]]
+; CHECK:       bb.peel.newph:
+; CHECK-NEXT:    br label [[BB1:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    br i1 [[C]], label [[BB1]], label [[BB4_LOOPEXIT:%.*]], [[LOOP0:!llvm.loop !.*]]
+; CHECK:       bb4.loopexit:
+; CHECK-NEXT:    [[TMP5_PH:%.*]] = phi i32 [ 8, [[BB1]] ]
+; CHECK-NEXT:    br label [[BB4]]
+; CHECK:       bb4:
+; CHECK-NEXT:    [[TMP5:%.*]] = phi i32 [ 0, [[BB1_PEEL]] ], [ 8, [[BB2_PEEL]] ], [ [[TMP5_PH]], [[BB4_LOOPEXIT]] ]
+; CHECK-NEXT:    [[TMP6:%.*]] = call i64 (...) @llvm.experimental.deoptimize.i64(i32 10) [ "deopt"() ]
+; CHECK-NEXT:    ret i64 [[TMP6]]
+;
+bb:
+  br label %bb1
+
+bb1:                                              ; preds = %bb2, %bb
+  %tmp = phi i32 [ 8, %bb2 ], [ 0, %bb ]
+  br i1 %c, label %bb2, label %bb4
+
+bb2:                                              ; preds = %bb1
+  %tmp3 = icmp slt i32 %tmp, 9
+  br i1 %tmp3, label %bb1, label %bb4
+
+bb4:                                              ; preds = %bb2, %bb1
+  %tmp5 = phi i32 [ 8, %bb2 ], [ %tmp, %bb1 ]
+  %tmp6 = call i64 (...) @llvm.experimental.deoptimize.i64(i32 10) [ "deopt"() ]
+  ret i64 %tmp6
+}
+
+declare i64 @llvm.experimental.deoptimize.i64(...)