[llvm-branch-commits] [llvm] [PGO] Load branch uniformity profiles and mark availability (PR #221629)
Yaxun Liu via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Mon Sep 21 20:55:31 PDT 2026
https://github.com/yxsamliu updated https://github.com/llvm/llvm-project/pull/221629
>From 760103188c6081b3bfd202d30210962375244129 Mon Sep 17 00:00:00 2001
From: "Yaxun (Sam) Liu" <yaxun.liu at amd.com>
Date: Mon, 7 Sep 2026 00:32:58 -0400
Subject: [PATCH] [PGO] Load branch uniformity profiles and mark availability
Uniformity profiles record block observations, but optimizations also
need information about branch decisions. A missing branch annotation
alone cannot tell them whether profile data is absent or no branches
are known to be uniform.
Mark functions when uniformity data is loaded and derive branch
annotations from blocks with a single conditional predecessor. Each
such block measures the lanes taking one outgoing edge. Mark a branch
as uniform only when all its instrumented outgoing edges are classified
as uniform, so one uniform edge cannot hide a divergent edge.
Replace existing branch annotations when loading the profile. This
clears an old uniform classification if the new data no longer supports
it. The function marker lets consumers use these results while keeping
their existing behavior when no uniformity profile is available.
---
.../Instrumentation/PGOInstrumentation.cpp | 21 ++++++++--
.../PGOInstrumentationTest.cpp | 41 +++++++++++++++----
2 files changed, 52 insertions(+), 10 deletions(-)
diff --git a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
index 6d3126bcb8ac7..9a0203a03d5c7 100644
--- a/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
+++ b/llvm/lib/Transforms/Instrumentation/PGOInstrumentation.cpp
@@ -1788,9 +1788,9 @@ void PGOUseFunc::setBlockUniformityAttribute() {
if (ProfileRecord.UniformityBits.empty())
return;
- // Annotate each uniform instrumented IR basic block so later codegen passes
- // (MachineFunction) can consume it without relying on fragile block numbering
- // heuristics.
+ // Mark the function as having uniformity profile, then annotate each uniform
+ // instrumented IR basic block so later codegen passes (MachineFunction) can
+ // consume it without relying on fragile block numbering heuristics.
// Metadata presence on a terminator means uniform; divergent blocks have no
// terminator metadata.
@@ -1799,14 +1799,29 @@ void PGOUseFunc::setBlockUniformityAttribute() {
LLVMContext &Ctx = F.getContext();
MDNode *UniformMD = MDNode::get(Ctx, {});
+ F.setMetadata(LLVMContext::MD_uniformity_profile, UniformMD);
+ DenseMap<CondBrInst *, bool> BranchUniformity;
for (size_t I = 0, E = InstrumentBBs.size(); I < E; ++I) {
BasicBlock *BB = InstrumentBBs[I];
if (!BB || !BB->getTerminator())
continue;
bool IsUniform = ProfileRecord.isBlockUniform(I);
+ // A counter placed in a block with a single conditional predecessor also
+ // measures the active lanes on that outgoing edge. Record the branch as
+ // uniform only when every instrumented outgoing edge is uniform.
+ if (BasicBlock *Pred = BB->getSinglePredecessor()) {
+ if (auto *Branch = dyn_cast<CondBrInst>(Pred->getTerminator())) {
+ auto It = BranchUniformity.try_emplace(Branch, true).first;
+ It->second &= IsUniform;
+ }
+ }
BB->getTerminator()->setMetadata(LLVMContext::MD_block_uniformity_profile,
IsUniform ? UniformMD : nullptr);
}
+ for (auto [Branch, IsUniform] : BranchUniformity) {
+ Branch->setMetadata(LLVMContext::MD_branch_uniformity_profile,
+ IsUniform ? UniformMD : nullptr);
+ }
LLVM_DEBUG({
dbgs() << "PGO: Set block uniformity profile for " << F.getName() << ": ";
diff --git a/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp b/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp
index 686410ac1468d..84a020aa0bbb6 100644
--- a/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp
+++ b/llvm/unittests/Transforms/Instrumentation/PGOInstrumentationTest.cpp
@@ -27,6 +27,7 @@ namespace {
using namespace llvm;
using testing::_;
+using ::testing::Combine;
using ::testing::DoDefault;
using ::testing::Invoke;
using ::testing::NotNull;
@@ -192,12 +193,15 @@ TEST_P(PGOInstrumentationGenTest, Instrumented) {
EXPECT_FALSE(IRInstrVar->isDeclaration());
}
-struct PGOInstrumentationUseTest : Test, WithParamInterface<bool> {};
+struct PGOInstrumentationUseTest : Test,
+ WithParamInterface<std::tuple<bool, bool>> {
+};
-INSTANTIATE_TEST_SUITE_P(ExistingMetadata, PGOInstrumentationUseTest,
- Values(false, true));
+INSTANTIATE_TEST_SUITE_P(ProfileAvailability, PGOInstrumentationUseTest,
+ Combine(Values(false, true), Values(false, true)));
-TEST_P(PGOInstrumentationUseTest, BlockUniformityMetadataUsesPresence) {
+TEST_P(PGOInstrumentationUseTest, UniformityMetadataUsesPresence) {
+ const auto [HasUniformityProfile, HasExistingMetadata] = GetParam();
static constexpr StringRef Code = R"(
define i32 @f(i1 %cond) {
entry:
@@ -259,7 +263,8 @@ TEST_P(PGOInstrumentationUseTest, BlockUniformityMetadataUsesPresence) {
std::string ProfileName = getIRPGOFuncName(*GenFunction);
NamedInstrProfRecord Record(ProfileName, FunctionHash,
std::vector<uint64_t>(NumCounters, 10));
- Record.UniformityBits = {static_cast<uint8_t>(UniformityMask)};
+ if (HasUniformityProfile)
+ Record.UniformityBits = {static_cast<uint8_t>(UniformityMask)};
InstrProfWriter Writer;
ASSERT_THAT_ERROR(Writer.mergeProfileKind(InstrProfKind::IRInstrumentation),
@@ -276,17 +281,38 @@ TEST_P(PGOInstrumentationUseTest, BlockUniformityMetadataUsesPresence) {
ASSERT_THAT(UseModule, NotNull());
Function *UseFunction = UseModule->getFunction("f");
ASSERT_THAT(UseFunction, NotNull());
- if (GetParam()) {
+ if (HasExistingMetadata) {
MDNode *UniformMD = MDNode::get(Context, {});
+ UseFunction->setMetadata(LLVMContext::MD_uniformity_profile, UniformMD);
for (BasicBlock &BB : *UseFunction)
BB.getTerminator()->setMetadata(
LLVMContext::MD_block_uniformity_profile, UniformMD);
+ UseFunction->getEntryBlock().getTerminator()->setMetadata(
+ LLVMContext::MD_branch_uniformity_profile, UniformMD);
}
ModulePassManager UseMPM;
UseMPM.addPass(PGOInstrumentationUse("/profile.profdata", "", false, FS));
UseMPM.run(*UseModule, MAM);
EXPECT_FALSE(verifyModule(*UseModule, &errs()));
+ MDNode *FunctionMD =
+ UseFunction->getMetadata(LLVMContext::MD_uniformity_profile);
+ // Missing uniformity data preserves existing hints. Observed nonuniformity
+ // replaces positive hints covered by the new profile.
+ EXPECT_EQ(FunctionMD != nullptr,
+ HasUniformityProfile || HasExistingMetadata);
+ if (FunctionMD)
+ EXPECT_EQ(FunctionMD->getNumOperands(), 0u);
+
+ auto *Branch =
+ cast<CondBrInst>(UseFunction->getEntryBlock().getTerminator());
+ MDNode *BranchMD =
+ Branch->getMetadata(LLVMContext::MD_branch_uniformity_profile);
+ EXPECT_EQ(BranchMD != nullptr,
+ HasUniformityProfile ? UniformityMask == 3 : HasExistingMetadata);
+ if (BranchMD)
+ EXPECT_EQ(BranchMD->getNumOperands(), 0u);
+
for (unsigned I = 0; I < NumCounters; ++I) {
BasicBlock *BB = nullptr;
for (BasicBlock &Candidate : *UseFunction)
@@ -296,7 +322,8 @@ TEST_P(PGOInstrumentationUseTest, BlockUniformityMetadataUsesPresence) {
MDNode *MD = BB->getTerminator()->getMetadata(
LLVMContext::MD_block_uniformity_profile);
- bool IsUniform = (UniformityMask & (1u << I)) != 0;
+ bool IsUniform = HasUniformityProfile ? (UniformityMask & (1u << I)) != 0
+ : HasExistingMetadata;
EXPECT_EQ(MD != nullptr, IsUniform);
if (MD)
EXPECT_EQ(MD->getNumOperands(), 0u);
More information about the llvm-branch-commits
mailing list