[LTO][AllocToken] Support AllocToken instrumentation in backend

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"
@@ -1451,6 +1452,7 @@ ModulePassManager
 PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
                                              ThinOrFullLTOPhase LTOPhase) {
   const bool LTOPreLink = isLTOPreLink(LTOPhase);
+  const bool LTOPostLink = isLTOPostLink(LTOPhase);
   ModulePassManager MPM;
 
   // Run partial inlining pass to partially inline functions that have
@@ -1615,6 +1617,9 @@ PassBuilder::buildModuleOptimizationPipeline(OptimizationLevel Level,
   MPM.addPass(createModuleToFunctionPassAdaptor(std::move(OptimizePM),
                                                 PTO.EagerlyInvalidateAnalyses));
 
+  if (LTOPostLink)
+    MPM.addPass(AllocTokenPass());
+
   invokeOptimizerLastEPCallbacks(MPM, Level, LTOPhase);
 
   // Split out cold code. Splitting is done late to avoid hiding context from
@@ -1853,6 +1858,7 @@ ModulePassManager PassBuilder::buildThinLTODefaultPipeline(
     MPM.addPass(LowerTypeTestsPass(nullptr, nullptr,
                                    lowertypetests::DropTestKind::Assume));
     MPM.addPass(buildCoroWrapper(ThinOrFullLTOPhase::ThinLTOPostLink));
+    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.
@@ -1913,6 +1919,7 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
                                    lowertypetests::DropTestKind::Assume));
 
     MPM.addPass(buildCoroWrapper(ThinOrFullLTOPhase::FullLTOPostLink));
+    MPM.addPass(AllocTokenPass());
 
     invokeFullLinkTimeOptimizationLastEPCallbacks(MPM, Level);
 
@@ -2000,6 +2007,7 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
                                    lowertypetests::DropTestKind::Assume));
 
     MPM.addPass(buildCoroWrapper(ThinOrFullLTOPhase::FullLTOPostLink));
+    MPM.addPass(AllocTokenPass());
 
     invokeFullLinkTimeOptimizationLastEPCallbacks(MPM, Level);
 
