Skip to content

[CodeGen] Preserve branch weights from PGO profile during instruction selection at -O0 #161620

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 7 commits into from
Oct 9, 2025
Merged

[CodeGen] Preserve branch weights from PGO profile during instruction selection at -O0 #161620

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c55f594
Calculate edge weights in MIR if PGO is available in O0
grigorypas Oct 1, 2025
6e4ebda
Register analysis only if PGO is present
grigorypas Oct 2, 2025
0af0751
Keep only pgo-kind option
grigorypas Oct 6, 2025
67050dd
Turn on PGO passes only if PGO enums are set to use
grigorypas Oct 6, 2025
293b095
Addressed comments
grigorypas Oct 7, 2025
a4773d9
Changed PGOOptions check for CSPGO
grigorypas Oct 7, 2025
b870ddd
Changed function name
grigorypas Oct 9, 2025
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
32 changes: 24 additions & 8 deletions llvm/lib/CodeGen/SelectionDAG/SelectionDAGISel.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -229,6 +229,19 @@ static bool dontUseFastISelFor(const Function &Fn) {
});
}

static bool maintainPGOProfile(const TargetMachine &TM,
CodeGenOptLevel OptLevel) {
if (OptLevel != CodeGenOptLevel::None)
return true;
if (TM.getPGOOption()) {
const PGOOptions &Options = *TM.getPGOOption();
return Options.Action == PGOOptions::PGOAction::IRUse ||
Options.Action == PGOOptions::PGOAction::SampleUse ||
Options.CSAction == PGOOptions::CSPGOAction::CSIRUse;
}
return false;
}

