[AllocToken] Introduce AllocToken instrumentation pass

llvm/docs/LangRef.rst

@@ -2427,6 +2427,9 @@ For example:
     if the attributed function is called during invocation of a function
     attributed with ``sanitize_realtime``.
     This attribute is incompatible with the ``sanitize_realtime`` attribute.
+``sanitize_alloc_token``
+    This attributes indicates that implicit allocation token instrumentation
+    is enabled for this function.
 ``speculative_load_hardening``
     This attribute indicates that
     `Speculative Load Hardening <https://llvm.org/docs/SpeculativeLoadHardening.html>`_

llvm/docs/ReleaseNotes.md

@@ -166,6 +166,10 @@ Changes to Sanitizers
 Other Changes
 -------------
 
+* Introduces the `AllocToken` pass, an instrumentation pass designed to provide
+  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/Bitcode/LLVMBitCodes.h

@@ -800,6 +800,7 @@ enum AttributeKindCodes {
   ATTR_KIND_SANITIZE_TYPE = 101,
   ATTR_KIND_CAPTURES = 102,
   ATTR_KIND_DEAD_ON_RETURN = 103,
+  ATTR_KIND_SANITIZE_ALLOC_TOKEN = 104,
 };
 
 enum ComdatSelectionKindCodes {

llvm/include/llvm/IR/Attributes.td

@@ -342,6 +342,9 @@ def SanitizeRealtime : EnumAttr<"sanitize_realtime", IntersectPreserve, [FnAttr]
 /// during a real-time sanitized function (see `sanitize_realtime`).
 def SanitizeRealtimeBlocking : EnumAttr<"sanitize_realtime_blocking", IntersectPreserve, [FnAttr]>;
 
+/// Allocation token instrumentation is on.
+def SanitizeAllocToken : EnumAttr<"sanitize_alloc_token", IntersectPreserve, [FnAttr]>;
+
 /// Speculative Load Hardening is enabled.
 ///
 /// Note that this uses the default compatibility (always compatible during

llvm/include/llvm/IR/FixedMetadataKinds.def

@@ -55,3 +55,4 @@ LLVM_FIXED_MD_KIND(MD_mmra, "mmra", 40)
 LLVM_FIXED_MD_KIND(MD_noalias_addrspace, "noalias.addrspace", 41)
 LLVM_FIXED_MD_KIND(MD_callee_type, "callee_type", 42)
 LLVM_FIXED_MD_KIND(MD_nofree, "nofree", 43)
+LLVM_FIXED_MD_KIND(MD_alloc_token_hint, "alloc_token_hint", 44)

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/Bitcode/Reader/BitcodeReader.cpp

@@ -2203,6 +2203,8 @@ static Attribute::AttrKind getAttrFromCode(uint64_t Code) {
     return Attribute::SanitizeRealtime;
   case bitc::ATTR_KIND_SANITIZE_REALTIME_BLOCKING:
     return Attribute::SanitizeRealtimeBlocking;
+  case bitc::ATTR_KIND_SANITIZE_ALLOC_TOKEN:
+    return Attribute::SanitizeAllocToken;
   case bitc::ATTR_KIND_SPECULATIVE_LOAD_HARDENING:
     return Attribute::SpeculativeLoadHardening;
   case bitc::ATTR_KIND_SWIFT_ERROR:

llvm/lib/Bitcode/Writer/BitcodeWriter.cpp

@@ -883,6 +883,8 @@ static uint64_t getAttrKindEncoding(Attribute::AttrKind Kind) {
     return bitc::ATTR_KIND_STRUCT_RET;
   case Attribute::SanitizeAddress:
     return bitc::ATTR_KIND_SANITIZE_ADDRESS;
+  case Attribute::SanitizeAllocToken:
+    return bitc::ATTR_KIND_SANITIZE_ALLOC_TOKEN;
   case Attribute::SanitizeHWAddress:
     return bitc::ATTR_KIND_SANITIZE_HWADDRESS;
   case Attribute::SanitizeThread:

llvm/lib/Passes/PassBuilder.cpp

@@ -237,6 +237,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

@@ -124,6 +124,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())