[llvm] [PGO] Implement PGO counter promotion for atomic updates (PR #202487)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Jun 9 17:47:39 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Wael Yehia (w2yehia)
<details>
<summary>Changes</summary>
---
Full diff: https://github.com/llvm/llvm-project/pull/202487.diff
2 Files Affected:
- (modified) llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp (+73-3)
- (added) llvm/test/Transforms/PGOProfile/atomic_counter_promote.ll (+61)
``````````diff
diff --git a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
index dabd495cddd49..4cde2d2809b3e 100644
--- a/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
+++ b/llvm/lib/Transforms/Instrumentation/InstrProfiling.cpp
@@ -15,6 +15,7 @@
#include "llvm/Transforms/Instrumentation/InstrProfiling.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/ScopeExit.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.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);
}
diff --git a/llvm/test/Transforms/PGOProfile/atomic_counter_promote.ll b/llvm/test/Transforms/PGOProfile/atomic_counter_promote.ll
new file mode 100644
index 0000000000000..28c83eb1a2aa5
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/atomic_counter_promote.ll
@@ -0,0 +1,61 @@
+; RUN: opt < %s -passes=instrprof -atomic-counter-promote -S | FileCheck %s
+
+; CHECK: define i32 @foo(i32 %n) {
+; CHECK: entry:
+; CHECK: atomicrmw add {{.*}}ptr @__profc_foo
+;
+; CHECK: for.cond.for.cond.cleanup_crit_edge:
+; CHECK-NOT: br
+; CHECK: atomicrmw add {{.*}}ptr @__profc_foo
+; CHECK-NEXT: atomicrmw add {{.*}}ptr @__profc_foo
+
+ at __profn_foo = private constant [3 x i8] c"foo"
+
+define i32 @foo(i32 %n) {
+entry:
+ call void @llvm.instrprof.increment(ptr @__profn_foo, i64 1124680652598534200, i32 3, i32 2)
+ %cmp16 = icmp slt i32 0, %n
+ br i1 %cmp16, label %for.cond1.preheader.lr.ph, label %for.cond.cleanup
+
+for.cond1.preheader.lr.ph:
+ br label %for.cond1.preheader
+
+for.cond1.preheader:
+ %i.018 = phi i32 [ 0, %for.cond1.preheader.lr.ph ], [ %inc6, %for.cond.cleanup3 ]
+ %x.017 = phi i32 [ 0, %for.cond1.preheader.lr.ph ], [ %x.1.lcssa, %for.cond.cleanup3 ]
+ %cmp213 = icmp slt i32 0, %n
+ br i1 %cmp213, label %for.body4.lr.ph, label %for.cond.cleanup3
+
+for.body4.lr.ph:
+ br label %for.body4
+
+for.cond.for.cond.cleanup_crit_edge:
+ %split19 = phi i32 [ %x.1.lcssa, %for.cond.cleanup3 ]
+ br label %for.cond.cleanup
+
+for.cond.cleanup:
+ %x.0.lcssa = phi i32 [ %split19, %for.cond.for.cond.cleanup_crit_edge ], [ 0, %entry ]
+ ret i32 %x.0.lcssa
+
+for.cond1.for.cond.cleanup3_crit_edge:
+ %split = phi i32 [ %add, %for.body4 ]
+ br label %for.cond.cleanup3
+
+for.cond.cleanup3:
+ %x.1.lcssa = phi i32 [ %split, %for.cond1.for.cond.cleanup3_crit_edge ], [ %x.017, %for.cond1.preheader ]
+ call void @llvm.instrprof.increment(ptr @__profn_foo, i64 1124680652598534200, i32 3, i32 1)
+ %inc6 = add nuw nsw i32 %i.018, 1
+ %cmp = icmp slt i32 %inc6, %n
+ br i1 %cmp, label %for.cond1.preheader, label %for.cond.for.cond.cleanup_crit_edge
+
+for.body4:
+ %j.015 = phi i32 [ 0, %for.body4.lr.ph ], [ %inc, %for.body4 ]
+ %x.114 = phi i32 [ %x.017, %for.body4.lr.ph ], [ %add, %for.body4 ]
+ call void @llvm.instrprof.increment(ptr @__profn_foo, i64 1124680652598534200, i32 3, i32 0)
+ %add = add nsw i32 %x.114, %j.015
+ %inc = add nuw nsw i32 %j.015, 1
+ %cmp2 = icmp slt i32 %inc, %n
+ br i1 %cmp2, label %for.body4, label %for.cond1.for.cond.cleanup3_crit_edge
+}
+
+declare void @llvm.instrprof.increment(ptr, i64, i32, i32)
``````````
</details>
https://github.com/llvm/llvm-project/pull/202487
More information about the llvm-commits
mailing list