[LV] Vectorizer should adjust trip count in profile information

[ebrevnov]

Summary

Vectorized loop processes VFxUF number of elements in one iteration thus total number of iterations decreases proportionally. In addition epilog loop may not have more than VFxUF - 1 iterations. This patch updates profile information accordingly.

[ebrevnov]

Minor test update

[ebrevnov]

ping

[DaniilSuchkov]

LGTM

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4053

Nit: VFxUF - 1

4059–4061

Style: it is usually advised to turn such conditions into early exits, it would reduce required indentation and slightly improve readability.

4080

Maybe introduce a new variable for this value (like EpilogueTakenCount)? Right now it's a bit surprising that OrigSomething is being changed. Same goes to OrigFallThroughCount.

[ebrevnov]

Minor fixes as requested by reviewer.

[ebrevnov]

I realized that current implementation has a flaw and and we should take into account that actual number of iterations is one greater than back edge taken count. In addition I believe that current structuring of calculations is easier for understanding.

[DaniilSuchkov]

LGTM

[Ayal]

Adding a few comments. Would be good to generalize and apply also to loop unroll (and jam).

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4062

OrigFallThroughCount can still be either the exit count or the continue-to-next-iteration count, according to the code below. Wait to test if its zero until we know what it stands for?

4074

Better use distinct names, e.g., OrigExitCount and OrigBackedgeTakenCount, than continue to call them Taken and FallThrough. Perhaps use Weight instead of Count, to denote total profile frequencies, as the latter is used elsewhere to denote the actual per-invocation TripCount.

4077

bel[l]ow

4079

How about "OrigAverageTripCount"?

Explanation about its computation:
OrigAverageTripCount = (number of times header block was executed) / (number of times header was reached from pre-header == number of times latch exited)
== (OrigTakenCount + OrigFallThroughCount) / OrigFallThroughCount
== OrigTakenCount / OrigFallThroughCount + 1.

4081

How about VecAverageTripCount = OrigAverageTripCount / (VF * UF);

4085

Just to clarify, maintaining branch frequencies through optimizations is best-effort and imprecise - a total weight that does not divide VF*UF implies that the trip count of at-least one invocation did not divide VF*UF, not necessarily all of them; w/o considering also the distribution of trip counts in addition to their sum.

Setting PRIterCount = 0 and VecAverageTipCount = round(OrigAverageTripCount / (VF*UV)) when Cost->foldTailByMasking() is probably the best that can be done. The former is redundant given that it applies to dead code, and the latter should perhaps apply to all cases, in general.

4089

There's also the special case of requiresScalarEpiloque() where 0 < PEIterCount <= VF*UF for each invocation of the loop, and hence the average is also strictly positive FWIW. But best keep the approximation general instead of trying to improve it, given general lack of information.

4096

This assumes the number of times the vector loop will be reached is equal to the number of times the original scalar loop was reached (OrigFallThrougCount). This holds is Cost->foldTailByMasking(), but otherwise invocations whose trip count < VF*UF will bypass the vector loop (and also == VF*UF if requireScalarEpilogue()), plus other run time guards.

4105

Similar to above comment, invocations whose trip count divides VF*UF will bypass the scalar remainder loop (w/o foldTailByMasking nor requireScalarEpilogue), so in general PEFallThroughCount <= OrigFallThroughCount.

llvm/test/Transforms/LoopVectorize/check-prof-info.ll

4

May want to also check with UF>1.

[ebrevnov]

Addressed Ayal's comments

[ebrevnov]

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4062

Good catch. Thanks!

4074

I fixed names. But I don't see reasons to use different variables here (if this is what you meant)

4077

fixed

4079

Ok. Turned your explanation to a comment.

4089

Agree. Let me remove this special case then.

4096

Please note that VecFallThrough is zero initially and set to OrigFallThrougCount only if vector loop is expected to be executed (VecIterCount > 0)

4105

Same explanation as for the above.

llvm/test/Transforms/LoopVectorize/check-prof-info.ll

4

I replaced masked case since we don't do anything special for it now.

[ebrevnov]

ping @Ayal

[Ayal]

Still think it would be better to provide this as a standalone function in Transforms/Utils/LoopUtils, for potential benefit of loop unroll (and jam) passes in addition to LV. Having agreed to ignore foldTail and requiresScalarEpilog, there's nothing vectorization-specific to do here. There's still an issue though with the fact that LV may use the scalar loop for both the remaining TC%(VF*UF) iterations when running the vector loop, and for all TC iterations when runtime guards bypass the vector loop. In absence of information, each such guard could be assigned 0.5 probability, or one could be aggressively optimistic and hope vector loop is always reached. In any case this deserves a comment.

Suggesting further variable name changes for the three Orig, Unrolled, and Remainder Loops, each having a LoopEntry==LoopExit edge weight, a Backedge weight, a HeaderBlock weight, and an AverageTripCount. The actual weights are recorded as TrueVal and FalseVal of the latch branches.

Patch needs to be clang-format'ed.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4051

"is less ... than original TC" >> "is smaller than the original TC by a factor of VFxUF"

4058

OrigTakenWeight >> OrigBackedgeTakenWeight or OrigBackedgeWeight ?
OrigExitWeight >> OrigLoopExitWeight?
May help to also set OrigLoopEntryWeight = OrigLoopExitWeight?

4061

OrigBackBranchI >> OrigLoopLatchBranch ?

4064