@@ -2234,6 +2242,7 @@ PassBuilder::buildLTODefaultPipeline(OptimizationLevel Level,
     MPM.addPass(CGProfilePass(/*InLTOPostLink=*/true));
 
   MPM.addPass(CoroCleanupPass());
+  MPM.addPass(AllocTokenPass());
 
   invokeFullLinkTimeOptimizationLastEPCallbacks(MPM, Level);
 
@@ -2351,6 +2360,9 @@ PassBuilder::buildO0DefaultPipeline(OptimizationLevel Level,
 
   MPM.addPass(buildCoroWrapper(Phase));
 
+  if (isLTOPostLink(Phase))
+    MPM.addPass(AllocTokenPass());
+
   invokeOptimizerLastEPCallbacks(MPM, Level, Phase);
 
   if (isLTOPreLink(Phase))

llvm/lib/Transforms/Instrumentation/AllocToken.cpp

@@ -234,20 +234,39 @@ class TypeHashPointerSplitMode : public TypeHashMode {
   }
 };
 
-// Apply opt overrides.
-AllocTokenOptions transformOptionsFromCl(AllocTokenOptions Opts) {
+// Apply opt overrides and module flags.
+static AllocTokenOptions resolveOptions(AllocTokenOptions Opts,
+                                        const Module &M) {
   if (!Opts.MaxTokens.has_value())
     Opts.MaxTokens = ClMaxTokens;
   Opts.FastABI |= ClFastABI;
   Opts.Extended |= ClExtended;
+
+  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();
+
   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)) {
     switch (Options.Mode) {

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

@@ -0,0 +1,30 @@
+; Test that all supported module flags are retrieved correctly.
+;
+; RUN: opt < %s -passes='inferattrs,alloc-token' -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)
+declare ptr @malloc(i64)
+declare ptr @my_malloc(i64)
+
+define void @test() sanitize_alloc_token {
+; CHECK-LABEL: define void @test(
+; CHECK: call ptr @__alloc_token_0_malloc(i64 8)
+; CHECK: call ptr @__alloc_token_1__Znwm(i64 8)
+; CHECK: call ptr @__alloc_token_2_malloc(i64 8)
+; CHECK: call ptr @__alloc_token_0_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.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
+}

llvm/test/Other/new-pm-O0-defaults.ll

@@ -45,6 +45,10 @@
 ; CHECK-PRE-LINK-NEXT: Running pass: NameAnonGlobalPass
 ; CHECK-THINLTO: Running pass: LowerTypeTestsPass
 ; CHECK-THINLTO-NEXT: Running pass: CoroConditionalWrapper
+; CHECK-THINLTO-NEXT: Running pass: AllocTokenPass
+; CHECK-THINLTO-NEXT: Running analysis: InnerAnalysisManagerProxy
+; CHECK-THINLTO-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis
+; CHECK-THINLTO-NEXT: Running analysis: TargetLibraryAnalysis
 ; CHECK-THINLTO-NEXT: Running pass: EliminateAvailableExternallyPass
 ; CHECK-THINLTO-NEXT: Running pass: GlobalDCEPass
 ; CHECK-LTO: Running pass: CrossDSOCFIPass on [module]
@@ -53,10 +57,12 @@
 ; CHECK-LTO-NEXT: Running pass: LowerTypeTestsPass
 ; CHECK-LTO-NEXT: Running pass: LowerTypeTestsPass
 ; CHECK-LTO-NEXT: CoroConditionalWrapper
+; CHECK-LTO-NEXT: Running pass: AllocTokenPass
+; CHECK-LTO-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis
+; CHECK-LTO-NEXT: Running analysis: TargetLibraryAnalysis
 ; CHECK-CORO-NEXT: Running pass: AnnotationRemarksPass
 ; CHECK-CORO-NEXT: Running analysis: TargetLibraryAnalysis
 ; CHECK-LTO-NEXT: Running pass: AnnotationRemarksPass
-; CHECK-LTO-NEXT: Running analysis: TargetLibraryAnalysis
 ; CHECK-NEXT: Running pass: PrintModulePass
 
 ; Make sure we get the IR back out without changes when we print the module.

llvm/test/Other/new-pm-lto-defaults.ll

@@ -163,6 +163,8 @@
 ; CHECK-O23SZ-NEXT: Running pass: CGProfilePass
 ; CHECK-O1-NEXT: Running pass: CoroConditionalWrapper
 ; CHECK-O23SZ-NEXT: Running pass: CoroCleanupPass
+; CHECK-O-NEXT: Running pass: AllocTokenPass
+; CHECK-O1-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis
 ; CHECK-EP-NEXT: Running pass: NoOpModulePass
 ; CHECK-O-NEXT: Running pass: AnnotationRemarksPass on foo
 ; CHECK-O-NEXT: Running pass: PrintModulePass

llvm/test/Other/new-pm-thinlto-postlink-defaults.ll

@@ -203,6 +203,7 @@
 ; CHECK-POSTLINK-O-NEXT: Running pass: DivRemPairsPass
 ; CHECK-POSTLINK-O-NEXT: Running pass: TailCallElimPass
 ; CHECK-POSTLINK-O-NEXT: Running pass: SimplifyCFGPass
+; CHECK-POSTLINK-O-NEXT: Running pass: AllocTokenPass
 ; CHECK-POST-EP-OPT-LAST-NEXT: Running pass: NoOpModulePass
 ; CHECK-POSTLINK-O-NEXT: Running pass: GlobalDCEPass
 ; CHECK-POSTLINK-O-NEXT: Running pass: ConstantMergePass

llvm/test/Other/new-pm-thinlto-postlink-pgo-defaults.ll

@@ -188,6 +188,7 @@
 ; CHECK-O-NEXT: Running pass: DivRemPairsPass
 ; CHECK-O-NEXT: Running pass: TailCallElimPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
+; CHECK-O-NEXT: Running pass: AllocTokenPass
 ; CHECK-O-NEXT: Running pass: GlobalDCEPass
 ; CHECK-O-NEXT: Running pass: ConstantMergePass
 ; CHECK-O-NEXT: Running pass: CGProfilePass

llvm/test/Other/new-pm-thinlto-postlink-samplepgo-defaults.ll

@@ -197,6 +197,7 @@
 ; CHECK-O-NEXT: Running pass: DivRemPairsPass
 ; CHECK-O-NEXT: Running pass: TailCallElimPass
 ; CHECK-O-NEXT: Running pass: SimplifyCFGPass
+; CHECK-O-NEXT: Running pass: AllocTokenPass
 ; CHECK-O-NEXT: Running pass: GlobalDCEPass
 ; CHECK-O-NEXT: Running pass: ConstantMergePass
 ; CHECK-O-NEXT: Running pass: CGProfilePass