LoopUnrollRuntime: Add weights to all branches

[MatzeB]

Summary

Make sure every conditional branch constructed by LoopUnrollRuntime code sets branch weights.

• Add new 1:127 weights for the conditional jumps checking whether the whole (unrolled) loop should be skipped in the generated prolog or epilog code.
• Remove updateLatchBranchWeightsForRemainderLoop function and just add weights immediately when constructing the relevant branches. This leads to simpler code and makes the code more obvious as every call to CreateCondBr now has a BranchWeights parameter.
• Rework formula for epilogue latch weights, to assume equal distribution of remainders and remove assert (as I was able to reach this code when forcing small unroll factors on the commandline).

[MatzeB]

I guess we may need a discussion about how to express these "nearly always taken" / "nearly never taken" branch weights. The hardcoded 1:127 ratio in the current patch is pretty arbitrary.

Looking around I see LikelyBranchWeight in llvm/Transforms/Scalar/LowerExpectIntrinsic.cpp but it's private to the pass and a warning advises to use TargetTransformInfo::getPredictableBranchThreshold(). Intuitively it feels wrong to me to have branch_weights change based on the selected target here and I am not sure I follow why LowerExpectIntrinsic gets to be an exception and not use getPredictableBranchThreshold() then...

CC @lebedev.ri @spatel

Some history in

• https://github.com/llvm/llvm-project/commit/e3a470162738871bba982416748ae5f5e3572947
• D98898

[MatzeB]

Maybe we should expose the "LikelyBranchWeight" of LowerExpectIntrinsic.cpp in an internal LLVM header? So frontends are still forced to use the llvm.expect abstractions while LLVM-internal code can do the quicker thing and add likely branch weights directly...

[mtrofin]

nits.

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

785–791

should this logic be factored in a getBranchWeights(Latch, Count) - something like that?

[mtrofin]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

789

one thought about the 1:127 ratio: how about we had a MDBuilder::createUnrolledLoopSkipWeights() doing the setting - more readable, API captures the intent, etc?

[hoy]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

789

Can this be computed based on the loop preheader count and the main loop trip count? IIUC, the block execution count of the preheader indicates how many times the loop is entered from outside the loop. Thus that the loop trip count divided by the preheader count indicates an average trip count per loop execution. If that is greater than the unrolling factor, then the unrolled loop should be always executed?

[MatzeB]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

785–791

and put it where? I find that having the branch weight decisions right next to the code creating the branch helps readability of the code as you don't have to jump around.

789

how about we had a MDBuilder::createUnrolledLoopSkipWeights()

As above I think this hurts readability as it moves the logic away into another file. Also if we decide this being the preferred method, then I'd rather see a mass-refactoring of relevant places rather than starting with one-offs in diffs.

Can this be computed based on the loop preheader count and the main loop trip count?

I don't see how. The old preheader (if it existed at all) will give you the ratio of zero and non-zero trip counts. But this here needs to catch the cases of trip-counts being smaller than the unroll factor...

[mtrofin]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

789

how about we had a MDBuilder::createUnrolledLoopSkipWeights()

As above I think this hurts readability as it moves the logic away into another file. Also if we decide this being the preferred method, then I'd rather see a mass-refactoring of relevant places rather than starting with one-offs in diffs.

I can see your way for the above, but here, the way I see it is that the choice of numbers and the intent behind them is not immediately obvious, so an API would make that clear. At minimum the values could be a constant that's reused? (maybe the ArrayRef overload of createBranchWeights would take that constant or something like that)

If this set of weights is this way elsewhere, doing a refactoring after makes sense, but (maybe I searched superficially) can't seem to find them other than in your recent patches on this topic, that's why I was thinking this may be the opportune moment - wdyt?

[hoy]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

789

I don't see how. The old preheader (if it existed at all) will give you the ratio of zero and non-zero trip counts. But this here needs to catch the cases of trip-counts being smaller than the unroll factor...

The old preheader can be just some block dominating the loop header, for the sake of block execution count.

Right, the branch is comparing trip count per each loop entry with the unroll factor. But to simulate that probability at compile time, I feel we can use the aggregated loop trip count (i.e., profile count) divided by the preheader count to compare with the unroll factor. Let me know if this feels wrong.

[hoy]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

789

Talked offline. The average trip count is already used to check if loop unroll is beneficial. Here we are pretty sure the unrolled loop body is going to be run, so giving a 99.9...% ratio is reasonable.

[MatzeB]

factor out branch weight constants.

[wenlei]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

61

nit: in the comment also explain why it's unlikely for loop to not enter unrolled version at all. Like does the unroll only happen with known trip count (which is bigger than unroll factor)?

62

It's good to have constants extracted out, but when we do that, we can try to make it meaningful when someone just look at the constants without looking at the uses.

If {1,127} means *low* probability for not entering unrolled loop at all, it'd be confusing for the same {1,127} to also mean *high* probability for epilog loop to be executed.

Easiest way to fix is to add comment for each individual weight..

310

nit: you mean [0, count)?

399–402

I know this is just refactoring from updateLatchBranchWeightsForRemainderLoop, but I'm wondering if the assumption is reasonable? If we assume linear distribution for loop trip counts, then the average trip count for the remainder should be Count / 2, rather than Count - 1?

[MatzeB]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

61

in the comment also explain why it's unlikely for loop to not enter unrolled version at all. Like does the unroll only happen with known trip count (which is bigger than unroll factor)?

The header deals with the amount of iterations not fitting into the unroll factor. i.e. if the loop requires 43 iterations and you have an unroll factor of 4, then the header will do 3 iterations and we enter the unrolled loop to process the remaining 40 iterations (that are now 10 unrolled ones).

