[llvm] [LoopPeel] Remove known trip count restriction when peeling last. (PR #140792)

[Florian Hahn via llvm-commits]

================
@@ -1080,11 +1086,40 @@ bool llvm::peelLoop(Loop *L, unsigned PeelCount, bool PeelLast, LoopInfo *LI,
     for (PHINode &P : Exit->phis())
       ExitValues[&P] = P.getIncomingValueForBlock(Latch);
 
+    const SCEV *BTC = SE->getBackedgeTakenCount(L);
+
     InsertTop = SplitEdge(Latch, Exit, &DT, LI);
     InsertBot = SplitBlock(InsertTop, InsertTop->getTerminator(), &DT, LI);
 
     InsertTop->setName(Exit->getName() + ".peel.begin");
     InsertBot->setName(Exit->getName() + ".peel.next");
+    NewPreHeader = nullptr;
+
+    // If the original loop may only execute a single iteration we need to
+    // insert a trip count check and skip the peeled loop if necessary.
+    if (!SE->isKnownPredicate(CmpInst::ICMP_UGT, BTC,
+                              SE->getZero(BTC->getType()))) {
+      NewPreHeader = SplitEdge(PreHeader, Header, &DT, LI);
+      SCEVExpander Expander(*SE, Latch->getDataLayout(), "loop-peel");
+
+      BranchInst *PreHeaderBR = cast<BranchInst>(PreHeader->getTerminator());
+      Value *BTCValue =
+          Expander.expandCodeFor(BTC, BTC->getType(), PreHeaderBR);
+      IRBuilder<> B(PreHeaderBR);
+      Value *Cond =
+          B.CreateICmpNE(BTCValue, ConstantInt::get(BTCValue->getType(), 0));
+      B.CreateCondBr(Cond, NewPreHeader, InsertTop);
+      PreHeaderBR->eraseFromParent();
+
+      // PreHeader now dominates InsertTop.
+      DT.changeImmediateDominator(InsertTop, PreHeader);
+
+      // If we branch from PreHeader to InsertTop, we are guaranteed to execute
+      // the peeled iteration, so the exit values from the original loop are
+      // dead. Use poison for them.
+      for (auto &PN : InsertTop->phis())
----------------
fhahn wrote:

The phis here are the LCSSA phis exiting the original loop. But the uses of the exit phis get replaced with new exit values later on. I update the code there to just remove the unused phi nodes, then there's nothing to do here.

https://github.com/llvm/llvm-project/pull/140792