Skip to content

[SimplifyCFG][PGO] Reuse existing setBranchWeights #160629

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Oct 1, 2025
Merged

[SimplifyCFG][PGO] Reuse existing setBranchWeights #160629

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
7 changes: 3 additions & 4 deletions llvm/include/llvm/IR/Instructions.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -32,6 +32,7 @@
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Intrinsics.h"
#include "llvm/IR/OperandTraits.h"
#include "llvm/IR/ProfDataUtils.h"
#include "llvm/IR/Use.h"
#include "llvm/IR/User.h"
#include "llvm/Support/AtomicOrdering.h"
Expand Down Expand Up @@ -3536,8 +3537,6 @@ class SwitchInstProfUpdateWrapper {
bool Changed = false;

protected:
LLVM_ABI MDNode *buildProfBranchWeightsMD();

LLVM_ABI void init();

public:
Expand All @@ -3549,8 +3548,8 @@ class SwitchInstProfUpdateWrapper {
SwitchInstProfUpdateWrapper(SwitchInst &SI) : SI(SI) { init(); }

~SwitchInstProfUpdateWrapper() {
if (Changed)
SI.setMetadata(LLVMContext::MD_prof, buildProfBranchWeightsMD());
if (Changed && Weights.has_value() && Weights->size() >= 2)
setBranchWeights(SI, Weights.value(), /*IsExpected=*/false);
}

/// Delegate the call to the underlying SwitchInst::removeCase() and remove
Expand Down
8 changes: 7 additions & 1 deletion llvm/include/llvm/IR/ProfDataUtils.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -145,7 +145,13 @@ LLVM_ABI bool extractProfTotalWeight(const Instruction &I,
/// \param Weights an array of weights to set on instruction I.
/// \param IsExpected were these weights added from an llvm.expect* intrinsic.
LLVM_ABI void setBranchWeights(Instruction &I, ArrayRef<uint32_t> Weights,
bool IsExpected);
bool IsExpected, bool ElideAllZero = false);

/// Variant of `setBranchWeights` where the `Weights` will be fit first to
/// uint32_t by shifting right.
LLVM_ABI void setFittedBranchWeights(Instruction &I, ArrayRef<uint64_t> Weights,
bool IsExpected,
bool ElideAllZero = false);

/// downscale the given weights preserving the ratio. If the maximum value is
/// not already known and not provided via \param KnownMaxCount , it will be
Expand Down
17 changes: 0 additions & 17 deletions llvm/lib/IR/Instructions.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4141,23 +4141,6 @@ void SwitchInst::growOperands() {
growHungoffUses(ReservedSpace);
}

MDNode *SwitchInstProfUpdateWrapper::buildProfBranchWeightsMD() {
assert(Changed && "called only if metadata has changed");

if (!Weights)
return nullptr;

assert(SI.getNumSuccessors() == Weights->size() &&
"num of prof branch_weights must accord with num of successors");

bool AllZeroes = all_of(*Weights, [](uint32_t W) { return W == 0; });

if (AllZeroes || Weights->size() < 2)
return nullptr;

return MDBuilder(SI.getParent()->getContext()).createBranchWeights(*Weights);
}

void SwitchInstProfUpdateWrapper::init() {
MDNode *ProfileData = getBranchWeightMDNode(SI);
if (!ProfileData)
Expand Down
38 changes: 36 additions & 2 deletions llvm/lib/IR/ProfDataUtils.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,13 +12,15 @@

#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"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/LLVMContext.h"
#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Metadata.h"
#include "llvm/Support/CommandLine.h"

using namespace llvm;

Expand Down Expand Up @@ -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());
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

nit: SmallVector has a constructor that accepts a size_t, so this can be combined with the previous line.

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

that resizes. I wanted reserve, so I can use push_back

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";
Expand Down Expand Up @@ -282,12 +305,23 @@ bool hasExplicitlyUnknownBranchWeights(const Instruction &I) {
}

void setBranchWeights(Instruction &I, ArrayRef<uint32_t> Weights,
bool IsExpected) {
bool IsExpected, bool ElideAllZero) {
if ((ElideAllZeroBranchWeights && ElideAllZero) &&
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should be || or && ? With &&, the zeros can not be elided (and default behavior is changed).

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

ElideAllZeroBranchWeights is true by default. It's only false when profcheck is enabled.

Then, the SimplifyCFG callers, which were the place where zero-elision was happening, call the new API with ElideAllZero = true

Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should the default parameter to be true instead of 'false'?

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

should the default parameter to be true instead of 'false'?

no, because the behavior elsewhere was to not elide. Only SimplifyCFG had an eliding variant (I mean, there could be other passes with local variants, and we'll address accordingly)

Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I could rename the compiler flag to ElideAllZeroBranchWeightsIfRequested for clarity.

llvm::all_of(Weights, [](uint32_t V) { return V == 0; })) {
I.setMetadata(LLVMContext::MD_prof, nullptr);
return;
}

MDBuilder MDB(I.getContext());
MDNode *BranchWeights = MDB.createBranchWeights(Weights, IsExpected);
I.setMetadata(LLVMContext::MD_prof, BranchWeights);
}

void setFittedBranchWeights(Instruction &I, ArrayRef<uint64_t> Weights,
bool IsExpected, bool ElideAllZero) {
setBranchWeights(I, fitWeights(Weights), IsExpected, ElideAllZero);
}

SmallVector<uint32_t> downscaleWeights(ArrayRef<uint64_t> Weights,
std::optional<uint64_t> KnownMaxCount) {
uint64_t MaxCount = KnownMaxCount.has_value() ? KnownMaxCount.value()
Expand Down
8 changes: 5 additions & 3 deletions llvm/lib/Transforms/IPO/SampleProfile.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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) {
Expand All @@ -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);
}
}
}
Expand Down
4 changes: 2 additions & 2 deletions llvm/lib/Transforms/Instrumentation/IndirectCallPromotion.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -672,8 +672,8 @@ CallBase &llvm::pgo::promoteIndirectCall(CallBase &CB, Function *DirectCallee,
createBranchWeights(CB.getContext(), Count, TotalCount - Count));

if (AttachProfToDirectCall)
setBranchWeights(NewInst, {static_cast<uint32_t>(Count)},
/*IsExpected=*/false);
setFittedBranchWeights(NewInst, {Count},
/*IsExpected=*/false);

using namespace ore;

Expand Down
Loading