[llvm] [TailCallElim] Optionally suppress tail call elim for cold calls (PR #209642)

Teresa Johnson via llvm-commits llvm-commits at lists.llvm.org
Fri Jul 17 14:59:20 PDT 2026


https://github.com/teresajohnson updated https://github.com/llvm/llvm-project/pull/209642

>From 1c43cf19c432e8f7401b51c8604cfb1c1c09b732 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Tue, 14 Jul 2026 15:01:16 -0700
Subject: [PATCH 1/5] [TailCallElim] Optionally suppress tail call elim for
 cold calls

Since eliminated tail calls result in missing frames when collecting
various types of profiles, provide an option for disabling the
elimination for cold non-musttail calls. By default the new handling is
disabled.
---
 .../Scalar/TailRecursionElimination.cpp       | 82 +++++++++++++++----
 .../disable-tail-call-elim-for-cold-calls.ll  | 70 ++++++++++++++++
 2 files changed, 138 insertions(+), 14 deletions(-)
 create mode 100644 llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 3c38171292718..ef87b605a8947 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -60,6 +60,7 @@
 #include "llvm/Analysis/Loads.h"
 #include "llvm/Analysis/OptimizationRemarkEmitter.h"
 #include "llvm/Analysis/PostDominators.h"
+#include "llvm/Analysis/ProfileSummaryInfo.h"
 #include "llvm/Analysis/TargetTransformInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/CFG.h"
@@ -89,12 +90,50 @@ using namespace llvm;
 STATISTIC(NumEliminated, "Number of tail calls removed");
 STATISTIC(NumRetDuped,   "Number of return duplicated");
 STATISTIC(NumAccumAdded, "Number of accumulators introduced");
+STATISTIC(NumTREPreventedCold,
+          "Number of tail calls/recursions prevented due to cold calling "
+          "convention or attribute");
 
 static cl::opt<bool> ForceDisableBFI(
     "tre-disable-entrycount-recompute", cl::init(false), cl::Hidden,
     cl::desc("Force disabling recomputing of function entry count, on "
              "successful tail recursion elimination."));
 
+static cl::opt<bool> DisableTailCallElimForColdCalls(
+    "disable-tail-call-elim-for-cold-calls", cl::Hidden, cl::init(false),
+    cl::desc("Disable tail call elimination and optimization for cold calls or "
+             "in cold functions"));
+
+static bool shouldDisableTailCallsForCold(const CallBase *CB,
+                                          const Function *Caller,
+                                          const ProfileSummaryInfo *PSI,
+                                          BlockFrequencyInfo *BFI) {
+  if (!DisableTailCallElimForColdCalls)
+    return false;
+
+  if (CB && CB->isMustTailCall())
+    return false;
+
+  if (CB && (CB->hasFnAttr(Attribute::Cold) ||
+             CB->getCallingConv() == CallingConv::Cold))
+    return true;
+
+  if (Caller && (Caller->hasFnAttribute(Attribute::Cold) ||
+                 Caller->getCallingConv() == CallingConv::Cold))
+    return true;
+
+  if (!PSI || !PSI->hasProfileSummary())
+    return false;
+
+  if (Caller && PSI->isFunctionEntryCold(Caller))
+    return true;
+
+  if (CB && BFI && PSI->isColdCallSite(*CB, BFI))
+    return true;
+
+  return false;
+}
+
 /// Scan the specified function for alloca instructions.
 /// If it contains any dynamic allocas, returns false.
 static bool canTRE(Function &F) {
@@ -194,7 +233,8 @@ struct AllocaDerivedValueTracker {
 };
 } // namespace
 
-static bool markTails(Function &F, OptimizationRemarkEmitter *ORE) {
+static bool markTails(Function &F, OptimizationRemarkEmitter *ORE,
+                      ProfileSummaryInfo *PSI, BlockFrequencyInfo *BFI) {
   if (F.callsFunctionThatReturnsTwice())
     return false;
 
@@ -258,10 +298,15 @@ static bool markTails(Function &F, OptimizationRemarkEmitter *ORE) {
 
       // Special-case operand bundles "clang.arc.attachedcall", "ptrauth", and
       // "kcfi".
-      bool IsNoTail = CI->isNoTailCall() ||
-                      CI->hasOperandBundlesOtherThan(
-                          {LLVMContext::OB_clang_arc_attachedcall,
-                           LLVMContext::OB_ptrauth, LLVMContext::OB_kcfi});
+      bool IsNoTail =
+          CI->isNoTailCall() ||
+          shouldDisableTailCallsForCold(CI, &F, PSI, BFI) ||
+          CI->hasOperandBundlesOtherThan(
+              {LLVMContext::OB_clang_arc_attachedcall, LLVMContext::OB_ptrauth,
+               LLVMContext::OB_kcfi});
+      if (!CI->isNoTailCall() &&
+          shouldDisableTailCallsForCold(CI, &F, PSI, BFI))
+        ++NumTREPreventedCold;
 
       if (!IsNoTail && CI->doesNotAccessMemory()) {
         // A call to a readnone function whose arguments are all things computed
@@ -407,6 +452,7 @@ class TailRecursionEliminator {
   OptimizationRemarkEmitter *ORE;
   DomTreeUpdater &DTU;
   BlockFrequencyInfo *const BFI;
+  ProfileSummaryInfo *const PSI;
   const uint64_t OrigEntryBBFreq;
   const uint64_t OrigEntryCount;
 
@@ -438,8 +484,9 @@ class TailRecursionEliminator {
 
   TailRecursionEliminator(Function &F, const TargetTransformInfo *TTI,
                           AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
-                          DomTreeUpdater &DTU, BlockFrequencyInfo *BFI)
-      : F(F), TTI(TTI), AA(AA), ORE(ORE), DTU(DTU), BFI(BFI),
+                          DomTreeUpdater &DTU, BlockFrequencyInfo *BFI,
+                          ProfileSummaryInfo *PSI)
+      : F(F), TTI(TTI), AA(AA), ORE(ORE), DTU(DTU), BFI(BFI), PSI(PSI),
         OrigEntryBBFreq(
             BFI ? BFI->getBlockFreq(&F.getEntryBlock()).getFrequency() : 0U),
         OrigEntryCount(F.getEntryCount() ? *F.getEntryCount() : 0) {
@@ -471,7 +518,8 @@ class TailRecursionEliminator {
 public:
   static bool eliminate(Function &F, const TargetTransformInfo *TTI,
                         AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
-                        DomTreeUpdater &DTU, BlockFrequencyInfo *BFI);
+                        DomTreeUpdater &DTU, BlockFrequencyInfo *BFI,
+                        ProfileSummaryInfo *PSI);
 };
 } // namespace
 
@@ -497,7 +545,8 @@ CallInst *TailRecursionEliminator::findTRECandidate(BasicBlock *BB) {
 
   assert((!CI->isTailCall() || !CI->isNoTailCall()) &&
          "Incompatible call site attributes(Tail,NoTail)");
-  if (!CI->isTailCall())
+  if (!CI->isTailCall() ||
+      shouldDisableTailCallsForCold(CI, &F, PSI, BFI))
     return nullptr;
 
   // As a special case, detect code like this:
@@ -906,12 +955,13 @@ bool TailRecursionEliminator::eliminate(Function &F,
                                         AliasAnalysis *AA,
                                         OptimizationRemarkEmitter *ORE,
                                         DomTreeUpdater &DTU,
-                                        BlockFrequencyInfo *BFI) {
+                                        BlockFrequencyInfo *BFI,
+                                        ProfileSummaryInfo *PSI) {
   if (F.getFnAttribute("disable-tail-calls").getValueAsBool())
     return false;
 
   bool MadeChange = false;
-  MadeChange |= markTails(F, ORE);
+  MadeChange |= markTails(F, ORE, PSI, BFI);
 
   // 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...
@@ -922,7 +972,7 @@ bool TailRecursionEliminator::eliminate(Function &F,
     return MadeChange;
 
   // Change any tail recursive calls to loops.
-  TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU, BFI);
+  TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU, BFI, PSI);
 
   for (BasicBlock &BB : F)
     MadeChange |= TRE.processBlock(BB);
@@ -961,11 +1011,13 @@ struct TailCallElim : public FunctionPass {
     // UpdateStrategy to Lazy if we find it profitable later.
     DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Eager);
 
+    auto *PSIWP = getAnalysisIfAvailable<ProfileSummaryInfoWrapperPass>();
+    auto *PSI = PSIWP ? &PSIWP->getPSI() : nullptr;
     return TailRecursionEliminator::eliminate(
         F, &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F),
         &getAnalysis<AAResultsWrapperPass>().getAAResults(),
         &getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE(), DTU,
-        /*BFI=*/nullptr);
+        /*BFI=*/nullptr, PSI);
   }
 };
 } // namespace
@@ -995,6 +1047,8 @@ PreservedAnalyses TailCallElimPass::run(Function &F,
                F.getEntryCount().has_value() && *F.getEntryCount())
                   ? &AM.getResult<BlockFrequencyAnalysis>(F)
                   : nullptr;
+  auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
+  auto *PSI = MAMProxy.getCachedResult<ProfileSummaryAnalysis>(*F.getParent());
   auto &ORE = AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
   auto *DT = AM.getCachedResult<DominatorTreeAnalysis>(F);
   auto *PDT = AM.getCachedResult<PostDominatorTreeAnalysis>(F);
@@ -1003,7 +1057,7 @@ PreservedAnalyses TailCallElimPass::run(Function &F,
   // UpdateStrategy to Lazy if we find it profitable later.
   DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Eager);
   bool Changed =
-      TailRecursionEliminator::eliminate(F, &TTI, &AA, &ORE, DTU, BFI);
+      TailRecursionEliminator::eliminate(F, &TTI, &AA, &ORE, DTU, BFI, PSI);
 
   if (!Changed)
     return PreservedAnalyses::all();
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
new file mode 100644
index 0000000000000..203ed22cbe6aa
--- /dev/null
+++ b/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll
@@ -0,0 +1,70 @@
+; RUN: opt < %s -passes=tailcallelim -disable-tail-call-elim-for-cold-calls=true -S | FileCheck %s --check-prefixes=CHECK,DISABLED
+; RUN: opt < %s -passes=tailcallelim -disable-tail-call-elim-for-cold-calls=false -S | FileCheck %s --check-prefixes=CHECK,ENABLED
+
+declare void @cold_callee() cold
+declare void @normal_callee()
+declare coldcc void @coldcc_callee()
+
+; Check that a call to a cold callee is not marked as tail when the flag is enabled.
+define void @test_cold_callee() {
+; CHECK-LABEL: @test_cold_callee(
+; DISABLED: call void @cold_callee()
+; ENABLED: tail call void @cold_callee()
+  call void @cold_callee()
+  ret void
+}
+
+; Check that a call to a callee with coldcc is not marked as tail when the flag is enabled.
+define void @test_coldcc_callee() {
+; CHECK-LABEL: @test_coldcc_callee(
+; DISABLED: call coldcc void @coldcc_callee()
+; ENABLED: tail call coldcc void @coldcc_callee()
+  call coldcc void @coldcc_callee()
+  ret void
+}
+
+; Check that a call inside a cold enclosing function is not marked as tail when the flag is enabled.
+define void @test_cold_caller() cold {
+; CHECK-LABEL: @test_cold_caller(
+; DISABLED: call void @normal_callee()
+; ENABLED: tail call void @normal_callee()
+  call void @normal_callee()
+  ret void
+}
+
+; Check that a callsite with coldcc is not marked as tail when the flag is enabled.
+define void @test_coldcc_callsite() {
+; CHECK-LABEL: @test_coldcc_callsite(
+; DISABLED: call coldcc void @normal_callee()
+; ENABLED: tail call coldcc void @normal_callee()
+  call coldcc void @normal_callee()
+  ret void
+}
+
+; Check that mandatory musttail calls are never disabled even when calling a cold function and the flag is true.
+define void @test_musttail_cold_callee() {
+; CHECK-LABEL: @test_musttail_cold_callee(
+; CHECK: musttail call void @cold_callee()
+  musttail call void @cold_callee()
+  ret void
+}
+
+; Check that tail recursion elimination to a loop header is disabled in a cold function when the flag is enabled.
+define i32 @test_recursive_cold_caller(i32 %X) cold {
+; CHECK-LABEL: @test_recursive_cold_caller(
+; DISABLED: call i32 @test_recursive_cold_caller(
+; DISABLED-NOT: br label %tailrecurse
+; ENABLED: tailrecurse:
+; ENABLED-NOT: call i32 @test_recursive_cold_caller(
+entry:
+  %cmp = icmp eq i32 %X, 0
+  br i1 %cmp, label %if.then, label %if.else
+
+if.then:
+  ret i32 0
+
+if.else:
+  %dec = sub i32 %X, 1
+  %call = call i32 @test_recursive_cold_caller(i32 %dec)
+  ret i32 %call
+}

>From 53e59bf316f9794ce4edb0c79196026658331f00 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Wed, 15 Jul 2026 14:38:57 -0700
Subject: [PATCH 2/5] Address comments

---
 .../Scalar/TailRecursionElimination.cpp       | 21 +++++++------------
 1 file changed, 8 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index ef87b605a8947..d7eb09d2173a0 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -91,8 +91,8 @@ STATISTIC(NumEliminated, "Number of tail calls removed");
 STATISTIC(NumRetDuped,   "Number of return duplicated");
 STATISTIC(NumAccumAdded, "Number of accumulators introduced");
 STATISTIC(NumTREPreventedCold,
-          "Number of tail calls/recursions prevented due to cold calling "
-          "convention or attribute");
+          "Number of tail calls/recursion eliminations prevented due to cold "
+          "calling convention or attribute");
 
 static cl::opt<bool> ForceDisableBFI(
     "tre-disable-entrycount-recompute", cl::init(false), cl::Hidden,
@@ -125,9 +125,6 @@ static bool shouldDisableTailCallsForCold(const CallBase *CB,
   if (!PSI || !PSI->hasProfileSummary())
     return false;
 
-  if (Caller && PSI->isFunctionEntryCold(Caller))
-    return true;
-
   if (CB && BFI && PSI->isColdCallSite(*CB, BFI))
     return true;
 
@@ -298,14 +295,12 @@ static bool markTails(Function &F, OptimizationRemarkEmitter *ORE,
 
       // Special-case operand bundles "clang.arc.attachedcall", "ptrauth", and
       // "kcfi".
-      bool IsNoTail =
-          CI->isNoTailCall() ||
-          shouldDisableTailCallsForCold(CI, &F, PSI, BFI) ||
-          CI->hasOperandBundlesOtherThan(
-              {LLVMContext::OB_clang_arc_attachedcall, LLVMContext::OB_ptrauth,
-               LLVMContext::OB_kcfi});
-      if (!CI->isNoTailCall() &&
-          shouldDisableTailCallsForCold(CI, &F, PSI, BFI))
+      bool DisableForCold = shouldDisableTailCallsForCold(CI, &F, PSI, BFI);
+      bool IsNoTail = CI->isNoTailCall() || DisableForCold ||
+                      CI->hasOperandBundlesOtherThan(
+                          {LLVMContext::OB_clang_arc_attachedcall,
+                           LLVMContext::OB_ptrauth, LLVMContext::OB_kcfi});
+      if (!CI->isNoTailCall() && DisableForCold)
         ++NumTREPreventedCold;
 
       if (!IsNoTail && CI->doesNotAccessMemory()) {

>From 0b16e76780fa544f68e389d9d1b3c89b42711e2e Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Fri, 17 Jul 2026 11:55:42 -0700
Subject: [PATCH 3/5] clang format

---
 .../Scalar/TailRecursionElimination.cpp        | 18 +++++-------------
 .../disable-tail-call-elim-for-cold-calls.ll   |  5 ++---
 2 files changed, 7 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index d7eb09d2173a0..03f8c2ffc721e 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -118,10 +118,6 @@ static bool shouldDisableTailCallsForCold(const CallBase *CB,
              CB->getCallingConv() == CallingConv::Cold))
     return true;
 
-  if (Caller && (Caller->hasFnAttribute(Attribute::Cold) ||
-                 Caller->getCallingConv() == CallingConv::Cold))
-    return true;
-
   if (!PSI || !PSI->hasProfileSummary())
     return false;
 
@@ -540,8 +536,7 @@ 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, &F, PSI, BFI))
     return nullptr;
 
   // As a special case, detect code like this:
@@ -945,13 +940,10 @@ bool TailRecursionEliminator::processBlock(BasicBlock &BB) {
   return false;
 }
 
-bool TailRecursionEliminator::eliminate(Function &F,
-                                        const TargetTransformInfo *TTI,
-                                        AliasAnalysis *AA,
-                                        OptimizationRemarkEmitter *ORE,
-                                        DomTreeUpdater &DTU,
-                                        BlockFrequencyInfo *BFI,
-                                        ProfileSummaryInfo *PSI) {
+bool TailRecursionEliminator::eliminate(
+    Function &F, const TargetTransformInfo *TTI, AliasAnalysis *AA,
+    OptimizationRemarkEmitter *ORE, DomTreeUpdater &DTU,
+    BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI) {
   if (F.getFnAttribute("disable-tail-calls").getValueAsBool())
     return false;
 
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 203ed22cbe6aa..12ad104779211 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
@@ -23,11 +23,10 @@ define void @test_coldcc_callee() {
   ret void
 }
 
-; Check that a call inside a cold enclosing function is not marked as tail when the flag is enabled.
+; Check that a call inside a cold enclosing function is marked as tail because we only suppress for cold callsites/callees (not cold callers).
 define void @test_cold_caller() cold {
 ; CHECK-LABEL: @test_cold_caller(
-; DISABLED: call void @normal_callee()
-; ENABLED: tail call void @normal_callee()
+; CHECK: tail call void @normal_callee()
   call void @normal_callee()
   ret void
 }

>From ffa798b5a731d7e2f9937e70d9002667ed243620 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Fri, 17 Jul 2026 12:47:01 -0700
Subject: [PATCH 4/5] Ensure BFI is available so profile coldness checks work

---
 .../Scalar/TailRecursionElimination.cpp       | 17 ++++---
 .../disable-tail-call-elim-for-cold-calls.ll  | 44 +++++++++++++++++--
 2 files changed, 50 insertions(+), 11 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 03f8c2ffc721e..5fd17e43c9371 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -118,10 +118,15 @@ static bool shouldDisableTailCallsForCold(const CallBase *CB,
              CB->getCallingConv() == CallingConv::Cold))
     return true;
 
+  if (Caller && (Caller->hasFnAttribute(Attribute::Cold) ||
+                 Caller->getCallingConv() == CallingConv::Cold))
+    return true;
+
   if (!PSI || !PSI->hasProfileSummary())
     return false;
 
-  if (CB && BFI && PSI->isColdCallSite(*CB, BFI))
+  if (CB && BFI &&
+      (PSI->isColdCallSite(*CB, BFI) || PSI->isColdBlock(CB->getParent(), BFI)))
     return true;
 
   return false;
@@ -483,10 +488,9 @@ class TailRecursionEliminator {
         OrigEntryCount(F.getEntryCount() ? *F.getEntryCount() : 0) {
     if (BFI) {
       // The assert is meant as API documentation for the caller.
-      assert((OrigEntryCount != 0 && OrigEntryBBFreq != 0) &&
-             "If a BFI was provided, the function should have both an entry "
-             "count that is non-zero and an entry basic block with a non-zero "
-             "frequency.");
+      assert(OrigEntryBBFreq != 0 &&
+             "If a BFI was provided, the function should have an entry "
+             "basic block with a non-zero frequency.");
     }
   }
 
@@ -1030,8 +1034,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 = (!ForceDisableBFI && UpdateFunctionEntryCount &&
-               F.getEntryCount().has_value() && *F.getEntryCount())
+  auto *BFI = (!ForceDisableBFI && F.getEntryCount().has_value())
                   ? &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.ll b/llvm/test/Transforms/TailCallElim/disable-tail-call-elim-for-cold-calls.ll
index 12ad104779211..62c00bf4114f8 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
@@ -1,5 +1,5 @@
-; RUN: opt < %s -passes=tailcallelim -disable-tail-call-elim-for-cold-calls=true -S | FileCheck %s --check-prefixes=CHECK,DISABLED
-; RUN: opt < %s -passes=tailcallelim -disable-tail-call-elim-for-cold-calls=false -S | FileCheck %s --check-prefixes=CHECK,ENABLED
+; 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 @cold_callee() cold
 declare void @normal_callee()
@@ -23,10 +23,20 @@ define void @test_coldcc_callee() {
   ret void
 }
 
-; Check that a call inside a cold enclosing function is marked as tail because we only suppress for cold callsites/callees (not cold callers).
+; Check that a call inside a cold enclosing function is not marked as tail when the flag is enabled.
 define void @test_cold_caller() cold {
 ; CHECK-LABEL: @test_cold_caller(
-; CHECK: tail call void @normal_callee()
+; DISABLED: call void @normal_callee()
+; ENABLED: tail call void @normal_callee()
+  call void @normal_callee()
+  ret void
+}
+
+; Check that a call inside an enclosing function with cold calling convention is not marked as tail when the flag is enabled.
+define coldcc void @test_coldcc_caller() {
+; CHECK-LABEL: @test_coldcc_caller(
+; DISABLED: call void @normal_callee()
+; ENABLED: tail call void @normal_callee()
   call void @normal_callee()
   ret void
 }
@@ -67,3 +77,29 @@ if.else:
   %call = call i32 @test_recursive_cold_caller(i32 %dec)
   ret i32 %call
 }
+
+; Check that a call inside a basic block verified cold via ProfileSummary and BFI (function_entry_count = 0) is not marked as tail when the flag is enabled.
+define void @test_profile_cold_block() !prof !14 {
+; CHECK-LABEL: @test_profile_cold_block(
+; DISABLED: call void @normal_callee()
+; ENABLED: tail call void @normal_callee()
+  call void @normal_callee()
+  ret void
+}
+
+!llvm.module.flags = !{!0}
+!0 = !{i32 1, !"ProfileSummary", !1}
+!1 = !{!2, !3, !4, !5, !6, !7, !8, !9}
+!2 = !{!"ProfileFormat", !"InstrProf"}
+!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}

>From 7caa480bf6e34cf0e118fff8047c7a6321302710 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Fri, 17 Jul 2026 14:58:01 -0700
Subject: [PATCH 5/5] Moved UpdateFunctionEntryCount and ForceDisableBFI
 (renamed to DisableEntryCountRecompute to match intent), to actually guard
 the function entry count update, since we compute BFI in more cases.

---
 .../Scalar/TailRecursionElimination.cpp       | 28 +++++++++++--------
 1 file changed, 16 insertions(+), 12 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 5fd17e43c9371..4ce2f08524baa 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -94,7 +94,7 @@ STATISTIC(NumTREPreventedCold,
           "Number of tail calls/recursion eliminations prevented due to cold "
           "calling convention or attribute");
 
-static cl::opt<bool> ForceDisableBFI(
+static cl::opt<bool> DisableEntryCountRecompute(
     "tre-disable-entrycount-recompute", cl::init(false), cl::Hidden,
     cl::desc("Force disabling recomputing of function entry count, on "
              "successful tail recursion elimination."));
@@ -449,6 +449,7 @@ class TailRecursionEliminator {
   DomTreeUpdater &DTU;
   BlockFrequencyInfo *const BFI;
   ProfileSummaryInfo *const PSI;
+  const bool UpdateFunctionEntryCount;
   const uint64_t OrigEntryBBFreq;
   const uint64_t OrigEntryCount;
 
@@ -481,8 +482,10 @@ class TailRecursionEliminator {
   TailRecursionEliminator(Function &F, const TargetTransformInfo *TTI,
                           AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
                           DomTreeUpdater &DTU, BlockFrequencyInfo *BFI,
-                          ProfileSummaryInfo *PSI)
+                          ProfileSummaryInfo *PSI,
+                          bool UpdateFunctionEntryCount)
       : F(F), TTI(TTI), AA(AA), ORE(ORE), DTU(DTU), BFI(BFI), PSI(PSI),
+        UpdateFunctionEntryCount(UpdateFunctionEntryCount),
         OrigEntryBBFreq(
             BFI ? BFI->getBlockFreq(&F.getEntryBlock()).getFrequency() : 0U),
         OrigEntryCount(F.getEntryCount() ? *F.getEntryCount() : 0) {
@@ -514,7 +517,7 @@ class TailRecursionEliminator {
   static bool eliminate(Function &F, const TargetTransformInfo *TTI,
                         AliasAnalysis *AA, OptimizationRemarkEmitter *ORE,
                         DomTreeUpdater &DTU, BlockFrequencyInfo *BFI,
-                        ProfileSummaryInfo *PSI);
+                        ProfileSummaryInfo *PSI, bool UpdateFunctionEntryCount);
 };
 } // namespace
 
@@ -803,7 +806,8 @@ bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
   CI->eraseFromParent();   // Remove call.
   DTU.applyUpdates({{DominatorTree::Insert, BB, HeaderBB}});
   ++NumEliminated;
-  if (OrigEntryBBFreq) {
+  if (!DisableEntryCountRecompute && UpdateFunctionEntryCount &&
+      OrigEntryBBFreq) {
     assert(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
@@ -947,7 +951,8 @@ bool TailRecursionEliminator::processBlock(BasicBlock &BB) {
 bool TailRecursionEliminator::eliminate(
     Function &F, const TargetTransformInfo *TTI, AliasAnalysis *AA,
     OptimizationRemarkEmitter *ORE, DomTreeUpdater &DTU,
-    BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI) {
+    BlockFrequencyInfo *BFI, ProfileSummaryInfo *PSI,
+    bool UpdateFunctionEntryCount) {
   if (F.getFnAttribute("disable-tail-calls").getValueAsBool())
     return false;
 
@@ -963,7 +968,8 @@ bool TailRecursionEliminator::eliminate(
     return MadeChange;
 
   // Change any tail recursive calls to loops.
-  TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU, BFI, PSI);
+  TailRecursionEliminator TRE(F, TTI, AA, ORE, DTU, BFI, PSI,
+                              UpdateFunctionEntryCount);
 
   for (BasicBlock &BB : F)
     MadeChange |= TRE.processBlock(BB);
@@ -1002,13 +1008,11 @@ struct TailCallElim : public FunctionPass {
     // UpdateStrategy to Lazy if we find it profitable later.
     DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Eager);
 
-    auto *PSIWP = getAnalysisIfAvailable<ProfileSummaryInfoWrapperPass>();
-    auto *PSI = PSIWP ? &PSIWP->getPSI() : nullptr;
     return TailRecursionEliminator::eliminate(
         F, &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F),
         &getAnalysis<AAResultsWrapperPass>().getAAResults(),
         &getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE(), DTU,
-        /*BFI=*/nullptr, PSI);
+        /*BFI=*/nullptr, /*PSI=*/nullptr, /*UpdateFunctionEntryCount=*/false);
   }
 };
 } // namespace
@@ -1034,7 +1038,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 = (!ForceDisableBFI && F.getEntryCount().has_value())
+  auto *BFI = F.getEntryCount().has_value()
                   ? &AM.getResult<BlockFrequencyAnalysis>(F)
                   : nullptr;
   auto &MAMProxy = AM.getResult<ModuleAnalysisManagerFunctionProxy>(F);
@@ -1046,8 +1050,8 @@ PreservedAnalyses TailCallElimPass::run(Function &F,
   // UpdateStrategy based on some test results. It is feasible to switch the
   // UpdateStrategy to Lazy if we find it profitable later.
   DomTreeUpdater DTU(DT, PDT, DomTreeUpdater::UpdateStrategy::Eager);
-  bool Changed =
-      TailRecursionEliminator::eliminate(F, &TTI, &AA, &ORE, DTU, BFI, PSI);
+  bool Changed = TailRecursionEliminator::eliminate(
+      F, &TTI, &AA, &ORE, DTU, BFI, PSI, UpdateFunctionEntryCount);
 
   if (!Changed)
     return PreservedAnalyses::all();



More information about the llvm-commits mailing list