However I think this is the wrong place to explain how this pass works. It's just a constant used in the algorithm. I'd rather not try to write explanations of how the algorithm works here, we should have that at the beginning of the file, in the functions performing the transformations and policy questions you hint at are in LoopUnrollPass.cpp this file here is mostly about mechanisms.

I have a feeling what you are really asking about here is some reasoning why "127" is a good value. And honestly I don't have a deeper analysis behind this, except for looking at some small toy examples. This should be clearly better than the default 50:50 split we end up when not adding anything, so I hope we can leave tuning for later (if necessary).

62

If {1,127} means *low* probability for not entering unrolled loop at all, it'd be confusing for the same {1,127} to also mean *high* probability for epilog loop to be executed.

This is about entering a loop epilogue, not the unrolled loop. And you are reading this correctly in that it should be a *low* probability.

The comment already states this. The code either produces a prologue or an epilogue to deal with the extra iterations that don't fit the unrolled loop. Though again this seems like the wrong place to explain the intricacies of this pass, it's just some constants...

[MatzeB]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

399–402

CC @hoy who originally added this (D83187).

I think it would be (Count / 2) - 1 for linear distribution. Guess I can change it as part of this patch...

[MatzeB]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

399–402

Given there is a pre-header to rule out zero-trip counts and this being a loop backedge, the average should be (Count - 2) / 2.

[wenlei]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

61

The header deals with the amount of iterations not fitting into the unroll factor. i.e. if the loop requires 43 iterations and you have an unroll factor of 4, then the header will do 3 iterations and we enter the unrolled loop to process the remaining 40 iterations (that are now 10 unrolled ones).

This is not answering the question, instead this is explaining how the pass works.

However I think this is the wrong place to explain how this pass works.

We do not need to explain how the pass works in order to explain why specific values are chosen. :) We can assume some knowledge of how the pass works from the readers.

I have a feeling what you are really asking about here is some reasoning why "127" is a good value.

Partially. But what I'm really asking is to put an explanation for the value you chose in the comment. Just saying this is generally unlikely, and it's a ball park estimate, not yet tuned can also be helpful (so people know they have more freedom to change it).

I actually think we should put something like Hongtao's comment here.

The average trip count is already used to check if loop unroll is beneficial. Here we are pretty sure the unrolled loop body is going to be run, so giving a 99.9...% ratio is reasonable.

62

This is about entering a loop epilogue, not the unrolled loop.

For the unrolled loop, I was referring to the UnrolledLoopHeaderWeights above which shares the same value {1,127}. Sorry for causing confusion.

And you are reading this correctly in that it should be a *low* probability.

No, I read it as it's a high probability to enter loop epilog. Given that we have 1 / UF chance of loop trip count being exact multiples of UF, we enter loop epilog with a probability of (UF - 1) / UF. Am I missing something here?

And btw, if I'm missing something obvious here, the answer doesn't need to go into comments.

Though again this seems like the wrong place to explain the intricacies of this pass, it's just some constants...

I guess I care less where this is explained. But I think there needs to be an explanation for the values chosen somewhere. Currently there is no explanation.

However, given the constants are now extracted out, I think the most natural way is to have the explanation go along with the constants/values.

We don't need explain the intricacies of this pass, but IMHO we need enough context for people to understand why these values were chosen, and how confident original author was about the chosen value. For such explanation, we can assume certain knowledge of the pass. Staring at some values and scratching our heads just isn't a good experience, so I'd err on the side of over-communicate when it comes to chosen constants.

[MatzeB]

Rework formula for prolog/epilog loop-backedge weights as suggested by @wenlei .

[MatzeB]

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

62

Sorry if this is confusing, the test here is whether the loop trip count is so small that we skip the unrolled loop and enter the epilogue immediatey. Trying to make the comment more clear now...

[MatzeB]

Tweak comments for weight constants.

[wenlei]

lgtm, thanks!

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

399–402

Makes sense. I didn't consider the fact that it's a backedge.

llvm/lib/Transforms/Utils/LoopUnrollRuntime.cpp

 static cl::opt<bool> UnrollRuntimeMultiExit(
     "unroll-runtime-multi-exit", cl::init(false), cl::Hidden,
     cl::desc("Allow runtime unrolling for loops with multiple exits, when "
              "epilog is generated"));
 static cl::opt<bool> UnrollRuntimeOtherExitPredictable(
     "unroll-runtime-other-exit-predictable", cl::init(false), cl::Hidden,
     cl::desc("Assume the non latch exit block to be predictable"));
 
+// Probability that the loop trip count is so small that after the prolog
+// we do not enter the unrolled loop at all.
+static const uint32_t UnrolledLoopHeaderWeights[] = {1, 127};
+// Probability that the epilogue loop will be executed at all.
+static const uint32_t EpilogHeaderWeights[] = {1, 127};
+
 /// Connect the unrolling prolog code to the original loop.
 /// The unrolling prolog code contains code to execute the
 /// 'extra' iterations if the run-time trip count modulo the
 /// unroll count is non-zero.
 ///
 /// This function performs the following:
 /// - Create PHI nodes at prolog end block to combine values
 ///   that exit the prolog code and jump around the prolog.
