[MemProf] Update hot/cold information after importing

[tejohnson]

Summary

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.

[davidxl]

lgtm

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

 ;; how the ThinLTO backend behavior is controlled.
 ; RUN: llvm-dis %t.out.index.bc -o - | FileCheck %s --check-prefix=CHECK-INDEX-ON
 ;; Flags are printed in decimal, but this corresponds to 0x161, and 0x100 is
 ;; the value indicating -supports-hot-cold-new was enabled.
 ; CHECK-INDEX-ON: flags: 353
 
 ; 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"
 
 ;; Next check without -supports-hot-cold-new, we should not perform
 ;; context disambiguation, and we should strip memprof metadata and
 ;; attributes before optimization during the distributed backend.
 ; RUN: llvm-lto2 run %t.o -save-temps \
 ; RUN:  -thinlto-distributed-indexes \
 ; RUN:  -r=%t.o,main,plx \
 ; RUN:  -r=%t.o,_Znam, \
 ; RUN:  -o %t.out
 
 ;; Ensure that the index file reflects not having -supports-hot-cold-new.
 ; RUN: llvm-dis %t.out.index.bc -o - | FileCheck %s --check-prefix=CHECK-INDEX-OFF
 ;; Flags are printed in decimal, but this corresponds to 0x61, without 0x100 set.
 ; CHECK-INDEX-OFF: flags: 97
 
 ; 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"
 
 source_filename = "thinlto-distributed-supports-hot-cold-new.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 {

llvm/lib/LTO/LTOBackend.cpp

Error lto::thinBackend(const Config &Conf, unsigned Task, AddStreamFn AddStream,
   if (!DiagFileOrErr)
     return DiagFileOrErr.takeError();
   auto DiagnosticOutputFile = std::move(*DiagFileOrErr);
 
   // Set the partial sample profile ratio in the profile summary module flag of
   // the module, if applicable.
   Mod.setPartialSampleProfileRatio(CombinedIndex);
 
-  updateMemProfAttributes(Mod, CombinedIndex);
-
   updatePublicTypeTestCalls(Mod, CombinedIndex.withWholeProgramVisibility());
 
   if (Conf.CodeGenOnly) {
     codegen(Conf, TM.get(), AddStream, Task, Mod, CombinedIndex);
     return finalizeOptimizationRemarks(std::move(DiagnosticOutputFile));
   }
 
   if (Conf.PreOptModuleHook && !Conf.PreOptModuleHook(Task, Mod))
auto ModuleLoader = [&](StringRef Identifier) {
     return MOrErr;
   };
 
   FunctionImporter Importer(CombinedIndex, ModuleLoader,
                             ClearDSOLocalOnDeclarations);
   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));
 
   return OptimizeAndCodegen(Mod, TM.get(), std::move(DiagnosticOutputFile));
 }
 
 BitcodeModule *lto::findThinLTOModule(MutableArrayRef<BitcodeModule> BMs) {
   if (ThinLTOAssumeMerged && BMs.size() == 1)

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

 ;; Test that passing -supports-hot-cold-new to the LTO link allows context
 ;; disambiguation to proceed, and also prevents memprof metadata and attributes
 ;; 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"
 
 define i32 @main() #0 {
 entry:
   %call = call ptr @_Znam(i64 0), !memprof !0, !callsite !5
   %call2 = call ptr @_Znam(i64 0) #1
   ret i32 0
 }
 
+define void @bar() {
+  call void @foo()
+  ret void
+}
+
+declare void @foo()
+
 declare ptr @_Znam(i64)
 
 attributes #0 = { noinline optnone }
 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}
+
+;--- 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}