[llvm] [MemProf] Handle empty stack context during ThinLTO cloning (PR #81008)

[Teresa Johnson via llvm-commits]

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

Fix for assert after PR#78264.

Handle the case where the MIB context is empty after skipping the
callsite context, because the callsite context is actually longer than
the MIB context. Presumably this happened as a result of inlining, but
in theory the metadata should have been replaced with an attribute in
that case. Need to investigate why this is occuring, but for now handle
this gracefully to fix the build regression.


>From e081b81d55bc9c5817296a0c076f42b847cdca18 Mon Sep 17 00:00:00 2001
From: Teresa Johnson <tejohnson at google.com>
Date: Wed, 7 Feb 2024 08:04:06 -0800
Subject: [PATCH] [MemProf] Handle empty stack context during ThinLTO cloning

Fix for assert after PR#78264.

Handle the case where the MIB context is empty after skipping the
callsite context, because the callsite context is actually longer than
the MIB context. Presumably this happened as a result of inlining, but
in theory the metadata should have been replaced with an attribute in
that case. Need to investigate why this is occuring, but for now handle
this gracefully to fix the build regression.
---
 .../IPO/MemProfContextDisambiguation.cpp      |  6 +++-
 llvm/test/ThinLTO/X86/memprof-import-fix.ll   | 36 +++++++++++++++++++
 2 files changed, 41 insertions(+), 1 deletion(-)
 create mode 100644 llvm/test/ThinLTO/X86/memprof-import-fix.ll

diff --git a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
index d81239dd49e4c..323f33fcc862d 100644
--- a/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
+++ b/llvm/lib/Transforms/IPO/MemProfContextDisambiguation.cpp
@@ -3475,7 +3475,11 @@ bool MemProfContextDisambiguation::applyImport(Module &M) {
             auto ContextIterBegin =
                 StackContext.beginAfterSharedPrefix(CallsiteContext);
             // Skip the checking on the first iteration.
-            uint64_t LastStackContextId = *ContextIterBegin == 0 ? 1 : 0;
+            uint64_t LastStackContextId =
+                (ContextIterBegin != StackContext.end() &&
+                 *ContextIterBegin == 0)
+                    ? 1
+                    : 0;
             for (auto ContextIter = ContextIterBegin;
                  ContextIter != StackContext.end(); ++ContextIter) {
               // If this is a direct recursion, simply skip the duplicate
diff --git a/llvm/test/ThinLTO/X86/memprof-import-fix.ll b/llvm/test/ThinLTO/X86/memprof-import-fix.ll
new file mode 100644
index 0000000000000..aeb54057976a3
--- /dev/null
+++ b/llvm/test/ThinLTO/X86/memprof-import-fix.ll
@@ -0,0 +1,36 @@
+;; Test to ensure that importing of cloning decisions does not assert when
+;; the callsite context is longer than the MIB context.
+;; FIXME: Presumably this happened as a result of inlining, but in theory the
+;; metadata should have been replaced with an attribute in that case. Need to
+;; investigate why this is occuring.
+
+; RUN: opt -thinlto-bc %s >%t.o
+; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
+; RUN:	-supports-hot-cold-new \
+; RUN:	-r=%t.o,main,plx \
+; RUN:	-r=%t.o,_Znam, \
+; RUN:	-save-temps \
+; RUN:	-o %t.out
+
+; RUN: llvm-dis %t.out.1.4.opt.bc -o - | FileCheck %s --check-prefix=IR
+;; The call to new is not changed in this case.
+; IR: call ptr @_Znam(i64 0)
+
+source_filename = "memprof-import-fix.ll"
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+target triple = "x86_64-unknown-linux-gnu"
+
+define i32 @main() #0 {
+entry:
+  %call = call ptr @_Znam(i64 0), !memprof !0, !callsite !3
+  ret i32 0
+}
+
+declare ptr @_Znam(i64)
+
+attributes #0 = { noinline optnone }
+
+!0 = !{!1}
+!1 = !{!2, !"notcold"}
+!2 = !{i64 9086428284934609951}
+!3 = !{i64 9086428284934609951, i64 -5964873800580613432}