diff --git a/clang/test/CodeGen/thinlto-distributed-supports-hot-cold-new.ll b/clang/test/CodeGen/thinlto-distributed-supports-hot-cold-new.ll
--- a/clang/test/CodeGen/thinlto-distributed-supports-hot-cold-new.ll
+++ b/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"
 
diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -565,8 +565,6 @@
   // the module, if applicable.
   Mod.setPartialSampleProfileRatio(CombinedIndex);
 
-  updateMemProfAttributes(Mod, CombinedIndex);
-
   updatePublicTypeTestCalls(Mod, CombinedIndex.withWholeProgramVisibility());
 
   if (Conf.CodeGenOnly) {
@@ -653,6 +651,9 @@
   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));
 
diff --git a/llvm/test/ThinLTO/X86/memprof-supports-hot-cold-new.ll b/llvm/test/ThinLTO/X86/memprof-supports-hot-cold-new.ll
--- a/llvm/test/ThinLTO/X86/memprof-supports-hot-cold-new.ll
+++ b/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 @@
   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 @@
 !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}