[llvm] [TailCallElim] Improve cold function detection for SamplePGO (PR #218153)
via llvm-commits
llvm-commits at lists.llvm.org
Sat Aug 22 10:43:52 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Teresa Johnson (teresajohnson)
<details>
<summary>Changes</summary>
Use isFunctionColdInCallGraph rather than checking the function entry
count to determine whether the function is cold (when disabling tail
call elimination for cold calls), to better handle SamplePGO where the
function entry count may be missing or inaccurate.
The coldness check is computed once per function and passed down
to avoid repeatedly invoking isFunctionColdInCallGraph which checks
coldness across all basic blocks.
Also added test coverage for missing and partial profile counts, as well
as SamplePGO profiles.
---
Full diff: https://github.com/llvm/llvm-project/pull/218153.diff
3 Files Affected:
- (modified) llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp (+35-43)
- (added) llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls-samplepgo.ll (+111)
- (modified) llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll (+23)
``````````diff
diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index b3df27433a9f5..e9a53e1df59b6 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -104,41 +104,17 @@ static cl::opt<bool> DisableTailCallElimForColdCalls(
cl::desc("Disable tail call elimination and optimization for cold calls or "
"in cold functions"));
+/// Return true if tail call elimination should be disabled for the given call
+/// instruction due to coldness. \p DisableColdTailCalls indicates whether the
+/// DisableTailCallElimForColdCalls flag is enabled and the enclosing function
+/// was determined to be cold overall (via attributes, calling convention, or
+/// profile). Note that mandatory tail calls (musttail) are never disabled.
static bool shouldDisableTailCallsForCold(const CallBase *CB,
- const Function *Caller,
- const ProfileSummaryInfo *PSI,
- BlockFrequencyInfo *BFI) {
- if (!DisableTailCallElimForColdCalls)
- return false;
-
+ bool DisableColdTailCalls) {
if (CB && CB->isMustTailCall())
return false;
- if (Caller && (Caller->hasFnAttribute(Attribute::Cold) ||
- Caller->getCallingConv() == CallingConv::Cold))
- return true;
-
- if (!PSI || !PSI->hasProfileSummary())
- return false;
-
- // We require both the function entry and the call site/block/callee to be
- // cold.
- // 1. Checking that the function entry is cold ensures we don't disable tail
- // call elimination in hot functions (with calls on cold conditional
- // paths), which would force stack frame setup and teardown on hot paths.
- // 2. Checking that the call site/block/callee is also cold ensures that if a
- // function has a cold entry count but contains a hot loop, we don't
- // disable tail call elimination for calls within that hot loop.
- if (Caller && PSI->isFunctionEntryCold(Caller) && CB) {
- if (CB->hasFnAttr(Attribute::Cold) ||
- CB->getCallingConv() == CallingConv::Cold)
- return true;
- if (BFI && (PSI->isColdCallSite(*CB, BFI) ||
- PSI->isColdBlock(CB->getParent(), BFI)))
- return true;
- }
-
- return false;
+ return DisableColdTailCalls;
}
/// Scan the specified function for alloca instructions.
@@ -241,7 +217,7 @@ struct AllocaDerivedValueTracker {
} // namespace
static bool markTails(Function &F, OptimizationRemarkEmitter *ORE,
- ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI) {
+ bool DisableColdTailCalls) {
if (F.callsFunctionThatReturnsTwice())
return false;
@@ -305,7 +281,8 @@ static bool markTails(Function &F, OptimizationRemarkEmitter *ORE,
// Special-case operand bundles "clang.arc.attachedcall", "ptrauth", and
// "kcfi".
- bool DisableForCold = shouldDisableTailCallsForCold(CI, &F, PSI, BFI);
+ bool DisableForCold =
+ shouldDisableTailCallsForCold(CI, DisableColdTailCalls);
bool IsNoTail = CI->isNoTailCall() || DisableForCold ||
CI->hasOperandBundlesOtherThan(
{LLVMContext::OB_clang_arc_attachedcall,
@@ -546,7 +523,7 @@ class TailRecursionEliminator {
OptimizationRemarkEmitter *ORE;
DomTreeUpdater &DTU;
BlockFrequencyInfo *const BFI;
- ProfileSummaryInfo *const PSI;
+ const bool DisableColdTailCalls;
const bool UpdateFunctionEntryCount;
const uint64_t OrigEntryBBFreq;
const uint64_t OrigEntryCount;
@@ -582,9 +559,10 @@ class TailRecursionEliminator {
TailRecursionEliminator(Function &F, const TargetTransformInfo *TTI,
AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
DomTreeUpdater &DTU, BlockFrequencyInfo *BFI,
- ProfileSummaryInfo *PSI,
+ bool DisableColdTailCalls,
bool UpdateFunctionEntryCount)
- : F(F), TTI(TTI), AA(AA), ORE(ORE), DTU(DTU), BFI(BFI), PSI(PSI),
+ : F(F), TTI(TTI), AA(AA), ORE(ORE), DTU(DTU), BFI(BFI),
+ DisableColdTailCalls(DisableColdTailCalls),
UpdateFunctionEntryCount(UpdateFunctionEntryCount),
OrigEntryBBFreq(
BFI ? BFI->getBlockFreq(&F.getEntryBlock()).getFrequency() : 0U),
@@ -643,7 +621,8 @@ CallInst *TailRecursionEliminator::findTRECandidate(BasicBlock *BB) {
assert((!CI->isTailCall() || !CI->isNoTailCall()) &&
"Incompatible call site attributes(Tail,NoTail)");
- if (!CI->isTailCall() || shouldDisableTailCallsForCold(CI, &F, PSI, BFI))
+ if (!CI->isTailCall() ||
+ shouldDisableTailCallsForCold(CI, DisableColdTailCalls))
return nullptr;
// As a special case, detect code like this:
@@ -912,8 +891,7 @@ bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
DTU.applyUpdates({{DominatorTree::Insert, BB, HeaderBB}});
++NumEliminated;
if (!DisableEntryCountRecompute && UpdateFunctionEntryCount &&
- OrigEntryBBFreq) {
- assert(F.getEntryCount().has_value());
+ OrigEntryBBFreq && F.getEntryCount().has_value()) {
// This pass is not expected to remove BBs, only add an entry BB. For that
// reason, and because the BB here isn't the new entry BB, the BFI lookup is
// expected to succeed.
@@ -1076,8 +1054,22 @@ bool TailRecursionEliminator::eliminate(
if (F.getFnAttribute("disable-tail-calls").getValueAsBool())
return false;
+ // We only disable tail calls in functions that are cold overall.
+ // 1. Checking that the function is cold ensures we don't disable tail call
+ // elimination in hot functions (with calls on cold conditional paths),
+ // which would force stack frame setup and teardown on hot paths.
+ // 2. Using isFunctionColdInCallGraph ensures all basic blocks in the function
+ // are cold, preventing disabling tail call elimination for calls within a
+ // hot loop inside a function with a cold entry count, and also handles
+ // SamplePGO where entry counts may be missing or inaccurate.
+ bool DisableColdTailCalls =
+ DisableTailCallElimForColdCalls &&
+ (F.hasFnAttribute(Attribute::Cold) ||
+ F.getCallingConv() == CallingConv::Cold ||
+ (PSI && BFI && PSI->isFunctionColdInCallGraph(&F, *BFI)));
+
bool MadeChange = false;
- MadeChange |= markTails(F, ORE, PSI, BFI);
+ MadeChange |= markTails(F, ORE, DisableColdTailCalls);
// If this function is a varargs function, we won't be able to PHI the args
// right, so don't even try to convert it...
@@ -1088,8 +1080,8 @@ bool TailRecursionEliminator::eliminate(
return MadeChange;
// Change any tail recursive calls to loops.
- TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU, BFI, PSI,
- UpdateFunctionEntryCount);
+ TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU, BFI,
+ DisableColdTailCalls, UpdateFunctionEntryCount);
for (BasicBlock &BB : F)
MadeChange |= TRE.processBlock(BB);
@@ -1158,7 +1150,7 @@ PreservedAnalyses TailCallElimPass::run(Function &F,
// This must come first. It needs the 2 analyses, meaning, if it came after
// the lines asking for the cached result, should they be nullptr (which, in
// the case of the PDT, is likely), updates to the trees would be missed.
- auto *BFI = F.getEntryCount().has_value()
+ auto *BFI = (F.getEntryCount().has_value() || DisableTailCallElimForColdCalls)
? &AM.getResult<BlockFrequencyAnalysis>(F)
: nullptr;
auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
diff --git a/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls-samplepgo.ll b/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls-samplepgo.ll
new file mode 100644
index 0000000000000..2b995ed4f0bed
--- /dev/null
+++ b/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls-samplepgo.ll
@@ -0,0 +1,111 @@
+; RUN: opt < %s -passes="require<profile-summary>,tailcallelim" -disable-tail-call-elim-for-cold-calls=true -S | FileCheck %s --check-prefixes=CHECK,DISABLED
+; RUN: opt < %s -passes="require<profile-summary>,tailcallelim" -disable-tail-call-elim-for-cold-calls=false -S | FileCheck %s --check-prefixes=CHECK,ENABLED
+
+declare void @normal_callee()
+declare void @cold_callee() cold
+
+; Check that in SamplePGO, a cold function (with cold entry and cold blocks) has tail call elimination disabled when the flag is enabled.
+define void @test_sample_cold_function() !prof !14 {
+; CHECK-LABEL: @test_sample_cold_function(
+; DISABLED: call void @normal_callee()
+; ENABLED: tail call void @normal_callee()
+ call void @normal_callee()
+ ret void
+}
+
+; Check that in SamplePGO, a function with hot entry count IS marked as tail.
+define void @test_sample_hot_function(i1 %cond) !prof !15 {
+; CHECK-LABEL: @test_sample_hot_function(
+; CHECK: tail call void @normal_callee()
+entry:
+ br i1 %cond, label %if.then, label %if.else, !prof !16
+
+if.then:
+ ret void
+
+if.else:
+ call void @normal_callee()
+ ret void
+}
+
+; Check that in SamplePGO, a function with missing entry count but cold block weights has tail call elimination disabled when the flag is enabled.
+define void @test_sample_missing_entry_cold_blocks(i1 %cond) {
+; CHECK-LABEL: @test_sample_missing_entry_cold_blocks(
+; DISABLED: call void @normal_callee()
+; ENABLED: tail call void @normal_callee()
+entry:
+ br i1 %cond, label %if.then, label %if.else, !prof !19
+
+if.then:
+ ret void
+
+if.else:
+ call void @normal_callee()
+ ret void
+}
+
+; Check that in SamplePGO, a function with missing entry count and missing block weights (unprofiled) IS marked as tail.
+define void @test_sample_missing_counts() {
+; CHECK-LABEL: @test_sample_missing_counts(
+; CHECK: tail call void @normal_callee()
+ call void @normal_callee()
+ ret void
+}
+
+; Check that in SamplePGO, a function with missing entry count and hot block weights IS marked as tail.
+define void @test_sample_missing_entry_hot_blocks(i1 %cond) {
+; CHECK-LABEL: @test_sample_missing_entry_hot_blocks(
+; CHECK: tail call void @normal_callee()
+entry:
+ br i1 %cond, label %if.then, label %if.else, !prof !16
+
+if.then:
+ ret void
+
+if.else:
+ call void @normal_callee()
+ ret void
+}
+
+; Check that in SamplePGO, a function with partial profile (one block has cold weights, one block is unprofiled) IS marked as tail.
+define void @test_sample_partial_profile(i1 %cond1, i1 %cond2) {
+; CHECK-LABEL: @test_sample_partial_profile(
+; CHECK: tail call void @normal_callee()
+entry:
+ br i1 %cond1, label %bb1, label %bb2, !prof !19
+
+bb1:
+ br i1 %cond2, label %exit, label %unprofiled_block
+
+unprofiled_block:
+ call void @normal_callee()
+ ret void
+
+bb2:
+ ret void
+
+exit:
+ ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"SampleProfile"}
+!3 = !{!"TotalCount", i64 10000}
+!4 = !{!"MaxCount", i64 10}
+!5 = !{!"MaxInternalCount", i64 1}
+!6 = !{!"MaxFunctionCount", i64 1000}
+!7 = !{!"NumCounts", i64 3}
+!8 = !{!"NumFunctions", i64 3}
+!9 = !{!"DetailedSummary", !10}
+!10 = !{!11, !12, !13}
+!11 = !{i32 10000, i64 100, i32 1}
+!12 = !{i32 999000, i64 100, i32 1}
+!13 = !{i32 999999, i64 1, i32 2}
+!14 = !{!"function_entry_count", i64 0}
+!15 = !{!"function_entry_count", i64 1000}
+!16 = !{!"branch_weights", i32 1000, i32 0}
+!17 = !{!"function_entry_count", i64 1}
+!18 = !{!"branch_weights", i32 1000, i32 1}
+!19 = !{!"branch_weights", i32 1, i32 0}
diff --git a/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll b/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll
index ee48b050175f5..2da30bb403368 100644
--- a/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll
+++ b/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll
@@ -137,6 +137,29 @@ exit:
ret void
}
+; Check that a call inside a function with missing profile count IS marked as tail.
+define void @test_missing_profile() {
+; CHECK-LABEL: @test_missing_profile(
+; CHECK: tail call void @normal_callee()
+ call void @normal_callee()
+ ret void
+}
+
+; Check that a call inside a function with missing entry count but with branch weights IS marked as tail, even though the block with the call is cold.
+define void @test_missing_entry_with_branch_weights(i1 %cond) {
+; CHECK-LABEL: @test_missing_entry_with_branch_weights(
+; CHECK: tail call void @normal_callee()
+entry:
+ br i1 %cond, label %if.then, label %if.else, !prof !16
+
+if.then:
+ ret void
+
+if.else:
+ call void @normal_callee()
+ ret void
+}
+
!llvm.module.flags = !{!0}
!0 = !{i32 1, !"ProfileSummary", !1}
!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
``````````
</details>
https://github.com/llvm/llvm-project/pull/218153
More information about the llvm-commits
mailing list