[Loop Peeling] Add support for peeling of loops with multiple exits

[skatkov]

Summary

This patch modifies the loop peeling transformation so that
it does not expect that there is only one loop exit from latch.

It modifies only transformation. Update of branch weights remains
only for exit from latch.

The motivation is that in follow-up patch I plan to enable loop peeling for
loops with multiple exits but only if other exits then from latch one goes to
block with call to deopt.

For now this patch is NFC.

[fhahn]

include/llvm/Analysis/LoopInfo.h

282

This is not directly related to the peeling changes. I think it would be better to split this off in a separate patch & also add a test to the loopInfo unit tests with differing const-ness.

lib/Transforms/Utils/LoopUnroll.cpp

407

Unless I am missing something, it looks like Succ here is outside the loop, which would mean that the latch would be an *exiting* block and Succ an *exit* block?

412

IIRC TripCount and TripMultiple would not be set for multi-exit loops (in LoopUnrollPass.cpp), but MaxTripCount would. But maybe I am missing the reasoning behind choosing the trip count of one of the exits?

[skatkov]

include/llvm/Analysis/LoopInfo.h

282

https://reviews.llvm.org/D64060

lib/Transforms/Utils/LoopUnroll.cpp

407

Ma bad. It should be latch. Will update a patch.

412

It is not only for multi exit but also for single exit.
We support other exits only for deopts, so no need to worry about them.

[fhahn]

lib/Transforms/Utils/LoopUnrollPeel.cpp

521–525

Maybe assign LatchBR->getNumSuccessors() to a variable, to avoid re-evaluating it?

556–564

Can you use a more descriptive variable name here? Same at line 620

557

Can you use It.second->phis() instead?

[skatkov]

rebase + changes according to review.

[fhahn]

lib/Transforms/Utils/LoopUnroll.cpp

412

We support other exits only for deopts, so no need to worry about them.

Right, I missed that restriction. I think it is fine to have this restriction to start with, but it would be good to explicitly mention that here in a comment and in the commit message.

[skatkov]

lib/Transforms/Utils/LoopUnroll.cpp

412

To be honest, I' mot sure what exactly you mean. This patch does not modify the guard to allow loops with multiple exits. with this patch we still support only loops with one exit on latch.

The patch does modification to not directly depend on the condition that there is only one exit.
So this particular line change is done because L->getExitingBlock() does not work for multiple exits but we are sure this exiting block is latch and it will be true both for single and multiple exits if others are calls to deopt.

I will update D63923 to comment here about other exits.

[fhahn]

lib/Transforms/Utils/LoopUnroll.cpp

412

I meant that I think we were missing an explanation here, for why we would always pick the latch (instead of the getExitingBlock()). Anyways, I missed that in a recent update, you added a comment where you use getLoopLatch, which makes it clear. Thanks

[fhahn]

LGTM, thanks.

I assume @reames marked it as requiring changes due to my comments, but it might be good to wait till early next week, in case he has additional comments.

[reames]

In D63921#1571190, @fhahn wrote:

LGTM, thanks.

I assume @reames marked it as requiring changes due to my comments, but it might be good to wait till early next week, in case he has additional comments.

Nope, purely off your previous comments, don't block on me!

[reames]

Just to clear conflicting phab state.

include/llvm/Analysis/LoopInfo.h

#endif
   /// This assumes that loop exits are in canonical form, i.e. all exits are
   /// dedicated exits.
   void getUniqueExitBlocks(SmallVectorImpl<BlockT *> &ExitBlocks) const;
 
   /// If getUniqueExitBlocks would return exactly one block, return that block.
   /// Otherwise return null.
   BlockT *getUniqueExitBlock() const;
 
-  /// Edge type.
-  typedef std::pair<const BlockT *, const BlockT *> Edge;
-
   /// Return all pairs of (_inside_block_,_outside_block_).
