[llvm] [CodeGen] Preserve pseudo-probe context when tail merging (PR #219502)
Kunal Pathak via llvm-commits
llvm-commits at lists.llvm.org
Fri Aug 28 08:39:57 PDT 2026
https://github.com/kunalspathak created https://github.com/llvm/llvm-project/pull/219502
BranchFolder normally ignores debug locations when deciding whether machine instructions are identical. That allowed it to merge `PSEUDO_PROBE `instructions and calls carrying packed pseudo-probe discriminators even when they came from different inline contexts. Merging their debug locations then kept only the common outer ancestry, which moved profile attribution to the wrong callsite.
Keep tail merging enabled, but stop matching at a probe-sensitive instruction unless both instructions have the same source and inline context. Call probes also require the same packed discriminator. Tails from the same probe context remain eligible for merging.
Add MIR coverage for block probes and call probes, including checks that same-context tails still merge, under both the legacy and new pass managers.
Assisted by: Claude
>From 1c08e59ac0774b2620d2101ff45ee6e31a67cb81 Mon Sep 17 00:00:00 2001
From: Kunal Pathak <kupathak at meta.com>
Date: Tue, 18 Aug 2026 17:34:40 -0700
Subject: [PATCH] [CodeGen] Preserve pseudo-probe context when tail merging
BranchFolder normally ignores debug locations when deciding whether machine
instructions are identical. That allowed it to merge PSEUDO_PROBE instructions
and calls carrying packed pseudo-probe discriminators even when they came from
different inline contexts. Merging their debug locations then kept only the
common outer ancestry, which moved profile attribution to the wrong callsite.
Keep tail merging enabled, but stop matching at a probe-sensitive instruction
unless both instructions have the same source and inline context. Call probes
also require the same packed discriminator. Tails from the same probe context
remain eligible for merging.
Add MIR coverage for block probes and call probes, including checks that
same-context tails still merge, under both the legacy and new pass managers.
Test Plan:
- ninja -j8 llc FileCheck
- Run both RUN lines in pseudo-probe-tail-merging.mir
- Verify the C++ EH reproducer retains both full cleanup contexts
---
llvm/lib/CodeGen/BranchFolding.cpp | 29 ++++
.../CodeGen/X86/pseudo-probe-tail-merging.mir | 142 ++++++++++++++++++
2 files changed, 171 insertions(+)
create mode 100644 llvm/test/CodeGen/X86/pseudo-probe-tail-merging.mir
diff --git a/llvm/lib/CodeGen/BranchFolding.cpp b/llvm/lib/CodeGen/BranchFolding.cpp
index 807f2971f6d6f..5145d25799306 100644
--- a/llvm/lib/CodeGen/BranchFolding.cpp
+++ b/llvm/lib/CodeGen/BranchFolding.cpp
@@ -374,6 +374,34 @@ static bool countsAsInstruction(const MachineInstr &MI) {
return !(MI.isDebugInstr() || MI.isCFIInstruction());
}
+static bool isPseudoProbeSensitiveInstruction(const MachineInstr &MI) {
+ if (MI.isPseudoProbe())
+ return true;
+ if (!MI.isCall())
+ return false;
+ const DILocation *DL = MI.getDebugLoc();
+ return DL && DILocation::isPseudoProbeDiscriminator(DL->getDiscriminator());
+}
+
+static bool haveSamePseudoProbeContext(const MachineInstr &MI1,
+ const MachineInstr &MI2) {
+ bool IsSensitive1 = isPseudoProbeSensitiveInstruction(MI1);
+ bool IsSensitive2 = isPseudoProbeSensitiveInstruction(MI2);
+ if (!IsSensitive1 && !IsSensitive2)
+ return true;
+ if (IsSensitive1 != IsSensitive2)
+ return false;
+
+ DebugLoc DL1 = MI1.getDebugLoc();
+ DebugLoc DL2 = MI2.getDebugLoc();
+ if (!DL1.isSameSourceLocation(DL2))
+ return false;
+
+ // A call probe's identity is encoded in its discriminator, which
+ // isSameSourceLocation intentionally ignores.
+ return !MI1.isCall() || DL1->getDiscriminator() == DL2->getDiscriminator();
+}
+
/// Iterate backwards from the given iterator \p I, towards the beginning of the
/// block. If a MI satisfying 'countsAsInstruction' is found, return an iterator
/// pointing to that MI. If no such MI is found, return the end iterator.
@@ -408,6 +436,7 @@ static unsigned ComputeCommonTailLength(MachineBasicBlock *MBB1,
if (MBBI1 == MBB1->end() || MBBI2 == MBB2->end())
break;
if (!MBBI1->isIdenticalTo(*MBBI2) ||
+ !haveSamePseudoProbeContext(*MBBI1, *MBBI2) ||
// FIXME: This check is dubious. It's used to get around a problem where
// people incorrectly expect inline asm directives to remain in the same
// relative order. This is untenable because normal compiler
diff --git a/llvm/test/CodeGen/X86/pseudo-probe-tail-merging.mir b/llvm/test/CodeGen/X86/pseudo-probe-tail-merging.mir
new file mode 100644
index 0000000000000..68ad8953cd4f9
--- /dev/null
+++ b/llvm/test/CodeGen/X86/pseudo-probe-tail-merging.mir
@@ -0,0 +1,142 @@
+# RUN: llc -mtriple=x86_64-- -run-pass=branch-folder -o - %s | FileCheck %s
+# RUN: llc -mtriple=x86_64-- -passes="require<profile-summary>,function(machine-function(branch-folder<enable-tail-merge>))" -o - %s | FileCheck %s
+
+--- |
+ target triple = "x86_64-unknown-unknown"
+
+ declare void @callee()
+
+ define void @probe_different_context() !dbg !4 {
+ ret void
+ }
+
+ define void @probe_same_context() !dbg !10 {
+ ret void
+ }
+
+ define void @call_different_context() !dbg !14 {
+ ret void
+ }
+
+ define void @call_same_context() !dbg !21 {
+ ret void
+ }
+
+ !llvm.dbg.cu = !{!0}
+ !llvm.module.flags = !{!2}
+
+ !0 = distinct !DICompileUnit(language: DW_LANG_C, file: !1, producer: "branch-folder test", isOptimized: true, runtimeVersion: 0, emissionKind: LineTablesOnly)
+ !1 = !DIFile(filename: "pseudo-probe-tail-merging.cpp", directory: "/")
+ !2 = !{i32 2, !"Debug Info Version", i32 3}
+ !3 = !DISubroutineType(types: !{})
+
+ !4 = distinct !DISubprogram(name: "probe_different_context", scope: !1, file: !1, line: 1, type: !3, scopeLine: 1, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !5 = distinct !DISubprogram(name: "inlined_probe", scope: !1, file: !1, line: 2, type: !3, scopeLine: 2, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !6 = distinct !DILocation(line: 10, column: 1, scope: !4)
+ !7 = distinct !DILocation(line: 20, column: 1, scope: !4)
+ !8 = !DILocation(line: 2, column: 1, scope: !5, inlinedAt: !6)
+ !9 = !DILocation(line: 2, column: 1, scope: !5, inlinedAt: !7)
+
+ !10 = distinct !DISubprogram(name: "probe_same_context", scope: !1, file: !1, line: 30, type: !3, scopeLine: 30, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !11 = distinct !DISubprogram(name: "inlined_same_probe", scope: !1, file: !1, line: 31, type: !3, scopeLine: 31, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !12 = distinct !DILocation(line: 32, column: 1, scope: !10)
+ !13 = !DILocation(line: 31, column: 1, scope: !11, inlinedAt: !12)
+
+ !14 = distinct !DISubprogram(name: "call_different_context", scope: !1, file: !1, line: 40, type: !3, scopeLine: 40, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !15 = distinct !DISubprogram(name: "inlined_call", scope: !1, file: !1, line: 41, type: !3, scopeLine: 41, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !16 = !DILexicalBlockFile(scope: !15, file: !1, discriminator: 455082007)
+ !17 = distinct !DILocation(line: 42, column: 1, scope: !14)
+ !18 = distinct !DILocation(line: 43, column: 1, scope: !14)
+ !19 = !DILocation(line: 41, column: 1, scope: !16, inlinedAt: !17)
+ !20 = !DILocation(line: 41, column: 1, scope: !16, inlinedAt: !18)
+
+ !21 = distinct !DISubprogram(name: "call_same_context", scope: !1, file: !1, line: 50, type: !3, scopeLine: 50, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !22 = distinct !DISubprogram(name: "inlined_same_call", scope: !1, file: !1, line: 51, type: !3, scopeLine: 51, spFlags: DISPFlagDefinition | DISPFlagOptimized, unit: !0)
+ !23 = !DILexicalBlockFile(scope: !22, file: !1, discriminator: 455082007)
+ !24 = distinct !DILocation(line: 52, column: 1, scope: !21)
+ !25 = !DILocation(line: 51, column: 1, scope: !23, inlinedAt: !24)
+...
+
+---
+name: probe_different_context
+body: |
+ ; CHECK-LABEL: name: probe_different_context
+ ; CHECK-COUNT-2: PSEUDO_PROBE 1, 1, 0, 0
+ bb.0:
+ liveins: $al
+ TEST8rr killed $al, killed $al, implicit-def $eflags
+ JCC_1 %bb.2, 5, implicit killed $eflags
+
+ bb.1:
+ PSEUDO_PROBE 1, 1, 0, 0, debug-location !8
+ NOOP
+ RET 0
+
+ bb.2:
+ PSEUDO_PROBE 1, 1, 0, 0, debug-location !9
+ NOOP
+ RET 0
+...
+
+---
+name: probe_same_context
+body: |
+ ; CHECK-LABEL: name: probe_same_context
+ ; CHECK-COUNT-1: PSEUDO_PROBE 2, 1, 0, 0
+ bb.0:
+ liveins: $al
+ TEST8rr killed $al, killed $al, implicit-def $eflags
+ JCC_1 %bb.2, 5, implicit killed $eflags
+
+ bb.1:
+ PSEUDO_PROBE 2, 1, 0, 0, debug-location !13
+ NOOP
+ RET 0
+
+ bb.2:
+ PSEUDO_PROBE 2, 1, 0, 0, debug-location !13
+ NOOP
+ RET 0
+...
+
+---
+name: call_different_context
+body: |
+ ; CHECK-LABEL: name: call_different_context
+ ; CHECK-COUNT-2: CALL64pcrel32 @callee
+ bb.0:
+ liveins: $al
+ TEST8rr killed $al, killed $al, implicit-def $eflags
+ JCC_1 %bb.2, 5, implicit killed $eflags
+
+ bb.1:
+ CALL64pcrel32 @callee, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, debug-location !19
+ NOOP
+ RET 0
+
+ bb.2:
+ CALL64pcrel32 @callee, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, debug-location !20
+ NOOP
+ RET 0
+...
+
+---
+name: call_same_context
+body: |
+ ; CHECK-LABEL: name: call_same_context
+ ; CHECK-COUNT-1: CALL64pcrel32 @callee
+ bb.0:
+ liveins: $al
+ TEST8rr killed $al, killed $al, implicit-def $eflags
+ JCC_1 %bb.2, 5, implicit killed $eflags
+
+ bb.1:
+ CALL64pcrel32 @callee, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, debug-location !25
+ NOOP
+ RET 0
+
+ bb.2:
+ CALL64pcrel32 @callee, csr_64, implicit $rsp, implicit $ssp, implicit-def $rsp, implicit-def $ssp, debug-location !25
+ NOOP
+ RET 0
+...
More information about the llvm-commits
mailing list