static void ConnectProlog(Loop *L, Value *BECount, unsigned Count,
   // then (BECount + 1) cannot unsigned-overflow.
   Value *BrLoopExit =
       B.CreateICmpULT(BECount, ConstantInt::get(BECount->getType(), Count - 1));
   // Split the exit to maintain loop canonicalization guarantees
   SmallVector<BasicBlock *, 4> Preds(predecessors(OriginalLoopLatchExit));
   SplitBlockPredecessors(OriginalLoopLatchExit, Preds, ".unr-lcssa", DT, LI,
                          nullptr, PreserveLCSSA);
   // Add the branch to the exit block (around the unrolled loop)
-  B.CreateCondBr(BrLoopExit, OriginalLoopLatchExit, NewPreHeader);
+  MDNode *BranchWeights = nullptr;
+  if (hasBranchWeightMD(*Latch->getTerminator())) {
+    // Assume loop is nearly always entered.
+    MDBuilder MDB(B.getContext());
+    BranchWeights = MDB.createBranchWeights(UnrolledLoopHeaderWeights);
+  }
+  B.CreateCondBr(BrLoopExit, OriginalLoopLatchExit, NewPreHeader,
+                 BranchWeights);
   InsertPt->eraseFromParent();
   if (DT) {
     auto *NewDom = DT->findNearestCommonDominator(OriginalLoopLatchExit,
                                                   PrologExit);
     DT->changeImmediateDominator(OriginalLoopLatchExit, NewDom);
   }
 }
 
 /// Connect the unrolling epilog code to the original loop.
 /// The unrolling epilog code contains code to execute the
 /// 'extra' iterations if the run-time trip count modulo the
 /// unroll count is non-zero.
 ///
 /// This function performs the following:
 /// - Update PHI nodes at the unrolling loop exit and epilog loop exit
 /// - Create PHI nodes at the unrolling loop exit to combine
 ///   values that exit the unrolling loop code and jump around it.
 /// - Update PHI operands in the epilog loop by the new PHI nodes
 /// - Branch around the epilog loop if extra iters (ModVal) is zero.
 ///
 static void ConnectEpilog(Loop *L, Value *ModVal, BasicBlock *NewExit,
                           BasicBlock *Exit, BasicBlock *PreHeader,
                           BasicBlock *EpilogPreHeader, BasicBlock *NewPreHeader,
                           ValueToValueMapTy &VMap, DominatorTree *DT,
-                          LoopInfo *LI, bool PreserveLCSSA,
-                          ScalarEvolution &SE) {
+                          LoopInfo *LI, bool PreserveLCSSA, ScalarEvolution &SE,
+                          unsigned Count) {
   BasicBlock *Latch = L->getLoopLatch();
   assert(Latch && "Loop must have a latch");
   BasicBlock *EpilogLatch = cast<BasicBlock>(VMap[Latch]);
 
   // Loop structure should be the following:
   //
   // PreHeader
   // NewPreHeader
static void ConnectEpilog(Loop *L, Value *ModVal, BasicBlock *NewExit,
   IRBuilder<> B(InsertPt);
   Value *BrLoopExit = B.CreateIsNotNull(ModVal, "lcmp.mod");
   assert(Exit && "Loop must have a single exit block only");
   // Split the epilogue exit to maintain loop canonicalization guarantees
   SmallVector<BasicBlock*, 4> Preds(predecessors(Exit));
   SplitBlockPredecessors(Exit, Preds, ".epilog-lcssa", DT, LI, nullptr,
                          PreserveLCSSA);
   // Add the branch to the exit block (around the unrolling loop)
-  B.CreateCondBr(BrLoopExit, EpilogPreHeader, Exit);
+  MDNode *BranchWeights = nullptr;
+  if (hasBranchWeightMD(*Latch->getTerminator())) {
+    // Assume equal distribution in interval [0;Count).
+    MDBuilder MDB(B.getContext());
+    BranchWeights = MDB.createBranchWeights(1, Count - 1);
+  }
+  B.CreateCondBr(BrLoopExit, EpilogPreHeader, Exit, BranchWeights);
   InsertPt->eraseFromParent();
   if (DT) {
     auto *NewDom = DT->findNearestCommonDominator(Exit, NewExit);
     DT->changeImmediateDominator(Exit, NewDom);
   }
 
   // Split the main loop exit to maintain canonicalization guarantees.
   SmallVector<BasicBlock*, 4> NewExitPreds{Latch};
   SplitBlockPredecessors(NewExit, NewExitPreds, ".loopexit", DT, LI, nullptr,
                          PreserveLCSSA);
 }
 
 /// Create a clone of the blocks in a loop and connect them together. A new
 /// loop will be created including all cloned blocks, and the iterator of the
 /// new loop switched to count NewIter down to 0.
 /// The cloned blocks should be inserted between InsertTop and InsertBot.
 /// InsertTop should be new preheader, InsertBot new loop exit.
 /// Returns the new cloned loop that is created.
 static Loop *
 CloneLoopBlocks(Loop *L, Value *NewIter, const bool UseEpilogRemainder,
                 const bool UnrollRemainder,
                 BasicBlock *InsertTop,
                 BasicBlock *InsertBot, BasicBlock *Preheader,
-                std::vector<BasicBlock *> &NewBlocks, LoopBlocksDFS &LoopBlocks,
-                ValueToValueMapTy &VMap, DominatorTree *DT, LoopInfo *LI) {
+                             std::vector<BasicBlock *> &NewBlocks,
+                             LoopBlocksDFS &LoopBlocks, ValueToValueMapTy &VMap,
+                             DominatorTree *DT, LoopInfo *LI, unsigned Count) {
   StringRef suffix = UseEpilogRemainder ? "epil" : "prol";
   BasicBlock *Header = L->getHeader();
   BasicBlock *Latch = L->getLoopLatch();
   Function *F = Header->getParent();
   LoopBlocksDFS::RPOIterator BlockBegin = LoopBlocks.beginRPO();
   LoopBlocksDFS::RPOIterator BlockEnd = LoopBlocks.endRPO();
   Loop *ParentLoop = L->getParentLoop();
   NewLoopsMap NewLoops;
if (Latch == *BB) {
       IRBuilder<> Builder(LatchBR);
       PHINode *NewIdx = PHINode::Create(NewIter->getType(), 2,
                                         suffix + ".iter",
                                         FirstLoopBB->getFirstNonPHI());
       auto *Zero = ConstantInt::get(NewIdx->getType(), 0);
       auto *One = ConstantInt::get(NewIdx->getType(), 1);
       Value *IdxNext = Builder.CreateAdd(NewIdx, One, NewIdx->getName() + ".next");
       Value *IdxCmp = Builder.CreateICmpNE(IdxNext, NewIter, NewIdx->getName() + ".cmp");
-      Builder.CreateCondBr(IdxCmp, FirstLoopBB, InsertBot);
+      MDNode *BranchWeights = nullptr;
+      uint64_t OrigExitWeight;
+      uint64_t OrigBackedgeWeight;
+      if (extractBranchWeights(*LatchBR, OrigExitWeight, OrigBackedgeWeight)) {
+        if (LatchBR->getSuccessor(0) == L->getHeader())
+          std::swap(OrigExitWeight, OrigBackedgeWeight);
+        // Assign the maximum possible trip count as the back edge weight for
+        // the remainder loop if the original loop comes with a branch weight.
+        assert(Count > 1 && "unroll factor must be 2 or more");
+        uint64_t BackEdgeWeight = (Count - 1) * OrigExitWeight;
+        MDBuilder MDB(Builder.getContext());
+        BranchWeights = MDB.createBranchWeights(BackEdgeWeight, OrigExitWeight);
+      }
+      Builder.CreateCondBr(IdxCmp, FirstLoopBB, InsertBot, BranchWeights);
       NewIdx->addIncoming(Zero, InsertTop);
       NewIdx->addIncoming(IdxNext, NewBB);
       LatchBR->eraseFromParent();
     }
   }
 
   // Change the incoming values to the ones defined in the preheader or
   // cloned loop.
return (OtherExits.size() == 1 &&
           (UnrollRuntimeOtherExitPredictable ||
            OtherExits[0]->getPostdominatingDeoptimizeCall()));
   // TODO: These can be fine-tuned further to consider code size or deopt states
   // that are captured by the deoptimize exit block.
   // Also, we can extend this to support more cases, if we actually
   // know of kinds of multiexit loops that would benefit from unrolling.
 }
 
