[llvm] [MemProf] Fix an assertion when writing distributed index for aliasee (PR #122946)

[Teresa Johnson via llvm-commits]

https://github.com/teresajohnson created https://github.com/llvm/llvm-project/pull/122946

The ThinLTO index bitcode writer uses a helper forEachSummary to manage
preparation and writing of summaries needed for each distributed index
file. For alias summaries, it invokes the provided callback for the
aliasee as well, as we at least need to produce a value id for the
alias's summary. However, all summary generation for the aliasee itself
should be skipped on calls when IsAliasee is true. We invoke the
callback again if that value's summary is to be written as well.

We were asserting in debug mode when invoking collectMemProfCallStacks,
because a given stack id index was not in the StackIdIndicesToIndex
map. It was not added because the forEachSummary invocation that records
these ids in the map (invoked from the IndexBitcodeWriter constructor)
was correctly skipping this handling when invoked for aliasees. We need
the same guard in the invocation that calls collectMemProfCallStacks.

Note that this doesn't cause any real problems in a non-asserts build
as the missing map lookup will return the default 0 value from the map,
which isn't used since we don't actually write the corresponding
summary.


>From adad23b50069b1f5a64883f2672c389443aa97cf Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Tue, 14 Jan 2025 10:35:23 -0800
Subject: [PATCH] [MemProf] Fix an assertion when writing distributed index for
 aliasee

The ThinLTO index bitcode writer uses a helper forEachSummary to manage
preparation and writing of summaries needed for each distributed index
file. For alias summaries, it invokes the provided callback for the
aliasee as well, as we at least need to produce a value id for the
alias's summary. However, all summary generation for the aliasee itself
should be skipped on calls when IsAliasee is true. We invoke the
callback again if that value's summary is to be written as well.

We were asserting in debug mode when invoking collectMemProfCallStacks,
because a given stack id index was not in the StackIdIndicesToIndex
map. It was not added because the forEachSummary invocation that records
these ids in the map (invoked from the IndexBitcodeWriter constructor)
was correctly skipping this handling when invoked for aliasees. We need
the same guard in the invocation that calls collectMemProfCallStacks.

Note that this doesn't cause any real problems in a non-asserts build
as the missing map lookup will return the default 0 value from the map,
which isn't used since we don't actually write the corresponding
summary.
---
 llvm/lib/Bitcode/Writer/BitcodeWriter.cpp         | 8 ++++++++
 llvm/test/ThinLTO/X86/memprof_direct_recursion.ll | 9 +++++++--
 2 files changed, 15 insertions(+), 2 deletions(-)

diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index 94d3afa6c1e332..31c96400dd0fe5 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
@@ -494,6 +494,9 @@ class IndexBitcodeWriter : public BitcodeWriterBase {
     // are currently saved in the index in terms of GUID.
     forEachSummary([&](GVInfo I, bool IsAliasee) {
       GUIDToValueIdMap[I.first] = ++GlobalValueId;
+      // If this is invoked for an aliasee, we want to record the above mapping,
+      // but not the information needed for its summary entry (if the aliasee is
+      // to be imported, we will invoke this separately with IsAliasee=false).
       if (IsAliasee)
         return;
       auto *FS = dyn_cast<FunctionSummary>(I.second);
@@ -4847,6 +4850,11 @@ void IndexBitcodeWriter::writeCombinedGlobalValueSummary() {
   // radix tree array are identified based on this order.
   MapVector<CallStackId, llvm::SmallVector<LinearFrameId>> CallStacks;
   forEachSummary([&](GVInfo I, bool IsAliasee) {
+    // Don't collect this when invoked for an aliasee, as it is not needed for
+    // the alias summary. If the aliasee is to be imported, we will invoke this
+    // separately with IsAliasee=false.
+    if (IsAliasee)
+      return;
     GlobalValueSummary *S = I.second;
     assert(S);
     auto *FS = dyn_cast<FunctionSummary>(S);
diff --git a/llvm/test/ThinLTO/X86/memprof_direct_recursion.ll b/llvm/test/ThinLTO/X86/memprof_direct_recursion.ll
index 102ee64d4638d8..63139cacd8fbab 100644
--- a/llvm/test/ThinLTO/X86/memprof_direct_recursion.ll
+++ b/llvm/test/ThinLTO/X86/memprof_direct_recursion.ll
@@ -34,6 +34,7 @@
 ; RUN:  -supports-hot-cold-new \
 ; RUN:  -thinlto-distributed-indexes \
 ; RUN:  -r=%t/b.o,_Z3fooi,plx \
+; RUN:  -r=%t/b.o,aliasee,plx \
 ; RUN:  -r=%t/b.o,a \
 ; RUN:  -r=%t/b.o,b \
 ; RUN:  -r=%t/b.o,_Znam \
@@ -65,11 +66,15 @@ source_filename = "b.cpp"
 target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
 target triple = "x86_64-unknown-linux-gnu"
 
+;; Make sure the distributed summary bitcode writing succeeds when the memprof
+;; metadata is in an aliasee.
+ at _Z3fooi = alias void (), ptr @aliasee
+
 @a = external local_unnamed_addr global ptr, align 8
 @b = external local_unnamed_addr global i32, align 4
 
 ; Function Attrs: mustprogress uwtable
-define dso_local void @_Z3fooi(i32 noundef %0) local_unnamed_addr #0 !dbg !9 {
+define dso_local void @aliasee(i32 noundef %0) local_unnamed_addr #0 !dbg !9 {
   br label %2, !dbg !12
 
 2:                                                ; preds = %7, %1
@@ -222,4 +227,4 @@ attributes #1 = { "no-trapping-math"="true" "stack-protector-buffer-size"="8" "t
 !19 = !DILocation(line: 7, column: 5, scope: !9)
 !20 = !{i64 8256520048276991898}
 !21 = !DILocation(line: 8, column: 5, scope: !9)
-!22 = !DISubprogram(name: "foo", linkageName: "_Z3fooi", scope: !1, file: !1, line: 1, type: !10, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)
\ No newline at end of file
+!22 = !DISubprogram(name: "foo", linkageName: "_Z3fooi", scope: !1, file: !1, line: 1, type: !10, flags: DIFlagPrototyped, spFlags: DISPFlagOptimized)