[llvm] [PGO] Implement PGO counter promotion for atomic updates (PR #202487)
Wael Yehia via llvm-commits
llvm-commits at lists.llvm.org
Mon Jun 8 19:50:03 PDT 2026
https://github.com/w2yehia created https://github.com/llvm/llvm-project/pull/202487
None
>From f62c74203f17a17af92e5af689d78609c1d479c9 Mon Sep 17 00:00:00 2001
From: Wael Yehia <wyehia at ca.ibm.com>
Date: Wed, 3 Jun 2026 16:24:51 -0400
Subject: [PATCH] WIP atomic promotion
---
.../Instrumentation/InstrProfiling.cpp | 76 ++++++++++++++++++-
1 file changed, 73 insertions(+), 3 deletions(-)
diff --git a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
index dabd495cddd49..da4990507a0fc 100644
--- a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
+++ b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
@@ -14,6 +14,7 @@
#include "llvm/Transforms/Instrumentation/InstrProfiling.h"
#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
@@ -36,6 +37,7 @@
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/GlobalVariable.h"
#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instruction.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
@@ -113,6 +115,16 @@ cl::opt<bool> AtomicCounterUpdateAll(
cl::desc("Make all profile counter updates atomic (for testing only)"),
cl::init(false));
+cl::opt<bool> AtomicCounterPromote(
+ "atomic-counter-promote",
+ cl::desc("Atomize profile counter updates and promote where possible"),
+ cl::init(false));
+
+cl::opt<bool> SanityCheck(
+ "atomic-counter-promote-check",
+ cl::desc("Check that all profile counter updates were made atomic"),
+ cl::init(false));
+
cl::opt<bool> AtomicCounterUpdatePromoted(
"atomic-counter-update-promoted",
cl::desc("Do counter update using atomic fetch add "
@@ -221,6 +233,20 @@ static SampledInstrumentationConfig getSampledInstrumentationConfig() {
using LoadStorePair = std::pair<Instruction *, Instruction *>;
+static void makeAtomic(Instruction *Load, Instruction *Store) {
+ // assert the load and store are accessing the same memory?
+ auto *Addition = dyn_cast<BinaryOperator>(Store->getOperand(0));
+ assert(Addition && Addition->getOpcode() == Instruction::BinaryOps::Add);
+ auto *Addend = Addition->getOperand(1);
+
+ IRBuilder<> Builder(Load);
+ Builder.CreateAtomicRMW(AtomicRMWInst::Add, Store->getOperand(1), Addend,
+ MaybeAlign(), AtomicOrdering::Monotonic);
+ Store->eraseFromParent();
+ Addition->eraseFromParent();
+ Load->eraseFromParent();
+}
+
static uint64_t getIntModuleFlagOrZero(const Module &M, StringRef Flag) {
auto *MD = dyn_cast_or_null<ConstantAsMetadata>(M.getModuleFlag(Flag));
if (!MD)
@@ -461,6 +487,12 @@ class PGOCounterPromoterHelper : public LoadAndStorePromoter {
Builder.CreateAtomicRMW(AtomicRMWInst::Add, Addr, LiveInValue,
MaybeAlign(),
AtomicOrdering::SequentiallyConsistent);
+ // Generate the relaxed atomic RMW if we've asked for it and no more
+ // promotion is possible.
+ else if (AtomicCounterPromote &&
+ (!IterativeCounterPromotion || !LI.getLoopFor(ExitBlock)))
+ Builder.CreateAtomicRMW(AtomicRMWInst::Add, Addr, LiveInValue,
+ MaybeAlign(), AtomicOrdering::Monotonic);
else {
LoadInst *OldVal = Builder.CreateLoad(Ty, Addr, "pgocount.promoted");
auto *NewVal = Builder.CreateAdd(OldVal, LiveInValue);
@@ -515,6 +547,20 @@ class PGOCounterPromoter {
}
bool run(int64_t *NumPromoted) {
+ // In this function we examine loop L and other parameters to decide what
+ // candidates are promotable. Once we've promoted what we can, we convert
+ // all remaining candidates to use atomics. This requires that promoted
+ // candidates are set to nullptr in the LoopToCandidates[&L] array.
+ llvm::scope_exit Cleanup([&]() {
+ if (!AtomicCounterPromote)
+ return;
+ for (auto &Cand : LoopToCandidates[&L]) {
+ if (Cand.first != nullptr && Cand.second != nullptr)
+ makeAtomic(Cand.first, Cand.second);
+ Cand = {nullptr, nullptr};
+ }
+ });
+
// Skip 'infinite' loops:
if (ExitBlocks.size() == 0)
return false;
@@ -536,7 +582,6 @@ class PGOCounterPromoter {
unsigned Promoted = 0;
for (auto &Cand : LoopToCandidates[&L]) {
-
SmallVector<PHINode *, 4> NewPHIs;
SSAUpdater SSA(&NewPHIs);
Value *InitVal = ConstantInt::get(Cand.first->getType(), 0);
@@ -558,6 +603,7 @@ class PGOCounterPromoter {
L.getLoopPreheader(), ExitBlocks,
InsertPts, LoopToCandidates, LI);
Promoter.run(SmallVector<Instruction *, 2>({Cand.first, Cand.second}));
+ Cand = {nullptr, nullptr};
Promoted++;
if (Promoted >= MaxProm)
break;
@@ -861,10 +907,27 @@ bool InstrLowerer::isSamplingEnabled() const {
bool InstrLowerer::isCounterPromotionEnabled() const {
if (DoCounterPromotion.getNumOccurrences() > 0)
return DoCounterPromotion;
-
+ if (AtomicCounterPromote)
+ return true;
return Options.DoCounterPromotion;
}
+static void doAtomicPromotionCheck(Function *F) {
+ for (const llvm::Instruction &I : llvm::instructions(F)) {
+ const Value *Addr = nullptr;
+ if (const LoadInst *LI = dyn_cast<LoadInst>(&I))
+ Addr = LI->getOperand(0);
+ else if (const StoreInst *LI = dyn_cast<StoreInst>(&I))
+ Addr = LI->getOperand(1);
+
+ if (Addr && Addr->stripInBoundsOffsets()->getName().starts_with(
+ getInstrProfCountersVarPrefix())) {
+ LLVM_DEBUG(dbgs() << "Missed candidate: "; I.dump());
+ assert(false && "Candidate load/store not converted to atomic");
+ }
+ }
+}
+
void InstrLowerer::promoteCounterLoadStores(Function *F) {
if (!isCounterPromotionEnabled())
return;
@@ -885,8 +948,11 @@ void InstrLowerer::promoteCounterLoadStores(Function *F) {
auto *CounterStore = LoadStore.second;
BasicBlock *BB = CounterLoad->getParent();
Loop *ParentLoop = LI.getLoopFor(BB);
- if (!ParentLoop)
+ if (!ParentLoop) {
+ if (AtomicCounterPromote)
+ makeAtomic(CounterLoad, CounterStore);
continue;
+ }
LoopPromotionCandidates[ParentLoop].emplace_back(CounterLoad, CounterStore);
}
@@ -898,6 +964,9 @@ void InstrLowerer::promoteCounterLoadStores(Function *F) {
PGOCounterPromoter Promoter(LoopPromotionCandidates, *Loop, LI, BFI.get());
Promoter.run(&TotalCountersPromoted);
}
+
+ if (AtomicCounterPromote && SanityCheck)
+ doAtomicPromotionCheck(F);
}
static bool needsRuntimeHookUnconditionally(const Triple &TT) {
@@ -1212,6 +1281,7 @@ void InstrLowerer::lowerIncrement(InstrProfIncrementInst *Inc) {
Value *Load = Builder.CreateLoad(IncStep->getType(), Addr, "pgocount");
auto *Count = Builder.CreateAdd(Load, Inc->getStep());
auto *Store = Builder.CreateStore(Count, Addr);
+
if (isCounterPromotionEnabled())
PromotionCandidates.emplace_back(cast<Instruction>(Load), Store);
}
More information about the llvm-commits
mailing list