-// Assign the maximum possible trip count as the back edge weight for the
-// remainder loop if the original loop comes with a branch weight.
-static void updateLatchBranchWeightsForRemainderLoop(Loop *OrigLoop,
-                                                     Loop *RemainderLoop,
-                                                     uint64_t UnrollFactor) {
-  uint64_t TrueWeight, FalseWeight;
-  BranchInst *LatchBR =
-      cast<BranchInst>(OrigLoop->getLoopLatch()->getTerminator());
-  if (!extractBranchWeights(*LatchBR, TrueWeight, FalseWeight))
-    return;
-  uint64_t ExitWeight = LatchBR->getSuccessor(0) == OrigLoop->getHeader()
-                            ? FalseWeight
-                            : TrueWeight;
-  assert(UnrollFactor > 1);
-  uint64_t BackEdgeWeight = (UnrollFactor - 1) * ExitWeight;
-  BasicBlock *Header = RemainderLoop->getHeader();
-  BasicBlock *Latch = RemainderLoop->getLoopLatch();
-  auto *RemainderLatchBR = cast<BranchInst>(Latch->getTerminator());
-  unsigned HeaderIdx = (RemainderLatchBR->getSuccessor(0) == Header ? 0 : 1);
-  MDBuilder MDB(RemainderLatchBR->getContext());
-  MDNode *WeightNode =
-    HeaderIdx ? MDB.createBranchWeights(ExitWeight, BackEdgeWeight)
-                : MDB.createBranchWeights(BackEdgeWeight, ExitWeight);
-  RemainderLatchBR->setMetadata(LLVMContext::MD_prof, WeightNode);
-}
-
 /// Calculate ModVal = (BECount + 1) % Count on the abstract integer domain
 /// accounting for the possibility of unsigned overflow in the 2s complement
 /// domain. Preconditions:
 /// 1) TripCount = BECount + 1 (allowing overflow)
 /// 2) Log2(Count) <= BitWidth(BECount)
 static Value *CreateTripRemainder(IRBuilder<> &B, Value *BECount,
                                   Value *TripCount, unsigned Count) {
   // Note that TripCount is BECount + 1.
bool llvm::UnrollRuntimeLoopRemainder(
   Value *BranchVal =
       UseEpilogRemainder ? B.CreateICmpULT(BECount,
                                            ConstantInt::get(BECount->getType(),
                                                             Count - 1)) :
                            B.CreateIsNotNull(ModVal, "lcmp.mod");
   BasicBlock *RemainderLoop = UseEpilogRemainder ? NewExit : PrologPreHeader;
   BasicBlock *UnrollingLoop = UseEpilogRemainder ? NewPreHeader : PrologExit;
   // Branch to either remainder (extra iterations) loop or unrolling loop.
-  B.CreateCondBr(BranchVal, RemainderLoop, UnrollingLoop);
+  MDNode *BranchWeights = nullptr;
+  if (hasBranchWeightMD(*Latch->getTerminator())) {
+    // Assume loop is nearly always entered.
+    MDBuilder MDB(B.getContext());
+    BranchWeights = MDB.createBranchWeights(EpilogHeaderWeights);
+  }
+  B.CreateCondBr(BranchVal, RemainderLoop, UnrollingLoop, BranchWeights);
   PreHeaderBR->eraseFromParent();
   if (DT) {
     if (UseEpilogRemainder)
       DT->changeImmediateDominator(NewExit, PreHeader);
     else
       DT->changeImmediateDominator(PrologExit, PreHeader);
   }
   Function *F = Header->getParent();
bool llvm::UnrollRuntimeLoopRemainder(
 
   // Clone all the basic blocks in the loop. If Count is 2, we don't clone
   // the loop, otherwise we create a cloned loop to execute the extra
   // iterations. This function adds the appropriate CFG connections.
   BasicBlock *InsertBot = UseEpilogRemainder ? LatchExit : PrologExit;
   BasicBlock *InsertTop = UseEpilogRemainder ? EpilogPreHeader : PrologPreHeader;
   Loop *remainderLoop = CloneLoopBlocks(
       L, ModVal, UseEpilogRemainder, UnrollRemainder, InsertTop, InsertBot,
-      NewPreHeader, NewBlocks, LoopBlocks, VMap, DT, LI);
-
-  // Assign the maximum possible trip count as the back edge weight for the
-  // remainder loop if the original loop comes with a branch weight.
-  if (remainderLoop && !UnrollRemainder)
-    updateLatchBranchWeightsForRemainderLoop(L, remainderLoop, Count);
+      NewPreHeader, NewBlocks, LoopBlocks, VMap, DT, LI, Count);
 
   // Insert the cloned blocks into the function.
   F->splice(InsertBot->getIterator(), F, NewBlocks[0]->getIterator(), F->end());
 
   // Now the loop blocks are cloned and the other exiting blocks from the
   // remainder are connected to the original Loop's exit blocks. The remaining
   // work is to update the phi nodes in the original loop, and take in the
   // values from the cloned region.
for (Instruction &I : *BB) {
                        RF_NoModuleLevelChanges | RF_IgnoreMissingLocals);
     }
   }
 
   if (UseEpilogRemainder) {
     // Connect the epilog code to the original loop and update the
     // PHI functions.
     ConnectEpilog(L, ModVal, NewExit, LatchExit, PreHeader, EpilogPreHeader,
-                  NewPreHeader, VMap, DT, LI, PreserveLCSSA, *SE);
+                  NewPreHeader, VMap, DT, LI, PreserveLCSSA, *SE, Count);
 
     // Update counter in loop for unrolling.
     // Use an incrementing IV.  Pre-incr/post-incr is backedge/trip count.
     // Subtle: TestVal can be 0 if we wrapped when computing the trip count,
     // thus we must compare the post-increment (wrapping) value.
     IRBuilder<> B2(NewPreHeader->getTerminator());
     Value *TestVal = B2.CreateSub(TripCount, ModVal, "unroll_iter");
     BranchInst *LatchBR = cast<BranchInst>(Latch->getTerminator());

