[LTO][AllocToken] Support AllocToken instrumentation in backend

clang/lib/CodeGen/BackendUtil.cpp

@@ -239,7 +239,8 @@ static AllocTokenOptions getAllocTokenOptions(const LangOptions &LangOpts,
   AllocTokenOptions Opts;
   if (LangOpts.AllocTokenMode)
     Opts.Mode = *LangOpts.AllocTokenMode;
-  Opts.MaxTokens = LangOpts.AllocTokenMax;
+  if (LangOpts.AllocTokenMax)
+    Opts.MaxTokens = *LangOpts.AllocTokenMax;
   Opts.Extended = CGOpts.SanitizeAllocTokenExtended;
   Opts.FastABI = CGOpts.SanitizeAllocTokenFastABI;
   return Opts;

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

@@ -33,11 +33,10 @@
 // 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
 
@@ -49,11 +48,10 @@
 // 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
 

llvm/include/llvm/Transforms/Instrumentation/AllocToken.h

@@ -25,7 +25,7 @@ class Module;
 
 struct AllocTokenOptions {
   AllocTokenMode Mode = DefaultAllocTokenMode;
-  std::optional<uint64_t> MaxTokens;
+  uint64_t MaxTokens = 0;
   bool FastABI = false;
   bool Extended = false;
   AllocTokenOptions() = default;

llvm/lib/Passes/PassBuilderPipelines.cpp

@@ -73,6 +73,7 @@
 #include "llvm/Transforms/IPO/SampleProfileProbe.h"
 #include "llvm/Transforms/IPO/WholeProgramDevirt.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
+#include "llvm/Transforms/Instrumentation/AllocToken.h"
 #include "llvm/Transforms/Instrumentation/CGProfile.h"
 #include "llvm/Transforms/Instrumentation/ControlHeightReduction.h"
 #include "llvm/Transforms/Instrumentation/InstrProfiling.h"
@@ -1615,6 +1616,11 @@ PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(OptimizePM),
                                                 PTO.EagerlyInvalidateAnalyses));
 
+  // AllocToken transforms heap allocation calls; this needs to run late after
+  // other allocation call transformations (such as those in InstCombine).
+  if (!LTOPreLink)
+    MPM.addPass(AllocTokenPass());
+
   invokeOptimizerLastEPCallbacks(MPM, Level, LTOPhase);
 
   // Split out cold code. Splitting is done late to avoid hiding context from
@@ -1853,6 +1859,11 @@ ModulePassManager PassBuilder::buildThinLTODefaultPipeline(
     MPM.addPass(LowerTypeTestsPass(nullptr, nullptr,
                                    lowertypetests::DropTestKind::Assume));
     MPM.addPass(buildCoroWrapper(ThinOrFullLTOPhase::ThinLTOPostLink));
+
+    // AllocToken transforms heap allocation calls; this needs to run late after
+    // other allocation call transformations (such as those in InstCombine).
+    MPM.addPass(AllocTokenPass());
+
     // Drop available_externally and unreferenced globals. This is necessary
     // with ThinLTO in order to avoid leaving undefined references to dead
     // globals in the object file.
@@ -1914,6 +1925,10 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
 
     MPM.addPass(buildCoroWrapper(ThinOrFullLTOPhase::FullLTOPostLink));
 
+    // AllocToken transforms heap allocation calls; this needs to run late after
+    // other allocation call transformations (such as those in InstCombine).
+    MPM.addPass(AllocTokenPass());
+
     invokeFullLinkTimeOptimizationLastEPCallbacks(MPM, Level);
 
     // Emit annotation remarks.
@@ -2001,6 +2016,10 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
 
     MPM.addPass(buildCoroWrapper(ThinOrFullLTOPhase::FullLTOPostLink));
 
+    // AllocToken transforms heap allocation calls; this needs to run late after
+    // other allocation call transformations (such as those in InstCombine).
+    MPM.addPass(AllocTokenPass());
+
     invokeFullLinkTimeOptimizationLastEPCallbacks(MPM, Level);
 
     // Emit annotation remarks.
@@ -2235,6 +2254,10 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
 
   MPM.addPass(CoroCleanupPass());
 
+  // AllocToken transforms heap allocation calls; this needs to run late after
+  // other allocation call transformations (such as those in InstCombine).
+  MPM.addPass(AllocTokenPass());
+
   invokeFullLinkTimeOptimizationLastEPCallbacks(MPM, Level);
 
   // Emit annotation remarks.
@@ -2351,6 +2374,11 @@ PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
 
   MPM.addPass(buildCoroWrapper(Phase));
 
+  // AllocToken transforms heap allocation calls; this needs to run late after
+  // other allocation call transformations (such as those in InstCombine).
+  if (!isLTOPreLink(Phase))
+    MPM.addPass(AllocTokenPass());
+
   invokeOptimizerLastEPCallbacks(MPM, Level, Phase);
 
   if (isLTOPreLink(Phase))

llvm/lib/Transforms/Instrumentation/AllocToken.cpp

@@ -234,34 +234,53 @@ class TypeHashPointerSplitMode : public TypeHashMode {
   }
 };
 