-  void getExitEdges(SmallVectorImpl<Edge> &ExitEdges) const;
+  template <class ParamBlockT>
+  void getExitEdges(SmallVectorImpl<std::pair<ParamBlockT *, ParamBlockT *> > &
+                        ExitEdges) const {
+    // Assert that ParamBlockT and BlockT are the same +/- const qualifier.
+    static_assert(
+        std::is_same<typename std::remove_const<BlockT>::type,
+                     typename std::remove_const<ParamBlockT>::type>::value,
+        "not same!");
+    assert(!isInvalid() && "Loop not in a valid state!");
+    for (const auto BB : blocks())
+      for (const auto &Succ : children<BlockT *>(BB))
+        if (!contains(Succ))
+          // Not in current loop? It must be an exit block.
+          ExitEdges.emplace_back(BB, Succ);
+  }
 
   /// If there is a preheader for this loop, return it. A loop has a preheader
   /// if there is only one edge to the header of the loop from outside of the
   /// loop. If this is the case, the block branching to the header of the loop
   /// is the preheader node.
   ///
   /// This method returns null if there is no preheader for the loop.
   BlockT *getLoopPreheader() const;

include/llvm/Analysis/LoopInfoImpl.h

 BlockT *LoopBase<BlockT, LoopT>::getUniqueExitBlock() const {
   SmallVector<BlockT *, 8> UniqueExitBlocks;
   getUniqueExitBlocks(UniqueExitBlocks);
   if (UniqueExitBlocks.size() == 1)
     return UniqueExitBlocks[0];
   return nullptr;
 }
 
-/// getExitEdges - Return all pairs of (_inside_block_,_outside_block_).
-template <class BlockT, class LoopT>
-void LoopBase<BlockT, LoopT>::getExitEdges(
-    SmallVectorImpl<Edge> &ExitEdges) const {
-  assert(!isInvalid() && "Loop not in a valid state!");
-  for (const auto BB : blocks())
-    for (const auto &Succ : children<BlockT *>(BB))
-      if (!contains(Succ))
-        // Not in current loop? It must be an exit block.
-        ExitEdges.emplace_back(BB, Succ);
-}
-
 /// getLoopPreheader - If there is a preheader for this loop, return it.  A
 /// loop has a preheader if there is only one edge to the header of the loop
 /// from outside of the loop and it is legal to hoist instructions into the
 /// predecessor. If this is the case, the block branching to the header of the
 /// loop is the preheader node.
 ///
 /// This method returns null if there is no preheader for the loop.
 ///

lib/Transforms/Utils/LoopUnroll.cpp