llvm/test/Transforms/LoopUnroll/runtime-exit-phi-scev-invalidation.ll

 ; CHECK-NEXT:    [[TMP1:%.*]] = sub i64 [[TMP0]], [[X]]
 ; CHECK-NEXT:    br label [[OUTER_HEADER:%.*]]
 ; CHECK:       outer.header:
 ; CHECK-NEXT:    [[OUTER_P:%.*]] = phi i32 [ 0, [[BB:%.*]] ], [ [[L_1_LCSSA:%.*]], [[OUTER_LATCH:%.*]] ]
 ; CHECK-NEXT:    [[TMP2:%.*]] = freeze i64 [[TMP1]]
 ; CHECK-NEXT:    [[TMP3:%.*]] = add i64 [[TMP2]], -1
 ; CHECK-NEXT:    [[XTRAITER:%.*]] = and i64 [[TMP2]], 7
 ; CHECK-NEXT:    [[LCMP_MOD:%.*]] = icmp ne i64 [[XTRAITER]], 0
-; CHECK-NEXT:    br i1 [[LCMP_MOD]], label [[INNER_1_HEADER_PROL_PREHEADER:%.*]], label [[INNER_1_HEADER_PROL_LOOPEXIT:%.*]]
+; CHECK-NEXT:    br i1 [[LCMP_MOD]], label [[INNER_1_HEADER_PROL_PREHEADER:%.*]], label [[INNER_1_HEADER_PROL_LOOPEXIT:%.*]], !prof [[PROF3:![0-9]+]]
 ; CHECK:       inner.1.header.prol.preheader:
 ; CHECK-NEXT:    br label [[INNER_1_HEADER_PROL:%.*]]
 ; CHECK:       inner.1.header.prol:
 ; CHECK-NEXT:    [[INNER_1_IV_PROL:%.*]] = phi i64 [ [[X]], [[INNER_1_HEADER_PROL_PREHEADER]] ], [ [[INNER_1_IV_NEXT_PROL:%.*]], [[INNER_1_LATCH_PROL:%.*]] ]
 ; CHECK-NEXT:    [[PROL_ITER:%.*]] = phi i64 [ 0, [[INNER_1_HEADER_PROL_PREHEADER]] ], [ [[PROL_ITER_NEXT:%.*]], [[INNER_1_LATCH_PROL]] ]
 ; CHECK-NEXT:    [[CMP_1_PROL:%.*]] = icmp sgt i32 [[OUTER_P]], 0
 ; CHECK-NEXT:    br i1 [[CMP_1_PROL]], label [[EXIT_DEOPT_LOOPEXIT1:%.*]], label [[INNER_1_LATCH_PROL]]
 ; CHECK:       inner.1.latch.prol:
 ; CHECK-NEXT:    [[L_1_PROL:%.*]] = load i32, ptr [[SRC:%.*]], align 4
 ; CHECK-NEXT:    store i32 [[L_1_PROL]], ptr [[DST:%.*]], align 8
 ; CHECK-NEXT:    [[INNER_1_IV_NEXT_PROL]] = add i64 [[INNER_1_IV_PROL]], 1
 ; CHECK-NEXT:    [[CMP_2_PROL:%.*]] = icmp sgt i64 [[INNER_1_IV_PROL]], 0
 ; CHECK-NEXT:    [[PROL_ITER_NEXT]] = add i64 [[PROL_ITER]], 1
 ; CHECK-NEXT:    [[PROL_ITER_CMP:%.*]] = icmp ne i64 [[PROL_ITER_NEXT]], [[XTRAITER]]
