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

clang/lib/CodeGen/BackendUtil.cpp

@@ -66,7 +66,6 @@
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizerOptions.h"
-#include "llvm/Transforms/Instrumentation/AllocToken.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
 #include "llvm/Transforms/Instrumentation/DataFlowSanitizer.h"
 #include "llvm/Transforms/Instrumentation/GCOVProfiler.h"
@@ -234,18 +233,6 @@ class EmitAssemblyHelper {
 };
 } // namespace
 
-static AllocTokenOptions getAllocTokenOptions(const LangOptions &LangOpts,
-                                              const CodeGenOptions &CGOpts) {
-  AllocTokenOptions Opts;
-  if (LangOpts.AllocTokenMode)
-    Opts.Mode = *LangOpts.AllocTokenMode;
-  if (LangOpts.AllocTokenMax)
-    Opts.MaxTokens = *LangOpts.AllocTokenMax;
-  Opts.Extended = CGOpts.SanitizeAllocTokenExtended;
-  Opts.FastABI = CGOpts.SanitizeAllocTokenFastABI;
-  return Opts;
-}
-
 static SanitizerCoverageOptions
 getSancovOptsFromCGOpts(const CodeGenOptions &CGOpts) {
   SanitizerCoverageOptions Opts;
@@ -874,23 +861,6 @@ static void addSanitizers(const Triple &TargetTriple,
   }
 }
 
-static void addAllocTokenPass(const Triple &TargetTriple,
-                              const CodeGenOptions &CodeGenOpts,
-                              const LangOptions &LangOpts, PassBuilder &PB) {
-  PB.registerOptimizerLastEPCallback([&](ModulePassManager &MPM,
-                                         OptimizationLevel Level,
-                                         ThinOrFullLTOPhase) {
-    if (Level == OptimizationLevel::O0 &&
-        LangOpts.Sanitize.has(SanitizerKind::AllocToken)) {
-      // The default pass builder only infers libcall function attrs when
-      // optimizing, so we insert it here because we need it for accurate
-      // memory allocation function detection with -fsanitize=alloc-token.
-      MPM.addPass(InferFunctionAttrsPass());
-    }
-    MPM.addPass(AllocTokenPass(getAllocTokenOptions(LangOpts, CodeGenOpts)));
-  });
-}
-
 void EmitAssemblyHelper::RunOptimizationPipeline(
     BackendAction Action, std::unique_ptr<raw_pwrite_stream> &OS,
     std::unique_ptr<llvm::ToolOutputFile> &ThinLinkOS, BackendConsumer *BC) {
@@ -1140,12 +1110,23 @@ void EmitAssemblyHelper::RunOptimizationPipeline(
         FPM.addPass(BoundsCheckingPass(Options));
       });
 
-    // Don't add sanitizers if we are here from ThinLTO PostLink. That already
-    // done on PreLink stage.
     if (!IsThinLTOPostLink) {
+      // Most sanitizers only run during PreLink stage.
       addSanitizers(TargetTriple, CodeGenOpts, LangOpts, PB);
       addKCFIPass(TargetTriple, LangOpts, PB);
-      addAllocTokenPass(TargetTriple, CodeGenOpts, LangOpts, PB);
+
+      PB.registerPipelineStartEPCallback(
+          [&](ModulePassManager &MPM, OptimizationLevel Level) {
+            if (Level == OptimizationLevel::O0 &&
+                LangOpts.Sanitize.has(SanitizerKind::AllocToken)) {
+              // With the default O0 pipeline, LibFunc attrs are not inferred,
+              // so we insert it here because we need it for accurate memory
+              // allocation function detection with -fsanitize=alloc-token.
+              // Note: This could also be added to the default O0 pipeline, but
+              // has a non-trivial effect on generated IR size (attributes).
+              MPM.addPass(InferFunctionAttrsPass());
+            }
+          });
     }
 
     if (std::optional<GCOVOptions> Options =

clang/lib/CodeGen/CodeGenModule.cpp

@@ -1635,6 +1635,22 @@ void CodeGenModule::EmitBackendOptionsMetadata(
     getModule().addModuleFlag(llvm::Module::Min, "SmallDataLimit",
                               CodeGenOpts.SmallDataLimit);
   }
+
+  // Set AllocToken configuration for backend pipeline.
+  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-lower.c

@@ -1,16 +1,20 @@
 // Test optimization pipelines do not interfere with AllocToken lowering, and we
 // pass on function attributes correctly.
 //
-// RUN: %clang_cc1     -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s
-// RUN: %clang_cc1 -O1 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s
-// RUN: %clang_cc1 -O2 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s
+// RUN: %clang_cc1     -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,DEFAULT
+// RUN: %clang_cc1 -O1 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,DEFAULT
+// RUN: %clang_cc1 -O2 -fsanitize=alloc-token -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,DEFAULT
+// RUN: %clang_cc1     -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,FASTABI
+// RUN: %clang_cc1 -O1 -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,FASTABI
+// RUN: %clang_cc1 -O2 -fsanitize=alloc-token -fsanitize-alloc-token-fast-abi -triple x86_64-linux-gnu -emit-llvm %s -o - | FileCheck %s --check-prefixes=CHECK,FASTABI
 
 typedef __typeof(sizeof(int)) size_t;
 
 void *malloc(size_t size);
 
 // CHECK-LABEL: @test_malloc(
-// CHECK: call{{.*}} ptr @__alloc_token_malloc(i64 noundef 4, i64 2689373973731826898){{.*}} !alloc_token [[META_INT:![0-9]+]]
+// DEFAULT: call{{.*}} ptr @__alloc_token_malloc(i64 noundef 4, i64 2689373973731826898){{.*}} !alloc_token [[META_INT:![0-9]+]]
+// FASTABI: call{{.*}} ptr @__alloc_token_2689373973731826898_malloc(i64 noundef 4){{.*}} !alloc_token [[META_INT:![0-9]+]]
 void *test_malloc() {
   return malloc(sizeof(int));
 }
@@ -26,6 +30,7 @@ void *no_sanitize_malloc(size_t size) __attribute__((no_sanitize("alloc-token"))
 // allocator will only implement standard allocation functions.
 void *nonstandard_malloc(size_t size) __attribute__((malloc));
 // CHECK-LABEL: @test_nonlibcall_malloc(
+// CHECK-NOT: __alloc_token_
 // CHECK: call{{.*}} ptr @nonstandard_malloc(i64 noundef 4){{.*}} !alloc_token [[META_INT]]
 void *test_nonlibcall_malloc() {
   return nonstandard_malloc(sizeof(int));

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,7 +32,6 @@
 // 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 pass: CanonicalizeAliasesPass
 // CHECK-FULL-O0-NEXT: Running pass: NameAnonGlobalPass
 // CHECK-FULL-O0-NEXT: Running pass: AnnotationRemarksPass
@@ -47,7 +46,6 @@
 // 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 pass: CanonicalizeAliasesPass
 // CHECK-THIN-O0-NEXT: Running pass: NameAnonGlobalPass
 // CHECK-THIN-O0-NEXT: Running pass: AnnotationRemarksPass