[llvm] [MemProf] Attach value profile metadata to the IR using CalleeGuids. (PR #141164)

[Snehasish Kumar via llvm-commits]

https://github.com/snehasish created https://github.com/llvm/llvm-project/pull/141164

Use the newly introduced CalleeGuids in CallSiteInfo to annotate the IR
where necessary with value profile metadata. Use a synthetic count of 1
since we don't have actual counts in the profile collection.

>From 2858bad04e66ce04ba0270da6d34af088b616df2 Mon Sep 17 00:00:00 2001
From: Snehasish Kumar <snehasishk at google.com>
Date: Thu, 22 May 2025 17:00:14 -0700
Subject: [PATCH] [MemProf] Attach value profile metadata to the IR using
 CalleeGuids.

Use the newly introduced CalleeGuids in CallSiteInfo to annotate the IR
where necessary with value profile metadata. Use a synthetic count of 1
since we don't have actual counts in the profile collection.
---
 .../Instrumentation/MemProfiler.cpp           | 35 ++++++++++++++++++
 .../memprof_annotate_indirect_call_yaml.test  | 37 +++++++++++++++++++
 2 files changed, 72 insertions(+)
 create mode 100644 llvm/test/Transforms/PGOProfile/memprof_annotate_indirect_call_yaml.test

diff --git a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
index 6538311571529..29206781f1743 100644
--- a/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
+++ b/llvm/lib/Transforms/Instrumentation/MemProfiler.cpp
@@ -1209,6 +1209,41 @@ readMemprof(Module &M, Function &F, IndexedInstrProfReader *MemProfReader,
                                                InlinedCallStack)) {
           NumOfMemProfMatchedCallSites++;
           addCallsiteMetadata(I, InlinedCallStack, Ctx);
+
+          // Check if this is an indirect call and we have GUID information
+          // from CallSiteInfo to attach value profile metadata
+          if (!CalledFunction) {
+            // This is an indirect call, look for CallSites with matching stacks
+            // that have CalleeGuids information
+            for (auto &CS : MemProfRec->CallSites) {
+              if (!CS.CalleeGuids.empty() && stackFrameIncludesInlinedCallStack(
+                                                 CS.Frames, InlinedCallStack)) {
+                // Create value profile data from the CalleeGuids
+                SmallVector<InstrProfValueData, 4> VDs;
+                uint64_t TotalCount = 0;
+
+                for (GlobalValue::GUID CalleeGUID : CS.CalleeGuids) {
+                  // For MemProf, we don't have actual call counts, so we assign
+                  // a weight of 1 to each potential target. This provides the
+                  // information needed for indirect call promotion without
+                  // specific count data.
+                  InstrProfValueData VD;
+                  VD.Value = CalleeGUID;
+                  VD.Count = 1; // Weight for ICP decision making
+                  VDs.push_back(VD);
+                  TotalCount += VD.Count;
+                }
+
+                if (!VDs.empty()) {
+                  // Attach value profile metadata for indirect call targets
+                  annotateValueSite(M, I, VDs, TotalCount,
+                                    IPVK_IndirectCallTarget, VDs.size());
+                }
+                break;
+              }
+            }
+          }
+
           // Only need to find one with a matching call stack and add a single
           // callsite metadata.
 
diff --git a/llvm/test/Transforms/PGOProfile/memprof_annotate_indirect_call_yaml.test b/llvm/test/Transforms/PGOProfile/memprof_annotate_indirect_call_yaml.test
new file mode 100644
index 0000000000000..dcea8f14f5fe9
--- /dev/null
+++ b/llvm/test/Transforms/PGOProfile/memprof_annotate_indirect_call_yaml.test
@@ -0,0 +1,37 @@
+; Make sure that we can ingest the MemProf profile in YAML with CalleeGuids
+; and annotate an indirect call with value profile metadata.
+
+; RUN: split-file %s %t
+; RUN: llvm-profdata merge --memprof-version=4 %t/memprof_annotate_indirect_call_yaml.yaml -o %t/memprof_annotate_indirect_call_yaml.memprofdata
+; RUN: opt < %t/memprof_annotate_indirect_call_yaml.ll -passes='memprof-use<profile-filename=%t/memprof_annotate_indirect_call_yaml.memprofdata>' -S 2>&1 | FileCheck %s
+
+;--- memprof_annotate_indirect_call_yaml.yaml
+---
+HeapProfileRecords:
+  - GUID:            _Z3barv
+    AllocSites:      []
+    CallSites:
+      - Frames:
+          - { Function: _Z3barv, LineOffset: 3, Column: 5, IsInlineFrame: false }
+        CalleeGuids:   [0x123456789abcdef0, 0x23456789abcdef01]
+...
+;--- memprof_annotate_indirect_call_yaml.ll
+define dso_local void @_Z3barv() !dbg !4 {
+entry:
+  %fp = alloca ptr, align 8
+  %0 = load ptr, ptr %fp, align 8
+  call void %0(), !dbg !5
+; CHECK: call void %0(), {{.*}} !prof ![[PROF:[0-9]+]]
+  ret void
+}
+
+; CHECK: ![[PROF]] = !{!"VP", i32 0, i64 2, i64 1311768467463790320, i64 1, i64 2541551405711093505, i64 1}
+
+!llvm.module.flags = !{!2, !3}
+
+!0 = distinct !DICompileUnit(language: DW_LANG_C_plus_plus_14, file: !1)
+!1 = !DIFile(filename: "t", directory: "/")
+!2 = !{i32 7, !"Dwarf Version", i32 5}
+!3 = !{i32 2, !"Debug Info Version", i32 3}
+!4 = distinct !DISubprogram(name: "bar", linkageName: "_Z3barv", scope: !1, file: !1, line: 1, unit: !0)
+!5 = !DILocation(line: 4, column: 5, scope: !4)