[LTO][AllocToken] Support AllocToken instrumentation in backend (#169358)

Unconditionally add AllocTokenPass to the optimization pipelines, and
ensure that it runs last in LTO backend pipelines. The latter ensures
that AllocToken instrumentation can be moved later in the LTO pipeline
to avoid interference with other optimizations (e.g. PGHO) and enable
late heap-allocation optimizations.

In preparation of removing AllocTokenPass being added by Clang, add
support for AllocTokenPass to read configuration options from LLVM
module flags.

To optimize given the pass is now runs unconditionally, only retrieve
TargetLibraryInfo and OptimizationRemarkEmitter when necessary.

Author: melver

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
+}