LoopUnrollResult llvm::UnrollLoop(Loop *L, UnrollLoopOptions ULO, LoopInfo *LI,
 
   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) {
-      BasicBlock *ExitingBlock = L->getExitingBlock();
+      BasicBlock *ExitingBlock = nullptr;
+      for (auto *Succ : successors(L->getLoopLatch()))
+        if (!L->contains(Succ)) {
+          ExitingBlock = Succ;
+          break;
+        }
       assert(ExitingBlock && "Loop without exiting block?");
       Preheader = L->getLoopPreheader();
       ULO.TripCount = SE->getSmallConstantTripCount(L, ExitingBlock);
       ULO.TripMultiple = SE->getSmallConstantTripMultiple(L, ExitingBlock);
     }
   }
 
   // Loops containing convergent instructions must have a count that divides
   // their TripMultiple.
   LLVM_DEBUG(
       {

lib/Transforms/Utils/LoopUnrollPeel.cpp

if (ExitWeight) {
     LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
   }
 }
 
 /// Clones the body of the loop L, putting it between \p InsertTop and \p
 /// InsertBot.
 /// \param IterNumber The serial number of the iteration currently being
 /// peeled off.
-/// \param Exit The exit block of the original loop.
+/// \param ExitEdges The exit edges of the original loop.
 /// \param[out] NewBlocks A list of the blocks in the newly created clone
 /// \param[out] VMap The value map between the loop and the new clone.
 /// \param LoopBlocks A helper for DFS-traversal of the loop.
 /// \param LVMap A value-map that maps instructions from the original loop to
 /// instructions in the last peeled-off iteration.
-static void cloneLoopBlocks(Loop *L, unsigned IterNumber, BasicBlock *InsertTop,
-                            BasicBlock *InsertBot, BasicBlock *Exit,
-                            SmallVectorImpl<BasicBlock *> &NewBlocks,
-                            LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap,
-                            ValueToValueMapTy &LVMap, DominatorTree *DT,
-                            LoopInfo *LI) {
+static void cloneLoopBlocks(
+    Loop *L, unsigned IterNumber, BasicBlock *InsertTop, BasicBlock *InsertBot,
+    SmallVectorImpl<std::pair<BasicBlock *, BasicBlock *> > &ExitEdges,
+    SmallVectorImpl<BasicBlock *> &NewBlocks, LoopBlocksDFS &LoopBlocks,
+    ValueToValueMapTy &VMap, ValueToValueMapTy &LVMap, DominatorTree *DT,
   BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   BasicBlock *PreHeader = L->getLoopPreheader();
 
   Function *F = Header->getParent();
   LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();
   LoopBlocksDFS::RPOIterator BlockEnd = LoopBlocks.endRPO();
   Loop *ParentLoop = L->getParentLoop();
static void cloneLoopBlocks(
 
   // 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);
+  for (unsigned idx = 0; idx < LatchBR->getNumSuccessors(); ++idx)
+    if (LatchBR->getSuccessor(idx) == Header) {
+      LatchBR->setSuccessor(idx, InsertBot);
+      break;
+    }
   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.
for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     }
     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
   // a value coming into the header.
-  for (BasicBlock::iterator I = Exit->begin(); isa<PHINode>(I); ++I) {
+  for (auto It : ExitEdges)
+    for (BasicBlock::iterator I = It.second->begin(); isa<PHINode>(I); ++I) {
       PHINode *PHI = cast<PHINode>(I);
-    Value *LatchVal = PHI->getIncomingValueForBlock(Latch);
+      Value *LatchVal = PHI->getIncomingValueForBlock(It.first);
       Instruction *LatchInst = dyn_cast<Instruction>(LatchVal);
       if (LatchInst && L->contains(LatchInst))
         LatchVal = VMap[LatchVal];
-    PHI->addIncoming(LatchVal, cast<BasicBlock>(VMap[Latch]));
+      PHI->addIncoming(LatchVal, cast<BasicBlock>(VMap[It.first]));
     }
 
   // 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;
 }
 
 /// Peel off the first \p PeelCount iterations of loop \p L.
bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   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();
-  BasicBlock *Exit = L->getUniqueExitBlock();
+  SmallVector<std::pair<BasicBlock *, BasicBlock *>, 4> ExitEdges;
+  L->getExitEdges(ExitEdges);
 
   Function *F = Header->getParent();
 
   // Set up all the necessary basic blocks. It is convenient to split the
   // preheader into 3 parts - two blocks to anchor the peeled copy of the loop
   // body, and a new preheader for the "real" loop.
 
   // Peeling the first iteration transforms.
for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     // weight is exactly the weight of the previous iteration's header.
     // FIXME: due to the way the distribution is constructed, we need a
     // guard here to make sure we don't end up with non-positive weights.
     if (ExitWeight < CurHeaderWeight)
       CurHeaderWeight -= ExitWeight;
     else
       CurHeaderWeight = 1;
 
-    cloneLoopBlocks(L, Iter, InsertTop, InsertBot, Exit,
-                    NewBlocks, LoopBlocks, VMap, LVMap, DT, LI);
+    cloneLoopBlocks(L, Iter, InsertTop, InsertBot, ExitEdges, NewBlocks,
+                    LoopBlocks, VMap, LVMap, DT, LI);
 
     // Remap to use values from the current iteration instead of the
     // previous one.
     remapInstructionsInBlocks(NewBlocks, VMap);
 
     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.
       if (Iter == 0)
-        DT->changeImmediateDominator(Exit, cast<BasicBlock>(LVMap[Latch]));
+        for (auto It : ExitEdges)
+          DT->changeImmediateDominator(It.second,
+                                       cast<BasicBlock>(LVMap[It.first]));
 #ifdef EXPENSIVE_CHECKS
       assert(DT->verify(DominatorTree::VerificationLevel::Fast));
 #endif
     }
 
     auto *LatchBRCopy = cast<BranchInst>(VMap[LatchBR]);
     updateBranchWeights(InsertBot, LatchBRCopy, Iter,
                         PeelCount, ExitWeight);