It seems clearer to call extractProfMetadata(TrueVal, FalseVal) and then set BackedgeTaken and LoopExit/Entry weights according to if (OrigLoopLatchBranch->getSuccessor(0) == OrigLoop->getHeader()), following LoopUtil's getLoopEstimatedTripCount(). Analogously for createBranchWeights(TrueVal, FalseVal).

In any case, better rename "IsTrueBackEdge*Loop".

4078

Patch needs to be clang-format'ed

4085

Instead of providing the explanation in a comment, seems better to implement the code this way, leaving the +1 for the compiler to optimize. I.e.,

const uint64_t OrigHeaderBlockWeight = OrigBackedgeTakenWeight + OrigLoopEntryWeight;
const unit64_t OrigAverageTripCount = OrigHeaderBlockWeight / OrigLoopEntryWeight;

4092

Better rename/expand "PE".
PEIterCount >> RemainderLoopAverageTripCount?

4096

In the general context, "Vec" >> "UnrolledLoop".
VecTakenCount >> UnrolledLoopBackedgeWeight
VecFallThrough >> UnrolledLoopExitWeight and/or UnrolledLoopEntryWeight

4098

How about

if (UnrolledLoopAverageTripCount > 0) {
  UnrolledLoopEntryWeight = OrigLoopEntryWeight;
  uint64_t UnrolledLoopHeaderWeight = UnrolledLoopAverageTripCount * UnrolledLoopEntryWeight; // Analogous to computing OrigLoopAverageTripCount from Header and Entry weights above.
  UnrolledLoopBackedgeWeight = UnrolledLoopHeaderWeight - UnrolledLoopEntryWeight;
}

leaving the -1 optimization to the compiler.

4104

PETakenCount >> RemainderLoopBackedgeWeight
PEFallThroughCount >> RemainderLoopExitWeight and/or RemainderLoopEntryWeight

[ebrevnov]

In D67905#1770563, @Ayal wrote:

Still think it would be better to provide this as a standalone function in Transforms/Utils/LoopUtils, for potential benefit of loop unroll (and jam) passes in addition to LV. Having agreed to ignore foldTail and requiresScalarEpilog, there's nothing vectorization-specific to do here. There's still an issue though with the fact that LV may use the scalar loop for both the remaining TC%(VF*UF) iterations when running the vector loop, and for all TC iterations when runtime guards bypass the vector loop. In absence of information, each such guard could be assigned 0.5 probability, or one could be aggressively optimistic and hope vector loop is always reached. In any case this deserves a comment.

Suggesting further variable name changes for the three Orig, Unrolled, and Remainder Loops, each having a LoopEntry==LoopExit edge weight, a Backedge weight, a HeaderBlock weight, and an AverageTripCount. The actual weights are recorded as TrueVal and FalseVal of the latch branches.

Patch needs to be clang-format'ed.

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

4064

Don't feel convinced. My point would be that extra variables and conditional reassignments make the code less readable. I think this is very subjective thing.

4098

That will make computations less stable to overflow. Personally I feel the way it's written today has the same level of complexity for understanding.

[ebrevnov]

Addressing issues raised by Ayal.

[ebrevnov]

Typo fixed.

[ebrevnov]

ping @Ayal

[Ayal]

llvm/include/llvm/Transforms/Utils/LoopUtils.h

373

The fact that OrigLoop is both the original loop containing the original profile weights, and acts as the RemainderLoop dedicated to leftover iterations, should be clarified. Alternatively, this utility can receive three loops: OrigLoop, UnrolledLoop and RemainderLoop, leaving it to the caller to decide if to pass OrigLoop also as RemainderLoop.

Would probably be clearer to start with UnrolledLoop receiving weights that reflect TC/UF iterations, and then OrigLoop which receives weights that reflect the remaining TC%UF iterations.

375

U[n]rolledLoop, several occurrences.

"\UF" >> "\p UF"

llvm/lib/Transforms/Utils/LoopUtils.cpp

1098

Worth commenting that OrigLoopEntryWeight also holds OrigLoopExitWeight, which is more clearly the weight associated with the (exit direction of the) latch branch.

1106

UnrolledBBI >> Unrolled[Loop]LatchBranch, as in OrigLoopLatchBranch.

As the names end up overflowing lines, can use Orig, Unrolled and Remainder to stand for OrigLoop, UnrolledLoop and RemainderLoop; i.e., taking "Loop" out.

1109

VecLoop >> UnrolledLoop

1119

Can drop the 'const', for consistency; these temporaries are obviously const's.

1122

Note that this is rounding down. Can add half of the denominator to the nominator before dividing in order to round more accurately; this is what getLoopEstimatedTripCount() does, but it seems to be off by 1 as it computes BackEdgeTakenWeight / LoopEntryWeight rounded to nearest, instead of HeaderBlockWeight / LoopEntryWeight rounded to nearest...

Simply call OrigAverageTripCount = getLoopEstimatedTripCount(OrigLoop)?

Perhaps having a setLoopEstimatedTripCount(Loop, EstimatedTripCount, EstimatedEntryWeight) would help fold the identical treatment of UnrolledLoop and RemainderLoop into one function, which also takes care of figuring out the True/False vs. Backedge/Exit directions?

1124

U[n]rolledAverageTripCount

1134

Seems slightly more logical to first set UnrolledLoopEntryWeight, and then using it set UnrolledLoopBackedgeWeight.

1143

