[llvm-branch-commits] [LTO] Support enabling AllocToken instrumentation in backend (PR #169358)

[via llvm-branch-commits]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-lto

Author: Marco Elver (melver)

<details>
<summary>Changes</summary>

Add support for running `AllocTokenPass` during the LTO backend phase,
controlled by a new internal option `-lto-alloc-token-mode`.

This is required to support running the pass after other heap allocation
optimizations (such as PGHO) in the LTO backend, avoiding interference
between them.

---

This change is part of the following series:
1. https://github.com/llvm/llvm-project/pull/169242
2. https://github.com/llvm/llvm-project/pull/169243
3. https://github.com/llvm/llvm-project/pull/169358
4. https://github.com/llvm/llvm-project/pull/169359
5. https://github.com/llvm/llvm-project/pull/169360

---
Full diff: https://github.com/llvm/llvm-project/pull/169358.diff


2 Files Affected:

- (modified) llvm/lib/LTO/LTOBackend.cpp (+33) 
- (added) llvm/test/LTO/X86/alloc-token.ll (+34) 


``````````diff
diff --git a/llvm/lib/LTO/LTOBackend.cpp b/llvm/lib/LTO/LTOBackend.cpp
index 93118becedbac..190cec4f5701c 100644
--- a/llvm/lib/LTO/LTOBackend.cpp
+++ b/llvm/lib/LTO/LTOBackend.cpp
@@ -31,6 +31,7 @@
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/Support/AllocToken.h"
 #include "llvm/Support/Error.h"
 #include "llvm/Support/FileSystem.h"
 #include "llvm/Support/MemoryBuffer.h"
@@ -42,6 +43,7 @@
 #include "llvm/Target/TargetMachine.h"
 #include "llvm/TargetParser/SubtargetFeature.h"
 #include "llvm/Transforms/IPO/WholeProgramDevirt.h"
+#include "llvm/Transforms/Instrumentation/AllocToken.h"
 #include "llvm/Transforms/Utils/FunctionImportUtils.h"
 #include "llvm/Transforms/Utils/SplitModule.h"
 #include <optional>
@@ -68,6 +70,10 @@ static cl::opt<LTOBitcodeEmbedding> EmbedBitcode(
                           "Embed post merge, but before optimizations")),
     cl::desc("Embed LLVM bitcode in object files produced by LTO"));
 
+static cl::opt<std::string> LTOAllocTokenMode(
+    "lto-alloc-token-mode", cl::init(""),
+    cl::desc("Enable AllocToken instrumentation during LTO with chosen mode"));
+
 static cl::opt<bool> ThinLTOAssumeMerged(
     "thinlto-assume-merged", cl::init(false),
     cl::desc("Assume the input has already undergone ThinLTO function "
@@ -198,6 +204,31 @@ static void RegisterPassPlugins(ArrayRef<std::string> PassPlugins,
   }
 }
 
+// Register instrumentation passes that need to run late in the pipeline; these
+// are non-optimization passes and need to run after most optimizations to avoid
+// interfering with them (e.g. PGHO) or to capture the final state of the code.
+static void registerBackendInstrumentation(PassBuilder &PB) {
+  if (!LTOAllocTokenMode.empty()) {
+    AllocTokenOptions Opts;
+    if (auto Mode = getAllocTokenModeFromString(LTOAllocTokenMode))
+      Opts.Mode = *Mode;
+    else
+      report_fatal_error("invalid lto-alloc-token-mode: " +
+                         Twine(LTOAllocTokenMode));
+
+    // ThinLTO backend
+    PB.registerOptimizerLastEPCallback(
+        [Opts](ModulePassManager &MPM, OptimizationLevel, ThinOrFullLTOPhase) {
+          MPM.addPass(AllocTokenPass(Opts));
+        });
+    // Full LTO backend
+    PB.registerFullLinkTimeOptimizationLastEPCallback(
+        [Opts](ModulePassManager &MPM, OptimizationLevel) {
+          MPM.addPass(AllocTokenPass(Opts));
+        });
+  }
+}
+
 static std::unique_ptr<TargetMachine>
 createTargetMachine(const Config &Conf, const Target *TheTarget, Module &M) {
   const Triple &TheTriple = M.getTargetTriple();
@@ -277,6 +308,8 @@ static void runNewPMPasses(const Config &Conf, Module &Mod, TargetMachine *TM,
 
   RegisterPassPlugins(Conf.PassPlugins, PB);
 
+  registerBackendInstrumentation(PB);
+
   std::unique_ptr<TargetLibraryInfoImpl> TLII(
       new TargetLibraryInfoImpl(TM->getTargetTriple()));
   if (Conf.Freestanding)
diff --git a/llvm/test/LTO/X86/alloc-token.ll b/llvm/test/LTO/X86/alloc-token.ll
new file mode 100644
index 0000000000000..873b7c94620bc
--- /dev/null
+++ b/llvm/test/LTO/X86/alloc-token.ll
@@ -0,0 +1,34 @@
+; RUN: llvm-as %s -o %t.bc
+;
+; RUN: llvm-lto2 run -lto-alloc-token-mode=default %t.bc -o %t.out \
+; RUN:   -r=%t.bc,main,plx \
+; RUN:   -r=%t.bc,_Znwm, \
+; RUN:   -r=%t.bc,sink,pl
+; RUN: llvm-objdump -d -r %t.out.0 | FileCheck %s --check-prefixes=CHECK,DEFAULT
+;
+; RUN: llvm-lto2 run -lto-alloc-token-mode=default -alloc-token-fast-abi -alloc-token-max=1 %t.bc -o %t.out \
+; RUN:   -r=%t.bc,main,plx \
+; RUN:   -r=%t.bc,_Znwm, \
+; RUN:   -r=%t.bc,sink,pl
+; RUN: llvm-objdump -d -r %t.out.0 | FileCheck %s --check-prefixes=CHECK,FASTABI
+
+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) #0
+
+ at sink = global ptr null
+
+; CHECK-LABEL: <main>:
+; CHECK: callq
+; DEFAULT-NEXT: R_X86_64_PLT32 __alloc_token__Znwm
+; FASTABI-NEXT: R_X86_64_PLT32 __alloc_token_0__Znwm
+define void @main() sanitize_alloc_token {
+  %call = call ptr @_Znwm(i64 8) #0, !alloc_token !0
+  store volatile ptr %call, ptr @sink
+  ret void
+}
+
+attributes #0 = { nobuiltin allocsize(0) }
+
+!0 = !{!"int", i1 0}

``````````

</details>


https://github.com/llvm/llvm-project/pull/169358