[llvm] [LoopRotate] Fix branch weights for rotated multi-exit loops (PR #202219)
Alok Kumar Sharma via llvm-commits
llvm-commits at lists.llvm.org
Mon Jun 8 05:44:01 PDT 2026
https://github.com/alokkrsharma updated https://github.com/llvm/llvm-project/pull/202219
>From 6cad892a211587c020bec7e670e0ff1049a65553 Mon Sep 17 00:00:00 2001
From: Alok Kumar Sharma <AlokKumar.Sharma at amd.com>
Date: Sun, 7 Jun 2026 22:02:08 +0530
Subject: [PATCH] [LoopRotate] Fix branch weights for rotated multi-exit loops
Loop rotation splits the header branch into a guard and latch.
updateBranchWeights assumed single-exit behavior, producing incorrect
weights for multi-exit loops.
Use BlockFrequencyInfo to recover loop entry count and derive correct
EnterWeight, scaling back to the original branch-weight form. Fall back
to the old heuristic when unavailable.
Apply only in the post-PGO path via a new flag. Also handle zero/never-
entered loops and drop invalid latch metadata.
---
.../llvm/Transforms/Scalar/LoopRotation.h | 4 +-
.../llvm/Transforms/Utils/LoopRotationUtils.h | 7 +-
llvm/lib/Passes/PassBuilder.cpp | 3 +
llvm/lib/Passes/PassBuilderPipelines.cpp | 9 +-
llvm/lib/Passes/PassRegistry.def | 5 +-
llvm/lib/Transforms/Scalar/LoopRotation.cpp | 29 +-
.../Transforms/Utils/LoopRotationUtils.cpp | 103 +++++-
llvm/test/Other/new-pm-print-pipeline.ll | 6 +-
...ded-preheader-multi-exit-branch-weights.ll | 74 ++++
.../LoopRotate/multi-exit-branch-weights.ll | 84 +++++
.../multi-exit-update-branch-weights.ll | 327 ++++++++++++++++++
.../never-entered-branch-weights.ll | 61 ++++
.../no-bfi-fallback-branch-weights.ll | 87 +++++
.../predecessor-no-prof-branch-weights.ll | 87 +++++
.../saturating-backedge-branch-weights.ll | 126 +++++++
.../scaled-preheader-branch-weights.ll | 63 ++++
.../LoopRotate/update-branch-weights.ll | 6 +-
17 files changed, 1053 insertions(+), 28 deletions(-)
create mode 100644 llvm/test/Transforms/LoopRotate/folded-preheader-multi-exit-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/multi-exit-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/multi-exit-update-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/never-entered-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/no-bfi-fallback-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/predecessor-no-prof-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/saturating-backedge-branch-weights.ll
create mode 100644 llvm/test/Transforms/LoopRotate/scaled-preheader-branch-weights.ll
diff --git a/llvm/include/llvm/Transforms/Scalar/LoopRotation.h b/llvm/include/llvm/Transforms/Scalar/LoopRotation.h
index bb17995116fec..3ab3f901ec5a4 100644
--- a/llvm/include/llvm/Transforms/Scalar/LoopRotation.h
+++ b/llvm/include/llvm/Transforms/Scalar/LoopRotation.h
@@ -25,7 +25,8 @@ class LoopRotatePass : public OptionalPassInfoMixin<LoopRotatePass> {
public:
LLVM_ABI LoopRotatePass(bool EnableHeaderDuplication = true,
bool PrepareForLTO = false,
- bool CheckExitCount = false);
+ bool CheckExitCount = false,
+ bool UpdateBranchWeights = false);
LLVM_ABI PreservedAnalyses run(Loop &L, LoopAnalysisManager &AM,
LoopStandardAnalysisResults &AR,
LPMUpdater &U);
@@ -38,6 +39,7 @@ class LoopRotatePass : public OptionalPassInfoMixin<LoopRotatePass> {
const bool EnableHeaderDuplication;
const bool PrepareForLTO;
const bool CheckExitCount;
+ const bool UpdateBranchWeights;
};
}
diff --git a/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h b/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h
index 18a1c4efcaab2..354e6cbdd90e0 100644
--- a/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h
+++ b/llvm/include/llvm/Transforms/Utils/LoopRotationUtils.h
@@ -18,6 +18,7 @@
namespace llvm {
class AssumptionCache;
+class BlockFrequencyInfo;
class DominatorTree;
class Loop;
class LoopInfo;
@@ -32,13 +33,17 @@ class TargetTransformInfo;
/// header. If the loop header's size exceeds the threshold, the loop rotation
/// will give up. The flag IsUtilMode controls the heuristic used in the
/// LoopRotation. If it is true, the profitability heuristic will be ignored.
+/// If \p BFI is provided, use it to recover loop entry counts when updating
+/// branch weights, ensuring correctness for multi-exit loops under both
+/// count-based and ratio-based profiles.
LLVM_ABI bool LoopRotation(Loop *L, LoopInfo *LI,
const TargetTransformInfo *TTI, AssumptionCache *AC,
DominatorTree *DT, ScalarEvolution *SE,
MemorySSAUpdater *MSSAU, const SimplifyQuery &SQ,
bool RotationOnly, unsigned Threshold,
bool IsUtilMode, bool PrepareForLTO = false,
- bool CheckExitCount = false);
+ bool CheckExitCount = false,
+ BlockFrequencyInfo *BFI = nullptr);
} // namespace llvm
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index f7db63ef8bf74..e0f6130b2c052 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1293,6 +1293,7 @@ struct LoopRotateOptions {
bool EnableHeaderDuplication = true;
bool PrepareForLTO = false;
bool CheckExitCount = false;
+ bool UpdateBranchWeights = false;
};
Expected<LoopRotateOptions> parseLoopRotateOptions(StringRef Params) {
@@ -1308,6 +1309,8 @@ Expected<LoopRotateOptions> parseLoopRotateOptions(StringRef Params) {
Result.PrepareForLTO = Enable;
} else if (ParamName == "check-exit-count") {
Result.CheckExitCount = Enable;
+ } else if (ParamName == "update-branch-weights") {
+ Result.UpdateBranchWeights = Enable;
} else {
return make_error<StringError>(
formatv("invalid LoopRotate pass parameter '{}'", ParamName).str(),
diff --git a/llvm/lib/Passes/PassBuilderPipelines.cpp b/llvm/lib/Passes/PassBuilderPipelines.cpp
index 9eea552fd263e..32b32dbd4e26e 100644
--- a/llvm/lib/Passes/PassBuilderPipelines.cpp
+++ b/llvm/lib/Passes/PassBuilderPipelines.cpp
@@ -883,9 +883,14 @@ void PassBuilder::addPostPGOLoopRotation(ModulePassManager &MPM,
OptimizationLevel Level) {
if (EnablePostPGOLoopRotation) {
// Disable header duplication in loop rotation at -Oz.
+ // Enable profile-based branch-weight updates so rotated multi-exit
+ // loops receive correct weights from the applied PGO profile.
MPM.addPass(createModuleToFunctionPassAdaptor(
- createFunctionToLoopPassAdaptor(LoopRotatePass(),
- /*UseMemorySSA=*/false),
+ createFunctionToLoopPassAdaptor(
+ LoopRotatePass(/*EnableHeaderDuplication=*/true,
+ /*PrepareForLTO=*/false, /*CheckExitCount=*/false,
+ /*UpdateBranchWeights=*/true),
+ /*UseMemorySSA=*/false),
PTO.EagerlyInvalidateAnalyses));
}
}
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 8f70c7cefa408..d22e9ac47c107 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -816,12 +816,13 @@ LOOP_PASS_WITH_PARAMS(
"loop-rotate", "LoopRotatePass",
[](LoopRotateOptions Params) {
return LoopRotatePass(Params.EnableHeaderDuplication, Params.PrepareForLTO,
- Params.CheckExitCount);
+ Params.CheckExitCount, Params.UpdateBranchWeights);
},
parseLoopRotateOptions,
"no-header-duplication;header-duplication;"
"no-prepare-for-lto;prepare-for-lto;"
- "no-check-exit-count;check-exit-count")
+ "no-check-exit-count;check-exit-count;"
+ "no-update-branch-weights;update-branch-weights")
LOOP_PASS_WITH_PARAMS(
"simple-loop-unswitch", "SimpleLoopUnswitchPass",
[](std::pair<bool, bool> Params) {
diff --git a/llvm/lib/Transforms/Scalar/LoopRotation.cpp b/llvm/lib/Transforms/Scalar/LoopRotation.cpp
index 2f2ce9e2782a1..379ce979f0855 100644
--- a/llvm/lib/Transforms/Scalar/LoopRotation.cpp
+++ b/llvm/lib/Transforms/Scalar/LoopRotation.cpp
@@ -12,6 +12,8 @@
#include "llvm/Transforms/Scalar/LoopRotation.h"
#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BranchProbabilityInfo.h"
#include "llvm/Analysis/InstructionSimplify.h"
#include "llvm/Analysis/LazyBlockFrequencyInfo.h"
#include "llvm/Analysis/LoopInfo.h"
@@ -48,9 +50,10 @@ static cl::opt<bool> EnableLoopHeaderDuplicationAtMinSize(
cl::desc("Enable loop header duplication even for minsize"));
LoopRotatePass::LoopRotatePass(bool EnableHeaderDuplication, bool PrepareForLTO,
- bool CheckExitCount)
+ bool CheckExitCount, bool UpdateBranchWeights)
: EnableHeaderDuplication(EnableHeaderDuplication),
- PrepareForLTO(PrepareForLTO), CheckExitCount(CheckExitCount) {}
+ PrepareForLTO(PrepareForLTO), CheckExitCount(CheckExitCount),
+ UpdateBranchWeights(UpdateBranchWeights) {}
void LoopRotatePass::printPipeline(
raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
@@ -67,7 +70,11 @@ void LoopRotatePass::printPipeline(
if (!CheckExitCount)
OS << "no-";
- OS << "check-exit-count";
+ OS << "check-exit-count;";
+
+ if (!UpdateBranchWeights)
+ OS << "no-";
+ OS << "update-branch-weights";
OS << ">";
}
@@ -89,10 +96,24 @@ PreservedAnalyses LoopRotatePass::run(Loop &L, LoopAnalysisManager &AM,
std::optional<MemorySSAUpdater> MSSAU;
if (AR.MSSA)
MSSAU = MemorySSAUpdater(AR.MSSA);
+
+ // On the post-PGO path, recover loop entry count from BFI so
+ // LoopRotation emits correct weights for multi-exit loops.
+ // Built locally since loop passes cannot safely access invalidatable
+ // function analyses.
+ BlockFrequencyInfo *BFI = nullptr;
+ std::optional<BlockFrequencyInfo> LocalBFI;
+ Function &F = *L.getHeader()->getParent();
+ if (UpdateBranchWeights && F.hasProfileData()) {
+ BranchProbabilityInfo BPI(F, AR.LI, &AR.TLI, &AR.DT);
+ LocalBFI.emplace(F, BPI, AR.LI);
+ BFI = &*LocalBFI;
+ }
+
bool Changed =
LoopRotation(&L, &AR.LI, &AR.TTI, &AR.AC, &AR.DT, &AR.SE,
MSSAU ? &*MSSAU : nullptr, SQ, false, Threshold, false,
- PrepareForLTO || PrepareForLTOOption, CheckExitCount);
+ PrepareForLTO || PrepareForLTOOption, CheckExitCount, BFI);
if (!Changed)
return PreservedAnalyses::all();
diff --git a/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp b/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
index c8bc5e4daeff3..6764211a0af8d 100644
--- a/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
+++ b/llvm/lib/Transforms/Utils/LoopRotationUtils.cpp
@@ -11,8 +11,10 @@
//===----------------------------------------------------------------------===//
#include "llvm/Transforms/Utils/LoopRotationUtils.h"
+#include "llvm/ADT/APInt.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/AssumptionCache.h"
+#include "llvm/Analysis/BlockFrequencyInfo.h"
#include "llvm/Analysis/CodeMetrics.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/InstructionSimplify.h"
@@ -35,6 +37,8 @@
#include "llvm/Transforms/Utils/Local.h"
#include "llvm/Transforms/Utils/SSAUpdater.h"
#include "llvm/Transforms/Utils/ValueMapper.h"
+#include <limits>
+#include <optional>
using namespace llvm;
#define DEBUG_TYPE "loop-rotate"
@@ -64,17 +68,18 @@ class LoopRotate {
bool IsUtilMode;
bool PrepareForLTO;
bool CheckExitCount;
+ BlockFrequencyInfo *BFI;
public:
LoopRotate(unsigned MaxHeaderSize, LoopInfo *LI,
const TargetTransformInfo *TTI, AssumptionCache *AC,
DominatorTree *DT, ScalarEvolution *SE, MemorySSAUpdater *MSSAU,
const SimplifyQuery &SQ, bool RotationOnly, bool IsUtilMode,
- bool PrepareForLTO, bool CheckExitCount)
+ bool PrepareForLTO, bool CheckExitCount, BlockFrequencyInfo *BFI)
: MaxHeaderSize(MaxHeaderSize), LI(LI), TTI(TTI), AC(AC), DT(DT), SE(SE),
MSSAU(MSSAU), SQ(SQ), RotationOnly(RotationOnly),
IsUtilMode(IsUtilMode), PrepareForLTO(PrepareForLTO),
- CheckExitCount(CheckExitCount) {}
+ CheckExitCount(CheckExitCount), BFI(BFI) {}
bool processLoop(Loop *L);
private:
@@ -210,8 +215,9 @@ static bool profitableToRotateLoopExitingLatch(Loop *L, ScalarEvolution *SE) {
}
static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
- bool HasConditionalPreHeader,
- bool SuccsSwapped) {
+ bool HasConditionalPreHeader, bool SuccsSwapped,
+ std::optional<uint64_t> PreHeaderEntries,
+ std::optional<uint64_t> HeaderEntries) {
MDNode *WeightMD = getBranchWeightMDNode(PreHeaderBI);
if (WeightMD == nullptr)
return;
@@ -232,6 +238,31 @@ static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
if (SuccsSwapped)
std::swap(OrigLoopExitWeight, OrigLoopBackedgeWeight);
+ // PreHeaderEntries, when present, holds the loop entry count from BFI.
+ // Scale it to the original header weight so it fits raw branch_weights,
+ // including ratio-based profiles.
+ std::optional<uint64_t> ScaledPreHeaderEntries;
+ if (PreHeaderEntries && HeaderEntries && *HeaderEntries != 0) {
+ APInt ScaledPE(128, *PreHeaderEntries);
+ APInt HeaderWeight(128, static_cast<uint64_t>(OrigLoopExitWeight) +
+ static_cast<uint64_t>(OrigLoopBackedgeWeight));
+ APInt HeaderCount(128, *HeaderEntries);
+ ScaledPE *= HeaderWeight;
+ ScaledPE = (ScaledPE + HeaderCount.lshr(1)).udiv(HeaderCount);
+ ScaledPreHeaderEntries = ScaledPE.getLimitedValue();
+ }
+
+ bool PreHeaderEntriesKnown = ScaledPreHeaderEntries.has_value();
+ uint64_t PE = PreHeaderEntriesKnown ? *ScaledPreHeaderEntries : 0;
+
+ // A valid profile requires x <= pe <= x + y.
+ // Otherwise, ignore pe and fall back to the single-exit estimate.
+ if (PreHeaderEntriesKnown &&
+ (PE < static_cast<uint64_t>(OrigLoopExitWeight) ||
+ PE > static_cast<uint64_t>(OrigLoopExitWeight) +
+ static_cast<uint64_t>(OrigLoopBackedgeWeight)))
+ PreHeaderEntriesKnown = false;
+
// Update branch weights. Consider the following edge-counts:
//
// | |-------- |
@@ -249,7 +280,8 @@ static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
//
// The following must hold:
// - x == x0 + x1 # counts to "exit" must stay the same.
- // - y0 == x - x0 == x1 # how often loop was entered at all.
+ // - y0 == pe - x0 # preheader budget; in single-exit loops pe == x,
+ // # so this reduces to the historical y0 == x1.
// - y1 == y - y0 # How often loop was repeated (after first iter.).
//
// We cannot generally deduce how often we had a zero-trip count loop so we
@@ -275,6 +307,13 @@ static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
if ((OrigLoopBackedgeWeight & HighBit) != 0 ||
(OrigLoopExitWeight & HighBit) != 0)
break;
+ // Keep pe on the same scale as the header weights.
+ if (PreHeaderEntriesKnown) {
+ if ((PE & HighBit) != 0)
+ PreHeaderEntriesKnown = false;
+ else
+ PE <<= 1;
+ }
OrigLoopBackedgeWeight <<= 1;
OrigLoopExitWeight <<= 1;
}
@@ -292,11 +331,30 @@ static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
if (OrigLoopExitWeight > OrigLoopBackedgeWeight)
OrigLoopBackedgeWeight = OrigLoopExitWeight;
}
+
+ // Raise ExitWeight0 so EnterWeight (y0 == pe - x0) stays within the
+ // original backedge weight.
+ if (PreHeaderEntriesKnown &&
+ PE > static_cast<uint64_t>(OrigLoopBackedgeWeight)) {
+ uint64_t MinExitWeight0 = PE - OrigLoopBackedgeWeight;
+ if (MinExitWeight0 > OrigLoopExitWeight)
+ MinExitWeight0 = OrigLoopExitWeight;
+ if (MinExitWeight0 > ExitWeight0)
+ ExitWeight0 = static_cast<uint32_t>(MinExitWeight0);
+ }
assert(OrigLoopExitWeight >= ExitWeight0 && "Bad branch weight");
ExitWeight1 = OrigLoopExitWeight - ExitWeight0;
- EnterWeight = ExitWeight1;
- assert(OrigLoopBackedgeWeight >= EnterWeight && "Bad branch weight");
- LoopBackWeight = OrigLoopBackedgeWeight - EnterWeight;
+ uint64_t EnterWeight64 = PreHeaderEntriesKnown
+ ? PE - static_cast<uint64_t>(ExitWeight0)
+ : static_cast<uint64_t>(ExitWeight1);
+ if (EnterWeight64 > std::numeric_limits<uint32_t>::max())
+ EnterWeight64 = std::numeric_limits<uint32_t>::max();
+ EnterWeight = static_cast<uint32_t>(EnterWeight64);
+ assert((OrigLoopBackedgeWeight >= EnterWeight || PreHeaderEntriesKnown) &&
+ "Bad branch weight");
+ LoopBackWeight = OrigLoopBackedgeWeight >= EnterWeight
+ ? OrigLoopBackedgeWeight - EnterWeight
+ : 0;
} else if (OrigLoopExitWeight == 0) {
if (OrigLoopBackedgeWeight == 0) {
// degenerate case... keep everything zero...
@@ -316,7 +374,13 @@ static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
} else {
// loop is never entered.
assert(OrigLoopBackedgeWeight == 0 && "remaining case is backedge zero");
- ExitWeight0 = 1;
+ if (PreHeaderEntriesKnown && PE > 0) {
+ uint64_t X0 =
+ std::min<uint64_t>(PE, std::numeric_limits<uint32_t>::max());
+ ExitWeight0 = static_cast<uint32_t>(X0);
+ } else {
+ ExitWeight0 = 1;
+ }
ExitWeight1 = 1;
EnterWeight = 0;
LoopBackWeight = 0;
@@ -326,7 +390,10 @@ static void updateBranchWeights(CondBrInst &PreHeaderBI, CondBrInst &LoopBI,
SuccsSwapped ? LoopBackWeight : ExitWeight1,
SuccsSwapped ? ExitWeight1 : LoopBackWeight,
};
- setBranchWeights(LoopBI, LoopBIWeights, /*IsExpected=*/false);
+ if (EnterWeight == 0 && LoopBackWeight == 0)
+ LoopBI.setMetadata(LLVMContext::MD_prof, nullptr);
+ else
+ setBranchWeights(LoopBI, LoopBIWeights, /*IsExpected=*/false);
if (HasConditionalPreHeader) {
const uint32_t PreHeaderBIWeights[] = {
SuccsSwapped ? EnterWeight : ExitWeight0,
@@ -427,6 +494,15 @@ bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
if (!OrigPreheader || !L->hasDedicatedExits())
return Rotated;
+ // Capture block counts before CFG changes. updateBranchWeights scales the
+ // preheader count to the original header branch-weight scale.
+ std::optional<uint64_t> PreHeaderEntries;
+ std::optional<uint64_t> HeaderEntries;
+ if (BFI) {
+ PreHeaderEntries = BFI->getBlockProfileCount(OrigPreheader);
+ HeaderEntries = BFI->getBlockProfileCount(OrigHeader);
+ }
+
// Anything ScalarEvolution may know about this loop or the PHI nodes
// in its header will soon be invalidated. We should also invalidate
// all outer loops because insertion and deletion of blocks that happens
@@ -756,7 +832,8 @@ bool LoopRotate::rotateLoop(Loop *L, bool SimplifiedLatch) {
!isa<ConstantInt>(Cond) ||
PHBI->getSuccessor(cast<ConstantInt>(Cond)->isZero()) != NewHeader;
- updateBranchWeights(*PHBI, *BI, HasConditionalPreHeader, BISuccsSwapped);
+ updateBranchWeights(*PHBI, *BI, HasConditionalPreHeader, BISuccsSwapped,
+ PreHeaderEntries, HeaderEntries);
if (HasConditionalPreHeader) {
// The conditional branch can't be folded, handle the general case.
@@ -976,8 +1053,8 @@ bool llvm::LoopRotation(Loop *L, LoopInfo *LI, const TargetTransformInfo *TTI,
const SimplifyQuery &SQ, bool RotationOnly = true,
unsigned Threshold = unsigned(-1),
bool IsUtilMode = true, bool PrepareForLTO,
- bool CheckExitCount) {
+ bool CheckExitCount, BlockFrequencyInfo *BFI) {
LoopRotate LR(Threshold, LI, TTI, AC, DT, SE, MSSAU, SQ, RotationOnly,
- IsUtilMode, PrepareForLTO, CheckExitCount);
+ IsUtilMode, PrepareForLTO, CheckExitCount, BFI);
return LR.processLoop(L);
}
diff --git a/llvm/test/Other/new-pm-print-pipeline.ll b/llvm/test/Other/new-pm-print-pipeline.ll
index 21da3eb33dd6d..55551a350db1f 100644
--- a/llvm/test/Other/new-pm-print-pipeline.ll
+++ b/llvm/test/Other/new-pm-print-pipeline.ll
@@ -4,7 +4,7 @@
; CHECK-0: function(adce),function(adce)
; RUN: opt -disable-output -disable-verify -print-pipeline-passes -passes='module(rpo-function-attrs,require<globals-aa>,function(float2int,lower-constant-intrinsics,loop(loop-rotate)),invalidate<globals-aa>)' < %s | FileCheck %s --match-full-lines --check-prefixes=CHECK-1
-; CHECK-1: rpo-function-attrs,require<globals-aa>,function(float2int,lower-constant-intrinsics,loop(loop-rotate<header-duplication;no-prepare-for-lto;no-check-exit-count>)),invalidate<globals-aa>
+; CHECK-1: rpo-function-attrs,require<globals-aa>,function(float2int,lower-constant-intrinsics,loop(loop-rotate<header-duplication;no-prepare-for-lto;no-check-exit-count;no-update-branch-weights>)),invalidate<globals-aa>
;; Test that we get ClassName printed when there is no ClassName to pass-name mapping (as is the case for the BitcodeWriterPass).
; RUN: opt -o /dev/null -disable-verify -print-pipeline-passes -passes='function(mem2reg)' < %s -disable-pipeline-verification | FileCheck %s --match-full-lines --check-prefixes=CHECK-3
@@ -66,7 +66,7 @@
;; Test that the loop-nest-pass lnicm is printed with the other loop-passes in the pipeline.
; RUN: opt -disable-output -disable-verify -print-pipeline-passes -passes='function(loop-mssa(licm,loop-rotate,loop-deletion,lnicm,loop-rotate))' < %s | FileCheck %s --match-full-lines --check-prefixes=CHECK-23
-; CHECK-23: function(loop-mssa(licm<allowspeculation>,loop-rotate<header-duplication;no-prepare-for-lto;no-check-exit-count>,loop-deletion,lnicm<allowspeculation>,loop-rotate<header-duplication;no-prepare-for-lto;no-check-exit-count>))
+; CHECK-23: function(loop-mssa(licm<allowspeculation>,loop-rotate<header-duplication;no-prepare-for-lto;no-check-exit-count;no-update-branch-weights>,loop-deletion,lnicm<allowspeculation>,loop-rotate<header-duplication;no-prepare-for-lto;no-check-exit-count;no-update-branch-weights>))
;; Test that -debugify and -check-debugify is printed correctly.
; RUN: opt -disable-output -disable-verify -print-pipeline-passes -passes='debugify,no-op-function,check-debugify' < %s | FileCheck %s --match-full-lines --check-prefixes=CHECK-24
@@ -110,7 +110,7 @@
; CHECK-32: cgscc(function<no-rerun>(no-op-function))
; RUN: opt -disable-output -disable-verify -print-pipeline-passes -passes='function(loop(loop-rotate<no-header-duplication;no-prepare-for-lto>))' < %s | FileCheck %s --match-full-lines --check-prefixes=CHECK-33
-; CHECK-33: function(loop(loop-rotate<no-header-duplication;no-prepare-for-lto;no-check-exit-count>))
+; CHECK-33: function(loop(loop-rotate<no-header-duplication;no-prepare-for-lto;no-check-exit-count;no-update-branch-weights>))
; RUN: opt -disable-output -disable-verify -print-pipeline-passes -passes='globaldce' < %s | FileCheck %s --match-full-lines --check-prefixes=CHECK-34
; CHECK-34: globaldce
diff --git a/llvm/test/Transforms/LoopRotate/folded-preheader-multi-exit-branch-weights.ll b/llvm/test/Transforms/LoopRotate/folded-preheader-multi-exit-branch-weights.ll
new file mode 100644
index 0000000000000..507d21c98e0d3
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/folded-preheader-multi-exit-branch-weights.ll
@@ -0,0 +1,74 @@
+; RUN: opt < %s -passes='loop-rotate<update-branch-weights>' -S | FileCheck %s
+;
+
+; The cloned preheader branch folds to an unconditional loop entry since the
+; initial IV cannot satisfy the exit condition. The loop remains multi-exit,
+; so the rotated latch must use the BFI-derived first-iteration count instead
+; of the single-exit estimate.
+;
+; Original profile scale:
+; entry -> header ; PreHeaderEntries = 210
+; header -> exit1 / body !prof {200, 9800}
+; body -> exit2 / latch !prof {10, 9790}
+; latch -> header
+;
+; Expected rotated latch (no guard, ExitWeight0 = 0):
+; ExitWeight1 = 200
+; EnterWeight = 210
+; LoopBackWeight = 9800 - 210 = 9590
+; LATCH !prof {200, 9590}
+
+define void @folded_preheader_multi_exit(ptr %p) !prof !14 {
+entry:
+ br label %header
+
+header: ; preds = %latch, %entry
+ %iv = phi i32 [ 0, %entry ], [ %iv.next, %latch ]
+ %hcmp = icmp eq i32 %iv, 42
+ br i1 %hcmp, label %exit1, label %body, !prof !15
+
+body: ; preds = %header
+ %addr = getelementptr i32, ptr %p, i32 %iv
+ %v = load i32, ptr %addr, align 4
+ %bcmp = icmp slt i32 %v, 0
+ br i1 %bcmp, label %exit2, label %latch, !prof !16
+
+latch: ; preds = %body
+ %iv.next = add i32 %iv, 1
+ br label %header
+
+exit1: ; preds = %header
+ ret void
+
+exit2: ; preds = %body
+ ret void
+
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 10800}
+!4 = !{!"MaxCount", i64 9800}
+!5 = !{!"MaxInternalCount", i64 9800}
+!6 = !{!"MaxFunctionCount", i64 210}
+!7 = !{!"NumCounts", i64 4}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 9800, i32 1}
+!12 = !{i32 999000, i64 210, i32 1}
+!13 = !{i32 999999, i64 200, i32 1}
+!14 = !{!"function_entry_count", i64 210}
+!15 = !{!"branch_weights", i32 200, i32 9800}
+!16 = !{!"branch_weights", i32 10, i32 9790}
+; CHECK-LABEL: define void @folded_preheader_multi_exit(
+
+; CHECK: entry:
+; CHECK-NEXT: br label %body
+
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %body, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 200, i32 9590}
diff --git a/llvm/test/Transforms/LoopRotate/multi-exit-branch-weights.ll b/llvm/test/Transforms/LoopRotate/multi-exit-branch-weights.ll
new file mode 100644
index 0000000000000..64a722b608202
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/multi-exit-branch-weights.ll
@@ -0,0 +1,84 @@
+; RUN: opt < %s -passes='loop-rotate<update-branch-weights>' -S | FileCheck %s
+;
+; Multi-exit loop where preheader entry count differs from the header’s
+; exit count.
+;
+; entry --(c0)--> ph !prof {210, 1}
+; ph -------> header
+; header -> exit1 / body !prof {200, 9800}
+; body -> exit2 / latch !prof {10, 9790} ; second exit
+; latch -------> header
+;
+; Entry counts are 10x the branch-weight scale. BFI reports
+; PreHeaderEntries = 2100 and HeaderEntries = 100000, which must be
+; scaled back to 210.
+;
+; Expected post-rotation:
+; ExitWeight0 + EnterWeight = 210
+; ExitWeight0 + ExitWeight1 = 200
+; EnterWeight + LoopBackWeight = 9800
+
+define void @f(ptr %p, i32 %start, i32 %limit, i32 %cond) !prof !14 {
+entry:
+ %c0 = icmp ne i32 %cond, 0
+ br i1 %c0, label %ph, label %ret, !prof !15
+
+ph: ; preds = %entry
+ br label %header
+
+header: ; preds = %latch, %ph
+ %iv = phi i32 [ %start, %ph ], [ %iv.next, %latch ]
+ ; loop-rotate cannot prove %start != %limit, so it must materialize a guard.
+ %hcmp = icmp eq i32 %iv, %limit
+ br i1 %hcmp, label %exit1, label %body, !prof !16
+
+body: ; preds = %header
+ %addr = getelementptr i32, ptr %p, i32 %iv
+ %v = load i32, ptr %addr, align 4
+ %bcmp = icmp slt i32 %v, 0
+ br i1 %bcmp, label %exit2, label %latch, !prof !17
+
+latch: ; preds = %body
+ %iv.next = add i32 %iv, 1
+ br label %header
+
+exit1: ; preds = %header
+ ret void
+
+exit2: ; preds = %body
+ ret void
+
+ret: ; preds = %entry
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 108000}
+!4 = !{!"MaxCount", i64 98000}
+!5 = !{!"MaxInternalCount", i64 98000}
+!6 = !{!"MaxFunctionCount", i64 2110}
+!7 = !{!"NumCounts", i64 4}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 98000, i32 1}
+!12 = !{i32 999000, i64 2110, i32 1}
+!13 = !{i32 999999, i64 2000, i32 1}
+!14 = !{!"function_entry_count", i64 2110}
+!15 = !{!"branch_weights", i32 210, i32 1}
+!16 = !{!"branch_weights", i32 200, i32 9800}
+!17 = !{!"branch_weights", i32 10, i32 9790}
+
+; CHECK-LABEL: define void @f(
+
+; CHECK: ph:
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}.lr.ph, !prof [[GUARD:![0-9]+]]
+
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %body, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[GUARD]] = !{!"branch_weights", i32 1, i32 209}
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 199, i32 9591}
diff --git a/llvm/test/Transforms/LoopRotate/multi-exit-update-branch-weights.ll b/llvm/test/Transforms/LoopRotate/multi-exit-update-branch-weights.ll
new file mode 100644
index 0000000000000..c3f0da37934b2
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/multi-exit-update-branch-weights.ll
@@ -0,0 +1,327 @@
+; RUN: opt < %s -passes='loop-rotate<update-branch-weights>' -S | FileCheck %s
+;
+; Multi-exit loop reduced from a real PGO trace (treeup.c::update_tree).
+; The loop has two exits (header and body). Loop-rotate previously assumed
+; preheader_entries == header_exit_count and derived EnterWeight from the
+; wrong edge, producing inconsistent weights on the rotated guard and latch.
+;
+; With the fix, EnterWeight is derived from the preheader entry count, so
+; rotated weights remain consistent with the original iteration profile.
+;
+; Original header (TEST):
+; br i1 %cmp18, label %CONTINUE, label %if.end20
+; !prof !51 ; {1977246, 82880387}
+; Preheader edge (RECURSION -> TEST.preheader): 31075547
+;
+; After rotation:
+; Guard (TEST.preheader):
+; !prof {1, 31075546}
+; Latch (if.end23):
+; !prof {1977245, 51804841}
+
+source_filename = "treeup.c"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+%struct.arc = type { i32, i64, ptr, ptr, i16, ptr, ptr, i64, i64 }
+%struct.node = type { i64, i32, ptr, ptr, ptr, ptr, ptr, ptr, ptr, ptr, i64, i64, i32, i32 }
+; Function Attrs: nounwind uwtable
+define dso_local void @update_tree(i64 noundef %cycle_ori, i64 noundef %new_orientation, i64 noundef %delta, i64 noundef %new_flow, ptr noundef %iplus, ptr noundef %jplus, ptr noundef captures(address) %iminus, ptr noundef captures(address) %jminus, ptr noundef readnone captures(
+address) %w, ptr noundef %bea, i64 noundef %sigma, i64 noundef %feas_tol) local_unnamed_addr #0 !prof !34 {
+entry:
+ %tail = getelementptr inbounds nuw %struct.arc, ptr %bea, i64 0, i32 2
+ %0 = load ptr, ptr %tail, align 8
+ %cmp = icmp eq ptr %0, %jplus
+ %cmp1 = icmp slt i64 %sigma, 0
+ %or.cond = and i1 %cmp1, %cmp
+ br i1 %or.cond, label %if.then, label %lor.lhs.false, !prof !45
+lor.lhs.false: ; preds = %entry
+ %cmp3 = icmp eq ptr %0, %iplus
+ %cmp5 = icmp sgt i64 %sigma, 0
+ %or.cond224 = and i1 %cmp5, %cmp3
+ br i1 %or.cond224, label %if.then, label %if.else, !prof !46
+if.then: ; preds = %lor.lhs.false, %entry
+ %cond = tail call i64 @llvm.abs.i64(i64 %sigma, i1 true)
+ br label %if.end
+if.else: ; preds = %lor.lhs.false
+ %cond12 = tail call i64 @llvm.abs.i64(i64 %sigma, i1 true)
+ %sub13 = sub nsw i64 0, %cond12
+ br label %if.end
+if.end: ; preds = %if.else, %if.then
+ %sigma.addr.0 = phi i64 [ %cond, %if.then ], [ %sub13, %if.else ]
+ %1 = load i64, ptr %iminus, align 8
+ %add = add nsw i64 %1, %sigma.addr.0
+ store i64 %add, ptr %iminus, align 8
+ br label %RECURSION
+RECURSION: ; preds = %ITERATION, %if.end
+ %father.0 = phi ptr [ %iminus, %if.end ], [ %temp.0, %ITERATION ]
+ %child = getelementptr inbounds nuw %struct.node, ptr %father.0, i64 0, i32 2
+ %2 = load ptr, ptr %child, align 8
+ %tobool.not = icmp eq ptr %2, null
+ br i1 %tobool.not, label %TEST.preheader, label %ITERATION, !prof !50
+TEST.preheader: ; preds = %RECURSION
+ br label %TEST
+ITERATION.loopexit: ; preds = %if.end20
+ %.lcssa = phi ptr [ %4, %if.end20 ]
+ br label %ITERATION
+ITERATION: ; preds = %ITERATION.loopexit, %RECURSION
+ %temp.0 = phi ptr [ %2, %RECURSION ], [ %.lcssa, %ITERATION.loopexit ]
+ %3 = load i64, ptr %temp.0, align 8
+ %add16 = add nsw i64 %3, %sigma.addr.0
+ store i64 %add16, ptr %temp.0, align 8
+ br label %RECURSION
+TEST: ; preds = %TEST.preheader, %if.end23
+ %father.1 = phi ptr [ %5, %if.end23 ], [ %father.0, %TEST.preheader ]
+ %cmp18 = icmp eq ptr %father.1, %iminus
+ br i1 %cmp18, label %CONTINUE, label %if.end20, !prof !51
+if.end20: ; preds = %TEST
+ %sibling = getelementptr inbounds nuw %struct.node, ptr %father.1, i64 0, i32 4
+ %4 = load ptr, ptr %sibling, align 8
+ %tobool21.not = icmp eq ptr %4, null
+ br i1 %tobool21.not, label %if.end23, label %ITERATION.loopexit, !prof !53
+if.end23: ; preds = %if.end20
+ %pred = getelementptr inbounds nuw %struct.node, ptr %father.1, i64 0, i32 3
+ %5 = load ptr, ptr %pred, align 8
+ br label %TEST
+CONTINUE: ; preds = %TEST
+ %pred24 = getelementptr inbounds nuw %struct.node, ptr %iplus, i64 0, i32 3
+ %6 = load ptr, ptr %pred24, align 8
+ %depth = getelementptr inbounds nuw %struct.node, ptr %iminus, i64 0, i32 11
+ %7 = load i64, ptr %depth, align 8
+ br label %while.cond
+while.cond: ; preds = %if.end61, %CONTINUE
+ %new_basic_arc.0 = phi ptr [ %bea, %CONTINUE ], [ %15, %if.end61 ]
+ %father.2 = phi ptr [ %6, %CONTINUE ], [ %17, %if.end61 ]
+ %temp.1 = phi ptr [ %iplus, %CONTINUE ], [ %father.2, %if.end61 ]
+ %new_pred.0 = phi ptr [ %jplus, %CONTINUE ], [ %temp.1, %if.end61 ]
+ %new_flow.addr.0 = phi i64 [ %new_flow, %CONTINUE ], [ %flow_temp.0, %if.end61 ]
+ %new_orientation.addr.0 = phi i64 [ %new_orientation, %CONTINUE ], [ %conv, %if.end61 ]
+ %new_depth.0 = phi i64 [ %7, %CONTINUE ], [ %sub68, %if.end61 ]
+ %cmp25.not = icmp eq ptr %temp.1, %jminus
+ br i1 %cmp25.not, label %while.end, label %while.body, !prof !56
+while.body: ; preds = %while.cond
+ %sibling26 = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 4
+ %8 = load ptr, ptr %sibling26, align 8
+ %tobool27.not = icmp eq ptr %8, null
+ br i1 %tobool27.not, label %if.end31, label %if.then28, !prof !57
+if.then28: ; preds = %while.body
+ %sibling_prev = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 5
+ %9 = load ptr, ptr %sibling_prev, align 8
+ %sibling_prev30 = getelementptr inbounds nuw %struct.node, ptr %8, i64 0, i32 5
+ store ptr %9, ptr %sibling_prev30, align 8
+ br label %if.end31
+if.end31: ; preds = %if.then28, %while.body
+ %sibling_prev32 = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 5
+ %10 = load ptr, ptr %sibling_prev32, align 8
+ %tobool33.not = icmp eq ptr %10, null
+ br i1 %tobool33.not, label %if.else38, label %if.then34, !prof !59
+if.then34: ; preds = %if.end31
+ %sibling37 = getelementptr inbounds nuw %struct.node, ptr %10, i64 0, i32 4
+ store ptr %8, ptr %sibling37, align 8
+ br label %if.end41
+if.else38: ; preds = %if.end31
+ %child40 = getelementptr inbounds nuw %struct.node, ptr %father.2, i64 0, i32 2
+ store ptr %8, ptr %child40, align 8
+ br label %if.end41
+if.end41: ; preds = %if.else38, %if.then34
+ %pred42 = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 3
+ store ptr %new_pred.0, ptr %pred42, align 8
+ %child43 = getelementptr inbounds nuw %struct.node, ptr %new_pred.0, i64 0, i32 2
+ %11 = load ptr, ptr %child43, align 8
+ store ptr %11, ptr %sibling26, align 8
+ %tobool46.not = icmp eq ptr %11, null
+ br i1 %tobool46.not, label %if.end50, label %if.then47, !prof !60
+if.then47: ; preds = %if.end41
+ %sibling_prev49 = getelementptr inbounds nuw %struct.node, ptr %11, i64 0, i32 5
+ store ptr %temp.1, ptr %sibling_prev49, align 8
+ br label %if.end50
+if.end50: ; preds = %if.then47, %if.end41
+ store ptr %temp.1, ptr %child43, align 8
+ store ptr null, ptr %sibling_prev32, align 8
+ %orientation = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 1
+ %12 = load i32, ptr %orientation, align 8
+ %tobool53.not = icmp eq i32 %12, 0
+ %conv = zext i1 %tobool53.not to i64
+ %cmp54 = icmp eq i64 %cycle_ori, %conv
+ br i1 %cmp54, label %if.then56, label %if.else58, !prof !62
+if.then56: ; preds = %if.end50
+ %flow = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 10
+ %13 = load i64, ptr %flow, align 8
+ %add57 = add nsw i64 %13, %delta
+ br label %if.end61
+if.else58: ; preds = %if.end50
+ %flow59 = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 10
+ %14 = load i64, ptr %flow59, align 8
+ %sub60 = sub nsw i64 %14, %delta
+ br label %if.end61
+if.end61: ; preds = %if.else58, %if.then56
+ %flow_temp.0 = phi i64 [ %add57, %if.then56 ], [ %sub60, %if.else58 ]
+ %basic_arc = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 6
+ %15 = load ptr, ptr %basic_arc, align 8
+ %depth62 = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 11
+ %16 = load i64, ptr %depth62, align 8
+ %conv63 = trunc i64 %new_orientation.addr.0 to i32
+ store i32 %conv63, ptr %orientation, align 8
+ %flow65 = getelementptr inbounds nuw %struct.node, ptr %temp.1, i64 0, i32 10
+ store i64 %new_flow.addr.0, ptr %flow65, align 8
+ store ptr %new_basic_arc.0, ptr %basic_arc, align 8
+ store i64 %new_depth.0, ptr %depth62, align 8
+ %sub68 = sub nsw i64 %7, %16
+ %pred69 = getelementptr inbounds nuw %struct.node, ptr %father.2, i64 0, i32 3
+ %17 = load ptr, ptr %pred69, align 8
+ br label %while.cond, !llvm.loop !65
+while.end: ; preds = %while.cond
+ %cmp70 = icmp sgt i64 %delta, %feas_tol
+ br i1 %cmp70, label %for.cond.preheader, label %for.cond110.preheader, !prof !67
+for.cond110.preheader: ; preds = %while.end
+ br label %for.cond110
+for.cond.preheader: ; preds = %while.end
+ br label %for.cond
+for.cond: ; preds = %for.cond.preheader, %for.inc
+ %temp.2 = phi ptr [ %22, %for.inc ], [ %jminus, %for.cond.preheader ]
+ %cmp73.not = icmp eq ptr %temp.2, %w
+ br i1 %cmp73.not, label %for.cond89.preheader, label %for.body, !prof !68
+for.cond89.preheader: ; preds = %for.cond
+ br label %for.cond89
+for.body: ; preds = %for.cond
+ %depth75 = getelementptr inbounds nuw %struct.node, ptr %temp.2, i64 0, i32 11
+ %18 = load i64, ptr %depth75, align 8
+ %sub76 = sub nsw i64 %18, %7
+ store i64 %sub76, ptr %depth75, align 8
+ %orientation77 = getelementptr inbounds nuw %struct.node, ptr %temp.2, i64 0, i32 1
+ %19 = load i32, ptr %orientation77, align 8
+ %conv78 = sext i32 %19 to i64
+ %cmp79.not = icmp eq i64 %cycle_ori, %conv78
+ br i1 %cmp79.not, label %if.else84, label %if.then81
+if.then81: ; preds = %for.body
+ %flow82 = getelementptr inbounds nuw %struct.node, ptr %temp.2, i64 0, i32 10
+ %20 = load i64, ptr %flow82, align 8
+ %add83 = add nsw i64 %20, %delta
+ store i64 %add83, ptr %flow82, align 8
+ br label %for.inc
+if.else84: ; preds = %for.body
+ %flow85 = getelementptr inbounds nuw %struct.node, ptr %temp.2, i64 0, i32 10
+ %21 = load i64, ptr %flow85, align 8
+ %sub86 = sub nsw i64 %21, %delta
+ store i64 %sub86, ptr %flow85, align 8
+ br label %for.inc
+for.inc: ; preds = %if.then81, %if.else84
+ %pred88 = getelementptr inbounds nuw %struct.node, ptr %temp.2, i64 0, i32 3
+ %22 = load ptr, ptr %pred88, align 8
+ br label %for.cond, !llvm.loop !69
+for.cond89: ; preds = %for.cond89.preheader, %for.inc106
+ %temp.3 = phi ptr [ %27, %for.inc106 ], [ %jplus, %for.cond89.preheader ]
+ %cmp90.not = icmp eq ptr %temp.3, %w
+ br i1 %cmp90.not, label %if.end128.loopexit, label %for.body92, !prof !68
+for.body92: ; preds = %for.cond89
+ %depth93 = getelementptr inbounds nuw %struct.node, ptr %temp.3, i64 0, i32 11
+ %23 = load i64, ptr %depth93, align 8
+ %add94 = add nsw i64 %23, %7
+ store i64 %add94, ptr %depth93, align 8
+ %orientation95 = getelementptr inbounds nuw %struct.node, ptr %temp.3, i64 0, i32 1
+ %24 = load i32, ptr %orientation95, align 8
+ %conv96 = sext i32 %24 to i64
+ %cmp97 = icmp eq i64 %cycle_ori, %conv96
+ br i1 %cmp97, label %if.then99, label %if.else102
+if.then99: ; preds = %for.body92
+ %flow100 = getelementptr inbounds nuw %struct.node, ptr %temp.3, i64 0, i32 10
+ %25 = load i64, ptr %flow100, align 8
+ %add101 = add nsw i64 %25, %delta
+ store i64 %add101, ptr %flow100, align 8
+ br label %for.inc106
+if.else102: ; preds = %for.body92
+ %flow103 = getelementptr inbounds nuw %struct.node, ptr %temp.3, i64 0, i32 10
+ %26 = load i64, ptr %flow103, align 8
+ %sub104 = sub nsw i64 %26, %delta
+ store i64 %sub104, ptr %flow103, align 8
+ br label %for.inc106
+for.inc106: ; preds = %if.then99, %if.else102
+ %pred107 = getelementptr inbounds nuw %struct.node, ptr %temp.3, i64 0, i32 3
+ %27 = load ptr, ptr %pred107, align 8
+ br label %for.cond89, !llvm.loop !70
+for.cond110: ; preds = %for.cond110.preheader, %for.body113
+ %temp.4 = phi ptr [ %29, %for.body113 ], [ %jminus, %for.cond110.preheader ]
+ %cmp111.not = icmp eq ptr %temp.4, %w
+ br i1 %cmp111.not, label %for.cond119.preheader, label %for.body113, !prof !71
+for.cond119.preheader: ; preds = %for.cond110
+ br label %for.cond119
+for.body113: ; preds = %for.cond110
+ %depth114 = getelementptr inbounds nuw %struct.node, ptr %temp.4, i64 0, i32 11
+ %28 = load i64, ptr %depth114, align 8
+ %sub115 = sub nsw i64 %28, %7
+ store i64 %sub115, ptr %depth114, align 8
+ %pred117 = getelementptr inbounds nuw %struct.node, ptr %temp.4, i64 0, i32 3
+ %29 = load ptr, ptr %pred117, align 8
+ br label %for.cond110, !llvm.loop !72
+for.cond119: ; preds = %for.cond119.preheader, %for.body122
+ %temp.5 = phi ptr [ %31, %for.body122 ], [ %jplus, %for.cond119.preheader ]
+ %cmp120.not = icmp eq ptr %temp.5, %w
+ br i1 %cmp120.not, label %if.end128.loopexit225, label %for.body122, !prof !73
+for.body122: ; preds = %for.cond119
+ %depth123 = getelementptr inbounds nuw %struct.node, ptr %temp.5, i64 0, i32 11
+ %30 = load i64, ptr %depth123, align 8
+ %add124 = add nsw i64 %30, %7
+ store i64 %add124, ptr %depth123, align 8
+ %pred126 = getelementptr inbounds nuw %struct.node, ptr %temp.5, i64 0, i32 3
+ %31 = load ptr, ptr %pred126, align 8
+ br label %for.cond119, !llvm.loop !74
+if.end128.loopexit: ; preds = %for.cond89
+ br label %if.end128
+if.end128.loopexit225: ; preds = %for.cond119
+ br label %if.end128
+if.end128: ; preds = %if.end128.loopexit225, %if.end128.loopexit
+ ret void
+}
+; Function Attrs: nocallback nofree nosync nounwind speculatable willreturn memory(none)
+declare i64 @llvm.abs.i64(i64, i1 immarg) #1
+
+!llvm.module.flags = !{!80}
+!80 = !{i32 1, !"ProfileSummary", !81}
+!81 = !{!82, !83, !84, !85, !86, !87, !88, !89}
+!82 = !{!"ProfileFormat", !"InstrProf"}
+!83 = !{!"TotalCount", i64 84857633}
+!84 = !{!"MaxCount", i64 82880387}
+!85 = !{!"MaxInternalCount", i64 82880387}
+!86 = !{!"MaxFunctionCount", i64 82880387}
+!87 = !{!"NumCounts", i64 64}
+!88 = !{!"NumFunctions", i64 1}
+!89 = !{!"DetailedSummary", !90}
+!90 = !{!91, !92, !93}
+!91 = !{i32 10000, i64 82880387, i32 1}
+!92 = !{i32 999000, i64 53782086, i32 1}
+!93 = !{i32 999999, i64 1977246, i32 1}
+
+!34 = !{!"function_entry_count", i64 1977246}
+!45 = !{!"branch_weights", i32 9915, i32 1967331}
+!46 = !{!"branch_weights", i32 6597, i32 1960734}
+!50 = !{!"branch_weights", i32 31075547, i32 53782086}
+!51 = !{!"branch_weights", i32 1977246, i32 82880387}
+!53 = !{!"branch_weights", i32 53782086, i32 29098301}
+!56 = !{!"branch_weights", i32 1977246, i32 2385079}
+!57 = !{!"branch_weights", i32 746895, i32 1638184}
+!59 = !{!"branch_weights", i32 791318, i32 1593761}
+!60 = !{!"branch_weights", i32 476345, i32 1908734}
+!62 = !{!"branch_weights", i32 240086, i32 2144993}
+!65 = distinct !{!65, !66}
+!66 = !{!"llvm.loop.mustprogress"}
+!67 = !{!"branch_weights", i32 16512, i32 1960734}
+!68 = !{!"branch_weights", i32 16512, i32 0}
+!69 = distinct !{!69, !66}
+!70 = distinct !{!70, !66}
+!71 = !{!"branch_weights", i32 1960734, i32 110585854}
+!72 = distinct !{!72, !66}
+!73 = !{!"branch_weights", i32 1960734, i32 119591748}
+!74 = distinct !{!74, !66}
+
+; CHECK-LABEL: define dso_local void @update_tree(
+
+; Rotated guard at the new preheader picks EnterWeight from the actual
+; preheader edge weight (31075547) rather than the header's loop-edge weight.
+; CHECK: TEST.preheader:
+; CHECK: br i1 %cmp{{[0-9]+}}, label %TEST.preheader.CONTINUE_crit_edge, label %if.end20.lr.ph, !prof [[GUARD:![0-9]+]]
+
+; Rotated latch carries the leftover of the original header weights.
+; CHECK: if.end23:
+; CHECK: br i1 %cmp{{[0-9]+}}, label %TEST.CONTINUE_crit_edge, label %if.end20, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[GUARD]] = !{!"branch_weights", i32 1, i32 31075546}
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 1977245, i32 51804841}
diff --git a/llvm/test/Transforms/LoopRotate/never-entered-branch-weights.ll b/llvm/test/Transforms/LoopRotate/never-entered-branch-weights.ll
new file mode 100644
index 0000000000000..48c0a476b457c
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/never-entered-branch-weights.ll
@@ -0,0 +1,61 @@
+; RUN: opt < %s -passes='loop-rotate<update-branch-weights>' -S | FileCheck %s
+;
+; Covers the "loop never entered" case in updateBranchWeights()
+; (OrigLoopBackedgeWeight == 0).
+;
+; entry --(c0)--> ph !prof {5, 1} ; PreHeaderEntries = 5
+; ph -------> header
+; header -> exit / latch !prof {5, 0} ; no backedge
+; latch -------> header
+;
+; Expected rotated header: !prof {5, 0}.
+
+define void @g(i32 %n, i32 %cond) !prof !14 {
+entry:
+ %c0 = icmp ne i32 %cond, 0
+ br i1 %c0, label %ph, label %ret, !prof !15
+
+ph: ; preds = %entry
+ br label %header
+
+header: ; preds = %latch, %ph
+ %iv = phi i32 [ 0, %ph ], [ %iv.next, %latch ]
+ %hcmp = icmp sge i32 %iv, %n
+ br i1 %hcmp, label %exit, label %latch, !prof !16
+
+latch: ; preds = %header
+ %iv.next = add i32 %iv, 1
+ br label %header
+
+exit: ; preds = %header
+ ret void
+
+ret: ; preds = %entry
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 6}
+!4 = !{!"MaxCount", i64 5}
+!5 = !{!"MaxInternalCount", i64 5}
+!6 = !{!"MaxFunctionCount", i64 6}
+!7 = !{!"NumCounts", i64 3}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 5, i32 1}
+!12 = !{i32 999000, i64 5, i32 1}
+!13 = !{i32 999999, i64 5, i32 1}
+!14 = !{!"function_entry_count", i64 6}
+!15 = !{!"branch_weights", i32 5, i32 1}
+!16 = !{!"branch_weights", i32 5, i32 0}
+
+; CHECK-LABEL: define void @g(
+
+; CHECK: br i1 %{{.*}}, label %exit, label %header, !prof [[GW:![0-9]+]]
+
+; CHECK-DAG: [[GW]] = !{!"branch_weights", i32 5, i32 0}
diff --git a/llvm/test/Transforms/LoopRotate/no-bfi-fallback-branch-weights.ll b/llvm/test/Transforms/LoopRotate/no-bfi-fallback-branch-weights.ll
new file mode 100644
index 0000000000000..0569fcf7b2f1d
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/no-bfi-fallback-branch-weights.ll
@@ -0,0 +1,87 @@
+; RUN: opt < %s -passes=loop-rotate -S | FileCheck %s
+;
+; Covers the fallback in updateBranchWeights() when BFI is unavailable.
+; In a plain loop-rotate run (no update-branch-weights flag), BFI is not
+; built, so PreHeaderEntries is unknown and the single-exit formula applies:
+;
+; EnterWeight = ExitWeight1
+; LoopBackWeight = OrigLoopBackedgeWeight - EnterWeight
+;
+; Original CFG:
+; entry --(c0)--> ph
+; ph -------> header
+; header -> exit / body !prof {200, 9800}
+; body -------> latch
+; latch -------> header
+;
+; Header weights are >= 127, so no rescaling is applied.
+;
+; Expected post-rotation:
+; ExitWeight0 = 1
+; ExitWeight1 = 199
+; EnterWeight = 199
+; LoopBackWeight = 9601
+;
+; GUARD !prof {1, 199}
+; LATCH !prof {199, 9601}
+
+define void @no_bfi(ptr %p, i32 %n, i32 %cond) !prof !14 {
+entry:
+ %c0 = icmp ne i32 %cond, 0
+ br i1 %c0, label %ph, label %ret
+
+ph: ; preds = %entry
+ br label %header
+
+header: ; preds = %latch, %ph
+ %iv = phi i32 [ 0, %ph ], [ %iv.next, %latch ]
+ %hcmp = icmp sge i32 %iv, %n
+ br i1 %hcmp, label %exit, label %body, !prof !15
+
+body: ; preds = %header
+ %addr = getelementptr i32, ptr %p, i32 %iv
+ store i32 %iv, ptr %addr, align 4
+ br label %latch
+
+latch: ; preds = %body
+ %iv.next = add i32 %iv, 1
+ br label %header
+
+exit: ; preds = %header
+ ret void
+
+ret: ; preds = %entry
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 10000}
+!4 = !{!"MaxCount", i64 9800}
+!5 = !{!"MaxInternalCount", i64 9800}
+!6 = !{!"MaxFunctionCount", i64 200}
+!7 = !{!"NumCounts", i64 3}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 9800, i32 1}
+!12 = !{i32 999000, i64 200, i32 1}
+!13 = !{i32 999999, i64 200, i32 1}
+!14 = !{!"function_entry_count", i64 201}
+!15 = !{!"branch_weights", i32 200, i32 9800}
+
+; CHECK-LABEL: define void @no_bfi(
+
+; Preheader-bypass guard: derived from the old single-exit formula because
+; BlockFrequencyInfo is not available to read PreHeaderEntries from.
+; CHECK: ph:
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}.lr.ph, !prof [[GUARD:![0-9]+]]
+
+; Rotated latch test.
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %body, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[GUARD]] = !{!"branch_weights", i32 1, i32 199}
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 199, i32 9601}
diff --git a/llvm/test/Transforms/LoopRotate/predecessor-no-prof-branch-weights.ll b/llvm/test/Transforms/LoopRotate/predecessor-no-prof-branch-weights.ll
new file mode 100644
index 0000000000000..f441524a5b53e
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/predecessor-no-prof-branch-weights.ll
@@ -0,0 +1,87 @@
+; RUN: opt < %s -passes=loop-rotate -S | FileCheck %s
+;
+; Covers the fallback in updateBranchWeights() when BFI is unavailable
+; (plain loop-rotate run). PreHeaderEntries cannot be derived, so the
+; single-exit formula is used:
+;
+; EnterWeight = ExitWeight1
+; LoopBackWeight = OrigLoopBackedgeWeight - EnterWeight
+;
+; Original CFG:
+; entry --(c0)--> ph
+; ph -------> header
+; header -> exit / body !prof {200, 9800}
+; body -------> latch
+; latch -------> header
+;
+; Header weights are >= 127, so no rescaling occurs.
+;
+; Expected post-rotation:
+; ExitWeight0 = 1
+; ExitWeight1 = 199
+; EnterWeight = 199
+; LoopBackWeight = 9601
+;
+; GUARD !prof {1, 199}
+; LATCH !prof {199, 9601}
+
+define void @no_bfi(ptr %p, i32 %n, i32 %cond) !prof !14 {
+entry:
+ %c0 = icmp ne i32 %cond, 0
+ br i1 %c0, label %ph, label %ret
+
+ph: ; preds = %entry
+ br label %header
+
+header: ; preds = %latch, %ph
+ %iv = phi i32 [ 0, %ph ], [ %iv.next, %latch ]
+ %hcmp = icmp sge i32 %iv, %n
+ br i1 %hcmp, label %exit, label %body, !prof !15
+
+body: ; preds = %header
+ %addr = getelementptr i32, ptr %p, i32 %iv
+ store i32 %iv, ptr %addr, align 4
+ br label %latch
+
+latch: ; preds = %body
+ %iv.next = add i32 %iv, 1
+ br label %header
+
+exit: ; preds = %header
+ ret void
+
+ret: ; preds = %entry
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 10000}
+!4 = !{!"MaxCount", i64 9800}
+!5 = !{!"MaxInternalCount", i64 9800}
+!6 = !{!"MaxFunctionCount", i64 200}
+!7 = !{!"NumCounts", i64 3}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 9800, i32 1}
+!12 = !{i32 999000, i64 200, i32 1}
+!13 = !{i32 999999, i64 200, i32 1}
+!14 = !{!"function_entry_count", i64 201}
+!15 = !{!"branch_weights", i32 200, i32 9800}
+
+; CHECK-LABEL: define void @no_bfi(
+
+; Preheader-bypass guard: derived from the old single-exit formula because
+; BlockFrequencyInfo is not available to read PreHeaderEntries from.
+; CHECK: ph:
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}.lr.ph, !prof [[GUARD:![0-9]+]]
+
+; Rotated latch test.
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %body, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[GUARD]] = !{!"branch_weights", i32 1, i32 199}
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 199, i32 9601}
diff --git a/llvm/test/Transforms/LoopRotate/saturating-backedge-branch-weights.ll b/llvm/test/Transforms/LoopRotate/saturating-backedge-branch-weights.ll
new file mode 100644
index 0000000000000..c56c130bb5fc9
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/saturating-backedge-branch-weights.ll
@@ -0,0 +1,126 @@
+; RUN: opt < %s -passes='loop-rotate<update-branch-weights>' -S | FileCheck %s
+;
+; Covers the per-block balance floor in updateBranchWeights() that protects
+; the saturating subtraction on LoopBackWeight:
+;
+; LoopBackWeight = OrigLoopBackedgeWeight >= EnterWeight
+; ? OrigLoopBackedgeWeight - EnterWeight : 0;
+;
+; Without the floor, an inconsistent profile where
+; EnterWeight = PreHeaderEntries - ExitWeight0 > OrigLoopBackedgeWeight
+; clamps LoopBackWeight to 0 but leaves the rotated body's incoming weight
+; (EnterWeight + LoopBackWeight) larger than its outgoing weight, breaking
+; per-block weight conservation. The floor lifts ExitWeight0 to
+; max(ExitWeight0, PreHeaderEntries - OrigLoopBackedgeWeight)
+; (capped at OrigLoopExitWeight) so EnterWeight stays within
+; OrigLoopBackedgeWeight and the saturating branch is never taken.
+;
+; CFG (weights shown as {taken, not-taken}; each block's incoming branch
+; weight equals the sum of its outgoing branch weights):
+;
+; entry --(c0)--> ph !prof { 999, 1 } ; PreHeaderEntries = 999
+; ph --------> header ; uncond, weight 999
+; header --(hcmp)--> exit1 / body !prof { 990, 10 }
+; body --(bcmp)--> exit2 / latch !prof { 9, 1 } ; second loop exit
+; latch ---------> header ; uncond, weight 1
+;
+; OrigLoopExitWeight (990) > OrigLoopBackedgeWeight (10), so the
+; "0-trip and 1-trip only" arm of the ExitWeight0 guess fires first:
+;
+; ExitWeight0(initial) = OrigLoopExitWeight - OrigLoopBackedgeWeight
+; = 990 - 10 = 980
+;
+; Floor (this fix): with PreHeaderEntries (999) > OrigBackedge (10),
+;
+; MinExitWeight0 = PreHeaderEntries - OrigLoopBackedgeWeight
+; = 999 - 10 = 989
+; ExitWeight0 = max(ExitWeight0, MinExitWeight0) = max(980, 989) = 989
+;
+; Remaining computation:
+;
+; ExitWeight1 = OrigLoopExitWeight - ExitWeight0 = 990 - 989 = 1
+; EnterWeight = PreHeaderEntries - ExitWeight0 = 999 - 989 = 10
+; LoopBackWeight = OrigBackedge - EnterWeight = 10 - 10 = 0
+; ; no longer saturated; EnterWeight == OrigBackedge.
+;
+; GUARD !prof { 989, 10 }
+; LATCH !prof { 1, 0 }
+;
+; Per-block conservation after rotation:
+; ph : in 999, out 989 + 10 = 999
+; .lr.ph: in 10, out 10
+; body : in 10 (.lr.ph) + 0 (latch) = 10, out 9 + 1 = 10
+; latch : in 1, out 1 + 0 = 1
+; exit1 : in 989 (ph) + 1 (latch) = 990
+; exit2 : in 9
+; ret : in 1
+; exits : 990 + 9 + 1 = 1000 = function_entry_count
+
+define void @sat_backedge(ptr %p, i32 %n, i32 %cond) !prof !14 {
+entry:
+ %c0 = icmp ne i32 %cond, 0
+ br i1 %c0, label %ph, label %ret, !prof !15
+
+ph: ; preds = %entry
+ br label %header
+
+header: ; preds = %latch, %ph
+ %iv = phi i32 [ 0, %ph ], [ %iv.next, %latch ]
+ %hcmp = icmp sge i32 %iv, %n
+ br i1 %hcmp, label %exit1, label %body, !prof !16
+
+body: ; preds = %header
+ %addr = getelementptr i32, ptr %p, i32 %iv
+ %v = load i32, ptr %addr, align 4
+ %bcmp = icmp slt i32 %v, 0
+ br i1 %bcmp, label %exit2, label %latch, !prof !17
+
+latch: ; preds = %body
+ %iv.next = add i32 %iv, 1
+ br label %header
+
+exit1: ; preds = %header
+ ret void
+
+exit2: ; preds = %body
+ ret void
+
+ret: ; preds = %entry
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!3 = !{!"TotalCount", i64 11000}
+!4 = !{!"MaxCount", i64 990}
+!5 = !{!"MaxInternalCount", i64 990}
+!6 = !{!"MaxFunctionCount", i64 1000}
+!7 = !{!"NumCounts", i64 4}
+!8 = !{!"NumFunctions", i64 1}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 990, i32 1}
+!12 = !{i32 999000, i64 1000, i32 1}
+!13 = !{i32 999999, i64 990, i32 1}
+!14 = !{!"function_entry_count", i64 1000}
+!15 = !{!"branch_weights", i32 999, i32 1}
+!16 = !{!"branch_weights", i32 990, i32 10}
+!17 = !{!"branch_weights", i32 9, i32 1}
+
+; CHECK-LABEL: define void @sat_backedge(
+
+; Guard at the new preheader: ExitWeight0 floored to
+; PreHeaderEntries - OrigBackedge = 999 - 10 = 989; EnterWeight = 10.
+; CHECK: ph:
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %{{.*}}.lr.ph, !prof [[GUARD:![0-9]+]]
+
+; Rotated latch: EnterWeight == OrigBackedge so LoopBackWeight = 0 without
+; tripping the saturating subtraction; the rotated body's incoming weight
+; (10) matches its outgoing weight (9 + 1).
+; CHECK: br i1 %{{.*}}, label %{{.*}}, label %body, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[GUARD]] = !{!"branch_weights", i32 989, i32 10}
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 1, i32 0}
diff --git a/llvm/test/Transforms/LoopRotate/scaled-preheader-branch-weights.ll b/llvm/test/Transforms/LoopRotate/scaled-preheader-branch-weights.ll
new file mode 100644
index 0000000000000..1b7f84214d7e2
--- /dev/null
+++ b/llvm/test/Transforms/LoopRotate/scaled-preheader-branch-weights.ll
@@ -0,0 +1,63 @@
+; RUN: opt < %s -passes='loop-rotate<update-branch-weights>' -S | FileCheck %s
+;
+; Covers interaction between PreHeaderEntries and the ZeroTripCountWeights
+; scale-up path. When header weights are small, updateBranchWeights scales
+; OrigLoopExitWeight and OrigLoopBackedgeWeight (here ×128) to represent the
+; zero-trip ratio. PreHeaderEntries must be scaled by the same factor; otherwise
+; the rotated guard mixes scaled header weights with unscaled entry counts.
+;
+; Original:
+; entry -> ph / ret !prof {1, 1} ; PreHeaderEntries = 1
+; header -> body/exit !prof {100, 1}
+;
+; After scaling (×128):
+; PreHeaderEntries = 128
+; OrigLoopExitWeight = 128
+; OrigLoopBackedgeWeight = 12800
+;
+; Expected post-rotation:
+; GUARD !prof {1, 127}
+; LATCH !prof {127, 12673}
+;
+; Per-loop flow is conserved after scaling.
+
+define void @small_scaled(ptr %p, i32 %n, i1 %c) !prof !0 {
+entry:
+ br i1 %c, label %ph, label %ret, !prof !1
+
+ph:
+ br label %header
+
+header:
+ %iv = phi i32 [ 0, %ph ], [ %next, %body ]
+ %cmp = icmp slt i32 %iv, %n
+ br i1 %cmp, label %body, label %exit, !prof !2
+
+body:
+ store volatile i32 %iv, ptr %p, align 4
+ %next = add i32 %iv, 1
+ br label %header
+
+exit:
+ ret void
+
+ret:
+ ret void
+}
+
+!0 = !{!"function_entry_count", i64 2}
+!1 = !{!"branch_weights", i32 1, i32 1}
+!2 = !{!"branch_weights", i32 100, i32 1}
+
+; CHECK-LABEL: define void @small_scaled(
+
+; The guard keeps the hot loop-entry edge hot after scaling PreHeaderEntries.
+; CHECK: ph:
+; CHECK: br i1 %{{.*}}, label %body.lr.ph, label %exit, !prof [[GUARD:![0-9]+]]
+
+; The latch uses the same scaled unit as the guard and original header weights.
+; CHECK: body:
+; CHECK: br i1 %{{.*}}, label %body, label %header.exit_crit_edge, !prof [[LATCH:![0-9]+]]
+
+; CHECK-DAG: [[GUARD]] = !{!"branch_weights", i32 127, i32 1}
+; CHECK-DAG: [[LATCH]] = !{!"branch_weights", i32 12673, i32 127}
diff --git a/llvm/test/Transforms/LoopRotate/update-branch-weights.ll b/llvm/test/Transforms/LoopRotate/update-branch-weights.ll
index 9a1f36ec5ff2b..2319d78943d76 100644
--- a/llvm/test/Transforms/LoopRotate/update-branch-weights.ll
+++ b/llvm/test/Transforms/LoopRotate/update-branch-weights.ll
@@ -187,7 +187,8 @@ loop_exit:
; IR: br i1 %cmp1, label %loop_exit, label %loop_body.lr.ph, !prof [[PROF_FUNC4_0:![0-9]+]]
; IR: loop_body:
-; IR: br i1 %cmp, label %loop_header.loop_exit_crit_edge, label %loop_body, !prof [[PROF_FUNC4_0]]
+; IR-NOT: !prof
+; IR: br i1 %cmp, label %loop_header.loop_exit_crit_edge, label %loop_body{{$}}
define void @func4_zero_branch_weight(i32 %n) !prof !3 {
entry:
@@ -212,7 +213,8 @@ loop_exit:
; IR: br i1 %cmp1, label %loop_exit, label %loop_body.lr.ph, !prof [[PROF_FUNC5_0:![0-9]+]]
; IR: loop_body:
-; IR: br i1 %cmp, label %loop_header.loop_exit_crit_edge, label %loop_body, !prof [[PROF_FUNC5_0]]
+; IR-NOT: !prof
+; IR: br i1 %cmp, label %loop_header.loop_exit_crit_edge, label %loop_body{{$}}
define void @func5_zero_branch_weight(i32 %n) !prof !3 {
entry:
More information about the llvm-commits
mailing list