[Clang][CodeGen] Remove explicit insertion of AllocToken pass

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);
+
+      // On ThinLTO or FullLTO pre-link, skip AllocTokenPass; it runs during the
+      // LTO backend compile phase to enable late heap-allocation optimizations
+      // to remain compatible with AllocToken instrumentation.
+      if (!PrepareForThinLTO && !PrepareForLTO)
+        addAllocTokenPass(TargetTriple, CodeGenOpts, LangOpts, PB);
     }
 
     if (std::optional<GCOVOptions> Options =

clang/lib/CodeGen/CodeGenModule.cpp

@@ -1635,6 +1635,25 @@ void CodeGenModule::EmitBackendOptionsMetadata(
     getModule().addModuleFlag(llvm::Module::Min, "SmallDataLimit",
                               CodeGenOpts.SmallDataLimit);
   }
+
+  if (LangOpts.AllocTokenMode) {
+    StringRef S = llvm::getAllocTokenModeAsString(*LangOpts.AllocTokenMode);
+    getModule().addModuleFlag(llvm::Module::Error, "alloc-token-mode",
+                              llvm::MDString::get(VMContext, S));
+  }
+
+  if (LangOpts.AllocTokenMax) {
+    getModule().addModuleFlag(
+        llvm::Module::Error, "alloc-token-max",
+        llvm::ConstantInt::get(llvm::Type::getInt64Ty(VMContext),
+                               *LangOpts.AllocTokenMax));
+  }
+
+  if (CodeGenOpts.SanitizeAllocTokenFastABI)
+    getModule().addModuleFlag(llvm::Module::Error, "alloc-token-fast-abi", 1);
+
+  if (CodeGenOpts.SanitizeAllocTokenExtended)
+    getModule().addModuleFlag(llvm::Module::Error, "alloc-token-extended", 1);
 }
 
 void CodeGenModule::UpdateCompletedType(const TagDecl *TD) {

clang/test/CodeGen/alloc-token-module-flags.c

@@ -0,0 +1,22 @@
+// RUN: %clang_cc1 -fsanitize=alloc-token -emit-llvm -o - %s | FileCheck %s --check-prefix=DEFAULT
+// RUN: %clang_cc1 -fsanitize=alloc-token -falloc-token-mode=increment -emit-llvm -o - %s | FileCheck %s --check-prefix=INCREMENT
+// RUN: %clang_cc1 -fsanitize=alloc-token -falloc-token-max=100 -emit-llvm -o - %s | FileCheck %s --check-prefix=MAX
+// RUN: %clang_cc1 -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi -emit-llvm -o - %s | FileCheck %s --check-prefix=FASTABI
+// RUN: %clang_cc1 -fsanitize=alloc-token -fsanitize-alloc-token-extended -emit-llvm -o - %s | FileCheck %s --check-prefix=EXTENDED
+
+// DEFAULT-NOT: !"alloc-token-mode"
+// DEFAULT-NOT: !"alloc-token-max"
+// DEFAULT-NOT: !"alloc-token-fast-abi"
+// DEFAULT-NOT: !"alloc-token-extended"
+
+// INCREMENT: !llvm.module.flags = !{{{.*}}![[FLAG:[0-9]+]]{{.*}}}
+// INCREMENT: ![[FLAG]] = !{i32 1, !"alloc-token-mode", !"increment"}
+
+// MAX: !llvm.module.flags = !{{{.*}}![[FLAG:[0-9]+]]{{.*}}}
+// MAX: ![[FLAG]] = !{i32 1, !"alloc-token-max", i64 100}
+
+// FASTABI: !llvm.module.flags = !{{{.*}}![[FLAG:[0-9]+]]{{.*}}}
+// FASTABI: ![[FLAG]] = !{i32 1, !"alloc-token-fast-abi", i32 1}
+
+// EXTENDED: !llvm.module.flags = !{{{.*}}![[FLAG:[0-9]+]]{{.*}}}
+// EXTENDED: ![[FLAG]] = !{i32 1, !"alloc-token-extended", i32 1}

clang/test/CodeGen/distributed-thin-lto/memprof-pgho.cpp

@@ -0,0 +1,67 @@
+// 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 --target=x86_64-linux-gnu -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 -triple x86_64-linux-gnu -O1 -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: %clang_cc1 -triple x86_64-linux-gnu -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 --target=x86_64-linux-gnu -O2 -flto=thin -g -fsanitize=alloc-token -falloc-token-max=32 -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 -triple x86_64-linux-gnu -O1 -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
+// RUN: %clang_cc1 -triple x86_64-linux-gnu -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 12){{.*}} !alloc_token
+char *foo() {
+  return new char[10];
+}
+
+int main() {
+  char *a = foo();
+  delete[] a;
+  return 0;
+}

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

@@ -32,12 +32,10 @@
 // CHECK-FULL-O0-NEXT: Running pass: AlwaysInlinerPass
 // CHECK-FULL-O0-NEXT: Running analysis: ProfileSummaryAnalysis
 // CHECK-FULL-O0-NEXT: Running pass: CoroConditionalWrapper
-// CHECK-FULL-O0-NEXT: Running pass: AllocTokenPass
-// CHECK-FULL-O0-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis
-// CHECK-FULL-O0-NEXT: Running analysis: TargetLibraryAnalysis
 // CHECK-FULL-O0-NEXT: Running pass: CanonicalizeAliasesPass
 // CHECK-FULL-O0-NEXT: Running pass: NameAnonGlobalPass
 // CHECK-FULL-O0-NEXT: Running pass: AnnotationRemarksPass
+// CHECK-FULL-O0-NEXT: Running analysis: TargetLibraryAnalysis
 // CHECK-FULL-O0-NEXT: Running pass: VerifierPass
 // CHECK-FULL-O0-NEXT: Running pass: BitcodeWriterPass
 
@@ -48,12 +46,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