[llvm] [LLVM/PGO] Preserve pseudo-probe discriminators in cloneWithBaseDiscriminator (PR #178267)

[via llvm-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-debuginfo

Author: Grigory Pastukhov (grigorypas)

<details>
<summary>Changes</summary>

In the O0 pipeline, `SampleProfileProbePass` runs before `AddDiscriminatorsPass`. The 
pseudo-probe pass stores probe IDs in the discriminator field using a special encoding 
(bits [2:0] = 0x7 as marker). When `AddDiscriminatorsPass` later processes instructions 
on the same source line, it calls `cloneWithBaseDiscriminator()` which overwrites these 
pseudo-probe discriminators with DWARF base discriminators, breaking sample profile 
matching.

This patch adds a check to `cloneWithBaseDiscriminator()` to detect and preserve 
pseudo-probe discriminators, similar to the existing check in 
`cloneByMultiplyingDuplicationFactor()`.

In optimized pipelines (O1+), AddDiscriminators runs before pseudo-probe instrumentation,
so this issue doesn't manifest. However, the fix makes the code robust regardless of 
pass ordering.


---
Full diff: https://github.com/llvm/llvm-project/pull/178267.diff


2 Files Affected:

- (modified) llvm/include/llvm/IR/DebugInfoMetadata.h (+6) 
- (added) llvm/test/Transforms/AddDiscriminators/pseudo-probe-discriminator.ll (+33) 


``````````diff
diff --git a/llvm/include/llvm/IR/DebugInfoMetadata.h b/llvm/include/llvm/IR/DebugInfoMetadata.h
index a41379b974069..220b6e28fd225 100644
--- a/llvm/include/llvm/IR/DebugInfoMetadata.h
+++ b/llvm/include/llvm/IR/DebugInfoMetadata.h
@@ -3010,6 +3010,12 @@ unsigned DILocation::getCopyIdentifier() const {
 
 std::optional<const DILocation *>
 DILocation::cloneWithBaseDiscriminator(unsigned D) const {
+  // Do not interfere with pseudo probes. Pseudo probe at a callsite uses
+  // the dwarf discriminator to store pseudo probe related information,
+  // such as the probe id.
+  if (isPseudoProbeDiscriminator(getDiscriminator()))
+    return this;
+
   unsigned BD, DF, CI;
 
   if (EnableFSDiscriminator) {
diff --git a/llvm/test/Transforms/AddDiscriminators/pseudo-probe-discriminator.ll b/llvm/test/Transforms/AddDiscriminators/pseudo-probe-discriminator.ll
new file mode 100644
index 0000000000000..14285793be153
--- /dev/null
+++ b/llvm/test/Transforms/AddDiscriminators/pseudo-probe-discriminator.ll
@@ -0,0 +1,33 @@
+; RUN: opt < %s -passes=add-discriminators -S | FileCheck %s
+
+; Test that AddDiscriminators preserves pseudo-probe discriminators.
+; Pseudo-probe discriminators use bits [2:0] = 0x7 as a marker.
+
+define void @foo() !dbg !4 {
+entry:
+  call void @bar(), !dbg !10
+  call void @bar(), !dbg !11
+  ret void
+}
+
+declare void @bar()
+
+!llvm.dbg.cu = !{!0}
+!llvm.module.flags = !{!7, !8}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus, file: !1, isOptimized: false, runtimeVersion: 0, emissionKind: NoDebug)
+!1 = !DIFile(filename: "test.cpp", directory: "/tmp")
+!4 = distinct !DISubprogram(name: "foo", scope: !1, file: !1, line: 1, type: !5, scopeLine: 1, unit: !0)
+!5 = !DISubroutineType(types: !6)
+!6 = !{null}
+!7 = !{i32 2, !"Dwarf Version", i32 4}
+!8 = !{i32 2, !"Debug Info Version", i32 3}
+
+; Two calls on the same line with pseudo-probe discriminators (bits [2:0] = 0x7)
+!9 = !DILexicalBlockFile(scope: !4, file: !1, discriminator: 455081999)
+!10 = !DILocation(line: 2, column: 3, scope: !9)
+!12 = !DILexicalBlockFile(scope: !4, file: !1, discriminator: 455082007)
+!11 = !DILocation(line: 2, column: 3, scope: !12)
+
+; CHECK-DAG: !DILexicalBlockFile({{.*}}discriminator: 455081999)
+; CHECK-DAG: !DILexicalBlockFile({{.*}}discriminator: 455082007)

``````````

</details>


https://github.com/llvm/llvm-project/pull/178267