[MemProf] Update hot/cold information after importing

The support added by D149215 to remove memprof metadata and attributes
if we don't link with an allocator supporting hot/cold operator new
interfaces did not update imported code. Move the update handling later
in the ThinLTO backend to just after importing, and update the test to
check this case.

Differential Revision: https://reviews.llvm.org/D150295

clang/test/CodeGen/thinlto-distributed-supports-hot-cold-new.ll

@@ -22,7 +22,7 @@
 
 ; RUN: %clang -target x86_64-unknown-linux-gnu -O2 -o %t1.o -x ir %t.o -c -fthinlto-index=%t.o.thinlto.bc -save-temps=obj
 
-; RUN: llvm-dis %t.s.0.preopt.bc -o - | FileCheck %s --check-prefix=CHECK-IR
+; RUN: llvm-dis %t.s.3.import.bc -o - | FileCheck %s --check-prefix=CHECK-IR
 ; CHECK-IR: !memprof {{.*}} !callsite
 ; CHECK-IR: "memprof"="cold"
 
@@ -42,7 +42,7 @@
 
 ; RUN: %clang -target x86_64-unknown-linux-gnu -O2 -o %t1.o -x ir %t.o -c -fthinlto-index=%t.o.thinlto.bc -save-temps=obj
 
-; RUN: llvm-dis %t.s.0.preopt.bc -o - | FileCheck %s \
+; RUN: llvm-dis %t.s.3.import.bc -o - | FileCheck %s \
 ; RUN: --implicit-check-not "!memprof" --implicit-check-not "!callsite" \
 ; RUN: --implicit-check-not "memprof"="cold"
 

llvm/lib/LTO/LTOBackend.cpp

@@ -565,8 +565,6 @@ Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   // the module, if applicable.
   Mod.setPartialSampleProfileRatio(CombinedIndex);
 
-  updateMemProfAttributes(Mod, CombinedIndex);
-
   updatePublicTypeTestCalls(Mod, CombinedIndex.withWholeProgramVisibility());
 
   if (Conf.CodeGenOnly) {
@@ -653,6 +651,9 @@ Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   if (Error Err = Importer.importFunctions(Mod, ImportList).takeError())
     return Err;
 
+  // Do this after any importing so that imported code is updated.
+  updateMemProfAttributes(Mod, CombinedIndex);
+
   if (Conf.PostImportModuleHook && !Conf.PostImportModuleHook(Task, Mod))
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
 

llvm/test/ThinLTO/X86/memprof-supports-hot-cold-new.ll

@@ -3,37 +3,53 @@
 ;; from being removed from the LTO backend, and vice versa without passing
 ;; -supports-hot-cold-new.
 
+; RUN: split-file %s %t
+
 ;; First check with -supports-hot-cold-new.
-; RUN: opt -thinlto-bc %s >%t.o
-; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
+; RUN: opt -thinlto-bc %t/main.ll >%t/main.o
+; RUN: opt -thinlto-bc %t/foo.ll >%t/foo.o
+; RUN: llvm-lto2 run %t/main.o %t/foo.o -enable-memprof-context-disambiguation \
 ; RUN:	-supports-hot-cold-new \
-; RUN:	-r=%t.o,main,plx \
-; RUN:	-r=%t.o,_Znam, \
+; RUN:	-r=%t/main.o,main,plx \
+; RUN:	-r=%t/main.o,bar,plx \
+; RUN:	-r=%t/main.o,foo, \
+; RUN:	-r=%t/main.o,_Znam, \
+; RUN:	-r=%t/foo.o,foo,plx \
+; RUN:	-r=%t/foo.o,_Znam, \
 ; RUN:	-memprof-dump-ccg \
 ; RUN:	 -save-temps \
 ; RUN:	-o %t.out 2>&1 | FileCheck %s --check-prefix=DUMP
 ; DUMP: Callsite Context Graph:
 
-; RUN: llvm-dis %t.out.1.0.preopt.bc -o - | FileCheck %s --check-prefix=IR
+; RUN: llvm-dis %t.out.1.3.import.bc -o - | FileCheck %s --check-prefix=IR
+; IR: @main()
+; IR: !memprof {{.*}} !callsite
+; IR: @_Znam(i64 0) #[[ATTR:[0-9]+]]
+; IR: @bar()
 ; IR: !memprof {{.*}} !callsite
-; IR: "memprof"="cold"
+; IR: @_Znam(i64 0) #[[ATTR:[0-9]+]]
+; IR: attributes #[[ATTR]] = { "memprof"="cold" }
 
 ;; Next check without -supports-hot-cold-new, we should not perform
 ;; context disambiguation, and we should strip memprof metadata and
 ;; attributes before optimization.
-; RUN: llvm-lto2 run %t.o -enable-memprof-context-disambiguation \
-; RUN:	-r=%t.o,main,plx \
-; RUN:	-r=%t.o,_Znam, \
+; RUN: llvm-lto2 run %t/main.o %t/foo.o -enable-memprof-context-disambiguation \
+; RUN:	-r=%t/main.o,main,plx \
+; RUN:	-r=%t/main.o,bar,plx \
+; RUN:	-r=%t/main.o,foo, \
+; RUN:	-r=%t/main.o,_Znam, \
+; RUN:	-r=%t/foo.o,foo,plx \
+; RUN:	-r=%t/foo.o,_Znam, \
 ; RUN:	-memprof-dump-ccg \
 ; RUN:	 -save-temps \
 ; RUN:	-o %t.out 2>&1 | FileCheck %s --allow-empty \
 ; RUN:  --implicit-check-not "Callsite Context Graph:"
 
-; RUN: llvm-dis %t.out.1.0.preopt.bc -o - | FileCheck %s \
+; RUN: llvm-dis %t.out.1.3.import.bc -o - | FileCheck %s \
 ; RUN: --implicit-check-not "!memprof" --implicit-check-not "!callsite" \
 ; RUN: --implicit-check-not "memprof"="cold"
 
-source_filename = "memprof-supports-hot-cold-new.ll"
+;--- main.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"
 
@@ -44,6 +60,13 @@ entry:
   ret i32 0
 }
 
+define void @bar() {
+  call void @foo()
+  ret void
+}
+
+declare void @foo()
+
 declare ptr @_Znam(i64)
 
 attributes #0 = { noinline optnone }
@@ -55,3 +78,25 @@ attributes #1 = { "memprof"="cold" }
 !3 = !{!4, !"cold"}
 !4 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 -3421689549917153178}
 !5 = !{i64 9086428284934609951}
+
+;--- foo.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 @foo() {
+entry:
+  %call = call ptr @_Znam(i64 0), !memprof !0, !callsite !5
+  %call2 = call ptr @_Znam(i64 0) #1
+  ret i32 0
+}
+
+declare ptr @_Znam(i64)
+
+attributes #1 = { "memprof"="cold" }
+
+!0 = !{!1, !3}
+!1 = !{!2, !"notcold"}
+!2 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 8632435727821051414}
+!3 = !{!4, !"cold"}
+!4 = !{i64 9086428284934609951, i64 -5964873800580613432, i64 2732490490862098848, i64 -3421689549917153178}
+!5 = !{i64 9086428284934609951}