ditto

1156

(This actually replaces the old profile metadata with the new one.)

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

118

Is this include still needed here?

3484

Comment below should start with a short sentence explaining that profile weights associated with the original loop are now distributed among the vector and scalar loops.

[ebrevnov]

llvm/lib/Transforms/Utils/LoopUtils.cpp

1098

IMHO instead of trying to clarify with a comment we better find self descriptive name for such a simple and commonly used thing. Strictly speaking OrigLoopEntryWeight != OrigLoopExitWeight. Do you find OrigBackEdgeExitWeight good enough?

1106

Removed "Loop" from most names to make them a little shorter.

1122

This "off by 1" stopped me from using it in the first place since that could be important in some cases. OK, let's reuse getLoopEstimatedTripCount. To be able to do that there are some changes to getLoopEstimatedTripCount.

[ebrevnov]

Updated as requested.

[ebrevnov]

Hi @Ayal. Thanks for you input. I fixed all places as you suggested. Please check.

[Ayal]

Thanks for making all the changes! More comments inline.

llvm/lib/Transforms/Utils/LoopUtils.cpp

691–692

Comment what this new function is for. Rename (see below)? Retain "Support loops ..." comment, added in D64553?

728

dyn_cast >> cast

Perhaps update above function to do here something like
BranchInst *LatchBR = getExpectedExitLoopLatchBranch(L)
checking if it returned nullptr or not?

746

Thanks for taking care of this fix to improve accuracy of estimated TC!
But doing so deserves a separate patch, and tests. Note that it also effects loop unrolling, i.e., its effects are beyond LV. This part can be introduced either before or after the part that teaches LV to maintain profiling info.

767

ditto (dyn_cast >> cast, ...)

768

Better check similar to above if (LatchBR->getSuccessor(0) != L->getHeader())

1094

"the \p UnrolledLoop \p RemainderLoop" >>
"\p UnrolledLoop and \p RemainderLoop"

1098

May also be worthwhile asserting that UF is positive (or greater than 1?)

1098

IMHO instead of trying to clarify with a comment we better find self descriptive name for such a simple and commonly used thing.

Definitely agree with (the preference of) finding self descriptive variable names.

Strictly speaking OrigLoopEntryWeight != OrigLoopExitWeight"

How so, given that OrigLoop has a single-entry, and an "expected" single-exit: there may be other "side/deopt" exits, but these are expected to have zero weight. E.g., when computing LoopHeaderWeight above, adding LoopEntryWeight (instead of LoopExitWeight) to BackEdgeTakenWeight seems more logical.

Do you find OrigBackEdgeExitWeight good enough?

Strictly speaking, a BackEdge cannot exit: it is an edge going from a latch block to a header block, and "BackEdgeTaken" is the number of times this edge is taken=traversed, which equals the number of times its latch branch is "taken" rather than "falls-thru" (if it's true direction points to the header). Perhaps LoopEntryExitWeight or LoopInvocationWeight could be used instead/in addition.

1099

are expected to be distinct

1122

This "off by 1" stopped me from using it in the first place since that could be important in some cases. OK, let's reuse getLoopEstimatedTripCount. To be able to do that there are some changes to getLoopEstimatedTripCount.

Indeed, best fix and reuse, in a dedicated patch as raised above, thereby isolating impact on such "some cases".

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

3493

"is not taken into account as unlikely case" >>
"is ignored, assigning all the weight to the vector loop, optimistically"

llvm/test/Transforms/LoopVectorize/check-prof-info.ll

7

Tests targeting x86 need to reside in LoopVectorize/X86

[ebrevnov]

One more round of updates.

[ebrevnov]

llvm/lib/Transforms/Utils/LoopUtils.cpp

691–692

The comment is related to get/setLoopEstimatedTripCount and still there...

728

I was thinking about that in the first place but didn't come up with a "good" enough name. I can go with that name if you like :-)

746

Ok.

1098

I think we better support 1 which could be used in some corner cases....

1098

How so, given that OrigLoop has a single-entry, and an "expected" single-exit: there may be other "side/deopt" exits, but these are expected to have zero weight. E.g., when computing LoopHeaderWeight above, adding LoopEntryWeight (instead of LoopExitWeight) to BackEdgeTakenWeight seems more logical.

For infinite loop entry is 1 while exit is 0 :-). I understand this is extreme we will never meet but still....

