[Clang][ThinLTO] Wire up AllocToken pass to run in ThinLTO backend

clang/lib/CodeGen/BackendUtil.cpp

@@ -1144,7 +1144,12 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
     if (!IsThinLTOPostLink) {
       addSanitizers(TargetTriple, CodeGenOpts, LangOpts, PB);
       addKCFIPass(TargetTriple, LangOpts, PB);
-      addAllocTokenPass(TargetTriple, CodeGenOpts, LangOpts, PB);
+
+      // If this is a ThinLTO pre-link compile, skip AllocTokenPass; it will be
+      // added during ThinLTO backend compile phase to enable optimizations such
+      // as MemProf to remain compatible with AllocToken instrumentation.
+      if (!CodeGenOpts.PrepareForThinLTO)
+        addAllocTokenPass(TargetTriple, CodeGenOpts, LangOpts, PB);
     }
 
     if (std::optional<GCOVOptions> Options =
@@ -1425,6 +1430,12 @@ runThinLTOBackend(CompilerInstance &CI, ModuleSummaryIndex *CombinedIndex,
   Conf.RemarksFormat = CGOpts.OptRecordFormat;
   Conf.SplitDwarfFile = CGOpts.SplitDwarfFile;
   Conf.SplitDwarfOutput = CGOpts.SplitDwarfOutput;
+  Conf.PassBuilderCallback = [&](PassBuilder &PB) {
+    // Skipped during pre-link phase to avoid instrumentation interfering with
+    // backend optimizations, and instead we run it as late as possible in the
+    // backend phase.
+    addAllocTokenPass(CI.getTarget().getTriple(), CGOpts, CI.getLangOpts(), PB);
+  };
   switch (Action) {
   case Backend_EmitNothing:
     Conf.PreCodeGenModuleHook = [](size_t Task, const llvm::Module &Mod) {

clang/test/CodeGen/lto-newpm-pipeline.c

@@ -48,12 +48,10 @@
 // CHECK-THIN-O0-NEXT: Running pass: AlwaysInlinerPass
 // CHECK-THIN-O0-NEXT: Running analysis: ProfileSummaryAnalysis
 // CHECK-THIN-O0-NEXT: Running pass: CoroConditionalWrapper
-// CHECK-THIN-O0-NEXT: Running pass: AllocTokenPass
-// CHECK-THIN-O0-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis
-// CHECK-THIN-O0-NEXT: Running analysis: TargetLibraryAnalysis
 // CHECK-THIN-O0-NEXT: Running pass: CanonicalizeAliasesPass
 // CHECK-THIN-O0-NEXT: Running pass: NameAnonGlobalPass
 // CHECK-THIN-O0-NEXT: Running pass: AnnotationRemarksPass
+// CHECK-THIN-O0-NEXT: Running analysis: TargetLibraryAnalysis
 // CHECK-THIN-O0-NEXT: Running pass: VerifierPass
 // CHECK-THIN-O0-NEXT: Running pass: ThinLTOBitcodeWriterPass
 

clang/test/CodeGen/memprof-pgho-thinlto.cpp

@@ -0,0 +1,65 @@
+// Test end-to-end ThinLTO optimization pipeline with PGHO, that it does not
+// interfere with other allocation instrumentation features.
+//
+// RUN: split-file %s %t
+// RUN: llvm-profdata merge %t/memprof.yaml -o %t/use.memprofdata
+//
+// RUN: %clangxx -O2 -flto=thin -g -fmemory-profile-use=%t/use.memprofdata %t/src.cpp -c -o %t.o
+// RUN: llvm-lto2 run %t.o -thinlto-distributed-indexes -supports-hot-cold-new -r=%t.o,main,plx -r=%t.o,_Z3foov,plx -r=%t.o,_Znam, -o %t.out
+// RUN: %clang_cc1 -O2 -x ir %t.o -fthinlto-index=%t.o.thinlto.bc -mllvm -optimize-hot-cold-new -emit-llvm -o - 2>&1 | FileCheck %s --check-prefixes=CHECK,DEFAULT
+//
+// RUN: %clangxx -O2 -flto=thin -g -fsanitize=alloc-token -fmemory-profile-use=%t/use.memprofdata %t/src.cpp -c -o %t.o
+// RUN: llvm-lto2 run %t.o -thinlto-distributed-indexes -supports-hot-cold-new -r=%t.o,main,plx -r=%t.o,_Z3foov,plx -r=%t.o,_Znam, -o %t.out
+// RUN: %clang_cc1 -O2 -x ir %t.o -fsanitize=alloc-token -fthinlto-index=%t.o.thinlto.bc -mllvm -optimize-hot-cold-new -emit-llvm -o - 2>&1 | FileCheck %s --check-prefixes=CHECK,ALLOCTOKEN
+
+//--- memprof.yaml
+---
+HeapProfileRecords:
+  - GUID: 0x7f8d88fcc70a347b
+    AllocSites:
+    - Callstack:
+      - { Function: 0x7f8d88fcc70a347b, LineOffset: 1, Column: 10, IsInlineFrame: false }
+      - { Function: 0xdb956436e78dd5fa, LineOffset: 1, Column: 13, IsInlineFrame: false }
+      MemInfoBlock:
+        AllocCount: 1
+        TotalAccessCount: 0
+        MinAccessCount: 0
+        MaxAccessCount: 0
+        TotalSize: 10
+        MinSize: 10
+        MaxSize: 10
+        AllocTimestamp: 100
+        DeallocTimestamp: 100
+        TotalLifetime: 100000
+        MinLifetime: 100000
+        MaxLifetime: 100000
+        AllocCpuId: 0
+        DeallocCpuId: 0
+        NumMigratedCpu: 0
+        NumLifetimeOverlaps: 0
+        NumSameAllocCpu: 0
+        NumSameDeallocCpu: 0
+        DataTypeId: 0
+        TotalAccessDensity: 0
+        MinAccessDensity: 0
+        MaxAccessDensity: 0
+        TotalLifetimeAccessDensity: 0
+        MinLifetimeAccessDensity: 0
+        MaxLifetimeAccessDensity: 0
+        AccessHistogramSize: 0
+        AccessHistogram: 0
+...
+
+//--- src.cpp
+// CHECK-LABEL: define{{.*}} ptr @_Z3foov()
+// DEFAULT:    call {{.*}} ptr @_Znam12__hot_cold_t(i64 10, i8 -128)
+// ALLOCTOKEN: call {{.*}} ptr @__alloc_token__Znam12__hot_cold_t(i64 10, i8 -128, i64 1538840549748785101){{.*}} !alloc_token
+char *foo() {
+  return new char[10];
+}
+
+int main() {
+  char *a = foo();
+  delete[] a;
+  return 0;
+}

llvm/include/llvm/LTO/Config.h

@@ -51,6 +51,8 @@ struct Config {
   std::vector<std::string> MAttrs;
   std::vector<std::string> MllvmArgs;
   std::vector<std::string> PassPlugins;
+  /// Callback to customize the PassBuilder.
+  std::function<void(PassBuilder &)> PassBuilderCallback;
   /// For adding passes that run right before codegen.
   std::function<void(legacy::PassManager &)> PreCodeGenPassesHook;
   std::optional<Reloc::Model> RelocModel = Reloc::PIC_;

llvm/lib/LTO/LTOBackend.cpp

@@ -275,6 +275,8 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
   SI.registerCallbacks(PIC, &MAM);
   PassBuilder PB(TM, Conf.PTO, PGOOpt, &PIC);
 
+  if (Conf.PassBuilderCallback)
+    Conf.PassBuilderCallback(PB);
   RegisterPassPlugins(Conf.PassPlugins, PB);
 
   std::unique_ptr<TargetLibraryInfoImpl> TLII(