-; CHECK-NEXT:    br i1 [[PROL_ITER_CMP]], label [[INNER_1_HEADER_PROL]], label [[INNER_1_HEADER_PROL_LOOPEXIT_UNR_LCSSA:%.*]], !prof [[PROF3:![0-9]+]], !llvm.loop [[LOOP4:![0-9]+]]
+; CHECK-NEXT:    br i1 [[PROL_ITER_CMP]], label [[INNER_1_HEADER_PROL]], label [[INNER_1_HEADER_PROL_LOOPEXIT_UNR_LCSSA:%.*]], !prof [[PROF4:![0-9]+]], !llvm.loop [[LOOP5:![0-9]+]]
 ; CHECK:       inner.1.header.prol.loopexit.unr-lcssa:
 ; CHECK-NEXT:    [[L_1_LCSSA_UNR_PH:%.*]] = phi i32 [ [[L_1_PROL]], [[INNER_1_LATCH_PROL]] ]
 ; CHECK-NEXT:    [[INNER_1_IV_UNR_PH:%.*]] = phi i64 [ [[INNER_1_IV_NEXT_PROL]], [[INNER_1_LATCH_PROL]] ]
 ; CHECK-NEXT:    br label [[INNER_1_HEADER_PROL_LOOPEXIT]]
 ; CHECK:       inner.1.header.prol.loopexit:
 ; CHECK-NEXT:    [[L_1_LCSSA_UNR:%.*]] = phi i32 [ undef, [[OUTER_HEADER]] ], [ [[L_1_LCSSA_UNR_PH]], [[INNER_1_HEADER_PROL_LOOPEXIT_UNR_LCSSA]] ]
 ; CHECK-NEXT:    [[INNER_1_IV_UNR:%.*]] = phi i64 [ [[X]], [[OUTER_HEADER]] ], [ [[INNER_1_IV_UNR_PH]], [[INNER_1_HEADER_PROL_LOOPEXIT_UNR_LCSSA]] ]
 ; CHECK-NEXT:    [[TMP4:%.*]] = icmp ult i64 [[TMP3]], 7
-; CHECK-NEXT:    br i1 [[TMP4]], label [[OUTER_MIDDLE:%.*]], label [[OUTER_HEADER_NEW:%.*]]
+; CHECK-NEXT:    br i1 [[TMP4]], label [[OUTER_MIDDLE:%.*]], label [[OUTER_HEADER_NEW:%.*]], !prof [[PROF3]]
 ; CHECK:       outer.header.new:
 ; CHECK-NEXT:    br label [[INNER_1_HEADER:%.*]]
 ; CHECK:       inner.1.header:
 ; CHECK-NEXT:    [[INNER_1_IV:%.*]] = phi i64 [ [[INNER_1_IV_UNR]], [[OUTER_HEADER_NEW]] ], [ [[INNER_1_IV_NEXT_7:%.*]], [[INNER_1_LATCH_7:%.*]] ]
 ; CHECK-NEXT:    [[CMP_1:%.*]] = icmp sgt i32 [[OUTER_P]], 0
 ; CHECK-NEXT:    br i1 [[CMP_1]], label [[EXIT_DEOPT_LOOPEXIT:%.*]], label [[INNER_1_LATCH:%.*]]
 ; CHECK:       inner.1.latch:
 ; CHECK-NEXT:    [[L_1:%.*]] = load i32, ptr [[SRC]], align 4

 ; CHECK-NEXT:    store i32 [[L_1_6]], ptr [[DST]], align 8
 ; CHECK-NEXT:    [[INNER_1_IV_NEXT_6:%.*]] = add i64 [[INNER_1_IV]], 7
 ; CHECK-NEXT:    br i1 false, label [[EXIT_DEOPT_LOOPEXIT]], label [[INNER_1_LATCH_7]]
 ; CHECK:       inner.1.latch.7:
 ; CHECK-NEXT:    [[L_1_7:%.*]] = load i32, ptr [[SRC]], align 4
 ; CHECK-NEXT:    store i32 [[L_1_7]], ptr [[DST]], align 8
 ; CHECK-NEXT:    [[INNER_1_IV_NEXT_7]] = add i64 [[INNER_1_IV]], 8
 ; CHECK-NEXT:    [[CMP_2_7:%.*]] = icmp sgt i64 [[INNER_1_IV_NEXT_6]], 0
-; CHECK-NEXT:    br i1 [[CMP_2_7]], label [[OUTER_MIDDLE_UNR_LCSSA:%.*]], label [[INNER_1_HEADER]], !prof [[PROF5:![0-9]+]]
+; CHECK-NEXT:    br i1 [[CMP_2_7]], label [[OUTER_MIDDLE_UNR_LCSSA:%.*]], label [[INNER_1_HEADER]], !prof [[PROF6:![0-9]+]]
 ; CHECK:       outer.middle.unr-lcssa:
 ; CHECK-NEXT:    [[L_1_LCSSA_PH:%.*]] = phi i32 [ [[L_1_7]], [[INNER_1_LATCH_7]] ]
 ; CHECK-NEXT:    br label [[OUTER_MIDDLE]]
 ; CHECK:       outer.middle:
 ; CHECK-NEXT:    [[L_1_LCSSA]] = phi i32 [ [[L_1_LCSSA_UNR]], [[INNER_1_HEADER_PROL_LOOPEXIT]] ], [ [[L_1_LCSSA_PH]], [[OUTER_MIDDLE_UNR_LCSSA]] ]
 ; CHECK-NEXT:    br label [[INNER_2:%.*]]
 ; CHECK:       inner.2:
 ; CHECK-NEXT:    [[INNER_2_IV:%.*]] = phi i32 [ [[L_1_LCSSA]], [[OUTER_MIDDLE]] ], [ [[INNER_2_IV_NEXT_2:%.*]], [[INNER_2]] ]