Strictly speaking, a BackEdge cannot exit: it is an edge going from a latch block to a header block, and "BackEdgeTaken" is the number of times this edge is taken=traversed, which equals the number of times its latch branch is "taken" rather than "falls-thru" (if it's true direction points to the header).

That's why I used TakenCount and FallThroughCount in the very first version what perfectly matches your description. I don't feel we are getting any better names with more iterations.... Perhaps LatchCycleWeight - number of times we go to loop header from the latch and LatchExitWeight - number of times we go to loop exit from the latch?

Perhaps LoopEntryExitWeight or LoopInvocationWeight could be used instead/in addition.

I think LoopEntryExitWeight may be confusing....
I think it makes sense to use EstimatedLoopInvocationWeight in conjunction with EstimatedTripCount as parameters to get/setEstimatedTripCount interface while LatchCycleWeight and LatchExitWeight in the implementation as they are little-bit more low level.

[ebrevnov]

Rebase

[ebrevnov]

Rebase

[ebrevnov]

ping @Ayal

[Ayal]

This looks good to me, thanks!

llvm/lib/Transforms/Utils/LoopUtils.cpp

775

The last part could call fixupBranchWeights() if moved here from Transforms/Utils/LoopUnrollPeel.cpp

llvm/test/Transforms/LoopVectorize/tripcount.ll

211

Following this, to clarify:
original loop has latchExitWeight=10 and backedgeTakenWeight=10,000, therefore estimatedBackedgeTakenCount=1,000 and estimatedTripCount=1,001.
Vectorizing by 4 produces
estimatedTripCounts of 1,001/4=250 and 1,001%4=1 for vectorized and remainder loops, respectively, therefore their
estimatedBackedgeTakenCounts are 249 and 0, and so the weights recorded with loop invocation weights of 10 are the above {10, 2490} and {10, 0}.

[ebrevnov]

llvm/lib/Transforms/Utils/LoopUtils.cpp

775

That would require to change implementation for fixupBranchWeights since it disregards to update when back edge taken count is zero.

llvm/test/Transforms/LoopVectorize/tripcount.ll

211

I will add this text to the test. I that what you wanted (just not sure :-))?

[ebrevnov]

Rebase

[Ayal]

llvm/lib/Transforms/Utils/LoopUtils.cpp

775

Agreed. Regarding disregarding to update when backedge taken weight is zero, note that in fixupBranchWeights(), BackedgeTakenWeight is called FallThroughWeight(?) "The weight of the edge from Latch to Header", and that

// FallThroughWeight is 0 means that there is no branch weights on original
// latch block or estimated trip count is zero.

Regarding the first meaning of 0, whoever calls fixupBranchWeights() should do so only if there were such weights on the original latch block, similar to the caller of setLoopEstimatedTripCount(). Regarding the second meaning of 0, it seem fixupBranchWeights() suffers from same +1 issue: estimating trip count to be zero when backedge taken weight is zero.

It would be good to fix and centralize the support for updating weights of loops, but such refactoring can be done as a separate follow-up patch, after landing this (accepted) patch.

llvm/include/llvm/Transforms/Utils/LoopUtils.h

 /// @}
 
 /// Set input string into loop metadata by keeping other values intact.
 /// If the string is already in loop metadata update value if it is
 /// different.
 void addStringMetadataToLoop(Loop *TheLoop, const char *MDString,
                              unsigned V = 0);
 
-/// Get a loop's estimated trip count based on branch weight metadata.
+/// Returns a loop's estimated trip count based on branch weight metadata.
+/// In addition if \p EstimatedLoopInvocationWeight is not null it is
+/// initialized with weight of loop's latch leading to the exit.
 /// Returns 0 when the count is estimated to be 0, or None when a meaningful
 /// estimate can not be made.
-Optional<unsigned> getLoopEstimatedTripCount(Loop *L);
+Optional<unsigned>
+getLoopEstimatedTripCount(Loop *L,
+                          unsigned *EstimatedLoopInvocationWeight = nullptr);
+
+/// Set a loop's branch weight metadata to reflect that loop has \p
+/// EstimatedTripCount iterations and \p EstimatedLoopInvocationWeight exits
+/// through latch. Returns true if metadata is successfully updated, false
+/// otherwise. Note that loop must have a latch block which controls loop exit
+/// in order to succeed.
+bool setLoopEstimatedTripCount(Loop *L, unsigned EstimatedTripCount,
+                               unsigned EstimatedLoopInvocationWeight);
 
 /// Check inner loop (L) backedge count is known to be invariant on all
 /// iterations of its outer loop. If the loop has no parent, this is trivially
 /// true.
 bool hasIterationCountInvariantInParent(Loop *L, ScalarEvolution &SE);
 
 /// Helper to consistently add the set of standard passes to a loop pass's \c
 /// AnalysisUsage.

 /// Returns true if \p S is defined and never is equal to signed/unsigned max.
 bool cannotBeMaxInLoop(const SCEV *S, const Loop *L, ScalarEvolution &SE,
                        bool Signed);
 
 /// Returns true if \p S is defined and never is equal to signed/unsigned min.
 bool cannotBeMinInLoop(const SCEV *S, const Loop *L, ScalarEvolution &SE,
                        bool Signed);
 
+/// Set weights for \p UnrolledLoop and \p RemainderLoop based on weights for
+/// \p OrigLoop and the following distribution of \p OrigLoop iteration among \p
+/// UnrolledLoop and \p RemainderLoop. \p UnrolledLoop receives weights that
+/// reflect TC/UF iterations, and \p RemainderLoop receives weights that reflect
+/// the remaining TC%UF iterations.
+///
+/// Note that \p OrigLoop may be equal to either \p UnrolledLoop or \p
+/// RemainderLoop in which case weights for \p OrigLoop are updated accordingly.
+/// Note also behavior is undefined if \p UnrolledLoop and \p RemainderLoop are
+/// equal. \p UF must be greater than zero.
+/// If \p OrigLoop has no profile info associated nothing happens.
+///
+/// This utility may be useful for such optimizations as unroller and
+/// vectorizer as it's typical transformation for them.
+void setProfileInfoAfterUnrolling(Loop *OrigLoop, Loop *UnrolledLoop,
+                                  Loop *RemainderLoop, uint64_t UF);
+
 } // end namespace llvm
 
 #endif // LLVM_TRANSFORMS_UTILS_LOOPUTILS_H