namespace llvm {

//===--------------------------------------------------------------------===//
Expand Down Expand Up @@ -390,6 +403,7 @@ SelectionDAGISel::~SelectionDAGISel() { delete CurDAG; }

void SelectionDAGISelLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
CodeGenOptLevel OptLevel = Selector->OptLevel;
bool RegisterPGOPasses = maintainPGOProfile(Selector->TM, Selector->OptLevel);
if (OptLevel != CodeGenOptLevel::None)
AU.addRequired<AAResultsWrapperPass>();
AU.addRequired<GCModuleInfo>();
Expand All @@ -398,15 +412,15 @@ void SelectionDAGISelLegacy::getAnalysisUsage(AnalysisUsage &AU) const {
AU.addRequired<TargetLibraryInfoWrapperPass>();
AU.addRequired<TargetTransformInfoWrapperPass>();
AU.addRequired<AssumptionCacheTracker>();
if (UseMBPI && OptLevel != CodeGenOptLevel::None)
AU.addRequired<BranchProbabilityInfoWrapperPass>();
if (UseMBPI && RegisterPGOPasses)
AU.addRequired<BranchProbabilityInfoWrapperPass>();
AU.addRequired<ProfileSummaryInfoWrapperPass>();
// AssignmentTrackingAnalysis only runs if assignment tracking is enabled for
// the module.
AU.addRequired<AssignmentTrackingAnalysis>();
AU.addPreserved<AssignmentTrackingAnalysis>();
if (OptLevel != CodeGenOptLevel::None)
LazyBlockFrequencyInfoPass::getLazyBFIAnalysisUsage(AU);
if (RegisterPGOPasses)
LazyBlockFrequencyInfoPass::getLazyBFIAnalysisUsage(AU);
MachineFunctionPass::getAnalysisUsage(AU);
}

Expand Down Expand Up @@ -459,6 +473,7 @@ void SelectionDAGISel::initializeAnalysisResults(
(void)MatchFilterFuncName;
#endif

bool RegisterPGOPasses = maintainPGOProfile(TM, OptLevel);
TII = MF->getSubtarget().getInstrInfo();
TLI = MF->getSubtarget().getTargetLowering();
RegInfo = &MF->getRegInfo();
Expand All @@ -469,7 +484,7 @@ void SelectionDAGISel::initializeAnalysisResults(
auto *PSI = MAMP.getCachedResult<ProfileSummaryAnalysis>(*Fn.getParent());
BlockFrequencyInfo *BFI = nullptr;
FAM.getResult<BlockFrequencyAnalysis>(Fn);
if (PSI && PSI->hasProfileSummary() && OptLevel != CodeGenOptLevel::None)
if (PSI && PSI->hasProfileSummary() && RegisterPGOPasses)
BFI = &FAM.getResult<BlockFrequencyAnalysis>(Fn);

FunctionVarLocs const *FnVarLocs = nullptr;
Expand All @@ -487,7 +502,7 @@ void SelectionDAGISel::initializeAnalysisResults(
// into account). That's unfortunate but OK because it just means we won't
// ask for passes that have been required anyway.

if (UseMBPI && OptLevel != CodeGenOptLevel::None)
if (UseMBPI && RegisterPGOPasses)
FuncInfo->BPI = &FAM.getResult<BranchProbabilityAnalysis>(Fn);
else
FuncInfo->BPI = nullptr;
Expand All @@ -513,6 +528,7 @@ void SelectionDAGISel::initializeAnalysisResults(MachineFunctionPass &MFP) {
(void)MatchFilterFuncName;
#endif

bool RegisterPGOPasses = maintainPGOProfile(TM, OptLevel);
TII = MF->getSubtarget().getInstrInfo();
TLI = MF->getSubtarget().getTargetLowering();
RegInfo = &MF->getRegInfo();
Expand All @@ -523,7 +539,7 @@ void SelectionDAGISel::initializeAnalysisResults(MachineFunctionPass &MFP) {
AC = &MFP.getAnalysis<AssumptionCacheTracker>().getAssumptionCache(Fn);
auto *PSI = &MFP.getAnalysis<ProfileSummaryInfoWrapperPass>().getPSI();
BlockFrequencyInfo *BFI = nullptr;
if (PSI && PSI->hasProfileSummary() && OptLevel != CodeGenOptLevel::None)
if (PSI && PSI->hasProfileSummary() && RegisterPGOPasses)
BFI = &MFP.getAnalysis<LazyBlockFrequencyInfoPass>().getBFI();

FunctionVarLocs const *FnVarLocs = nullptr;
Expand All @@ -544,7 +560,7 @@ void SelectionDAGISel::initializeAnalysisResults(MachineFunctionPass &MFP) {
// into account). That's unfortunate but OK because it just means we won't
// ask for passes that have been required anyway.

if (UseMBPI && OptLevel != CodeGenOptLevel::None)
if (UseMBPI && RegisterPGOPasses)
FuncInfo->BPI =
&MFP.getAnalysis<BranchProbabilityInfoWrapperPass>().getBPI();
else
Expand Down
49 changes: 49 additions & 0 deletions llvm/test/CodeGen/X86/pgo-profile-o0.ll
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,49 @@
; RUN: llc -mtriple=x86_64-- -O0 -pgo-kind=pgo-sample-use-pipeline -debug-pass=Structure %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=PASSES
; RUN: llc -mtriple=x86_64-- -O0 -pgo-kind=pgo-sample-use-pipeline -debug-only=branch-prob %s -o /dev/null 2>&1 | FileCheck %s --check-prefix=BRANCH_PROB
; RUN: llc -mtriple=x86_64-- -O0 -pgo-kind=pgo-sample-use-pipeline -stop-after=finalize-isel %s -o - | FileCheck %s --check-prefix=MIR

; REQUIRES: asserts

; This test verifies that PGO profile information (branch weights) is preserved
; during instruction selection at -O0.

; Test function with explicit branch weights from PGO.
define i32 @test_pgo_preservation(i32 %x) !prof !15 {
entry:
%cmp = icmp sgt i32 %x, 10
; This branch has bias: 97 taken vs 3 not taken
br i1 %cmp, label %if.then, label %if.else, !prof !16

if.then:
; Hot path - should have high frequency
%add = add nsw i32 %x, 100
br label %if.end

if.else:
; Cold path - should have low frequency
%sub = sub nsw i32 %x, 50
br label %if.end

if.end:
%result = phi i32 [ %add, %if.then ], [ %sub, %if.else ]
ret i32 %result
}

; Profile metadata with branch weights 97:3.
!15 = !{!"function_entry_count", i64 100}
!16 = !{!"branch_weights", i32 97, i32 3}

; Verify that Branch Probability Analysis runs at O0.
; PASSES: Branch Probability Analysis

; Verify that the branch probabilities reflect the exact profile data.
; BRANCH_PROB: ---- Branch Probability Info : test_pgo_preservation ----
; BRANCH_PROB: set edge entry -> 0 successor probability to {{.*}} = 97.00%
; BRANCH_PROB: set edge entry -> 1 successor probability to {{.*}} = 3.00%

; Verify that machine IR preserves the branch probabilities from profile data
; MIR: bb.0.entry:
; MIR-NEXT: successors: %bb.{{[0-9]+}}({{0x03d70a3d|0x7c28f5c3}}), %bb.{{[0-9]+}}({{0x7c28f5c3|0x03d70a3d}})
; The two successor probability values should be:
; - 0x7c28f5c3: approximately 97% (high probability successor)
; - 0x03d70a3d: approximately 3% (low probability successor)
40 changes: 40 additions & 0 deletions llvm/tools/llc/llc.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -44,6 +44,7 @@
#include "llvm/Support/FileSystem.h"
#include "llvm/Support/FormattedStream.h"
#include "llvm/Support/InitLLVM.h"
#include "llvm/Support/PGOOptions.h"
#include "llvm/Support/PluginLoader.h"
#include "llvm/Support/SourceMgr.h"
#include "llvm/Support/TargetSelect.h"
Expand Down Expand Up @@ -239,6 +240,39 @@ static cl::opt<RunPassOption, true, cl::parser<std::string>> RunPass(
cl::desc("Run compiler only for specified passes (comma separated list)"),
cl::value_desc("pass-name"), cl::location(RunPassOpt));

// PGO command line options
enum PGOKind {
NoPGO,
SampleUse,
};

static cl::opt<PGOKind>
PGOKindFlag("pgo-kind", cl::init(NoPGO), cl::Hidden,
cl::desc("The kind of profile guided optimization"),
cl::values(clEnumValN(NoPGO, "nopgo", "Do not use PGO."),
clEnumValN(SampleUse, "pgo-sample-use-pipeline",
"Use sampled profile to guide PGO.")));

// Function to set PGO options on TargetMachine based on command line flags.
static void setPGOOptions(TargetMachine &TM) {
std::optional<PGOOptions> PGOOpt;

switch (PGOKindFlag) {
case SampleUse:
// Use default values for other PGOOptions parameters. This parameter
// is used to test that PGO data is preserved at -O0.
PGOOpt = PGOOptions("", "", "", "", PGOOptions::SampleUse,
PGOOptions::NoCSAction);
break;
case NoPGO:
PGOOpt = std::nullopt;
break;
}

if (PGOOpt)
TM.setPGOOption(PGOOpt);
Copy link
Contributor

@apolloww apolloww Oct 3, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

If the changes here are just to cover the case of directly invoke llc on an IR module with profile information (like the test does), I think we can use a simple flag instead of creating a PGOOption that acts like a boolean switch. In most scenarios, we'd feed cpp source and profile to clang driver, which would set PGOOption correctly.

Copy link
Contributor

Choose a reason for hiding this comment

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

+1 on a simple flag, would it make sense to still use PGOOption, but only with a single value? something like --pgo-kind=any-pgo-pipeline

we can pass any of the PGO option.

if (PGOKindFlag == AnyPGO)
 TM.setPGOOption(PGOOptions("", "", "", "",PGOOptions::SampleUse /*Any PGO action*/,
                            PGOOptions::NoAction);

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I thought it might be useful in the future. There is no way to inject PGOOption into llc. However, there might be logic relying on those parameters in CodeGen in the future that would be easier to test. Maybe I can keep just pgo-kind cmd parameter for now?

Copy link
Member

Choose a reason for hiding this comment

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

I agree with the others. Adding all of this just for a possible future use case seems unreasonable, we can always add more if necessary.

Also +1 to using --pgo-kind with some arbitrary value for now.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

I kept only pgo-kind option with the only value pgo-sample-use-pipeline.

}

static int compileModule(char **, LLVMContext &);

[[noreturn]] static void reportError(Twine Msg, StringRef Filename = "") {
Expand Down Expand Up @@ -554,6 +588,9 @@ static int compileModule(char **argv, LLVMContext &Context) {
TheTriple, CPUStr, FeaturesStr, Options, RM, CM, OLvl));
assert(Target && "Could not allocate target machine!");

// Set PGO options based on command line flags
setPGOOptions(*Target);

return Target->createDataLayout().getStringRepresentation();
};
if (InputLanguage == "mir" ||
Expand Down Expand Up @@ -597,6 +634,9 @@ static int compileModule(char **argv, LLVMContext &Context) {
TheTriple, CPUStr, FeaturesStr, Options, RM, CM, OLvl));
assert(Target && "Could not allocate target machine!");

// Set PGO options based on command line flags
setPGOOptions(*Target);

// If we don't have a module then just exit now. We do this down
// here since the CPU/Feature help is underneath the target machine
// creation.
Expand Down