[llvm-branch-commits] [TailCallElim] Add profile annotations to return value selects (PR #202518)
via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Jun 9 00:08:52 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-llvm-transforms
Author: Aiden Grossman (boomanaiden154)
<details>
<summary>Changes</summary>
Tail call elimination in some cases can create selects on possible
return values conditioned on whether or not execution is currently in
what was a recursive call. That is equal to the probability with which
we recurse, which in turn can be computed from the block frequencies of
blocks that recurse and blocks that directly return.
---
Full diff: https://github.com/llvm/llvm-project/pull/202518.diff
3 Files Affected:
- (modified) llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp (+33)
- (added) llvm/test/Transforms/TailCallElim/return-value-select-pgo.ll (+121)
- (modified) llvm/utils/profcheck-xfail.txt (-6)
``````````diff
diff --git a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
index 98ba75d1d45ae..16e6a312eeda0 100644
--- a/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
+++ b/llvm/lib/Transforms/Scalar/TailRecursionElimination.cpp
@@ -54,6 +54,7 @@
#include "llvm/ADT/SmallPtrSet.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Analysis/BlockFrequencyInfo.h"
+#include "llvm/Analysis/BlockFrequencyInfoImpl.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/Analysis/GlobalsModRef.h"
#include "llvm/Analysis/InstructionSimplify.h"
@@ -73,9 +74,12 @@
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/Instructions.h"
#include "llvm/IR/IntrinsicInst.h"
+#include "llvm/IR/MDBuilder.h"
#include "llvm/IR/Module.h"
+#include "llvm/IR/ProfDataUtils.h"
#include "llvm/InitializePasses.h"
#include "llvm/Pass.h"
+#include "llvm/Support/BlockFrequency.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/raw_ostream.h"
@@ -426,6 +430,11 @@ class TailRecursionEliminator {
// to either propagate RetPN or select a new return value.
SmallVector<SelectInst *, 8> RetSelects;
+ // Keep track of the sum of frequencies of blocks that have calls eliminated
+ // so we can synthesize branch weights later that require information on
+ // recursion frequency.
+ uint64_t EliminateBlocksFrequencySum;
+
// The below are shared state needed when performing accumulator recursion.
// There values should be populated by insertAccumulator the first time we
// find an elimination that requires an accumulator.
@@ -680,6 +689,9 @@ bool TailRecursionEliminator::eliminateCall(CallInst *CI) {
BasicBlock *BB = Ret->getParent();
+ if (BFI)
+ EliminateBlocksFrequencySum += BFI->getBlockFreq(BB).getFrequency();
+
using namespace ore;
ORE->emit([&]() {
return OptimizationRemark(DEBUG_TYPE, "tailcall-recursion", CI)
@@ -858,6 +870,27 @@ void TailRecursionEliminator::cleanupAndFinalize() {
}
}
}
+
+ if (BFI) {
+ uint64_t BaseCaseBlocksFrequencySum = 0;
+ for (BasicBlock &BB : F)
+ if (isa<ReturnInst>(BB.getTerminator()))
+ BaseCaseBlocksFrequencySum += BFI->getBlockFreq(&BB).getFrequency();
+
+ BlockFrequencyInfoImplBase::Distribution Distribution;
+ Distribution.addExit(0, EliminateBlocksFrequencySum);
+ Distribution.addExit(1, BaseCaseBlocksFrequencySum);
+ if (Distribution.Total == 0)
+ return;
+ Distribution.normalize();
+ MDBuilder MDB(F.getContext());
+ MDNode *BranchWeights = MDB.createBranchWeights(
+ {static_cast<uint32_t>(Distribution.Weights[0].Amount),
+ static_cast<uint32_t>(Distribution.Weights[1].Amount)},
+ false);
+ for (SelectInst *SI : RetSelects)
+ SI->setMetadata(LLVMContext::MD_prof, BranchWeights);
+ }
}
}
diff --git a/llvm/test/Transforms/TailCallElim/return-value-select-pgo.ll b/llvm/test/Transforms/TailCallElim/return-value-select-pgo.ll
new file mode 100644
index 0000000000000..9302142226970
--- /dev/null
+++ b/llvm/test/Transforms/TailCallElim/return-value-select-pgo.ll
@@ -0,0 +1,121 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 6
+; RUN: opt < %s -passes=tailcallelim -S | FileCheck %s
+
+define i32 @test1(i32 %c) !prof !0 {
+; CHECK-LABEL: define i32 @test1(
+; CHECK-SAME: i32 [[C:%.*]]) !prof [[PROF0:![0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*]]:
+; CHECK-NEXT: br label %[[TAILRECURSE:.*]]
+; CHECK: [[TAILRECURSE]]:
+; CHECK-NEXT: [[C_TR:%.*]] = phi i32 [ [[C]], %[[ENTRY]] ], [ [[TMP:%.*]], %[[DEFAULT:.*]] ]
+; CHECK-NEXT: [[RET_TR:%.*]] = phi i32 [ poison, %[[ENTRY]] ], [ [[CURRENT_RET_TR:%.*]], %[[DEFAULT]] ]
+; CHECK-NEXT: [[RET_KNOWN_TR:%.*]] = phi i1 [ false, %[[ENTRY]] ], [ true, %[[DEFAULT]] ]
+; CHECK-NEXT: switch i32 [[C_TR]], label %[[DEFAULT]] [
+; CHECK-NEXT: i32 0, label %[[ONZERO:.*]]
+; CHECK-NEXT: ], !prof [[PROF1:![0-9]+]]
+; CHECK: [[ONZERO]]:
+; CHECK-NEXT: [[CURRENT_RET_TR1:%.*]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 0, !prof [[PROF2:![0-9]+]]
+; CHECK-NEXT: ret i32 [[CURRENT_RET_TR1]]
+; CHECK: [[DEFAULT]]:
+; CHECK-NEXT: [[TMP]] = add i32 [[C_TR]], -1
+; CHECK-NEXT: [[CURRENT_RET_TR]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 1, !prof [[PROF2]]
+; CHECK-NEXT: br label %[[TAILRECURSE]]
+;
+entry:
+ switch i32 %c, label %default [ i32 0, label %onzero ], !prof !1
+onzero:
+ ret i32 0
+default:
+ %tmp = add i32 %c, -1
+ %tmp.2 = call i32 @test1(i32 %tmp)
+ ret i32 1
+}
+
+define i32 @test2(i32 %c) !prof !0 {
+; CHECK-LABEL: define i32 @test2(
+; CHECK-SAME: i32 [[C:%.*]]) !prof [[PROF3:![0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*]]:
+; CHECK-NEXT: br label %[[TAILRECURSE:.*]]
+; CHECK: [[TAILRECURSE]]:
+; CHECK-NEXT: [[C_TR:%.*]] = phi i32 [ [[C]], %[[ENTRY]] ], [ [[TMP:%.*]], %[[DEFAULT:.*]] ]
+; CHECK-NEXT: [[RET_TR:%.*]] = phi i32 [ poison, %[[ENTRY]] ], [ [[CURRENT_RET_TR:%.*]], %[[DEFAULT]] ]
+; CHECK-NEXT: [[RET_KNOWN_TR:%.*]] = phi i1 [ false, %[[ENTRY]] ], [ true, %[[DEFAULT]] ]
+; CHECK-NEXT: switch i32 [[C_TR]], label %[[DEFAULT]] [
+; CHECK-NEXT: i32 0, label %[[ONZERO:.*]]
+; CHECK-NEXT: i32 1, label %[[ONONE:.*]]
+; CHECK-NEXT: ], !prof [[PROF4:![0-9]+]]
+; CHECK: [[ONZERO]]:
+; CHECK-NEXT: [[CURRENT_RET_TR1:%.*]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 0, !prof [[PROF5:![0-9]+]]
+; CHECK-NEXT: ret i32 [[CURRENT_RET_TR1]]
+; CHECK: [[ONONE]]:
+; CHECK-NEXT: [[CURRENT_RET_TR2:%.*]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 1, !prof [[PROF5]]
+; CHECK-NEXT: ret i32 [[CURRENT_RET_TR2]]
+; CHECK: [[DEFAULT]]:
+; CHECK-NEXT: [[TMP]] = add i32 [[C_TR]], -1
+; CHECK-NEXT: [[CURRENT_RET_TR]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 2, !prof [[PROF5]]
+; CHECK-NEXT: br label %[[TAILRECURSE]]
+;
+entry:
+ switch i32 %c, label %default [ i32 0, label %onzero i32 1, label %onone], !prof !2
+onzero:
+ ret i32 0
+onone:
+ ret i32 1
+default:
+ %tmp = add i32 %c, -1
+ %tmp.2 = call i32 @test2(i32 %tmp)
+ ret i32 2
+}
+
+define i32 @test3(i32 %c) !prof !0 {
+; CHECK-LABEL: define i32 @test3(
+; CHECK-SAME: i32 [[C:%.*]]) !prof [[PROF6:![0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*]]:
+; CHECK-NEXT: br label %[[TAILRECURSE:.*]]
+; CHECK: [[TAILRECURSE]]:
+; CHECK-NEXT: [[C_TR:%.*]] = phi i32 [ [[C]], %[[ENTRY]] ], [ [[TMP:%.*]], %[[ONONE:.*]] ], [ [[TMP_3:%.*]], %[[DEFAULT:.*]] ]
+; CHECK-NEXT: [[RET_TR:%.*]] = phi i32 [ poison, %[[ENTRY]] ], [ [[CURRENT_RET_TR:%.*]], %[[ONONE]] ], [ [[CURRENT_RET_TR1:%.*]], %[[DEFAULT]] ]
+; CHECK-NEXT: [[RET_KNOWN_TR:%.*]] = phi i1 [ false, %[[ENTRY]] ], [ true, %[[ONONE]] ], [ true, %[[DEFAULT]] ]
+; CHECK-NEXT: switch i32 [[C_TR]], label %[[DEFAULT]] [
+; CHECK-NEXT: i32 0, label %[[ONZERO:.*]]
+; CHECK-NEXT: i32 1, label %[[ONONE]]
+; CHECK-NEXT: ], !prof [[PROF4]]
+; CHECK: [[ONZERO]]:
+; CHECK-NEXT: [[CURRENT_RET_TR2:%.*]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 0, !prof [[PROF7:![0-9]+]]
+; CHECK-NEXT: ret i32 [[CURRENT_RET_TR2]]
+; CHECK: [[ONONE]]:
+; CHECK-NEXT: [[TMP]] = add i32 [[C_TR]], 1
+; CHECK-NEXT: [[CURRENT_RET_TR]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 1, !prof [[PROF7]]
+; CHECK-NEXT: br label %[[TAILRECURSE]]
+; CHECK: [[DEFAULT]]:
+; CHECK-NEXT: [[TMP_3]] = add i32 [[C_TR]], -2
+; CHECK-NEXT: [[CURRENT_RET_TR1]] = select i1 [[RET_KNOWN_TR]], i32 [[RET_TR]], i32 2, !prof [[PROF7]]
+; CHECK-NEXT: br label %[[TAILRECURSE]]
+;
+entry:
+ switch i32 %c, label %default [ i32 0, label %onzero i32 1, label %onone], !prof !2
+onzero:
+ ret i32 0
+onone:
+ %tmp = add i32 %c, 1
+ %tmp.2 = call i32 @test3(i32 %tmp)
+ ret i32 1
+default:
+ %tmp.3 = add i32 %c, -2
+ %tmp.4 = call i32 @test3(i32 %tmp.3)
+ ret i32 2
+}
+
+!0 = !{!"function_entry_count", i64 1000}
+!1 = !{!"branch_weights", i32 2, i32 3}
+!2 = !{!"branch_weights", i32 2, i32 3, i32 5}
+;.
+; CHECK: [[PROF0]] = !{!"function_entry_count", i64 600}
+; CHECK: [[PROF1]] = !{!"branch_weights", i32 2, i32 3}
+; CHECK: [[PROF2]] = !{!"branch_weights", i32 429496730, i32 644245095}
+; CHECK: [[PROF3]] = !{!"function_entry_count", i64 800}
+; CHECK: [[PROF4]] = !{!"branch_weights", i32 2, i32 3, i32 5}
+; CHECK: [[PROF5]] = !{!"branch_weights", i32 214748365, i32 858993459}
+; CHECK: [[PROF6]] = !{!"function_entry_count", i64 300}
+; CHECK: [[PROF7]] = !{!"branch_weights", i32 751619277, i32 322122547}
+;.
diff --git a/llvm/utils/profcheck-xfail.txt b/llvm/utils/profcheck-xfail.txt
index 78a19b0ae3db6..8130fa4d6f0f1 100644
--- a/llvm/utils/profcheck-xfail.txt
+++ b/llvm/utils/profcheck-xfail.txt
@@ -159,12 +159,6 @@ Transforms/ScalarizeMaskedMemIntrin/X86/expand-masked-expandload.ll
Transforms/ScalarizeMaskedMemIntrin/X86/expand-masked-load.ll
Transforms/ScalarizeMaskedMemIntrin/X86/expand-masked-store.ll
Transforms/StackProtector/cross-dso-cfi-stack-chk-fail.ll
-Transforms/TailCallElim/2010-06-26-MultipleReturnValues.ll
-Transforms/TailCallElim/accum_recursion.ll
-Transforms/TailCallElim/basic.ll
-Transforms/TailCallElim/debugloc.ll
-Transforms/TailCallElim/dropping_debugloc_acc_rec_inst_rnew.ll
-Transforms/TailCallElim/inf-recursion.ll
Transforms/Util/control-flow-hub-finalize-same-succ-crash.ll
Transforms/Util/libcalls-opt-remarks.ll
Transforms/Util/lowerswitch.ll
``````````
</details>
https://github.com/llvm/llvm-project/pull/202518
More information about the llvm-branch-commits
mailing list