llvm/lib/Transforms/Utils/LoopUtils.cpp

 #include "llvm/Analysis/ScalarEvolutionAliasAnalysis.h"
 #include "llvm/Analysis/ScalarEvolutionExpressions.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/DIBuilder.h"
 #include "llvm/IR/Dominators.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/MDBuilder.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/PatternMatch.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/InitializePasses.h"
 #include "llvm/Pass.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/KnownBits.h"
 #include "llvm/Transforms/Utils/BasicBlockUtils.h"
if (LI) {
     } else {
       Loop::iterator I = find(LI->begin(), LI->end(), L);
       assert(I != LI->end() && "Couldn't find loop");
       LI->removeLoop(I);
     }
     LI->destroy(L);
   }
 }
 
-Optional<unsigned> llvm::getLoopEstimatedTripCount(Loop *L) {
-  // Support loops with an exiting latch and other existing exists only
-  // deoptimize.
-
-  // Get the branch weights for the loop's backedge.
+/// Checks if \p L has single exit through latch block except possibly
+/// "deoptimizing" exits. Returns branch instruction terminating the loop
+/// latch if above check is successful, nullptr otherwise.
+static BranchInst *getExpectedExitLoopLatchBranch(Loop *L) {
   BasicBlock *Latch = L->getLoopLatch();
   if (!Latch)
-    return None;
+    return nullptr;
+
   BranchInst *LatchBR = dyn_cast<BranchInst>(Latch->getTerminator());
   if (!LatchBR || LatchBR->getNumSuccessors() != 2 || !L->isLoopExiting(Latch))
-    return None;
+    return nullptr;
 
   assert((LatchBR->getSuccessor(0) == L->getHeader() ||
           LatchBR->getSuccessor(1) == L->getHeader()) &&
          "At least one edge out of the latch must go to the header");
 
   SmallVector<BasicBlock *, 4> ExitBlocks;
   L->getUniqueNonLatchExitBlocks(ExitBlocks);
   if (any_of(ExitBlocks, [](const BasicBlock *EB) {
         return !EB->getTerminatingDeoptimizeCall();
       }))
+    return nullptr;
+
+  return LatchBR;
+}
+
+Optional<unsigned>
+llvm::getLoopEstimatedTripCount(Loop *L,
+                                unsigned *EstimatedLoopInvocationWeight) {
+  // Support loops with an exiting latch and other existing exists only
+  // deoptimize.
+  BranchInst *LatchBranch = getExpectedExitLoopLatchBranch(L);
+  if (!LatchBranch)
     return None;
 
   // To estimate the number of times the loop body was executed, we want to
   // know the number of times the backedge was taken, vs. the number of times
   // we exited the loop.
   uint64_t BackedgeTakenWeight, LatchExitWeight;
-  if (!LatchBR->extractProfMetadata(BackedgeTakenWeight, LatchExitWeight))
+  if (!LatchBranch->extractProfMetadata(BackedgeTakenWeight, LatchExitWeight))
     return None;
 
-  if (LatchBR->getSuccessor(0) != L->getHeader())
+  if (LatchBranch->getSuccessor(0) != L->getHeader())
     std::swap(BackedgeTakenWeight, LatchExitWeight);
 
+  if (EstimatedLoopInvocationWeight)
+    *EstimatedLoopInvocationWeight = LatchExitWeight;
+
   if (!BackedgeTakenWeight || !LatchExitWeight)
     return 0;
 
   // Trip count is a ratio of loop header block execution weight to loop exit
   // weight.
   uint64_t HeaderBlockWeight = BackedgeTakenWeight + LatchExitWeight;
   return llvm::divideNearest(HeaderBlockWeight, LatchExitWeight);
 }
 
+bool llvm::setLoopEstimatedTripCount(Loop *L, unsigned EstimatedTripCount,
+                                     unsigned EstimatedloopInvocationWeight) {
+  // Support loops with an exiting latch and other existing exists only
+  // deoptimize.
+  BranchInst *LatchBranch = getExpectedExitLoopLatchBranch(L);
+  if (!LatchBranch)
+    return false;
+
+  // Calculate taken and exit weights.
+  unsigned LatchExitWeight = 0;
+  unsigned BackedgeTakenWeight = 0;
+
+  if (EstimatedTripCount > 0) {
+    LatchExitWeight = EstimatedloopInvocationWeight;
+    BackedgeTakenWeight = (EstimatedTripCount - 1) * LatchExitWeight;
+  }
+
+  // Make a swap if back edge is taken when condition is "false".
+  if (LatchBranch->getSuccessor(0) != L->getHeader())
+    std::swap(BackedgeTakenWeight, LatchExitWeight);
+
+  MDBuilder MDB(LatchBranch->getContext());
+
+  // Set/Update profile metadata.
+  LatchBranch->setMetadata(
+      LLVMContext::MD_prof,
+      MDB.createBranchWeights(BackedgeTakenWeight, LatchExitWeight));
+
+  return true;
+}
+
 bool llvm::hasIterationCountInvariantInParent(Loop *InnerLoop,
                                               ScalarEvolution &SE) {
   Loop *OuterL = InnerLoop->getParentLoop();
   if (!OuterL)
     return true;
 
   // Get the backedge taken count for the inner loop
   BasicBlock *InnerLoopLatch = InnerLoop->getLoopLatch();
bool llvm::cannotBeMaxInLoop(const SCEV *S, const Loop *L, ScalarEvolution &SE,
   unsigned BitWidth = cast<IntegerType>(S->getType())->getBitWidth();
   APInt Max = Signed ? APInt::getSignedMaxValue(BitWidth) :
     APInt::getMaxValue(BitWidth);
   auto Predicate = Signed ? ICmpInst::ICMP_SLT : ICmpInst::ICMP_ULT;
   return SE.isAvailableAtLoopEntry(S, L) &&
          SE.isLoopEntryGuardedByCond(L, Predicate, S,
                                      SE.getConstant(Max));
 }
+
+/// Set weights for \p UnrolledLoop and \p RemainderLoop based on weights for
+/// \p OrigLoop.
+void llvm::setProfileInfoAfterUnrolling(Loop *OrigLoop, Loop *UnrolledLoop,
+                                        Loop *RemainderLoop, uint64_t UF) {
+  assert(UF > 0 && "Zero unrolled factor is not supported");
+  assert(UnrolledLoop != RemainderLoop &&
+         "Unrolled and Remainder loops are expected to distinct");
+
+  // Get number of iterations in the original scalar loop.
+  unsigned OrigLoopInvocationWeight = 0;
+  Optional<unsigned> OrigAverageTripCount =
+      getLoopEstimatedTripCount(OrigLoop, &OrigLoopInvocationWeight);
+  if (!OrigAverageTripCount)
+    return;
+
+  // Calculate number of iterations in unrolled loop.
+  unsigned UnrolledAverageTripCount = *OrigAverageTripCount / UF;
+  // Calculate number of iterations for remainder loop.
+  unsigned RemainderAverageTripCount = *OrigAverageTripCount % UF;
+
+  setLoopEstimatedTripCount(UnrolledLoop, UnrolledAverageTripCount,
+                            OrigLoopInvocationWeight);
+  setLoopEstimatedTripCount(RemainderLoop, RemainderAverageTripCount,
+                            OrigLoopInvocationWeight);
+}

llvm/lib/Transforms/Vectorize/LoopVectorize.cpp

 #include "llvm/IR/Function.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/InstrTypes.h"
 #include "llvm/IR/Instruction.h"
 #include "llvm/IR/Instructions.h"
 #include "llvm/IR/IntrinsicInst.h"
 #include "llvm/IR/Intrinsics.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Metadata.h"
 #include "llvm/IR/Module.h"
 #include "llvm/IR/Operator.h"
 #include "llvm/IR/Type.h"
 #include "llvm/IR/Use.h"
 #include "llvm/IR/User.h"
 #include "llvm/IR/Value.h"
 #include "llvm/IR/ValueHandle.h"
 #include "llvm/IR/Verifier.h"
fixupIVUsers(Entry.first, Entry.second,
                  IVEndValues[Entry.first], LoopMiddleBlock);
 
   fixLCSSAPHIs();
   for (Instruction *PI : PredicatedInstructions)
     sinkScalarOperands(&*PI);
 
   // Remove redundant induction instructions.
   cse(LoopVectorBody);
+
+  // Set/update profile weights for the vector and remainder loops as original
+  // loop iterations are now distributed among them. Note that original loop
+  // represented by LoopScalarBody becomes remainder loop after vectorization.
+  //
+  // For cases like foldTailByMasking() and requiresScalarEpiloque() we may
+  // end up getting slightly roughened result but that should be OK since
+  // profile is not inherently precise anyway. Note also possible bypass of
+  // vector code caused by legality checks is ignored, assigning all the weight
+  // to the vector loop, optimistically.
+  setProfileInfoAfterUnrolling(LI->getLoopFor(LoopScalarBody),
+                               LI->getLoopFor(LoopVectorBody),
+                               LI->getLoopFor(LoopScalarBody), VF * UF);
 }
 
 void InnerLoopVectorizer::fixCrossIterationPHIs() {
   // In order to support recurrences we need to be able to vectorize Phi nodes.
   // Phi nodes have cycles, so we need to vectorize them in two stages. This is
   // stage #2: We now need to fix the recurrences by adding incoming edges to
   // the currently empty PHI nodes. At this point every instruction in the
   // original loop is widened to a vector form so we can use them to construct
for (unsigned i = 0; i < NumIncomingValues; ++i) {
 
       // Scalar incoming value may need a broadcast
       Value *NewIncV = getOrCreateVectorValue(ScIncV, 0);
       NewPhi->addIncoming(NewIncV, NewPredBB);
     }
   }
 }
 
 void InnerLoopVectorizer::widenGEP(GetElementPtrInst *GEP, unsigned UF,
                                    unsigned VF, bool IsPtrLoopInvariant,
                                    SmallBitVector &IsIndexLoopInvariant) {
   // Construct a vector GEP by widening the operands of the scalar GEP as
   // necessary. We mark the vector GEP 'inbounds' if appropriate. A GEP
   // results in a vector of pointers when at least one operand of the GEP
   // is vector-typed. Thus, to keep the representation compact, we only use
   // vector-typed operands for loop-varying values.
 
   if (VF > 1 && IsPtrLoopInvariant && IsIndexLoopInvariant.all()) {
     // If we are vectorizing, but the GEP has only loop-invariant operands,
     // the GEP we build (by only using vector-typed operands for
     // loop-varying values) would be a scalar pointer. Thus, to ensure we
     // produce a vector of pointers, we need to either arbitrarily pick an
     // operand to broadcast, or broadcast a clone of the original GEP.
     // Here, we broadcast a clone of the original.
     //
     // TODO: If at some point we decide to scalarize instructions having
     //       loop-invariant operands, this special case will no longer be
     //       required. We would add the scalarization decision to
     //       collectLoopScalars() and teach getVectorValue() to broadcast
     //       the lane-zero scalar value.
     auto *Clone = Builder.Insert(GEP->clone());
     for (unsigned Part = 0; Part < UF; ++Part) {
       Value *EntryPart = Builder.CreateVectorSplat(VF, Clone);
       VectorLoopValueMap.setVectorValue(GEP, Part, EntryPart);
       addMetadata(EntryPart, GEP);
     }
   } else {
     // If the GEP has at least one loop-varying operand, we are sure to
     // produce a vector of pointers. But if we are only unrolling, we want
     // to produce a scalar GEP for each unroll part. Thus, the GEP we
     // produce with the code below will be scalar (if VF == 1) or vector
     // (otherwise). Note that for the unroll-only case, we still maintain
     // values in the vector mapping with initVector, as we do for other
     // instructions.
     for (unsigned Part = 0; Part < UF; ++Part) {
       // The pointer operand of the new GEP. If it's loop-invariant, we
       // won't broadcast it.
       auto *Ptr = IsPtrLoopInvariant
                       ? GEP->getPointerOperand()
                       : getOrCreateVectorValue(GEP->getPointerOperand(), Part);
 
       // Collect all the indices for the new GEP. If any index is
       // loop-invariant, we won't broadcast it.
       SmallVector<Value *, 4> Indices;
       for (auto Index : enumerate(GEP->indices())) {
         Value *User = Index.value().get();
         if (IsIndexLoopInvariant[Index.index()])
           Indices.push_back(User);
         else
           Indices.push_back(getOrCreateVectorValue(User, Part));
       }
 
       // Create the new GEP. Note that this GEP may be a scalar if VF == 1,
       // but it should be a vector, otherwise.
       auto *NewGEP =
           GEP->isInBounds()
               ? Builder.CreateInBoundsGEP(GEP->getSourceElementType(), Ptr,
                                           Indices)
               : Builder.CreateGEP(GEP->getSourceElementType(), Ptr, Indices);
       assert((VF == 1 || NewGEP->getType()->isVectorTy()) &&
              "NewGEP is not a pointer vector");

llvm/test/Transforms/LoopVectorize/check-prof-info.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt  -passes="print<block-freq>,loop-vectorize" -force-vector-width=4 -force-vector-interleave=1 -S < %s |  FileCheck %s
+; RUN: opt  -passes="print<block-freq>,loop-vectorize" -force-vector-width=4 -force-vector-interleave=4 -S < %s |  FileCheck %s -check-prefix=CHECK-MASKED
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+@a = dso_local global [1024 x i32] zeroinitializer, align 16
+@b = dso_local global [1024 x i32] zeroinitializer, align 16
+
+; Check correctness of profile info for vectorization without epilog.
+; Function Attrs: nofree norecurse nounwind uwtable
+define dso_local void @_Z3foov() local_unnamed_addr #0 {
+; CHECK-LABEL: @_Z3foov(
+; CHECK:  [[VECTOR_BODY:vector\.body]]:
+; CHECK:    br i1 [[TMP:%.*]], label [[MIDDLE_BLOCK:%.*]], label %[[VECTOR_BODY]], !prof [[LP1_255:\!.*]],
+; CHECK:  [[FOR_BODY:for\.body]]:
+; CHECK:    br i1 [[EXITCOND:%.*]], label [[FOR_END_LOOPEXIT:%.*]], label %[[FOR_BODY]], !prof [[LP0_0:\!.*]],
+; CHECK-MASKED:  [[VECTOR_BODY:vector\.body]]:
+; CHECK-MASKED:    br i1 [[TMP:%.*]], label [[MIDDLE_BLOCK:%.*]], label %[[VECTOR_BODY]], !prof [[LP1_63:\!.*]],
+; CHECK-MASKED:  [[FOR_BODY:for\.body]]:
+; CHECK-MASKED:    br i1 [[EXITCOND:%.*]], label [[FOR_END_LOOPEXIT:%.*]], label %[[FOR_BODY]], !prof [[LP0_0:\!.*]],
+;
+entry:
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+
+for.body:                                         ; preds = %for.body, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds [1024 x i32], [1024 x i32]* @b, i64 0, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4, !tbaa !2
+  %1 = trunc i64 %indvars.iv to i32
+  %mul = mul nsw i32 %0, %1
+  %arrayidx2 = getelementptr inbounds [1024 x i32], [1024 x i32]* @a, i64 0, i64 %indvars.iv
+  %2 = load i32, i32* %arrayidx2, align 4, !tbaa !2
+  %add = add nsw i32 %2, %mul
+  store i32 %add, i32* %arrayidx2, align 4, !tbaa !2
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv.next, 1024
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body, !prof !6
+}
+
+; Check correctness of profile info for vectorization with epilog.
+; Function Attrs: nofree norecurse nounwind uwtable
+define dso_local void @_Z3foo2v() local_unnamed_addr #0 {
+; CHECK-LABEL: @_Z3foo2v(
+; CHECK:  [[VECTOR_BODY:vector\.body]]:
+; CHECK:    br i1 [[TMP:%.*]], label [[MIDDLE_BLOCK:%.*]], label %[[VECTOR_BODY]], !prof [[LP1_255:\!.*]],
+; CHECK:  [[FOR_BODY:for\.body]]:
+; CHECK:    br i1 [[EXITCOND:%.*]], label [[FOR_END_LOOPEXIT:%.*]], label %[[FOR_BODY]], !prof [[LP1_2:\!.*]],
+; CHECK-MASKED:  [[VECTOR_BODY:vector\.body]]:
+; CHECK-MASKED:    br i1 [[TMP:%.*]], label [[MIDDLE_BLOCK:%.*]], label %[[VECTOR_BODY]], !prof [[LP1_63:\!.*]],
+; CHECK-MASKED:  [[FOR_BODY:for\.body]]:
+; CHECK-MASKED:    br i1 [[EXITCOND:%.*]], label [[FOR_END_LOOPEXIT:%.*]], label %[[FOR_BODY]], !prof [[LP1_2:\!.*]],
+;
+entry:
+  br label %for.body
+
+for.cond.cleanup:                                 ; preds = %for.body
+  ret void
+
+for.body:                                         ; preds = %for.body, %entry
+  %indvars.iv = phi i64 [ 0, %entry ], [ %indvars.iv.next, %for.body ]
+  %arrayidx = getelementptr inbounds [1024 x i32], [1024 x i32]* @b, i64 0, i64 %indvars.iv
+  %0 = load i32, i32* %arrayidx, align 4, !tbaa !2
+  %1 = trunc i64 %indvars.iv to i32
+  %mul = mul nsw i32 %0, %1
+  %arrayidx2 = getelementptr inbounds [1024 x i32], [1024 x i32]* @a, i64 0, i64 %indvars.iv
+  %2 = load i32, i32* %arrayidx2, align 4, !tbaa !2
+  %add = add nsw i32 %2, %mul
+  store i32 %add, i32* %arrayidx2, align 4, !tbaa !2
+  %indvars.iv.next = add nuw nsw i64 %indvars.iv, 1
+  %exitcond = icmp eq i64 %indvars.iv.next, 1027
+  br i1 %exitcond, label %for.cond.cleanup, label %for.body, !prof !7
+}
+
+attributes #0 = { "use-soft-float"="false" }
+
+!llvm.module.flags = !{!0}
+!llvm.ident = !{!1}
+
+; CHECK: [[LP1_255]] = !{!"branch_weights", i32 1, i32 255}
+; CHECK: [[LP0_0]] = !{!"branch_weights", i32 0, i32 0}
+; CHECK-MASKED: [[LP1_63]] = !{!"branch_weights", i32 1, i32 63}
+; CHECK-MASKED: [[LP0_0]] = !{!"branch_weights", i32 0, i32 0}
+; CHECK: [[LP1_2]] = !{!"branch_weights", i32 1, i32 2}
+
+!0 = !{i32 1, !"wchar_size", i32 4}
+!1 = !{!"clang version 10.0.0 (https://github.com/llvm/llvm-project c292b5b5e059e6ce3e6449e6827ef7e1037c21c4)"}
+!2 = !{!3, !3, i64 0}
+!3 = !{!"int", !4, i64 0}
+!4 = !{!"omnipotent char", !5, i64 0}
+!5 = !{!"Simple C++ TBAA"}
+!6 = !{!"branch_weights", i32 1, i32 1023}
+!7 = !{!"branch_weights", i32 1, i32 1026}

llvm/test/Transforms/LoopVectorize/tripcount.ll

for.end:                                          ; preds = %for.body
   ret i32 0
 }
 
 define i32 @foo_low_trip_count3(i1 %cond, i32 %bound) !prof !0 {
 ; The loop has low invocation count compare to the function invocation count,
 ; but has a high trip count per invocation. Vectorize it.
 
 ; CHECK-LABEL: @foo_low_trip_count3(
-; CHECK: vector.body:
-
+; CHECK:  [[VECTOR_BODY:vector\.body]]:
+; CHECK:    br i1 [[TMP9:%.*]], label [[MIDDLE_BLOCK:%.*]], label %[[VECTOR_BODY]], !prof [[LP3:\!.*]],
+; CHECK:  [[FOR_BODY:for\.body]]:
+; CHECK:    br i1 [[EXITCOND:%.*]], label [[FOR_END_LOOPEXIT:%.*]], label %[[FOR_BODY]], !prof [[LP6:\!.*]],
 entry:
   br i1 %cond, label %for.preheader, label %for.end, !prof !2
 
 for.preheader:
   br label %for.body
 
 for.body:                                         ; preds = %for.body, %entry
   %i.08 = phi i32 [ 0, %for.preheader ], [ %inc, %for.body ]
for.body:                                         ; preds = %for.body, %entry
   %inc = add nsw i32 %i.08, 1
   %exitcond = icmp slt i32 %i.08, 1000
   br i1 %exitcond, label %for.body, label %for.end, !prof !1
 
 for.end:                                          ; preds = %for.body
   ret i32 0
 }
 
+; CHECK: [[LP3]] = !{!"branch_weights", i32 10, i32 2490}
+; CHECK: [[LP6]] = !{!"branch_weights", i32 10, i32 0}
+
 !0 = !{!"function_entry_count", i64 100}
 !1 = !{!"branch_weights", i32 100, i32 0}
 !2 = !{!"branch_weights", i32 10, i32 90}
 !3 = !{!"branch_weights", i32 10, i32 10000}