-// Apply opt overrides.
-AllocTokenOptions transformOptionsFromCl(AllocTokenOptions Opts) {
-  if (!Opts.MaxTokens.has_value())
+// Apply opt overrides and module flags.
+static AllocTokenOptions resolveOptions(AllocTokenOptions Opts,
+                                        const Module &M) {
+  auto IntModuleFlagOrNull = [&](StringRef Key) {
+    return mdconst::extract_or_null<ConstantInt>(M.getModuleFlag(Key));
+  };
+
+  if (auto *S = dyn_cast_or_null<MDString>(M.getModuleFlag("alloc-token-mode")))
+    if (auto Mode = getAllocTokenModeFromString(S->getString()))
+      Opts.Mode = *Mode;
+  if (auto *Val = IntModuleFlagOrNull("alloc-token-max"))
+    Opts.MaxTokens = Val->getZExtValue();
+  if (auto *Val = IntModuleFlagOrNull("alloc-token-fast-abi"))
+    Opts.FastABI |= Val->isOne();
+  if (auto *Val = IntModuleFlagOrNull("alloc-token-extended"))
+    Opts.Extended |= Val->isOne();
+
+  // Allow overriding options from command line options.
+  if (ClMaxTokens.getNumOccurrences())
     Opts.MaxTokens = ClMaxTokens;
-  Opts.FastABI |= ClFastABI;
-  Opts.Extended |= ClExtended;
+  if (ClFastABI.getNumOccurrences())
+    Opts.FastABI = ClFastABI;
+  if (ClExtended.getNumOccurrences())
+    Opts.Extended = ClExtended;
+
   return Opts;
 }
 
 class AllocToken {
 public:
   explicit AllocToken(AllocTokenOptions Opts, Module &M,
                       ModuleAnalysisManager &MAM)
-      : Options(transformOptionsFromCl(std::move(Opts))), Mod(M),
+      : Options(resolveOptions(std::move(Opts), M)), Mod(M),
         FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
-        Mode(IncrementMode(*IntPtrTy, *Options.MaxTokens)) {
+        Mode(IncrementMode(*IntPtrTy, Options.MaxTokens)) {
     switch (Options.Mode) {
     case TokenMode::Increment:
       break;
     case TokenMode::Random:
-      Mode.emplace<RandomMode>(*IntPtrTy, *Options.MaxTokens,
+      Mode.emplace<RandomMode>(*IntPtrTy, Options.MaxTokens,
                                M.createRNG(DEBUG_TYPE));
       break;
     case TokenMode::TypeHash:
-      Mode.emplace<TypeHashMode>(*IntPtrTy, *Options.MaxTokens);
+      Mode.emplace<TypeHashMode>(*IntPtrTy, Options.MaxTokens);
       break;
     case TokenMode::TypeHashPointerSplit:
-      Mode.emplace<TypeHashPointerSplitMode>(*IntPtrTy, *Options.MaxTokens);
+      Mode.emplace<TypeHashPointerSplitMode>(*IntPtrTy, Options.MaxTokens);
       break;
     }
   }
@@ -318,8 +337,6 @@ bool AllocToken::instrumentFunction(Function &F) {
   if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)
     return false;
 
-  auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
-  auto &TLI = FAM.getResult<TargetLibraryAnalysis>(F);
   SmallVector<std::pair<CallBase *, LibFunc>, 4> AllocCalls;
   SmallVector<IntrinsicInst *, 4> IntrinsicInsts;
 
@@ -328,6 +345,10 @@ bool AllocToken::instrumentFunction(Function &F) {
       F.hasFnAttribute(Attribute::SanitizeAllocToken) &&
       !F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation);
 
+  // Get TLI only when required.
+  const TargetLibraryInfo *TLI =
+      InstrumentFunction ? &FAM.getResult<TargetLibraryAnalysis>(F) : nullptr;
+
   // Collect all allocation calls to avoid iterator invalidation.
   for (Instruction &I : instructions(F)) {
     // Collect all alloc_token_* intrinsics.
@@ -343,26 +364,28 @@ bool AllocToken::instrumentFunction(Function &F) {
     auto *CB = dyn_cast<CallBase>(&I);
     if (!CB)
       continue;
-    if (std::optional<LibFunc> Func = shouldInstrumentCall(*CB, TLI))
+    if (std::optional<LibFunc> Func = shouldInstrumentCall(*CB, *TLI))
       AllocCalls.emplace_back(CB, Func.value());
   }
 
+  // Return early to avoid unnecessarily instantiating the ORE.
+  if (AllocCalls.empty() && IntrinsicInsts.empty())
+    return false;
+
+  auto &ORE = FAM.getResult<OptimizationRemarkEmitterAnalysis>(F);
   bool Modified = false;
 
-  if (!AllocCalls.empty()) {
-    for (auto &[CB, Func] : AllocCalls)
-      Modified |= replaceAllocationCall(CB, Func, ORE, TLI);
-    if (Modified)
-      NumFunctionsModified++;
-  }
+  for (auto &[CB, Func] : AllocCalls)
+    Modified |= replaceAllocationCall(CB, Func, ORE, *TLI);
 
-  if (!IntrinsicInsts.empty()) {
-    for (auto *II : IntrinsicInsts)
-      replaceIntrinsicInst(II, ORE);
+  for (auto *II : IntrinsicInsts) {
+    replaceIntrinsicInst(II, ORE);
     Modified = true;
-    NumFunctionsModified++;
   }
 
+  if (Modified)
+    NumFunctionsModified++;
+
   return Modified;
 }
 

llvm/test/CodeGen/AArch64/print-pipeline-passes.ll

@@ -2,7 +2,7 @@
 ; RUN: opt -mtriple=aarch64 -S -passes='default<O2>' -print-pipeline-passes < %s | FileCheck %s
 
 ; CHECK: loop-idiom-vectorize
-; O0: {{^}}function(ee-instrument<>),always-inline,coro-cond(coro-early,cgscc(coro-split),coro-cleanup,globaldce),function(annotation-remarks),verify,print{{$}}
+; O0: {{^}}function(ee-instrument<>),always-inline,coro-cond(coro-early,cgscc(coro-split),coro-cleanup,globaldce),alloc-token,function(annotation-remarks),verify,print{{$}}
 
 define void @foo() {
 entry:

llvm/test/Instrumentation/AllocToken/hot-cold-new.ll

@@ -0,0 +1,20 @@
+; Manually add instcombine to ensure the hot/cold transformation happens before
+; the LTO pipeline. The default LTO pipeline includes MemProfRemoveInfo which
+; strips the memprof attributes unless the summary index indicates support.
+; RUN: opt < %s -passes='function(instcombine),thinlto<O2>' -optimize-hot-cold-new -S | FileCheck %s
+; RUN: opt < %s -passes='function(instcombine),lto<O2>' -optimize-hot-cold-new -S | FileCheck %s
+; RUN: opt < %s -passes='function(instcombine),alloc-token' -optimize-hot-cold-new -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+declare ptr @_Znwm(i64)
+
+define ptr @new_hot() sanitize_alloc_token {
+; CHECK-LABEL: @new_hot(
+; CHECK: call {{.*}} @__alloc_token__Znwm12__hot_cold_t(i64 10, i8 -2, i64 2689373973731826898){{.*}} !alloc_token
+  %ret = call ptr @_Znwm(i64 10) #0, !alloc_token !0
+  ret ptr %ret
+}
+
+attributes #0 = { builtin allocsize(0) "memprof"="hot" }
+!0 = !{!"int", i1 false}

llvm/test/Instrumentation/AllocToken/module-flags.ll

@@ -0,0 +1,35 @@
+; Test that all supported module flags are retrieved correctly.
+;
+; RUN: opt < %s -passes='inferattrs,alloc-token' -S | FileCheck %s --check-prefixes=CHECK,DEFAULT
+; RUN: opt < %s -passes='inferattrs,alloc-token' -alloc-token-max=2 -alloc-token-fast-abi=0 -alloc-token-extended=0 -S | FileCheck %s --check-prefixes=CHECK,OVERRIDE
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+declare ptr @_Znwm(i64)
+declare ptr @malloc(i64)
+declare ptr @my_malloc(i64)
+
+define void @test() sanitize_alloc_token {
+; CHECK-LABEL: define void @test(
+; DEFAULT: call ptr @__alloc_token_0_malloc(i64 8)
+; DEFAULT: call ptr @__alloc_token_1__Znwm(i64 8)
+; DEFAULT: call ptr @__alloc_token_2_malloc(i64 8)
+; DEFAULT: call ptr @__alloc_token_0_my_malloc(i64 8)
+; OVERRIDE: call ptr @__alloc_token_malloc(i64 8, i64 0)
+; OVERRIDE: call ptr @__alloc_token__Znwm(i64 8, i64 1)
+; OVERRIDE: call ptr @__alloc_token_malloc(i64 8, i64 0)
+; OVERRIDE: call ptr @my_malloc(i64 8)
+  %1 = call ptr @malloc(i64 8)
+  %2 = call ptr @_Znwm(i64 8)
+  %3 = call ptr @malloc(i64 8)
+  %4 = call ptr @my_malloc(i64 8), !alloc_token !0
+  ret void
+}
+
+!0 = !{!"int", i1 0}
+
+!llvm.module.flags = !{!1, !2, !3, !4}
+!1 = !{i32 1, !"alloc-token-mode", !"increment"}
+!2 = !{i32 1, !"alloc-token-max", i64 3}
+!3 = !{i32 1, !"alloc-token-fast-abi", i64 1}
+!4 = !{i32 1, !"alloc-token-extended", i64 1}

llvm/test/LTO/X86/alloc-token-hot-cold-new.ll

@@ -0,0 +1,25 @@
+; RUN: opt -module-summary -o %t.thin.bc %s
+; RUN: llvm-lto2 run %t.thin.bc -o %t.thin.out \
+; RUN:   -r=%t.thin.bc,main,plx \
+; RUN:   -r=%t.thin.bc,_Znwm, \
+; RUN:   -r=%t.thin.bc,sink,pl \
+; RUN:   -supports-hot-cold-new -optimize-hot-cold-new
+; RUN: llvm-objdump -d -r %t.thin.out.1 | FileCheck %s
+
+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"
+
+declare ptr @_Znwm(i64)
+
+@sink = global ptr null
+
+; CHECK-LABEL: <main>:
+; CHECK: callq
+; CHECK-NEXT: R_X86_64_PLT32 __alloc_token__Znwm12__hot_cold_t
+define void @main() sanitize_alloc_token {
+  %call = call ptr @_Znwm(i64 8) #0
+  store volatile ptr %call, ptr @sink
+  ret void
+}
+
+attributes #0 = { builtin allocsize(0) "memprof"="hot" }

llvm/test/LTO/X86/alloc-token.ll

@@ -0,0 +1,27 @@
+; --- Full LTO ---
+; RUN: llvm-as %s -o %t.bc
+; RUN: llvm-lto -exported-symbol=main -o %t.out %t.bc
+; RUN: llvm-objdump -d -r %t.out | FileCheck %s
+; --- ThinLTO ---
+; RUN: opt -module-summary -o %t.thin.bc %s
+; RUN: llvm-lto2 run %t.thin.bc -o %t.thin.out \
+; RUN:   -r=%t.thin.bc,main,plx \
+; RUN:   -r=%t.thin.bc,_Znwm, \
+; RUN:   -r=%t.thin.bc,sink,pl
+; RUN: llvm-objdump -d -r %t.thin.out.1 | FileCheck %s
+
+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"
+
+declare ptr @_Znwm(i64)
+
+@sink = global ptr null
+
+; CHECK-LABEL: <main>:
+; CHECK: callq
+; CHECK-NEXT: R_X86_64_PLT32 __alloc_token__Znwm
+define void @main() sanitize_alloc_token {
+  %call = call ptr @_Znwm(i64 8)
+  store volatile ptr %call, ptr @sink
+  ret void
+}