[clang] [llvm] [Clang] Refactor allocation type inference logic (PR #163636)
Marco Elver via llvm-commits
llvm-commits at lists.llvm.org
Wed Oct 22 06:04:06 PDT 2025
https://github.com/melver updated https://github.com/llvm/llvm-project/pull/163636
>From 2ea2825781ac00b10d2f173ff748696e3d07f19c Mon Sep 17 00:00:00 2001
From: Marco Elver <elver at google.com>
Date: Wed, 15 Oct 2025 23:28:58 +0200
Subject: [PATCH 1/3] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20ch?=
=?UTF-8?q?anges=20to=20main=20this=20commit=20is=20based=20on?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Created using spr 1.3.8-beta.1
[skip ci]
---
clang/docs/AllocToken.rst | 4 +-
clang/include/clang/Basic/CodeGenOptions.h | 4 -
clang/include/clang/Basic/LangOptions.h | 8 +
clang/include/clang/Driver/Options.td | 4 +
clang/lib/CodeGen/BackendUtil.cpp | 9 +-
clang/lib/Frontend/CompilerInvocation.cpp | 60 +++++--
clang/test/Driver/fsanitize-alloc-token.c | 11 ++
llvm/include/llvm/IR/Intrinsics.td | 8 +
llvm/include/llvm/Support/AllocToken.h | 62 ++++++++
.../Transforms/Instrumentation/AllocToken.h | 2 +
llvm/lib/Passes/PassBuilder.cpp | 32 ++++
llvm/lib/Passes/PassRegistry.def | 5 +-
llvm/lib/Support/AllocToken.cpp | 46 ++++++
llvm/lib/Support/CMakeLists.txt | 1 +
.../Transforms/Instrumentation/AllocToken.cpp | 148 ++++++++++--------
llvm/test/Instrumentation/AllocToken/basic.ll | 2 +-
.../Instrumentation/AllocToken/basic32.ll | 2 +-
llvm/test/Instrumentation/AllocToken/fast.ll | 2 +-
.../Instrumentation/AllocToken/intrinsic.ll | 32 ++++
.../Instrumentation/AllocToken/intrinsic32.ll | 32 ++++
.../test/Instrumentation/AllocToken/invoke.ll | 2 +-
.../Instrumentation/AllocToken/nonlibcalls.ll | 2 +-
.../AllocToken/typehashpointersplit.ll | 2 +-
23 files changed, 382 insertions(+), 98 deletions(-)
create mode 100644 llvm/include/llvm/Support/AllocToken.h
create mode 100644 llvm/lib/Support/AllocToken.cpp
create mode 100644 llvm/test/Instrumentation/AllocToken/intrinsic.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/intrinsic32.ll
diff --git a/clang/docs/AllocToken.rst b/clang/docs/AllocToken.rst
index bda84669456ce..b65e18ccfa967 100644
--- a/clang/docs/AllocToken.rst
+++ b/clang/docs/AllocToken.rst
@@ -37,8 +37,8 @@ The default mode to calculate tokens is:
pointers.
Other token ID assignment modes are supported, but they may be subject to
-change or removal. These may (experimentally) be selected with ``-mllvm
--alloc-token-mode=<mode>``:
+change or removal. These may (experimentally) be selected with ``-Xclang
+-falloc-token-mode=<mode>``:
* ``typehash``: This mode assigns a token ID based on the hash of the allocated
type's name.
diff --git a/clang/include/clang/Basic/CodeGenOptions.h b/clang/include/clang/Basic/CodeGenOptions.h
index cae06c3c9495a..5d5cf250b56b9 100644
--- a/clang/include/clang/Basic/CodeGenOptions.h
+++ b/clang/include/clang/Basic/CodeGenOptions.h
@@ -447,10 +447,6 @@ class CodeGenOptions : public CodeGenOptionsBase {
std::optional<double> AllowRuntimeCheckSkipHotCutoff;
- /// Maximum number of allocation tokens (0 = no max), nullopt if none set (use
- /// pass default).
- std::optional<uint64_t> AllocTokenMax;
-
/// List of backend command-line options for -fembed-bitcode.
std::vector<uint8_t> CmdArgs;
diff --git a/clang/include/clang/Basic/LangOptions.h b/clang/include/clang/Basic/LangOptions.h
index 41595ec2a060d..83becb73076f9 100644
--- a/clang/include/clang/Basic/LangOptions.h
+++ b/clang/include/clang/Basic/LangOptions.h
@@ -25,6 +25,7 @@
#include "llvm/ADT/FloatingPointMode.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/BinaryFormat/DXContainer.h"
+#include "llvm/Support/AllocToken.h"
#include "llvm/TargetParser/Triple.h"
#include <optional>
#include <string>
@@ -565,6 +566,13 @@ class LangOptions : public LangOptionsBase {
bool AtomicFineGrainedMemory = false;
bool AtomicIgnoreDenormalMode = false;
+ /// Maximum number of allocation tokens (0 = no max), nullopt if none set (use
+ /// target default).
+ std::optional<uint64_t> AllocTokenMax;
+
+ /// The allocation token mode.
+ std::optional<llvm::AllocTokenMode> AllocTokenMode;
+
LangOptions();
/// Set language defaults for the given input language and
diff --git a/clang/include/clang/Driver/Options.td b/clang/include/clang/Driver/Options.td
index 611b68e5281f0..370fedd545fec 100644
--- a/clang/include/clang/Driver/Options.td
+++ b/clang/include/clang/Driver/Options.td
@@ -2751,6 +2751,10 @@ def falloc_token_max_EQ : Joined<["-"], "falloc-token-max=">,
MetaVarName<"<N>">,
HelpText<"Limit to maximum N allocation tokens (0 = no max)">;
+def falloc_token_mode_EQ : Joined<["-"], "falloc-token-mode=">,
+ Group<f_Group>, Visibility<[CC1Option]>,
+ HelpText<"Set the allocation token mode (experimental)">;
+
def fallow_runtime_check_skip_hot_cutoff_EQ
: Joined<["-"], "fallow-runtime-check-skip-hot-cutoff=">,
Group<f_clang_Group>,
diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp
index f8e8086afc36f..23ad11ac9f792 100644
--- a/clang/lib/CodeGen/BackendUtil.cpp
+++ b/clang/lib/CodeGen/BackendUtil.cpp
@@ -234,9 +234,12 @@ class EmitAssemblyHelper {
};
} // namespace
-static AllocTokenOptions getAllocTokenOptions(const CodeGenOptions &CGOpts) {
+static AllocTokenOptions getAllocTokenOptions(const LangOptions &LangOpts,
+ const CodeGenOptions &CGOpts) {
AllocTokenOptions Opts;
- Opts.MaxTokens = CGOpts.AllocTokenMax;
+ if (LangOpts.AllocTokenMode)
+ Opts.Mode = *LangOpts.AllocTokenMode;
+ Opts.MaxTokens = LangOpts.AllocTokenMax;
Opts.Extended = CGOpts.SanitizeAllocTokenExtended;
Opts.FastABI = CGOpts.SanitizeAllocTokenFastABI;
return Opts;
@@ -808,7 +811,7 @@ static void addSanitizers(const Triple &TargetTriple,
// memory allocation function detection.
MPM.addPass(InferFunctionAttrsPass());
}
- MPM.addPass(AllocTokenPass(getAllocTokenOptions(CodeGenOpts)));
+ MPM.addPass(AllocTokenPass(getAllocTokenOptions(LangOpts, CodeGenOpts)));
}
};
if (ClSanitizeOnOptimizerEarlyEP) {
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 292adce8180bc..9ce1df728336e 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -1833,10 +1833,6 @@ void CompilerInvocationBase::GenerateCodeGenArgs(const CodeGenOptions &Opts,
serializeSanitizerKinds(Opts.SanitizeAnnotateDebugInfo))
GenerateArg(Consumer, OPT_fsanitize_annotate_debug_info_EQ, Sanitizer);
- if (Opts.AllocTokenMax)
- GenerateArg(Consumer, OPT_falloc_token_max_EQ,
- std::to_string(*Opts.AllocTokenMax));
-
if (!Opts.EmitVersionIdentMetadata)
GenerateArg(Consumer, OPT_Qn);
@@ -2350,15 +2346,6 @@ bool CompilerInvocation::ParseCodeGenArgs(CodeGenOptions &Opts, ArgList &Args,
}
}
- if (const auto *Arg = Args.getLastArg(options::OPT_falloc_token_max_EQ)) {
- StringRef S = Arg->getValue();
- uint64_t Value = 0;
- if (S.getAsInteger(0, Value))
- Diags.Report(diag::err_drv_invalid_value) << Arg->getAsString(Args) << S;
- else
- Opts.AllocTokenMax = Value;
- }
-
Opts.EmitVersionIdentMetadata = Args.hasFlag(OPT_Qy, OPT_Qn, true);
if (!LangOpts->CUDAIsDevice)
@@ -3966,6 +3953,29 @@ void CompilerInvocationBase::GenerateLangArgs(const LangOptions &Opts,
if (!Opts.RandstructSeed.empty())
GenerateArg(Consumer, OPT_frandomize_layout_seed_EQ, Opts.RandstructSeed);
+
+ if (Opts.AllocTokenMax)
+ GenerateArg(Consumer, OPT_falloc_token_max_EQ,
+ std::to_string(*Opts.AllocTokenMax));
+
+ if (Opts.AllocTokenMode) {
+ StringRef S;
+ switch (*Opts.AllocTokenMode) {
+ case llvm::AllocTokenMode::Increment:
+ S = "increment";
+ break;
+ case llvm::AllocTokenMode::Random:
+ S = "random";
+ break;
+ case llvm::AllocTokenMode::TypeHash:
+ S = "typehash";
+ break;
+ case llvm::AllocTokenMode::TypeHashPointerSplit:
+ S = "typehashpointersplit";
+ break;
+ }
+ GenerateArg(Consumer, OPT_falloc_token_mode_EQ, S);
+ }
}
bool CompilerInvocation::ParseLangArgs(LangOptions &Opts, ArgList &Args,
@@ -4544,6 +4554,30 @@ bool CompilerInvocation::ParseLangArgs(LangOptions &Opts, ArgList &Args,
if (const Arg *A = Args.getLastArg(OPT_frandomize_layout_seed_EQ))
Opts.RandstructSeed = A->getValue(0);
+ if (const auto *Arg = Args.getLastArg(options::OPT_falloc_token_max_EQ)) {
+ StringRef S = Arg->getValue();
+ uint64_t Value = 0;
+ if (S.getAsInteger(0, Value))
+ Diags.Report(diag::err_drv_invalid_value) << Arg->getAsString(Args) << S;
+ else
+ Opts.AllocTokenMax = Value;
+ }
+
+ if (const auto *Arg = Args.getLastArg(options::OPT_falloc_token_mode_EQ)) {
+ StringRef S = Arg->getValue();
+ auto Mode = llvm::StringSwitch<std::optional<llvm::AllocTokenMode>>(S)
+ .Case("increment", llvm::AllocTokenMode::Increment)
+ .Case("random", llvm::AllocTokenMode::Random)
+ .Case("typehash", llvm::AllocTokenMode::TypeHash)
+ .Case("typehashpointersplit",
+ llvm::AllocTokenMode::TypeHashPointerSplit)
+ .Default(std::nullopt);
+ if (Mode)
+ Opts.AllocTokenMode = Mode;
+ else
+ Diags.Report(diag::err_drv_invalid_value) << Arg->getAsString(Args) << S;
+ }
+
// Validate options for HLSL
if (Opts.HLSL) {
// TODO: Revisit restricting SPIR-V to logical once we've figured out how to
diff --git a/clang/test/Driver/fsanitize-alloc-token.c b/clang/test/Driver/fsanitize-alloc-token.c
index 2964f60c4f26f..6d8bda16dfb96 100644
--- a/clang/test/Driver/fsanitize-alloc-token.c
+++ b/clang/test/Driver/fsanitize-alloc-token.c
@@ -41,3 +41,14 @@
// CHECK-MAX: "-falloc-token-max=42"
// RUN: not %clang --target=x86_64-linux-gnu -fsanitize=alloc-token -falloc-token-max=-1 %s 2>&1 | FileCheck -check-prefix=CHECK-INVALID-MAX %s
// CHECK-INVALID-MAX: error: invalid value
+
+// RUN: %clang --target=x86_64-linux-gnu -Xclang -falloc-token-mode=increment %s -### 2>&1 | FileCheck -check-prefix=CHECK-MODE-INCREMENT %s
+// CHECK-MODE-INCREMENT: "-falloc-token-mode=increment"
+// RUN: %clang --target=x86_64-linux-gnu -Xclang -falloc-token-mode=random %s -### 2>&1 | FileCheck -check-prefix=CHECK-MODE-RANDOM %s
+// CHECK-MODE-RANDOM: "-falloc-token-mode=random"
+// RUN: %clang --target=x86_64-linux-gnu -Xclang -falloc-token-mode=typehash %s -### 2>&1 | FileCheck -check-prefix=CHECK-MODE-TYPEHASH %s
+// CHECK-MODE-TYPEHASH: "-falloc-token-mode=typehash"
+// RUN: %clang --target=x86_64-linux-gnu -Xclang -falloc-token-mode=typehashpointersplit %s -### 2>&1 | FileCheck -check-prefix=CHECK-MODE-TYPEHASHPTRSPLIT %s
+// CHECK-MODE-TYPEHASHPTRSPLIT: "-falloc-token-mode=typehashpointersplit"
+// RUN: not %clang --target=x86_64-linux-gnu -Xclang -falloc-token-mode=asdf %s 2>&1 | FileCheck -check-prefix=CHECK-INVALID-MODE %s
+// CHECK-INVALID-MODE: error: invalid value 'asdf'
diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td
index 8856eda250ed6..0c13b059c4cd0 100644
--- a/llvm/include/llvm/IR/Intrinsics.td
+++ b/llvm/include/llvm/IR/Intrinsics.td
@@ -2853,7 +2853,15 @@ def int_ptrauth_blend :
def int_ptrauth_sign_generic :
DefaultAttrsIntrinsic<[llvm_i64_ty], [llvm_i64_ty, llvm_i64_ty], [IntrNoMem]>;
+//===----------------- AllocToken Intrinsics ------------------------------===//
+
+// Return the token ID for the given !alloc_token metadata.
+def int_alloc_token_id :
+ DefaultAttrsIntrinsic<[llvm_anyint_ty], [llvm_metadata_ty],
+ [IntrNoMem, NoUndef<RetIndex>]>;
+
//===----------------------------------------------------------------------===//
+
//===------- Convergence Intrinsics ---------------------------------------===//
def int_experimental_convergence_entry
diff --git a/llvm/include/llvm/Support/AllocToken.h b/llvm/include/llvm/Support/AllocToken.h
new file mode 100644
index 0000000000000..48db026957443
--- /dev/null
+++ b/llvm/include/llvm/Support/AllocToken.h
@@ -0,0 +1,62 @@
+//===- llvm/Support/AllocToken.h - Allocation Token Calculation -----*- C++ -*//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Definition of AllocToken modes and shared calculation of stateless token IDs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_SUPPORT_ALLOCTOKEN_H
+#define LLVM_SUPPORT_ALLOCTOKEN_H
+
+#include "llvm/ADT/SmallString.h"
+#include <cstdint>
+#include <optional>
+
+namespace llvm {
+
+/// Modes for generating allocation token IDs.
+enum class AllocTokenMode {
+ /// Incrementally increasing token ID.
+ Increment,
+
+ /// Simple mode that returns a statically-assigned random token ID.
+ Random,
+
+ /// Token ID based on allocated type hash.
+ TypeHash,
+
+ /// Token ID based on allocated type hash, where the top half ID-space is
+ /// reserved for types that contain pointers and the bottom half for types
+ /// that do not contain pointers.
+ TypeHashPointerSplit,
+};
+
+/// The default allocation token mode.
+inline constexpr AllocTokenMode DefaultAllocTokenMode =
+ AllocTokenMode::TypeHashPointerSplit;
+
+/// Metadata about an allocation used to generate a token ID.
+struct AllocTokenMetadata {
+ SmallString<64> TypeName;
+ bool ContainsPointer;
+};
+
+/// Calculates stable allocation token ID. Returns std::nullopt for stateful
+/// modes that are only available in the AllocToken pass.
+///
+/// \param Mode The token generation mode.
+/// \param Metadata The metadata about the allocation.
+/// \param MaxTokens The maximum number of tokens (must not be 0)
+/// \return The calculated allocation token ID, or std::nullopt.
+std::optional<uint64_t> getAllocTokenHash(AllocTokenMode Mode,
+ const AllocTokenMetadata &Metadata,
+ uint64_t MaxTokens);
+
+} // end namespace llvm
+
+#endif // LLVM_SUPPORT_ALLOCTOKEN_H
diff --git a/llvm/include/llvm/Transforms/Instrumentation/AllocToken.h b/llvm/include/llvm/Transforms/Instrumentation/AllocToken.h
index b1391cb04302c..077703c214745 100644
--- a/llvm/include/llvm/Transforms/Instrumentation/AllocToken.h
+++ b/llvm/include/llvm/Transforms/Instrumentation/AllocToken.h
@@ -16,6 +16,7 @@
#include "llvm/IR/Analysis.h"
#include "llvm/IR/PassManager.h"
+#include "llvm/Support/AllocToken.h"
#include <optional>
namespace llvm {
@@ -23,6 +24,7 @@ namespace llvm {
class Module;
struct AllocTokenOptions {
+ AllocTokenMode Mode = DefaultAllocTokenMode;
std::optional<uint64_t> MaxTokens;
bool FastABI = false;
bool Extended = false;
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index 53cf0046bd858..c3522a38eb9e1 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1095,6 +1095,38 @@ Expected<MemorySanitizerOptions> parseMSanPassOptions(StringRef Params) {
return Result;
}
+Expected<AllocTokenOptions> parseAllocTokenPassOptions(StringRef Params) {
+ AllocTokenOptions Result;
+ while (!Params.empty()) {
+ StringRef ParamName;
+ std::tie(ParamName, Params) = Params.split(';');
+
+ if (ParamName.consume_front("mode=")) {
+ auto Mode = StringSwitch<std::optional<AllocTokenMode>>(ParamName)
+ .Case("increment", AllocTokenMode::Increment)
+ .Case("random", AllocTokenMode::Random)
+ .Case("typehash", AllocTokenMode::TypeHash)
+ .Case("typehashpointersplit",
+ AllocTokenMode::TypeHashPointerSplit)
+ .Default(std::nullopt);
+ if (Mode)
+ Result.Mode = *Mode;
+ else
+ return make_error<StringError>(
+ formatv("invalid argument to AllocToken pass mode "
+ "parameter: '{}'",
+ ParamName)
+ .str(),
+ inconvertibleErrorCode());
+ } else {
+ return make_error<StringError>(
+ formatv("invalid AllocToken pass parameter '{}'", ParamName).str(),
+ inconvertibleErrorCode());
+ }
+ }
+ return Result;
+}
+
/// Parser of parameters for SimplifyCFG pass.
Expected<SimplifyCFGOptions> parseSimplifyCFGOptions(StringRef Params) {
SimplifyCFGOptions Result;
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index 1b1652555cd28..fdd99a43e360d 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -125,7 +125,6 @@ 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())
@@ -181,6 +180,10 @@ MODULE_PASS("wholeprogramdevirt", WholeProgramDevirtPass())
#ifndef MODULE_PASS_WITH_PARAMS
#define MODULE_PASS_WITH_PARAMS(NAME, CLASS, CREATE_PASS, PARSER, PARAMS)
#endif
+MODULE_PASS_WITH_PARAMS(
+ "alloc-token", "AllocTokenPass",
+ [](AllocTokenOptions Opts) { return AllocTokenPass(Opts); },
+ parseAllocTokenPassOptions, "mode=<mode>")
MODULE_PASS_WITH_PARAMS(
"asan", "AddressSanitizerPass",
[](AddressSanitizerOptions Opts) { return AddressSanitizerPass(Opts); },
diff --git a/llvm/lib/Support/AllocToken.cpp b/llvm/lib/Support/AllocToken.cpp
new file mode 100644
index 0000000000000..6c6f80ac4997c
--- /dev/null
+++ b/llvm/lib/Support/AllocToken.cpp
@@ -0,0 +1,46 @@
+//===- AllocToken.cpp - Allocation Token Calculation ----------------------===//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+//
+// Definition of AllocToken modes and shared calculation of stateless token IDs.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/AllocToken.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/SipHash.h"
+
+namespace llvm {
+std::optional<uint64_t> getAllocTokenHash(AllocTokenMode Mode,
+ const AllocTokenMetadata &Metadata,
+ uint64_t MaxTokens) {
+ assert(MaxTokens && "Must provide concrete max tokens");
+
+ switch (Mode) {
+ case AllocTokenMode::Increment:
+ case AllocTokenMode::Random:
+ // Stateful modes cannot be implemented as a pure function.
+ return std::nullopt;
+
+ case AllocTokenMode::TypeHash: {
+ return getStableSipHash(Metadata.TypeName) % MaxTokens;
+ }
+
+ case AllocTokenMode::TypeHashPointerSplit: {
+ if (MaxTokens == 1)
+ return 0;
+ const uint64_t HalfTokens = MaxTokens / 2;
+ uint64_t Hash = getStableSipHash(Metadata.TypeName) % HalfTokens;
+ if (Metadata.ContainsPointer)
+ Hash += HalfTokens;
+ return Hash;
+ }
+ }
+
+ llvm_unreachable("");
+}
+} // namespace llvm
diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt
index 42b21b5e62029..f06b460a0113d 100644
--- a/llvm/lib/Support/CMakeLists.txt
+++ b/llvm/lib/Support/CMakeLists.txt
@@ -147,6 +147,7 @@ add_llvm_component_library(LLVMSupport
ARMBuildAttributes.cpp
AArch64AttributeParser.cpp
AArch64BuildAttributes.cpp
+ AllocToken.cpp
ARMAttributeParser.cpp
ARMWinEH.cpp
Allocator.cpp
diff --git a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
index 40720ae4b39ae..bfda56b1f746d 100644
--- a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
@@ -31,10 +31,12 @@
#include "llvm/IR/InstIterator.h"
#include "llvm/IR/InstrTypes.h"
#include "llvm/IR/Instructions.h"
+#include "llvm/IR/IntrinsicInst.h"
#include "llvm/IR/Metadata.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/IR/Type.h"
+#include "llvm/Support/AllocToken.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/Compiler.h"
@@ -53,47 +55,14 @@
#include <variant>
using namespace llvm;
+using TokenMode = AllocTokenMode;
#define DEBUG_TYPE "alloc-token"
namespace {
-//===--- Constants --------------------------------------------------------===//
-
-enum class TokenMode : unsigned {
- /// Incrementally increasing token ID.
- Increment = 0,
-
- /// Simple mode that returns a statically-assigned random token ID.
- Random = 1,
-
- /// Token ID based on allocated type hash.
- TypeHash = 2,
-
- /// Token ID based on allocated type hash, where the top half ID-space is
- /// reserved for types that contain pointers and the bottom half for types
- /// that do not contain pointers.
- TypeHashPointerSplit = 3,
-};
-
//===--- Command-line options ---------------------------------------------===//
-cl::opt<TokenMode> ClMode(
- "alloc-token-mode", cl::Hidden, cl::desc("Token assignment mode"),
- cl::init(TokenMode::TypeHashPointerSplit),
- cl::values(
- clEnumValN(TokenMode::Increment, "increment",
- "Incrementally increasing token ID"),
- clEnumValN(TokenMode::Random, "random",
- "Statically-assigned random token ID"),
- clEnumValN(TokenMode::TypeHash, "typehash",
- "Token ID based on allocated type hash"),
- clEnumValN(
- TokenMode::TypeHashPointerSplit, "typehashpointersplit",
- "Token ID based on allocated type hash, where the top half "
- "ID-space is reserved for types that contain pointers and the "
- "bottom half for types that do not contain pointers. ")));
-
cl::opt<std::string> ClFuncPrefix("alloc-token-prefix",
cl::desc("The allocation function prefix"),
cl::Hidden, cl::init("__alloc_token_"));
@@ -131,7 +100,7 @@ cl::opt<uint64_t> ClFallbackToken(
//===--- Statistics -------------------------------------------------------===//
-STATISTIC(NumFunctionsInstrumented, "Functions instrumented");
+STATISTIC(NumFunctionsModified, "Functions modified");
STATISTIC(NumAllocationsInstrumented, "Allocations instrumented");
//===----------------------------------------------------------------------===//
@@ -140,9 +109,19 @@ STATISTIC(NumAllocationsInstrumented, "Allocations instrumented");
///
/// Expected format is: !{<type-name>, <contains-pointer>}
MDNode *getAllocTokenMetadata(const CallBase &CB) {
- MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token);
- if (!Ret)
- return nullptr;
+ MDNode *Ret = nullptr;
+ if (auto *II = dyn_cast<IntrinsicInst>(&CB);
+ II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {
+ auto *MDV = cast<MetadataAsValue>(II->getArgOperand(0));
+ Ret = cast<MDNode>(MDV->getMetadata());
+ // If the intrinsic has an empty MDNode, type inference failed.
+ if (Ret->getNumOperands() == 0)
+ return nullptr;
+ } else {
+ Ret = CB.getMetadata(LLVMContext::MD_alloc_token);
+ if (!Ret)
+ return nullptr;
+ }
assert(Ret->getNumOperands() == 2 && "bad !alloc_token");
assert(isa<MDString>(Ret->getOperand(0)));
assert(isa<ConstantAsMetadata>(Ret->getOperand(1)));
@@ -206,22 +185,20 @@ class TypeHashMode : public ModeBase {
using ModeBase::ModeBase;
uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
- const auto [N, H] = getHash(CB, ORE);
- return N ? boundedToken(H) : H;
- }
-protected:
- std::pair<MDNode *, uint64_t> getHash(const CallBase &CB,
- OptimizationRemarkEmitter &ORE) {
if (MDNode *N = getAllocTokenMetadata(CB)) {
MDString *S = cast<MDString>(N->getOperand(0));
- return {N, getStableSipHash(S->getString())};
+ AllocTokenMetadata Metadata{S->getString(), containsPointer(N)};
+ if (auto Token =
+ getAllocTokenHash(TokenMode::TypeHash, Metadata, MaxTokens))
+ return *Token;
}
// Fallback.
remarkNoMetadata(CB, ORE);
- return {nullptr, ClFallbackToken};
+ return ClFallbackToken;
}
+protected:
/// Remark that there was no precise type information.
static void remarkNoMetadata(const CallBase &CB,
OptimizationRemarkEmitter &ORE) {
@@ -242,20 +219,18 @@ class TypeHashPointerSplitMode : public TypeHashMode {
using TypeHashMode::TypeHashMode;
uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
- if (MaxTokens == 1)
- return 0;
- const uint64_t HalfTokens = MaxTokens / 2;
- const auto [N, H] = getHash(CB, ORE);
- if (!N) {
- // Pick the fallback token (ClFallbackToken), which by default is 0,
- // meaning it'll fall into the pointer-less bucket. Override by setting
- // -alloc-token-fallback if that is the wrong choice.
- return H;
+ if (MDNode *N = getAllocTokenMetadata(CB)) {
+ MDString *S = cast<MDString>(N->getOperand(0));
+ AllocTokenMetadata Metadata{S->getString(), containsPointer(N)};
+ if (auto Token = getAllocTokenHash(TokenMode::TypeHashPointerSplit,
+ Metadata, MaxTokens))
+ return *Token;
}
- uint64_t Hash = H % HalfTokens; // base hash
- if (containsPointer(N))
- Hash += HalfTokens;
- return Hash;
+ // Pick the fallback token (ClFallbackToken), which by default is 0, meaning
+ // it'll fall into the pointer-less bucket. Override by setting
+ // -alloc-token-fallback if that is the wrong choice.
+ remarkNoMetadata(CB, ORE);
+ return ClFallbackToken;
}
};
@@ -275,7 +250,7 @@ class AllocToken {
: Options(transformOptionsFromCl(std::move(Opts))), Mod(M),
FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
Mode(IncrementMode(*IntPtrTy, *Options.MaxTokens)) {
- switch (ClMode.getValue()) {
+ switch (Options.Mode) {
case TokenMode::Increment:
break;
case TokenMode::Random:
@@ -315,6 +290,9 @@ class AllocToken {
FunctionCallee getTokenAllocFunction(const CallBase &CB, uint64_t TokenID,
LibFunc OriginalFunc);
+ /// Lower alloc_token_* intrinsics.
+ void replaceIntrinsicInst(IntrinsicInst *II, OptimizationRemarkEmitter &ORE);
+
/// Return the token ID from metadata in the call.
uint64_t getToken(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
return std::visit([&](auto &&Mode) { return Mode(CB, ORE); }, Mode);
@@ -336,21 +314,32 @@ bool AllocToken::instrumentFunction(Function &F) {
// Do not apply any instrumentation for naked functions.
if (F.hasFnAttribute(Attribute::Naked))
return false;
- if (F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation))
- return false;
// Don't touch available_externally functions, their actual body is elsewhere.
if (F.getLinkage() == GlobalValue::AvailableExternallyLinkage)
return false;
- // Only instrument functions that have the sanitize_alloc_token attribute.
- if (!F.hasFnAttribute(Attribute::SanitizeAllocToken))
- 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;
+
+ // Only instrument functions that have the sanitize_alloc_token attribute.
+ const bool InstrumentFunction =
+ F.hasFnAttribute(Attribute::SanitizeAllocToken) &&
+ !F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation);
// Collect all allocation calls to avoid iterator invalidation.
for (Instruction &I : instructions(F)) {
+ // Collect all alloc_token_* intrinsics.
+ if (auto *II = dyn_cast<IntrinsicInst>(&I);
+ II && II->getIntrinsicID() == Intrinsic::alloc_token_id) {
+ IntrinsicInsts.emplace_back(II);
+ continue;
+ }
+
+ if (!InstrumentFunction)
+ continue;
+
auto *CB = dyn_cast<CallBase>(&I);
if (!CB)
continue;
@@ -359,11 +348,22 @@ bool AllocToken::instrumentFunction(Function &F) {
}
bool Modified = false;
- for (auto &[CB, Func] : AllocCalls)
- Modified |= replaceAllocationCall(CB, Func, ORE, TLI);
- if (Modified)
- NumFunctionsInstrumented++;
+ if (!AllocCalls.empty()) {
+ for (auto &[CB, Func] : AllocCalls)
+ Modified |= replaceAllocationCall(CB, Func, ORE, TLI);
+ if (Modified)
+ NumFunctionsModified++;
+ }
+
+ if (!IntrinsicInsts.empty()) {
+ for (auto *II : IntrinsicInsts) {
+ replaceIntrinsicInst(II, ORE);
+ }
+ Modified = true;
+ NumFunctionsModified++;
+ }
+
return Modified;
}
@@ -528,6 +528,16 @@ FunctionCallee AllocToken::getTokenAllocFunction(const CallBase &CB,
return TokenAlloc;
}
+void AllocToken::replaceIntrinsicInst(IntrinsicInst *II,
+ OptimizationRemarkEmitter &ORE) {
+ assert(II->getIntrinsicID() == Intrinsic::alloc_token_id);
+
+ uint64_t TokenID = getToken(*II, ORE);
+ Value *V = ConstantInt::get(IntPtrTy, TokenID);
+ II->replaceAllUsesWith(V);
+ II->eraseFromParent();
+}
+
} // namespace
AllocTokenPass::AllocTokenPass(AllocTokenOptions Opts)
diff --git a/llvm/test/Instrumentation/AllocToken/basic.ll b/llvm/test/Instrumentation/AllocToken/basic.ll
index 099d37df264d6..0c34b1373cfa1 100644
--- a/llvm/test/Instrumentation/AllocToken/basic.ll
+++ b/llvm/test/Instrumentation/AllocToken/basic.ll
@@ -1,5 +1,5 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=increment -S | FileCheck %s
+; RUN: opt < %s -passes='inferattrs,alloc-token<mode=increment>' -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"
diff --git a/llvm/test/Instrumentation/AllocToken/basic32.ll b/llvm/test/Instrumentation/AllocToken/basic32.ll
index 944a452f4b4d7..52d1d1446f0ef 100644
--- a/llvm/test/Instrumentation/AllocToken/basic32.ll
+++ b/llvm/test/Instrumentation/AllocToken/basic32.ll
@@ -1,5 +1,5 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=increment -S | FileCheck %s
+; RUN: opt < %s -passes='inferattrs,alloc-token<mode=increment>' -S | FileCheck %s
target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:32:32-n8:16:32-S128"
diff --git a/llvm/test/Instrumentation/AllocToken/fast.ll b/llvm/test/Instrumentation/AllocToken/fast.ll
index 19a3ef6bb9ede..f6bf5ee0b7c97 100644
--- a/llvm/test/Instrumentation/AllocToken/fast.ll
+++ b/llvm/test/Instrumentation/AllocToken/fast.ll
@@ -1,5 +1,5 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=increment -alloc-token-fast-abi -alloc-token-max=3 -S | FileCheck %s
+; RUN: opt < %s -passes='inferattrs,alloc-token<mode=increment>' -alloc-token-fast-abi -alloc-token-max=3 -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"
diff --git a/llvm/test/Instrumentation/AllocToken/intrinsic.ll b/llvm/test/Instrumentation/AllocToken/intrinsic.ll
new file mode 100644
index 0000000000000..5c6f2f147b3d0
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/intrinsic.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Test that the alloc-token pass lowers the intrinsic to a constant token ID.
+;
+; RUN: opt < %s -passes='alloc-token<mode=typehashpointersplit>' -alloc-token-max=2 -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"
+target triple = "x86_64-unknown-linux-gnu"
+
+declare i64 @llvm.alloc.token.id.i64(metadata)
+
+define i64 @test_intrinsic_lowering() {
+; CHECK-LABEL: define i64 @test_intrinsic_lowering() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: ret i64 0
+;
+entry:
+ %token_no_ptr = call i64 @llvm.alloc.token.id.i64(metadata !0)
+ ret i64 %token_no_ptr
+}
+
+define i64 @test_intrinsic_lowering_ptr() {
+; CHECK-LABEL: define i64 @test_intrinsic_lowering_ptr() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: ret i64 1
+;
+entry:
+ %token_with_ptr = call i64 @llvm.alloc.token.id.i64(metadata !1)
+ ret i64 %token_with_ptr
+}
+
+!0 = !{!"NoPointerType", i1 false}
+!1 = !{!"PointerType", i1 true}
diff --git a/llvm/test/Instrumentation/AllocToken/intrinsic32.ll b/llvm/test/Instrumentation/AllocToken/intrinsic32.ll
new file mode 100644
index 0000000000000..15f7c258e2a5d
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/intrinsic32.ll
@@ -0,0 +1,32 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Test that the alloc-token pass lowers the intrinsic to a constant token ID.
+;
+; RUN: opt < %s -passes='alloc-token<mode=typehashpointersplit>' -alloc-token-max=2 -S | FileCheck %s
+
+target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:32:32-n8:16:32-S128"
+target triple = "i386-pc-linux-gnu"
+
+declare i32 @llvm.alloc.token.id.i32(metadata)
+
+define i32 @test_intrinsic_lowering() {
+; CHECK-LABEL: define i32 @test_intrinsic_lowering() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: ret i32 0
+;
+entry:
+ %token_no_ptr = call i32 @llvm.alloc.token.id.i32(metadata !0)
+ ret i32 %token_no_ptr
+}
+
+define i32 @test_intrinsic_lowering_ptr() {
+; CHECK-LABEL: define i32 @test_intrinsic_lowering_ptr() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: ret i32 1
+;
+entry:
+ %token_with_ptr = call i32 @llvm.alloc.token.id.i32(metadata !1)
+ ret i32 %token_with_ptr
+}
+
+!0 = !{!"NoPointerType", i1 false}
+!1 = !{!"PointerType", i1 true}
diff --git a/llvm/test/Instrumentation/AllocToken/invoke.ll b/llvm/test/Instrumentation/AllocToken/invoke.ll
index 347c99a2e8f8d..8e7ab3848dc05 100644
--- a/llvm/test/Instrumentation/AllocToken/invoke.ll
+++ b/llvm/test/Instrumentation/AllocToken/invoke.ll
@@ -1,5 +1,5 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=increment -S | FileCheck %s
+; RUN: opt < %s -passes='inferattrs,alloc-token<mode=increment>' -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"
diff --git a/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll b/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
index 19673da1bcfb6..45f573ee7b044 100644
--- a/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
+++ b/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
@@ -1,5 +1,5 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=increment -alloc-token-extended -S | FileCheck %s
+; RUN: opt < %s -passes='inferattrs,alloc-token<mode=increment>' -alloc-token-extended -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"
diff --git a/llvm/test/Instrumentation/AllocToken/typehashpointersplit.ll b/llvm/test/Instrumentation/AllocToken/typehashpointersplit.ll
index 1f776480c5b3a..4d1be5eac8cd2 100644
--- a/llvm/test/Instrumentation/AllocToken/typehashpointersplit.ll
+++ b/llvm/test/Instrumentation/AllocToken/typehashpointersplit.ll
@@ -1,5 +1,5 @@
; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=typehashpointersplit -alloc-token-max=2 -S | FileCheck %s
+; RUN: opt < %s -passes='inferattrs,alloc-token<mode=typehashpointersplit>' -alloc-token-max=2 -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"
>From 089297186dd868669f589a6b5b2e701053fd232c Mon Sep 17 00:00:00 2001
From: Marco Elver <elver at google.com>
Date: Wed, 15 Oct 2025 23:47:05 +0200
Subject: [PATCH 2/3] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20ch?=
=?UTF-8?q?anges=20introduced=20through=20rebase?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Created using spr 1.3.8-beta.1
[skip ci]
---
llvm/lib/Support/CMakeLists.txt | 2 +-
llvm/utils/gn/secondary/llvm/lib/Support/BUILD.gn | 1 +
2 files changed, 2 insertions(+), 1 deletion(-)
diff --git a/llvm/lib/Support/CMakeLists.txt b/llvm/lib/Support/CMakeLists.txt
index f06b460a0113d..671a5fe941cef 100644
--- a/llvm/lib/Support/CMakeLists.txt
+++ b/llvm/lib/Support/CMakeLists.txt
@@ -147,9 +147,9 @@ add_llvm_component_library(LLVMSupport
ARMBuildAttributes.cpp
AArch64AttributeParser.cpp
AArch64BuildAttributes.cpp
- AllocToken.cpp
ARMAttributeParser.cpp
ARMWinEH.cpp
+ AllocToken.cpp
Allocator.cpp
AutoConvert.cpp
Base64.cpp
diff --git a/llvm/utils/gn/secondary/llvm/lib/Support/BUILD.gn b/llvm/utils/gn/secondary/llvm/lib/Support/BUILD.gn
index 38ba4661daacc..df9ddf91f2c49 100644
--- a/llvm/utils/gn/secondary/llvm/lib/Support/BUILD.gn
+++ b/llvm/utils/gn/secondary/llvm/lib/Support/BUILD.gn
@@ -45,6 +45,7 @@ static_library("Support") {
"ARMAttributeParser.cpp",
"ARMBuildAttributes.cpp",
"ARMWinEH.cpp",
+ "AllocToken.cpp",
"Allocator.cpp",
"AutoConvert.cpp",
"BalancedPartitioning.cpp",
>From 76d33ee5c9a5d1932f7448b03b05dda093e103f4 Mon Sep 17 00:00:00 2001
From: Marco Elver <elver at google.com>
Date: Fri, 17 Oct 2025 20:14:33 +0200
Subject: [PATCH 3/3] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20ch?=
=?UTF-8?q?anges=20introduced=20through=20rebase?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Created using spr 1.3.8-beta.1
[skip ci]
---
clang/lib/Frontend/CompilerInvocation.cpp | 9 +----
llvm/include/llvm/Support/AllocToken.h | 12 +++++--
llvm/lib/Passes/PassBuilder.cpp | 9 +----
llvm/lib/Support/AllocToken.cpp | 35 +++++++++++++------
.../Transforms/Instrumentation/AllocToken.cpp | 12 +++----
5 files changed, 41 insertions(+), 36 deletions(-)
diff --git a/clang/lib/Frontend/CompilerInvocation.cpp b/clang/lib/Frontend/CompilerInvocation.cpp
index 9ce1df728336e..85ba85099500a 100644
--- a/clang/lib/Frontend/CompilerInvocation.cpp
+++ b/clang/lib/Frontend/CompilerInvocation.cpp
@@ -4565,14 +4565,7 @@ bool CompilerInvocation::ParseLangArgs(LangOptions &Opts, ArgList &Args,
if (const auto *Arg = Args.getLastArg(options::OPT_falloc_token_mode_EQ)) {
StringRef S = Arg->getValue();
- auto Mode = llvm::StringSwitch<std::optional<llvm::AllocTokenMode>>(S)
- .Case("increment", llvm::AllocTokenMode::Increment)
- .Case("random", llvm::AllocTokenMode::Random)
- .Case("typehash", llvm::AllocTokenMode::TypeHash)
- .Case("typehashpointersplit",
- llvm::AllocTokenMode::TypeHashPointerSplit)
- .Default(std::nullopt);
- if (Mode)
+ if (auto Mode = getAllocTokenModeFromString(S))
Opts.AllocTokenMode = Mode;
else
Diags.Report(diag::err_drv_invalid_value) << Arg->getAsString(Args) << S;
diff --git a/llvm/include/llvm/Support/AllocToken.h b/llvm/include/llvm/Support/AllocToken.h
index 48db026957443..e40d8163a9d7c 100644
--- a/llvm/include/llvm/Support/AllocToken.h
+++ b/llvm/include/llvm/Support/AllocToken.h
@@ -14,6 +14,7 @@
#define LLVM_SUPPORT_ALLOCTOKEN_H
#include "llvm/ADT/SmallString.h"
+#include "llvm/ADT/StringRef.h"
#include <cstdint>
#include <optional>
@@ -40,6 +41,11 @@ enum class AllocTokenMode {
inline constexpr AllocTokenMode DefaultAllocTokenMode =
AllocTokenMode::TypeHashPointerSplit;
+/// Returns the AllocTokenMode from its canonical string name; if an invalid
+/// name was provided returns nullopt.
+LLVM_ABI std::optional<AllocTokenMode>
+getAllocTokenModeFromString(StringRef Name);
+
/// Metadata about an allocation used to generate a token ID.
struct AllocTokenMetadata {
SmallString<64> TypeName;
@@ -53,9 +59,9 @@ struct AllocTokenMetadata {
/// \param Metadata The metadata about the allocation.
/// \param MaxTokens The maximum number of tokens (must not be 0)
/// \return The calculated allocation token ID, or std::nullopt.
-std::optional<uint64_t> getAllocTokenHash(AllocTokenMode Mode,
- const AllocTokenMetadata &Metadata,
- uint64_t MaxTokens);
+LLVM_ABI std::optional<uint64_t>
+getAllocToken(AllocTokenMode Mode, const AllocTokenMetadata &Metadata,
+ uint64_t MaxTokens);
} // end namespace llvm
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index c3522a38eb9e1..4cebb0bb32e4e 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/llvm/lib/Passes/PassBuilder.cpp
@@ -1102,14 +1102,7 @@ Expected<AllocTokenOptions> parseAllocTokenPassOptions(StringRef Params) {
std::tie(ParamName, Params) = Params.split(';');
if (ParamName.consume_front("mode=")) {
- auto Mode = StringSwitch<std::optional<AllocTokenMode>>(ParamName)
- .Case("increment", AllocTokenMode::Increment)
- .Case("random", AllocTokenMode::Random)
- .Case("typehash", AllocTokenMode::TypeHash)
- .Case("typehashpointersplit",
- AllocTokenMode::TypeHashPointerSplit)
- .Default(std::nullopt);
- if (Mode)
+ if (auto Mode = getAllocTokenModeFromString(ParamName))
Result.Mode = *Mode;
else
return make_error<StringError>(
diff --git a/llvm/lib/Support/AllocToken.cpp b/llvm/lib/Support/AllocToken.cpp
index 6c6f80ac4997c..8e9e89f0df353 100644
--- a/llvm/lib/Support/AllocToken.cpp
+++ b/llvm/lib/Support/AllocToken.cpp
@@ -11,14 +11,31 @@
//===----------------------------------------------------------------------===//
#include "llvm/Support/AllocToken.h"
+#include "llvm/ADT/StringSwitch.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/SipHash.h"
-namespace llvm {
-std::optional<uint64_t> getAllocTokenHash(AllocTokenMode Mode,
- const AllocTokenMetadata &Metadata,
- uint64_t MaxTokens) {
- assert(MaxTokens && "Must provide concrete max tokens");
+using namespace llvm;
+
+std::optional<AllocTokenMode>
+llvm::getAllocTokenModeFromString(StringRef Name) {
+ return StringSwitch<std::optional<AllocTokenMode>>(Name)
+ .Case("increment", AllocTokenMode::Increment)
+ .Case("random", AllocTokenMode::Random)
+ .Case("typehash", AllocTokenMode::TypeHash)
+ .Case("typehashpointersplit", AllocTokenMode::TypeHashPointerSplit)
+ .Default(std::nullopt);
+}
+
+static uint64_t getStableHash(const AllocTokenMetadata &Metadata,
+ uint64_t MaxTokens) {
+ return getStableSipHash(Metadata.TypeName) % MaxTokens;
+}
+
+std::optional<uint64_t> llvm::getAllocToken(AllocTokenMode Mode,
+ const AllocTokenMetadata &Metadata,
+ uint64_t MaxTokens) {
+ assert(MaxTokens && "Must provide non-zero max tokens");
switch (Mode) {
case AllocTokenMode::Increment:
@@ -26,15 +43,14 @@ std::optional<uint64_t> getAllocTokenHash(AllocTokenMode Mode,
// Stateful modes cannot be implemented as a pure function.
return std::nullopt;
- case AllocTokenMode::TypeHash: {
- return getStableSipHash(Metadata.TypeName) % MaxTokens;
- }
+ case AllocTokenMode::TypeHash:
+ return getStableHash(Metadata, MaxTokens);
case AllocTokenMode::TypeHashPointerSplit: {
if (MaxTokens == 1)
return 0;
const uint64_t HalfTokens = MaxTokens / 2;
- uint64_t Hash = getStableSipHash(Metadata.TypeName) % HalfTokens;
+ uint64_t Hash = getStableHash(Metadata, HalfTokens);
if (Metadata.ContainsPointer)
Hash += HalfTokens;
return Hash;
@@ -43,4 +59,3 @@ std::optional<uint64_t> getAllocTokenHash(AllocTokenMode Mode,
llvm_unreachable("");
}
-} // namespace llvm
diff --git a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
index bfda56b1f746d..8181e4ef1d74f 100644
--- a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
+++ b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
@@ -189,8 +189,7 @@ class TypeHashMode : public ModeBase {
if (MDNode *N = getAllocTokenMetadata(CB)) {
MDString *S = cast<MDString>(N->getOperand(0));
AllocTokenMetadata Metadata{S->getString(), containsPointer(N)};
- if (auto Token =
- getAllocTokenHash(TokenMode::TypeHash, Metadata, MaxTokens))
+ if (auto Token = getAllocToken(TokenMode::TypeHash, Metadata, MaxTokens))
return *Token;
}
// Fallback.
@@ -222,8 +221,8 @@ class TypeHashPointerSplitMode : public TypeHashMode {
if (MDNode *N = getAllocTokenMetadata(CB)) {
MDString *S = cast<MDString>(N->getOperand(0));
AllocTokenMetadata Metadata{S->getString(), containsPointer(N)};
- if (auto Token = getAllocTokenHash(TokenMode::TypeHashPointerSplit,
- Metadata, MaxTokens))
+ if (auto Token = getAllocToken(TokenMode::TypeHashPointerSplit, Metadata,
+ MaxTokens))
return *Token;
}
// Pick the fallback token (ClFallbackToken), which by default is 0, meaning
@@ -357,9 +356,8 @@ bool AllocToken::instrumentFunction(Function &F) {
}
if (!IntrinsicInsts.empty()) {
- for (auto *II : IntrinsicInsts) {
+ for (auto *II : IntrinsicInsts)
replaceIntrinsicInst(II, ORE);
- }
Modified = true;
NumFunctionsModified++;
}
@@ -381,7 +379,7 @@ AllocToken::shouldInstrumentCall(const CallBase &CB,
if (TLI.getLibFunc(*Callee, Func)) {
if (isInstrumentableLibFunc(Func, CB, TLI))
return Func;
- } else if (Options.Extended && getAllocTokenMetadata(CB)) {
+ } else if (Options.Extended && CB.getMetadata(LLVMContext::MD_alloc_token)) {
return NotLibFunc;
}
More information about the llvm-commits
mailing list