[llvm] [SimplifyCFG][profcheck] Fix artificially-failing `preserve-branchweights.ll` (PR #160629)
Mircea Trofin via llvm-commits
llvm-commits at lists.llvm.org
Mon Sep 29 12:34:30 PDT 2025
https://github.com/mtrofin updated https://github.com/llvm/llvm-project/pull/160629
>From e5a8d3f2f264ba242373ab9b94ba0a6a93b43976 Mon Sep 17 00:00:00 2001
From: Mircea Trofin <mtrofin at google.com>
Date: Thu, 25 Sep 2025 02:18:30 +0000
Subject: [PATCH] [SimplifyCFG][profcheck] Fix artificially-failing
`preserve-branchweights.ll`
---
llvm/include/llvm/IR/ProfDataUtils.h | 6 ++
llvm/lib/IR/ProfDataUtils.cpp | 39 +++++++-
llvm/lib/Transforms/IPO/SampleProfile.cpp | 8 +-
.../Instrumentation/IndirectCallPromotion.cpp | 2 +-
llvm/lib/Transforms/Utils/ProfileVerify.cpp | 5 +-
llvm/lib/Transforms/Utils/SimplifyCFG.cpp | 98 +++++--------------
llvm/utils/profcheck-xfail.txt | 1 -
7 files changed, 76 insertions(+), 83 deletions(-)
diff --git a/llvm/include/llvm/IR/ProfDataUtils.h b/llvm/include/llvm/IR/ProfDataUtils.h
index de9675f48c79b..d3c9a50c98234 100644
--- a/llvm/include/llvm/IR/ProfDataUtils.h
+++ b/llvm/include/llvm/IR/ProfDataUtils.h
@@ -147,6 +147,12 @@ LLVM_ABI bool extractProfTotalWeight(const Instruction &I,
LLVM_ABI void setBranchWeights(Instruction &I, ArrayRef<uint32_t> Weights,
bool IsExpected);
+LLVM_ABI void setBranchWeights(Instruction &I, ArrayRef<uint64_t> Weights,
+ bool IsExpected);
+
+LLVM_ABI void setBranchWeights(Instruction &I, uint64_t TrueWeight,
+ uint64_t FalseWeight, bool IsExpected);
+
/// downscale the given weights preserving the ratio. If the maximum value is
/// not already known and not provided via \param KnownMaxCount , it will be
/// obtained from \param Weights.
diff --git a/llvm/lib/IR/ProfDataUtils.cpp b/llvm/lib/IR/ProfDataUtils.cpp
index 5827292cee39b..5cdb20abbb0a8 100644
--- a/llvm/lib/IR/ProfDataUtils.cpp
+++ b/llvm/lib/IR/ProfDataUtils.cpp
@@ -12,6 +12,7 @@
#include "llvm/IR/ProfDataUtils.h"
+#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Function.h"
@@ -19,6 +20,7 @@
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Metadata.h"
+#include "llvm/Support/CommandLine.h"
using namespace llvm;
@@ -84,10 +86,31 @@ static void extractFromBranchWeightMD(const MDNode *ProfileData,
}
}
+/// Push the weights right to fit in uint32_t.
+static SmallVector<uint32_t> fitWeights(ArrayRef<uint64_t> Weights) {
+ SmallVector<uint32_t> Ret;
+ Ret.reserve(Weights.size());
+ uint64_t Max = *llvm::max_element(Weights);
+ if (Max > UINT_MAX) {
+ unsigned Offset = 32 - llvm::countl_zero(Max);
+ for (const uint64_t &Value : Weights)
+ Ret.push_back(static_cast<uint32_t>(Value >> Offset));
+ } else {
+ append_range(Ret, Weights);
+ }
+ return Ret;
+}
+
} // namespace
namespace llvm {
-
+cl::opt<bool> ElideAllZeroBranchWeights("elide-all-zero-branch-weights",
+#if defined(LLVM_ENABLE_PROFCHECK)
+ cl::init(false)
+#else
+ cl::init(true)
+#endif
+);
const char *MDProfLabels::BranchWeights = "branch_weights";
const char *MDProfLabels::ExpectedBranchWeights = "expected";
const char *MDProfLabels::ValueProfile = "VP";
@@ -276,11 +299,25 @@ bool hasExplicitlyUnknownBranchWeights(const Instruction &I) {
void setBranchWeights(Instruction &I, ArrayRef<uint32_t> Weights,
bool IsExpected) {
+ if (ElideAllZeroBranchWeights &&
+ llvm::all_of(Weights, [](uint32_t V) { return V == 0; }))
+ return;
MDBuilder MDB(I.getContext());
MDNode *BranchWeights = MDB.createBranchWeights(Weights, IsExpected);
I.setMetadata(LLVMContext::MD_prof, BranchWeights);
}
+void setBranchWeights(Instruction &I, ArrayRef<uint64_t> Weights,
+ bool IsExpected) {
+ setBranchWeights(I, fitWeights(Weights), IsExpected);
+}
+
+void setBranchWeights(Instruction &I, uint64_t TrueWeight, uint64_t FalseWeight,
+ bool IsExpected) {
+ return setBranchWeights(I, ArrayRef<uint64_t>{TrueWeight, FalseWeight},
+ IsExpected);
+}
+
SmallVector<uint32_t> downscaleWeights(ArrayRef<uint64_t> Weights,
std::optional<uint64_t> KnownMaxCount) {
uint64_t MaxCount = KnownMaxCount.has_value() ? KnownMaxCount.value()
diff --git a/llvm/lib/Transforms/IPO/SampleProfile.cpp b/llvm/lib/Transforms/IPO/SampleProfile.cpp
index 5bc7e34938127..99b8b88ebedbb 100644
--- a/llvm/lib/Transforms/IPO/SampleProfile.cpp
+++ b/llvm/lib/Transforms/IPO/SampleProfile.cpp
@@ -1664,8 +1664,9 @@ void SampleProfileLoader::generateMDProfMetadata(Function &F) {
else if (OverwriteExistingWeights)
I.setMetadata(LLVMContext::MD_prof, nullptr);
} else if (!isa<IntrinsicInst>(&I)) {
- setBranchWeights(I, {static_cast<uint32_t>(BlockWeights[BB])},
- /*IsExpected=*/false);
+ setBranchWeights(
+ I, ArrayRef<uint32_t>{static_cast<uint32_t>(BlockWeights[BB])},
+ /*IsExpected=*/false);
}
}
} else if (OverwriteExistingWeights || ProfileSampleBlockAccurate) {
@@ -1676,7 +1677,8 @@ void SampleProfileLoader::generateMDProfMetadata(Function &F) {
if (cast<CallBase>(I).isIndirectCall()) {
I.setMetadata(LLVMContext::MD_prof, nullptr);
} else {
- setBranchWeights(I, {uint32_t(0)}, /*IsExpected=*/false);
+ setBranchWeights(I, ArrayRef<uint32_t>{uint32_t(0)},
+ /*IsExpected=*/false);
}
}
}
diff --git a/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp b/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
index f451c2b471aa6..5d21bac1db4c8 100644
--- a/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
+++ b/llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
@@ -672,7 +672,7 @@ CallBase &llvm::pgo::promoteIndirectCall(CallBase &CB, Function *DirectCallee,
createBranchWeights(CB.getContext(), Count, TotalCount - Count));
if (AttachProfToDirectCall)
- setBranchWeights(NewInst, {static_cast<uint32_t>(Count)},
+ setBranchWeights(NewInst, ArrayRef<uint32_t>{static_cast<uint32_t>(Count)},
/*IsExpected=*/false);
using namespace ore;
diff --git a/llvm/lib/Transforms/Utils/ProfileVerify.cpp b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
index c578b4b839258..72d3dcba85dea 100644
--- a/llvm/lib/Transforms/Utils/ProfileVerify.cpp
+++ b/llvm/lib/Transforms/Utils/ProfileVerify.cpp
@@ -103,8 +103,9 @@ bool ProfileInjector::inject() {
if (AnnotateSelect) {
for (auto &I : BB) {
if (isa<SelectInst>(I) && !I.getMetadata(LLVMContext::MD_prof))
- setBranchWeights(I, {SelectTrueWeight, SelectFalseWeight},
- /*IsExpected=*/false);
+ setBranchWeights(
+ I, ArrayRef<uint32_t>{SelectTrueWeight, SelectFalseWeight},
+ /*IsExpected=*/false);
}
}
auto *Term = getTerminatorBenefitingFromMDProf(BB);
diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 216bdf4eb9efb..964d9345f575f 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -955,33 +955,6 @@ static bool valuesOverlap(std::vector<ValueEqualityComparisonCase> &C1,
return false;
}
-// Set branch weights on SwitchInst. This sets the metadata if there is at
-// least one non-zero weight.
-static void setBranchWeights(SwitchInst *SI, ArrayRef<uint32_t> Weights,
- bool IsExpected) {
- // Check that there is at least one non-zero weight. Otherwise, pass
- // nullptr to setMetadata which will erase the existing metadata.
- MDNode *N = nullptr;
- if (llvm::any_of(Weights, [](uint32_t W) { return W != 0; }))
- N = MDBuilder(SI->getParent()->getContext())
- .createBranchWeights(Weights, IsExpected);
- SI->setMetadata(LLVMContext::MD_prof, N);
-}
-
-// Similar to the above, but for branch and select instructions that take
-// exactly 2 weights.
-static void setBranchWeights(Instruction *I, uint32_t TrueWeight,
- uint32_t FalseWeight, bool IsExpected) {
- assert(isa<BranchInst>(I) || isa<SelectInst>(I));
- // Check that there is at least one non-zero weight. Otherwise, pass
- // nullptr to setMetadata which will erase the existing metadata.
- MDNode *N = nullptr;
- if (TrueWeight || FalseWeight)
- N = MDBuilder(I->getParent()->getContext())
- .createBranchWeights(TrueWeight, FalseWeight, IsExpected);
- I->setMetadata(LLVMContext::MD_prof, N);
-}
-
/// If TI is known to be a terminator instruction and its block is known to
/// only have a single predecessor block, check to see if that predecessor is
/// also a value comparison with the same value, and if that comparison
@@ -1181,16 +1154,6 @@ static void getBranchWeights(Instruction *TI,
}
}
-/// Keep halving the weights until all can fit in uint32_t.
-static void fitWeights(MutableArrayRef<uint64_t> Weights) {
- uint64_t Max = *llvm::max_element(Weights);
- if (Max > UINT_MAX) {
- unsigned Offset = 32 - llvm::countl_zero(Max);
- for (uint64_t &I : Weights)
- I >>= Offset;
- }
-}
-
static void cloneInstructionsIntoPredecessorBlockAndUpdateSSAUses(
BasicBlock *BB, BasicBlock *PredBlock, ValueToValueMapTy &VMap) {
Instruction *PTI = PredBlock->getTerminator();
@@ -1446,14 +1409,8 @@ bool SimplifyCFGOpt::performValueComparisonIntoPredecessorFolding(
for (ValueEqualityComparisonCase &V : PredCases)
NewSI->addCase(V.Value, V.Dest);
- if (PredHasWeights || SuccHasWeights) {
- // Halve the weights if any of them cannot fit in an uint32_t
- fitWeights(Weights);
-
- SmallVector<uint32_t, 8> MDWeights(Weights.begin(), Weights.end());
-
- setBranchWeights(NewSI, MDWeights, /*IsExpected=*/false);
- }
+ if (PredHasWeights || SuccHasWeights)
+ setBranchWeights(*NewSI, Weights, /*IsExpected=*/false);
eraseTerminatorAndDCECond(PTI);
@@ -4053,39 +4010,33 @@ static bool performBranchToCommonDestFolding(BranchInst *BI, BranchInst *PBI,
// Try to update branch weights.
uint64_t PredTrueWeight, PredFalseWeight, SuccTrueWeight, SuccFalseWeight;
- SmallVector<uint32_t, 2> MDWeights;
+ SmallVector<uint64_t, 2> MDWeights;
if (extractPredSuccWeights(PBI, BI, PredTrueWeight, PredFalseWeight,
SuccTrueWeight, SuccFalseWeight)) {
- SmallVector<uint64_t, 8> NewWeights;
if (PBI->getSuccessor(0) == BB) {
// PBI: br i1 %x, BB, FalseDest
// BI: br i1 %y, UniqueSucc, FalseDest
// TrueWeight is TrueWeight for PBI * TrueWeight for BI.
- NewWeights.push_back(PredTrueWeight * SuccTrueWeight);
+ MDWeights.push_back(PredTrueWeight * SuccTrueWeight);
// FalseWeight is FalseWeight for PBI * TotalWeight for BI +
// TrueWeight for PBI * FalseWeight for BI.
// We assume that total weights of a BranchInst can fit into 32 bits.
// Therefore, we will not have overflow using 64-bit arithmetic.
- NewWeights.push_back(PredFalseWeight *
- (SuccFalseWeight + SuccTrueWeight) +
- PredTrueWeight * SuccFalseWeight);
+ MDWeights.push_back(PredFalseWeight * (SuccFalseWeight + SuccTrueWeight) +
+ PredTrueWeight * SuccFalseWeight);
} else {
// PBI: br i1 %x, TrueDest, BB
// BI: br i1 %y, TrueDest, UniqueSucc
// TrueWeight is TrueWeight for PBI * TotalWeight for BI +
// FalseWeight for PBI * TrueWeight for BI.
- NewWeights.push_back(PredTrueWeight * (SuccFalseWeight + SuccTrueWeight) +
- PredFalseWeight * SuccTrueWeight);
+ MDWeights.push_back(PredTrueWeight * (SuccFalseWeight + SuccTrueWeight) +
+ PredFalseWeight * SuccTrueWeight);
// FalseWeight is FalseWeight for PBI * FalseWeight for BI.
- NewWeights.push_back(PredFalseWeight * SuccFalseWeight);
+ MDWeights.push_back(PredFalseWeight * SuccFalseWeight);
}
- // Halve the weights if any of them cannot fit in an uint32_t
- fitWeights(NewWeights);
-
- append_range(MDWeights, NewWeights);
- setBranchWeights(PBI, MDWeights[0], MDWeights[1], /*IsExpected=*/false);
+ setBranchWeights(*PBI, MDWeights, /*IsExpected=*/false);
// TODO: If BB is reachable from all paths through PredBlock, then we
// could replace PBI's branch probabilities with BI's.
@@ -4125,7 +4076,7 @@ static bool performBranchToCommonDestFolding(BranchInst *BI, BranchInst *PBI,
if (auto *SI = dyn_cast<SelectInst>(PBI->getCondition()))
if (!MDWeights.empty()) {
assert(isSelectInRoleOfConjunctionOrDisjunction(SI));
- setBranchWeights(SI, MDWeights[0], MDWeights[1],
+ setBranchWeights(*SI, MDWeights[0], MDWeights[1],
/*IsExpected=*/false);
}
@@ -4478,7 +4429,7 @@ static bool mergeConditionalStoreToAddress(
if (InvertQCond)
std::swap(QWeights[0], QWeights[1]);
auto CombinedWeights = getDisjunctionWeights(PWeights, QWeights);
- setBranchWeights(PostBB->getTerminator(), CombinedWeights[0],
+ setBranchWeights(*PostBB->getTerminator(), CombinedWeights[0],
CombinedWeights[1],
/*IsExpected=*/false);
}
@@ -4836,10 +4787,8 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
uint64_t NewWeights[2] = {PredCommon * (SuccCommon + SuccOther) +
PredOther * SuccCommon,
PredOther * SuccOther};
- // Halve the weights if any of them cannot fit in an uint32_t
- fitWeights(NewWeights);
- setBranchWeights(PBI, NewWeights[0], NewWeights[1], /*IsExpected=*/false);
+ setBranchWeights(*PBI, NewWeights, /*IsExpected=*/false);
// Cond may be a select instruction with the first operand set to "true", or
// the second to "false" (see how createLogicalOp works for `and` and `or`)
if (!ProfcheckDisableMetadataFixes)
@@ -4849,7 +4798,7 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
assert(dyn_cast<SelectInst>(SI)->getCondition() == PBICond);
// The corresponding probabilities are what was referred to above as
// PredCommon and PredOther.
- setBranchWeights(SI, PredCommon, PredOther,
+ setBranchWeights(*SI, PredCommon, PredOther,
/*IsExpected=*/false);
}
}
@@ -4876,7 +4825,7 @@ static bool SimplifyCondBranchToCondBranch(BranchInst *PBI, BranchInst *BI,
if (HasWeights) {
uint64_t TrueWeight = PBIOp ? PredFalseWeight : PredTrueWeight;
uint64_t FalseWeight = PBIOp ? PredTrueWeight : PredFalseWeight;
- setBranchWeights(NV, TrueWeight, FalseWeight,
+ setBranchWeights(*NV, TrueWeight, FalseWeight,
/*IsExpected=*/false);
}
}
@@ -4940,7 +4889,7 @@ bool SimplifyCFGOpt::simplifyTerminatorOnSelect(Instruction *OldTerm,
// Create a conditional branch sharing the condition of the select.
BranchInst *NewBI = Builder.CreateCondBr(Cond, TrueBB, FalseBB);
if (TrueWeight != FalseWeight)
- setBranchWeights(NewBI, TrueWeight, FalseWeight, /*IsExpected=*/false);
+ setBranchWeights(*NewBI, TrueWeight, FalseWeight, /*IsExpected=*/false);
}
} else if (KeepEdge1 && (KeepEdge2 || TrueBB == FalseBB)) {
// Neither of the selected blocks were successors, so this
@@ -5889,7 +5838,7 @@ bool SimplifyCFGOpt::turnSwitchRangeIntoICmp(SwitchInst *SI,
TrueWeight /= 2;
FalseWeight /= 2;
}
- setBranchWeights(NewBI, TrueWeight, FalseWeight, /*IsExpected=*/false);
+ setBranchWeights(*NewBI, TrueWeight, FalseWeight, /*IsExpected=*/false);
}
}
@@ -6364,7 +6313,7 @@ static Value *foldSwitchToSelect(const SwitchCaseResultVectorTy &ResultVector,
// BranchWeights. We want the probability and negative probability of
// Condition == SecondCase.
assert(BranchWeights.size() == 3);
- setBranchWeights(SI, BranchWeights[2],
+ setBranchWeights(*SI, BranchWeights[2],
BranchWeights[0] + BranchWeights[1],
/*IsExpected=*/false);
}
@@ -6381,7 +6330,7 @@ static Value *foldSwitchToSelect(const SwitchCaseResultVectorTy &ResultVector,
size_t FirstCasePos = (Condition != nullptr);
size_t SecondCasePos = FirstCasePos + 1;
uint32_t DefaultCase = (Condition != nullptr) ? BranchWeights[0] : 0;
- setBranchWeights(SI, BranchWeights[FirstCasePos],
+ setBranchWeights(*SI, BranchWeights[FirstCasePos],
DefaultCase + BranchWeights[SecondCasePos],
/*IsExpected=*/false);
}
@@ -6427,7 +6376,7 @@ static Value *foldSwitchToSelect(const SwitchCaseResultVectorTy &ResultVector,
// We know there's a Default case. We base the resulting branch
// weights off its probability.
assert(BranchWeights.size() >= 2);
- setBranchWeights(SI, accumulate(drop_begin(BranchWeights), 0),
+ setBranchWeights(*SI, accumulate(drop_begin(BranchWeights), 0),
BranchWeights[0], /*IsExpected=*/false);
}
return Ret;
@@ -6451,7 +6400,7 @@ static Value *foldSwitchToSelect(const SwitchCaseResultVectorTy &ResultVector,
Builder.CreateSelect(Cmp, ResultVector[0].first, DefaultResult);
if (auto *SI = dyn_cast<SelectInst>(Ret); SI && HasBranchWeights) {
assert(BranchWeights.size() >= 2);
- setBranchWeights(SI, accumulate(drop_begin(BranchWeights), 0),
+ setBranchWeights(*SI, accumulate(drop_begin(BranchWeights), 0),
BranchWeights[0], /*IsExpected=*/false);
}
return Ret;
@@ -6469,7 +6418,7 @@ static Value *foldSwitchToSelect(const SwitchCaseResultVectorTy &ResultVector,
Builder.CreateSelect(Cmp, ResultVector[0].first, DefaultResult);
if (auto *SI = dyn_cast<SelectInst>(Ret); SI && HasBranchWeights) {
assert(BranchWeights.size() >= 2);
- setBranchWeights(SI, accumulate(drop_begin(BranchWeights), 0),
+ setBranchWeights(*SI, accumulate(drop_begin(BranchWeights), 0),
BranchWeights[0], /*IsExpected=*/false);
}
return Ret;
@@ -8152,8 +8101,7 @@ static bool mergeNestedCondBranch(BranchInst *BI, DomTreeUpdater *DTU) {
if (HasWeight) {
uint64_t Weights[2] = {BBTWeight * BB1FWeight + BBFWeight * BB2TWeight,
BBTWeight * BB1TWeight + BBFWeight * BB2FWeight};
- fitWeights(Weights);
- setBranchWeights(BI, Weights[0], Weights[1], /*IsExpected=*/false);
+ setBranchWeights(*BI, Weights, /*IsExpected=*/false);
}
return true;
}
diff --git a/llvm/utils/profcheck-xfail.txt b/llvm/utils/profcheck-xfail.txt
index bd6627b5b6158..ca7665f03a9d6 100644
--- a/llvm/utils/profcheck-xfail.txt
+++ b/llvm/utils/profcheck-xfail.txt
@@ -1863,7 +1863,6 @@ Transforms/SimplifyCFG/merge-cond-stores.ll
Transforms/SimplifyCFG/multiple-phis.ll
Transforms/SimplifyCFG/PhiBlockMerge.ll
Transforms/SimplifyCFG/pr48641.ll
-Transforms/SimplifyCFG/preserve-branchweights.ll
Transforms/SimplifyCFG/preserve-store-alignment.ll
Transforms/SimplifyCFG/rangereduce.ll
Transforms/SimplifyCFG/RISCV/select-trunc-i64.ll
More information about the llvm-commits
mailing list