[AllocToken] Introduce AllocToken instrumentation pass

Introduce `AllocToken`, an instrumentation pass designed to provide
tokens to memory allocators enabling various heap organization
strategies, such as heap partitioning.

Initially, the pass instruments functions marked with a new attribute
`sanitize_alloc_token` by rewriting allocation calls to include a token
ID, appended as a function argument with the default ABI.

The design aims to provide a flexible framework for implementing
different token generation schemes. It currently supports the following
token modes:

- TypeHash (default): token IDs based on a hash of the allocated type
- Random: statically-assigned pseudo-random token IDs
- Increment: incrementing token IDs per TU

For the `TypeHash` mode introduce support for `!alloc_token` metadata:
the metadata can be attached to allocation calls to provide richer
semantic information to be consumed by the AllocToken pass.
Optimization remarks can be enabled to show where no metadata was
available.

An alternative "fast ABI" is provided, where instead of passing the
token ID as an argument (e.g., `__alloc_token_malloc(size, id)`), the
token ID is directly encoded into the name of the called function (e.g.,
`__alloc_token_0_malloc(size)`). Where the maximum tokens is small, this
offers more efficient instrumentation by avoiding the overhead of
passing an additional argument at each allocation site.

Link: https://discourse.llvm.org/t/rfc-a-framework-for-allocator-partitioning-hints/87434 [1]

Pull Request: llvm#156838

Author: melver

llvm/docs/ReleaseNotes.md

@@ -177,6 +177,10 @@ Changes to Sanitizers
 Other Changes
 -------------
 
+* Introduces the `AllocToken` pass, an instrumentation pass providing tokens to
+  memory allocators enabling various heap organization strategies, such as heap
+  partitioning.
+
 External Open Source Projects Using LLVM {{env.config.release}}
 ===============================================================
 

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

@@ -0,0 +1,46 @@
+//===- AllocToken.h - Allocation token instrumentation --------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the AllocTokenPass, an instrumentation pass that
+// replaces allocation calls with ones including an allocation token.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_TRANSFORMS_INSTRUMENTATION_ALLOCTOKEN_H
+#define LLVM_TRANSFORMS_INSTRUMENTATION_ALLOCTOKEN_H
+
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/PassManager.h"
+#include <optional>
+
+namespace llvm {
+
+class Module;
+
+struct AllocTokenOptions {
+  std::optional<uint64_t> MaxTokens;
+  bool FastABI = false;
+  bool Extended = false;
+  AllocTokenOptions() = default;
+};
+
+/// A module pass that rewrites heap allocations to use token-enabled
+/// allocation functions based on various source-level properties.
+class AllocTokenPass : public PassInfoMixin<AllocTokenPass> {
+public:
+  LLVM_ABI explicit AllocTokenPass(AllocTokenOptions Opts = {});
+  LLVM_ABI PreservedAnalyses run(Module &M, ModuleAnalysisManager &MAM);
+  static bool isRequired() { return true; }
+
+private:
+  const AllocTokenOptions Options;
+};
+
+} // namespace llvm
+
+#endif // LLVM_TRANSFORMS_INSTRUMENTATION_ALLOCTOKEN_H

llvm/lib/Passes/PassBuilder.cpp

@@ -240,6 +240,7 @@
 #include "llvm/Transforms/IPO/WholeProgramDevirt.h"
 #include "llvm/Transforms/InstCombine/InstCombine.h"
 #include "llvm/Transforms/Instrumentation/AddressSanitizer.h"
+#include "llvm/Transforms/Instrumentation/AllocToken.h"
 #include "llvm/Transforms/Instrumentation/BoundsChecking.h"
 #include "llvm/Transforms/Instrumentation/CGProfile.h"
 #include "llvm/Transforms/Instrumentation/ControlHeightReduction.h"

llvm/lib/Passes/PassRegistry.def

@@ -125,6 +125,7 @@ MODULE_PASS("openmp-opt", OpenMPOptPass())
 MODULE_PASS("openmp-opt-postlink",
             OpenMPOptPass(ThinOrFullLTOPhase::FullLTOPostLink))
 MODULE_PASS("partial-inliner", PartialInlinerPass())
+MODULE_PASS("alloc-token", AllocTokenPass())
 MODULE_PASS("pgo-icall-prom", PGOIndirectCallPromotion())
 MODULE_PASS("pgo-instr-gen", PGOInstrumentationGen())
 MODULE_PASS("pgo-instr-use", PGOInstrumentationUse())