[LoopPeel] Re-factor llvm::peelLoop method. NFC.

[skatkov]

Summary

Extract code dealing with branch weights in separate functions.

[fhahn]

LGTM, just 2 small nits.

lib/Transforms/Utils/LoopUnrollPeel.cpp

390 ↗

(On Diff #206996)

Using weight in the description for ExitWeight and # of times for CurHeaderWeigth seems inconsistent? (Also in some comments)

397 ↗

(On Diff #206996)

nit: no brackets required?

[reames]

Marking as accepted per previous reviewer (just so that search shows it as such)

[skatkov]

lib/Transforms/Utils/LoopUnrollPeel.cpp

390 ↗

(On Diff #206996)

Could you please elaborate a bit on this comment?
I did not change anything in the semantic, just moved the code in a separate functions.

So the weights still make sense...

397 ↗

(On Diff #206996)

Sure, will remove before landing.

[skatkov]

Hi Florian, I will land this change with comment update because I do not understand what exactly comment needs to be fixed.
Feel free to update this review with details what wrong in comment you see and I will land it as follow-up patch.

[fhahn]

lib/Transforms/Utils/LoopUnrollPeel.cpp

390 ↗

(On Diff #206996)

Yep, it was = a small thing. I just found it a bit inconsistent that ExitWeight and CurHeaderWeight are described differently. I think something like \param[out] ExitWeight The number of times the edge from Latch to Exit block is executed (or taken) would be more consistent with the description for CurHeaderWeight.

llvm/trunk/lib/Transforms/Utils/LoopUnrollPeel.cpp

if (PeeledHeaderWeight) {
     MDBuilder MDB(LatchBR->getContext());
     MDNode *WeightNode =
         HeaderIdx ? MDB.createBranchWeights(ExitWeight, FallThruWeight)
                   : MDB.createBranchWeights(FallThruWeight, ExitWeight);
     LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
   }
 }
 
+/// Initialize the weights.
+///
+/// \param Header The header block.
+/// \param LatchBR The latch branch.
+/// \param AvgIters The average number of iterations we expect the loop to have.
+/// \param[out] ExitWeight The weight of the edge from Latch to Exit block.
+/// \param[out] CurHeaderWeight The # of time the header is executed.
+static void initBranchWeights(BasicBlock *Header, BranchInst *LatchBR,
+                              unsigned AvgIters, uint64_t &ExitWeight,
+                              uint64_t &CurHeaderWeight) {
+  uint64_t TrueWeight, FalseWeight;
+  if (!LatchBR->extractProfMetadata(TrueWeight, FalseWeight))
+    return;
+  unsigned HeaderIdx = LatchBR->getSuccessor(0) == Header ? 0 : 1;
+  ExitWeight = HeaderIdx ? TrueWeight : FalseWeight;
+  // The # of times the loop body executes is the sum of the exit block
+  // weight and the # of times the backedges are taken.
+  CurHeaderWeight = TrueWeight + FalseWeight;
+}
+
+/// Update the weights of original Latch block after peeling off all iterations.
+///
+/// \param Header The header block.
+/// \param LatchBR The latch branch.
+/// \param ExitWeight The weight of the edge from Latch to Exit block.
+/// \param CurHeaderWeight The # of time the header is executed.
+static void fixupBranchWeights(BasicBlock *Header, BranchInst *LatchBR,
+                               uint64_t ExitWeight, uint64_t CurHeaderWeight) {
+  // Adjust the branch weights on the loop exit.
+  if (ExitWeight) {
+    // The backedge count is the difference of current header weight and
+    // current loop exit weight. If the current header weight is smaller than
+    // the current loop exit weight, we mark the loop backedge weight as 1.
+    uint64_t BackEdgeWeight = 0;
+    if (ExitWeight < CurHeaderWeight)
+      BackEdgeWeight = CurHeaderWeight - ExitWeight;
+    else
+      BackEdgeWeight = 1;
+    MDBuilder MDB(LatchBR->getContext());
+    unsigned HeaderIdx = LatchBR->getSuccessor(0) == Header ? 0 : 1;
+    MDNode *WeightNode =
+        HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)
+                  : MDB.createBranchWeights(BackEdgeWeight, 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[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.
bool llvm::peelLoop(Loop *L, unsigned PeelCount, LoopInfo *LI,
   NewPreHeader->setName(PreHeader->getName() + ".peel.newph");
 
   ValueToValueMapTy LVMap;
 
   // If we have branch weight information, we'll want to update it for the
   // newly created branches.
   BranchInst *LatchBR =
       cast<BranchInst>(cast<BasicBlock>(Latch)->getTerminator());
-  unsigned HeaderIdx = (LatchBR->getSuccessor(0) == Header ? 0 : 1);
-
-  uint64_t TrueWeight, FalseWeight;
   uint64_t ExitWeight = 0, CurHeaderWeight = 0;
-  if (LatchBR->extractProfMetadata(TrueWeight, FalseWeight)) {
-    ExitWeight = HeaderIdx ? TrueWeight : FalseWeight;
-    // The # of times the loop body executes is the sum of the exit block
-    // weight and the # of times the backedges are taken.
-    CurHeaderWeight = TrueWeight + FalseWeight;
-  }
+  initBranchWeights(Header, LatchBR, PeelCount, ExitWeight, CurHeaderWeight);
 
   // For each peeled-off iteration, make a copy of the loop.
   for (unsigned Iter = 0; Iter < PeelCount; ++Iter) {
     SmallVector<BasicBlock *, 8> NewBlocks;
     ValueToValueMapTy VMap;
 
     // Subtract the exit weight from the current header weight -- the exit
     // weight is exactly the weight of the previous iteration's header.
for (BasicBlock::iterator I = Header->begin(); isa<PHINode>(I); ++I) {
     Value *NewVal = PHI->getIncomingValueForBlock(Latch);
     Instruction *LatchInst = dyn_cast<Instruction>(NewVal);
     if (LatchInst && L->contains(LatchInst))
       NewVal = LVMap[LatchInst];
 
     PHI->setIncomingValueForBlock(NewPreHeader, NewVal);
   }
 
-  // Adjust the branch weights on the loop exit.
-  if (ExitWeight) {
-    // The backedge count is the difference of current header weight and
-    // current loop exit weight. If the current header weight is smaller than
-    // the current loop exit weight, we mark the loop backedge weight as 1.
-    uint64_t BackEdgeWeight = 0;
-    if (ExitWeight < CurHeaderWeight)
-      BackEdgeWeight = CurHeaderWeight - ExitWeight;
-    else
-      BackEdgeWeight = 1;
-    MDBuilder MDB(LatchBR->getContext());
-    MDNode *WeightNode =
-        HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)
-                  : MDB.createBranchWeights(BackEdgeWeight, ExitWeight);
-    LatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
-  }
+  fixupBranchWeights(Header, LatchBR, ExitWeight, CurHeaderWeight);
 
   if (Loop *ParentLoop = L->getParentLoop())
     L = ParentLoop;
 
   // We modified the loop, update SE.
   SE->forgetTopmostLoop(L);
 
   // FIXME: Incrementally update loop-simplify
   simplifyLoop(L, DT, LI, SE, AC, nullptr, PreserveLCSSA);
 
   NumPeeled++;
 
   return true;
 }