llvm/test/Transforms/LoopUnroll/runtime-loop.ll

 ; RUN: opt < %s -S -passes='require<opt-remark-emit>,loop(loop-unroll-full)' -unroll-runtime=true -unroll-runtime-epilog=false | FileCheck %s -check-prefixes=NOPROLOG,COMMON
 
 target datalayout = "e-p:64:64:64-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:64:64-f32:32:32-f64:64:64-v64:64:64-v128:128:128-a0:0:64-s0:64:64-f80:128:128-n8:16:32:64-S128"
 
 ; Tests for unrolling loops with run-time trip counts
 
 ; COMMON-LABEL: @test(
 
-; EPILOG: %xtraiter = and i32 %n
-; EPILOG:  %lcmp.mod = icmp ne i32 %xtraiter, 0
-; EPILOG:  br i1 %lcmp.mod, label %for.body.epil.preheader, label %for.end.loopexit
+; EPILOG: entry:
+; EPILOG:   br i1 %cmp1, label %for.end, label %for.body.preheader, !prof [[EPILOG_PROF_0:![0-9]+]]
+; EPILOG: for.body.preheader:
+; EPILOG:   %xtraiter = and i32 %n
+; EPILOG:   br i1 %1, label %for.end.loopexit.unr-lcssa, label %for.body.preheader.new, !prof [[EPILOG_PROF_1:![0-9]+]]
+
+; EPILOG: for.end.loopexit.unr-lcssa:
+; EPILOG:   %lcmp.mod = icmp ne i32 %xtraiter, 0
+; EPILOG:   br i1 %lcmp.mod, label %for.body.epil.preheader, label %for.end.loopexit, !prof [[EPILOG_PROF_2:![0-9]+]]
 
 ; NOEPILOG-NOT: %xtraiter = and i32 %n
 
-; PROLOG: %xtraiter = and i32 %n
-; PROLOG:  %lcmp.mod = icmp ne i32 %xtraiter, 0
-; PROLOG:  br i1 %lcmp.mod, label %for.body.prol.preheader, label %for.body.prol.loopexit
+; PROLOG: entry:
+; PROLOG:   br i1 %cmp1, label %for.end, label %for.body.preheader, !prof [[PROLOG_PROF_0:![0-9]+]]
+
+; PROLOG: for.body.preheader:
+; PROLOG:   %xtraiter = and i32 %n
+; PROLOG:   %lcmp.mod = icmp ne i32 %xtraiter, 0
+; PROLOG:   br i1 %lcmp.mod, label %for.body.prol.preheader, label %for.body.prol.loopexit, !prof [[PROLOG_PROF_1:![0-9]+]]
 
 ; NOPROLOG-NOT: %xtraiter = and i32 %n
 
 ; EPILOG: for.body.epil:
-; EPILOG: %indvars.iv.epil = phi i64 [ %indvars.iv.next.epil, %for.body.epil ],  [ %indvars.iv.unr, %for.body.epil.preheader ]
-; EPILOG:  %epil.iter.next = add i32 %epil.iter, 1
-; EPILOG:  %epil.iter.cmp = icmp ne i32 %epil.iter.next, %xtraiter
-; EPILOG:  br i1 %epil.iter.cmp, label %for.body.epil, label %for.end.loopexit.epilog-lcssa, !llvm.loop !0
+; EPILOG:   %indvars.iv.epil = phi i64 [ %indvars.iv.next.epil, %for.body.epil ],  [ %indvars.iv.unr, %for.body.epil.preheader ]
+; EPILOG:   %epil.iter.next = add i32 %epil.iter, 1
+; EPILOG:   %epil.iter.cmp = icmp ne i32 %epil.iter.next, %xtraiter
+; EPILOG:   br i1 %epil.iter.cmp, label %for.body.epil, label %for.end.loopexit.epilog-lcssa, !prof [[EPILOG_PROF_3:![0-9]+]], !llvm.loop [[EPILOG_LOOP:![0-9]+]]
 
 ; NOEPILOG: for.body:
 ; NOEPILOG-NOT: for.body.epil:
 
 ; PROLOG: for.body.prol:
-; PROLOG: %indvars.iv.prol = phi i64 [ %indvars.iv.next.prol, %for.body.prol ], [ 0, %for.body.prol.preheader ]
-; PROLOG:  %prol.iter.next = add i32 %prol.iter, 1
-; PROLOG:  %prol.iter.cmp = icmp ne i32 %prol.iter.next, %xtraiter
-; PROLOG:  br i1 %prol.iter.cmp, label %for.body.prol, label %for.body.prol.loopexit.unr-lcssa, !llvm.loop !0
+; PROLOG:   %indvars.iv.prol = phi i64 [ %indvars.iv.next.prol, %for.body.prol ], [ 0, %for.body.prol.preheader ]
+; PROLOG:   %prol.iter.next = add i32 %prol.iter, 1
+; PROLOG:   %prol.iter.cmp = icmp ne i32 %prol.iter.next, %xtraiter
+; PROLOG:   br i1 %prol.iter.cmp, label %for.body.prol, label %for.body.prol.loopexit.unr-lcssa, !prof [[PROLOG_PROF_2:![0-9]+]], !llvm.loop [[PROLOG_LOOP:![0-9]+]]
+
+; PROLOG: for.body.prol.loopexit:
+; PROLOG:   br i1 %2, label %for.end.loopexit, label %for.body.preheader.new, !prof [[PROLOG_PROF_1:![0-9]+]]
 
 ; NOPROLOG: for.body:
 ; NOPROLOG-NOT: for.body.prol:
 
 
-define i32 @test(ptr nocapture %a, i32 %n) nounwind uwtable readonly {
+define i32 @test(ptr nocapture %a, i32 %n) nounwind uwtable readonly !prof !2 {
 entry:
   %cmp1 = icmp eq i32 %n, 0
-  br i1 %cmp1, label %for.end, label %for.body
+  br i1 %cmp1, label %for.end, label %for.body, !prof !3
 
 for.body:                                         ; preds = %for.body, %entry
   %indvars.iv = phi i64 [ %indvars.iv.next, %for.body ], [ 0, %entry ]
   %sum.02 = phi i32 [ %add, %for.body ], [ 0, %entry ]
   %arrayidx = getelementptr inbounds i32, ptr %a, i64 %indvars.iv
   %0 = load i32, ptr %arrayidx, align 4
   %add = add nsw i32 %0, %sum.02
   %indvars.iv.next = add i64 %indvars.iv, 1
   %lftr.wideiv = trunc i64 %indvars.iv.next to i32
   %exitcond = icmp eq i32 %lftr.wideiv, %n
-  br i1 %exitcond, label %for.end, label %for.body
+  br i1 %exitcond, label %for.end, label %for.body, !prof !4
 
 for.end:                                          ; preds = %for.body, %entry
   %sum.0.lcssa = phi i32 [ 0, %entry ], [ %add, %for.body ]
   ret i32 %sum.0.lcssa
 }
 
 
 ; Still try to completely unroll loops with compile-time trip counts

  ret void
 
 exit2.loopexit:
   ret void
 }
 
 !0 = distinct !{!0, !1}
 !1 = !{!"llvm.loop.unroll.runtime.disable"}
+!2 = !{!"function_entry_count", i64 1}
+!3 = !{!"branch_weights", i32 1, i32 11}
+!4 = !{!"branch_weights", i32 1, i32 42}
 
 ; need to use LABEL here to separate function IR matching from metadata matching
 ; COMMON-LABEL: {{^}}!0 =
 
-; EPILOG-SAME: distinct !{!0, !1}
-; EPILOG: !1 = !{!"llvm.loop.unroll.disable"}
+; EPILOG: [[EPILOG_PROF_0]] = !{!"branch_weights", i32 1, i32 11}
+; EPILOG: [[EPILOG_PROF_1]] = !{!"branch_weights", i32 1, i32 127}
+; EPILOG: [[EPILOG_PROF_2]] = !{!"branch_weights", i32 1, i32 7}
+; EPILOG: [[EPILOG_PROF_3]] = !{!"branch_weights", i32 7, i32 1}
+
+; EPILOG: [[EPILOG_LOOP]] = distinct !{[[EPILOG_LOOP]], [[EPILOG_LOOP_1:![0-9]+]]}
+; EPILOG: [[EPILOG_LOOP_1]] = !{!"llvm.loop.unroll.disable"}
+
+; PROLOG: [[PROLOG_PROF_0]] = !{!"branch_weights", i32 1, i32 11}
+; PROLOG: [[PROLOG_PROF_1]] = !{!"branch_weights", i32 1, i32 127}
+; PROLOG: [[PROLOG_PROF_2]] = !{!"branch_weights", i32 7, i32 1}
 
-; PROLOG-SAME: distinct !{!0, !1}
-; PROLOG: !1 = !{!"llvm.loop.unroll.disable"}
+; PROLOG: distinct !{[[PROLOG_LOOP]], [[PROLOG_LOOP_1:![0-9]+]]}
+; PROLOG: [[PROLOG_LOOP_1]] = !{!"llvm.loop.unroll.disable"}

llvm/test/Transforms/LoopUnroll/unroll-heuristics-pgo.ll

 ; RUN: opt < %s -S -passes=loop-unroll -unroll-runtime -unroll-threshold=40 -unroll-max-percent-threshold-boost=100 | FileCheck %s
 
 @known_constant = internal unnamed_addr constant [9 x i32] [i32 0, i32 -1, i32 0, i32 -1, i32 5, i32 -1, i32 0, i32 -1, i32 0], align 16
 
 ; CHECK-LABEL: @bar_prof
 ; CHECK: loop:
-; CHECK: %mul = mul
-; CHECK: %mul.1 = mul
-; CHECK: %mul.2 = mul
-; CHECK: %mul.3 = mul
-; CHECK: br i1 %niter.ncmp.7, label %loop.end.unr-lcssa.loopexit, label %loop, !prof !1
+; CHECK:   %mul = mul
+; CHECK:   %mul.1 = mul
+; CHECK:   %mul.2 = mul
+; CHECK:   %mul.3 = mul
+; CHECK:   br i1 %niter.ncmp.7, label %loop.end.unr-lcssa.loopexit, label %loop, !prof [[PROF0:![0-9]+]]
 ; CHECK: loop.epil:
-; CHECK:   br i1 %epil.iter.cmp, label %loop.epil, label %loop.end.epilog-lcssa, !prof !2, !llvm.loop !3
+; CHECK:   br i1 %epil.iter.cmp, label %loop.epil, label %loop.end.epilog-lcssa, !prof [[PROF1:![0-9]+]], !llvm.loop {{![0-9]+}}
 define i32 @bar_prof(ptr noalias nocapture readonly %src, i64 %c) !prof !1 {
 entry:
   br label %loop
 
 loop:
   %iv = phi i64 [ 0, %entry ], [ %inc, %loop ]
   %r  = phi i32 [ 0, %entry ], [ %add, %loop ]
   %arrayidx = getelementptr inbounds i32, ptr %src, i64 %iv

 loop.end:
   %r.lcssa = phi i32 [ %r, %loop ]
   ret i32 %r.lcssa
 }
 
 !1 = !{!"function_entry_count", i64 1}
 !2 = !{!"branch_weights", i32 1, i32 1000}
 
-; CHECK: !1 = !{!"branch_weights", i32 1, i32 124}
-; CHECK: !2 = !{!"branch_weights", i32 7, i32 1}
+; CHECK: [[PROF0]] = !{!"branch_weights", i32 1, i32 124}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 7, i32 1}