[llvm] [Clang][CodeGen] Introduce the AllocToken SanitizerKind (PR #162098)
Marco Elver via llvm-commits
llvm-commits at lists.llvm.org
Tue Oct 7 05:56:12 PDT 2025
https://github.com/melver updated https://github.com/llvm/llvm-project/pull/162098
>From a77cb189f426829fed18edef7e91a43de6c8e1e6 Mon Sep 17 00:00:00 2001
From: Marco Elver <elver at google.com>
Date: Mon, 6 Oct 2025 16:55:46 +0200
Subject: [PATCH] =?UTF-8?q?[=F0=9D=98=80=F0=9D=97=BD=F0=9D=97=BF]=20change?=
=?UTF-8?q?s=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]
---
llvm/docs/LangRef.rst | 18 +
llvm/docs/ReleaseNotes.md | 4 +
llvm/include/llvm/Bitcode/LLVMBitCodes.h | 1 +
llvm/include/llvm/IR/Attributes.td | 3 +
llvm/include/llvm/IR/FixedMetadataKinds.def | 1 +
.../Transforms/Instrumentation/AllocToken.h | 46 ++
llvm/lib/Bitcode/Reader/BitcodeReader.cpp | 2 +
llvm/lib/Bitcode/Writer/BitcodeWriter.cpp | 2 +
llvm/lib/IR/Verifier.cpp | 10 +
llvm/lib/Passes/PassBuilder.cpp | 1 +
llvm/lib/Passes/PassRegistry.def | 1 +
.../Transforms/Instrumentation/AllocToken.cpp | 494 ++++++++++++++++++
.../Transforms/Instrumentation/CMakeLists.txt | 1 +
llvm/lib/Transforms/Utils/CodeExtractor.cpp | 1 +
llvm/lib/Transforms/Utils/Local.cpp | 7 +
llvm/test/Bitcode/attributes.ll | 6 +
llvm/test/Bitcode/compatibility.ll | 8 +-
llvm/test/Instrumentation/AllocToken/basic.ll | 98 ++++
.../Instrumentation/AllocToken/basic32.ll | 31 ++
.../AllocToken/extralibfuncs.ll | 44 ++
llvm/test/Instrumentation/AllocToken/fast.ll | 42 ++
.../test/Instrumentation/AllocToken/ignore.ll | 29 +
.../test/Instrumentation/AllocToken/invoke.ll | 123 +++++
.../Instrumentation/AllocToken/nonlibcalls.ll | 85 +++
.../test/Instrumentation/AllocToken/remark.ll | 38 ++
llvm/test/Transforms/Inline/attributes.ll | 42 ++
.../SimplifyCFG/merge-calls-alloc-token.ll | 104 ++++
llvm/utils/emacs/llvm-mode.el | 2 +-
.../lib/Transforms/Instrumentation/BUILD.gn | 1 +
llvm/utils/llvm.grm | 1 +
llvm/utils/vim/syntax/llvm.vim | 1 +
.../vscode/llvm/syntaxes/ll.tmLanguage.yaml | 1 +
32 files changed, 1245 insertions(+), 3 deletions(-)
create mode 100644 llvm/include/llvm/Transforms/Instrumentation/AllocToken.h
create mode 100644 llvm/lib/Transforms/Instrumentation/AllocToken.cpp
create mode 100644 llvm/test/Instrumentation/AllocToken/basic.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/basic32.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/extralibfuncs.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/fast.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/ignore.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/invoke.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
create mode 100644 llvm/test/Instrumentation/AllocToken/remark.ll
create mode 100644 llvm/test/Transforms/SimplifyCFG/merge-calls-alloc-token.ll
diff --git a/llvm/docs/LangRef.rst b/llvm/docs/LangRef.rst
index 8e863939781a2..633a62860dfc3 100644
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@@ -2527,6 +2527,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 attribute 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>`_
@@ -8548,6 +8551,21 @@ Example:
The ``nofree`` metadata indicates the memory pointed by the pointer will not be
freed after the attached instruction.
+'``alloc_token``' Metadata
+^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The ``alloc_token`` metadata may be attached to calls to memory allocation
+functions, and contains richer semantic information about the type of the
+allocation. This information is consumed by the ``alloc-token`` pass to
+instrument such calls with allocation token IDs.
+
+The metadata contains a string with the type of an allocation.
+
+.. code-block:: none
+
+ call ptr @malloc(i64 64), !alloc_token !0
+
+ !0 = !{!"<type-name>"}
Module Flags Metadata
=====================
diff --git a/llvm/docs/ReleaseNotes.md b/llvm/docs/ReleaseNotes.md
index 85c16b9c33f10..94c9c2cdbc34c 100644
--- a/llvm/docs/ReleaseNotes.md
+++ b/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}}
===============================================================
diff --git a/llvm/include/llvm/Bitcode/LLVMBitCodes.h b/llvm/include/llvm/Bitcode/LLVMBitCodes.h
index 1c7d3462b6bae..464f475098ec5 100644
--- a/llvm/include/llvm/Bitcode/LLVMBitCodes.h
+++ b/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 {
diff --git a/llvm/include/llvm/IR/Attributes.td b/llvm/include/llvm/IR/Attributes.td
index ef816fb86ed1d..8e7d9dcebfe2a 100644
--- a/llvm/include/llvm/IR/Attributes.td
+++ b/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
diff --git a/llvm/include/llvm/IR/FixedMetadataKinds.def b/llvm/include/llvm/IR/FixedMetadataKinds.def
index d09cc15d65ff6..86a4d79afc2de 100644
--- a/llvm/include/llvm/IR/FixedMetadataKinds.def
+++ b/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, "alloc_token", 44)
diff --git a/llvm/include/llvm/Transforms/Instrumentation/AllocToken.h b/llvm/include/llvm/Transforms/Instrumentation/AllocToken.h
new file mode 100644
index 0000000000000..b1391cb04302c
--- /dev/null
+++ b/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
diff --git a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp b/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
index 832aa9ff7ed3d..aaee1f0a7687c 100644
--- a/llvm/lib/Bitcode/Reader/BitcodeReader.cpp
+++ b/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:
diff --git a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp b/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
index c4070e1f44688..6d86809a5c282 100644
--- a/llvm/lib/Bitcode/Writer/BitcodeWriter.cpp
+++ b/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:
diff --git a/llvm/lib/IR/Verifier.cpp b/llvm/lib/IR/Verifier.cpp
index 8c03d6f809d50..a5c7558dce954 100644
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@@ -539,6 +539,7 @@ class Verifier : public InstVisitor<Verifier>, VerifierSupport {
void visitDIAssignIDMetadata(Instruction &I, MDNode *MD);
void visitMMRAMetadata(Instruction &I, MDNode *MD);
void visitAnnotationMetadata(MDNode *Annotation);
+ void visitAllocTokenMetadata(Instruction &I, MDNode *MD);
void visitAliasScopeMetadata(const MDNode *MD);
void visitAliasScopeListMetadata(const MDNode *MD);
void visitAccessGroupMetadata(const MDNode *MD);
@@ -5373,6 +5374,12 @@ void Verifier::visitAccessGroupMetadata(const MDNode *MD) {
}
}
+void Verifier::visitAllocTokenMetadata(Instruction &I, MDNode *MD) {
+ Check(isa<CallBase>(I), "!alloc_token should only exist on calls", &I);
+ Check(MD->getNumOperands() == 1, "!alloc_token must have 1 operand", MD);
+ Check(isa<MDString>(MD->getOperand(0)), "expected string", MD);
+}
+
/// verifyInstruction - Verify that an instruction is well formed.
///
void Verifier::visitInstruction(Instruction &I) {
@@ -5600,6 +5607,9 @@ void Verifier::visitInstruction(Instruction &I) {
if (MDNode *Annotation = I.getMetadata(LLVMContext::MD_annotation))
visitAnnotationMetadata(Annotation);
+ if (MDNode *MD = I.getMetadata(LLVMContext::MD_alloc_token))
+ visitAllocTokenMetadata(I, MD);
+
if (MDNode *N = I.getDebugLoc().getAsMDNode()) {
CheckDI(isa<DILocation>(N), "invalid !dbg metadata attachment", &I, N);
visitMDNode(*N, AreDebugLocsAllowed::Yes);
diff --git a/llvm/lib/Passes/PassBuilder.cpp b/llvm/lib/Passes/PassBuilder.cpp
index c234623caecf9..20dcde8fff0d4 100644
--- a/llvm/lib/Passes/PassBuilder.cpp
+++ b/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"
diff --git a/llvm/lib/Passes/PassRegistry.def b/llvm/lib/Passes/PassRegistry.def
index f0e7d36f78aab..6b9252fb55a71 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/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())
diff --git a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
new file mode 100644
index 0000000000000..782d5a162a314
--- /dev/null
+++ b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp
@@ -0,0 +1,494 @@
+//===- AllocToken.cpp - 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 implements AllocToken, an instrumentation pass that
+// replaces allocation calls with token-enabled versions.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Transforms/Instrumentation/AllocToken.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SmallPtrSet.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/Statistic.h"
+#include "llvm/ADT/StringExtras.h"
+#include "llvm/ADT/StringRef.h"
+#include "llvm/Analysis/MemoryBuiltins.h"
+#include "llvm/Analysis/OptimizationRemarkEmitter.h"
+#include "llvm/Analysis/TargetLibraryInfo.h"
+#include "llvm/IR/Analysis.h"
+#include "llvm/IR/Attributes.h"
+#include "llvm/IR/Constants.h"
+#include "llvm/IR/DerivedTypes.h"
+#include "llvm/IR/Function.h"
+#include "llvm/IR/GlobalValue.h"
+#include "llvm/IR/IRBuilder.h"
+#include "llvm/IR/InstIterator.h"
+#include "llvm/IR/InstrTypes.h"
+#include "llvm/IR/Instructions.h"
+#include "llvm/IR/Metadata.h"
+#include "llvm/IR/Module.h"
+#include "llvm/IR/PassManager.h"
+#include "llvm/IR/Type.h"
+#include "llvm/Support/Casting.h"
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Compiler.h"
+#include "llvm/Support/ErrorHandling.h"
+#include "llvm/Support/RandomNumberGenerator.h"
+#include "llvm/Support/SipHash.h"
+#include "llvm/Support/raw_ostream.h"
+#include <cassert>
+#include <cstddef>
+#include <cstdint>
+#include <limits>
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+#include <variant>
+
+using namespace llvm;
+
+#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,
+};
+
+//===--- Command-line options ---------------------------------------------===//
+
+cl::opt<TokenMode>
+ ClMode("alloc-token-mode", cl::Hidden, cl::desc("Token assignment mode"),
+ cl::init(TokenMode::TypeHash),
+ 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")));
+
+cl::opt<std::string> ClFuncPrefix("alloc-token-prefix",
+ cl::desc("The allocation function prefix"),
+ cl::Hidden, cl::init("__alloc_token_"));
+
+cl::opt<uint64_t> ClMaxTokens("alloc-token-max",
+ cl::desc("Maximum number of tokens (0 = no max)"),
+ cl::Hidden, cl::init(0));
+
+cl::opt<bool>
+ ClFastABI("alloc-token-fast-abi",
+ cl::desc("The token ID is encoded in the function name"),
+ cl::Hidden, cl::init(false));
+
+// Instrument libcalls only by default - compatible allocators only need to take
+// care of providing standard allocation functions. With extended coverage, also
+// instrument non-libcall allocation function calls with !alloc_token
+// metadata.
+cl::opt<bool>
+ ClExtended("alloc-token-extended",
+ cl::desc("Extend coverage to custom allocation functions"),
+ cl::Hidden, cl::init(false));
+
+// C++ defines ::operator new (and variants) as replaceable (vs. standard
+// library versions), which are nobuiltin, and are therefore not covered by
+// isAllocationFn(). Cover by default, as users of AllocToken are already
+// required to provide token-aware allocation functions (no defaults).
+cl::opt<bool> ClCoverReplaceableNew("alloc-token-cover-replaceable-new",
+ cl::desc("Cover replaceable operator new"),
+ cl::Hidden, cl::init(true));
+
+cl::opt<uint64_t> ClFallbackToken(
+ "alloc-token-fallback",
+ cl::desc("The default fallback token where none could be determined"),
+ cl::Hidden, cl::init(0));
+
+//===--- Statistics -------------------------------------------------------===//
+
+STATISTIC(NumFunctionsInstrumented, "Functions instrumented");
+STATISTIC(NumAllocationsInstrumented, "Allocations instrumented");
+
+//===----------------------------------------------------------------------===//
+
+/// Returns the !alloc_token metadata if available.
+///
+/// Expected format is: !{<type-name>}
+MDNode *getAllocTokenMetadata(const CallBase &CB) {
+ MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token);
+ if (!Ret)
+ return nullptr;
+ assert(Ret->getNumOperands() == 1 && "bad !alloc_token");
+ assert(isa<MDString>(Ret->getOperand(0)));
+ return Ret;
+}
+
+class ModeBase {
+public:
+ explicit ModeBase(const IntegerType &TokenTy, uint64_t MaxTokens)
+ : MaxTokens(MaxTokens ? MaxTokens : TokenTy.getBitMask()) {
+ assert(MaxTokens <= TokenTy.getBitMask());
+ }
+
+protected:
+ uint64_t boundedToken(uint64_t Val) const {
+ assert(MaxTokens != 0);
+ return Val % MaxTokens;
+ }
+
+ const uint64_t MaxTokens;
+};
+
+/// Implementation for TokenMode::Increment.
+class IncrementMode : public ModeBase {
+public:
+ using ModeBase::ModeBase;
+
+ uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &) {
+ return boundedToken(Counter++);
+ }
+
+private:
+ uint64_t Counter = 0;
+};
+
+/// Implementation for TokenMode::Random.
+class RandomMode : public ModeBase {
+public:
+ RandomMode(const IntegerType &TokenTy, uint64_t MaxTokens,
+ std::unique_ptr<RandomNumberGenerator> RNG)
+ : ModeBase(TokenTy, MaxTokens), RNG(std::move(RNG)) {}
+ uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &) {
+ return boundedToken((*RNG)());
+ }
+
+private:
+ std::unique_ptr<RandomNumberGenerator> RNG;
+};
+
+/// Implementation for TokenMode::TypeHash. The implementation ensures
+/// hashes are stable across different compiler invocations. Uses SipHash as the
+/// hash function.
+class TypeHashMode : public ModeBase {
+public:
+ using ModeBase::ModeBase;
+
+ uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
+ if (MDNode *N = getAllocTokenMetadata(CB)) {
+ MDString *S = cast<MDString>(N->getOperand(0));
+ return boundedToken(getStableSipHash(S->getString()));
+ }
+ remarkNoMetadata(CB, ORE);
+ return ClFallbackToken;
+ }
+
+ /// Remark that there was no precise type information.
+ static void remarkNoMetadata(const CallBase &CB,
+ OptimizationRemarkEmitter &ORE) {
+ ORE.emit([&] {
+ ore::NV FuncNV("Function", CB.getParent()->getParent());
+ const Function *Callee = CB.getCalledFunction();
+ ore::NV CalleeNV("Callee", Callee ? Callee->getName() : "<unknown>");
+ return OptimizationRemark(DEBUG_TYPE, "NoAllocToken", &CB)
+ << "Call to '" << CalleeNV << "' in '" << FuncNV
+ << "' without source-level type token";
+ });
+ }
+};
+
+// Apply opt overrides.
+AllocTokenOptions transformOptionsFromCl(AllocTokenOptions Opts) {
+ if (!Opts.MaxTokens.has_value())
+ Opts.MaxTokens = ClMaxTokens;
+ Opts.FastABI |= ClFastABI;
+ Opts.Extended |= ClExtended;
+ return Opts;
+}
+
+class AllocToken {
+public:
+ explicit AllocToken(AllocTokenOptions Opts, Module &M,
+ ModuleAnalysisManager &MAM)
+ : Options(transformOptionsFromCl(std::move(Opts))), Mod(M),
+ FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
+ Mode(IncrementMode(*IntPtrTy, *Options.MaxTokens)) {
+ switch (ClMode.getValue()) {
+ case TokenMode::Increment:
+ break;
+ case TokenMode::Random:
+ Mode.emplace<RandomMode>(*IntPtrTy, *Options.MaxTokens,
+ M.createRNG(DEBUG_TYPE));
+ break;
+ case TokenMode::TypeHash:
+ Mode.emplace<TypeHashMode>(*IntPtrTy, *Options.MaxTokens);
+ break;
+ }
+ }
+
+ bool instrumentFunction(Function &F);
+
+private:
+ /// Returns the LibFunc (or NotLibFunc) if this call should be instrumented.
+ std::optional<LibFunc>
+ shouldInstrumentCall(const CallBase &CB, const TargetLibraryInfo &TLI) const;
+
+ /// Returns true for functions that are eligible for instrumentation.
+ static bool isInstrumentableLibFunc(LibFunc Func, const CallBase &CB,
+ const TargetLibraryInfo &TLI);
+
+ /// Returns true for isAllocationFn() functions that we should ignore.
+ static bool ignoreInstrumentableLibFunc(LibFunc Func);
+
+ /// Replace a call/invoke with a call/invoke to the allocation function
+ /// with token ID.
+ bool replaceAllocationCall(CallBase *CB, LibFunc Func,
+ OptimizationRemarkEmitter &ORE,
+ const TargetLibraryInfo &TLI);
+
+ /// Return replacement function for a LibFunc that takes a token ID.
+ FunctionCallee getTokenAllocFunction(const CallBase &CB, uint64_t TokenID,
+ LibFunc OriginalFunc);
+
+ /// 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);
+ }
+
+ const AllocTokenOptions Options;
+ Module &Mod;
+ IntegerType *IntPtrTy = Mod.getDataLayout().getIntPtrType(Mod.getContext());
+ FunctionAnalysisManager &FAM;
+ // Cache for replacement functions.
+ DenseMap<std::pair<LibFunc, uint64_t>, FunctionCallee> TokenAllocFunctions;
+ // Selected mode.
+ std::variant<IncrementMode, RandomMode, TypeHashMode> Mode;
+};
+
+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;
+
+ // Collect all allocation calls to avoid iterator invalidation.
+ for (Instruction &I : instructions(F)) {
+ auto *CB = dyn_cast<CallBase>(&I);
+ if (!CB)
+ continue;
+ if (std::optional<LibFunc> Func = shouldInstrumentCall(*CB, TLI))
+ AllocCalls.emplace_back(CB, Func.value());
+ }
+
+ bool Modified = false;
+ for (auto &[CB, Func] : AllocCalls)
+ Modified |= replaceAllocationCall(CB, Func, ORE, TLI);
+
+ if (Modified)
+ NumFunctionsInstrumented++;
+ return Modified;
+}
+
+std::optional<LibFunc>
+AllocToken::shouldInstrumentCall(const CallBase &CB,
+ const TargetLibraryInfo &TLI) const {
+ const Function *Callee = CB.getCalledFunction();
+ if (!Callee)
+ return std::nullopt;
+
+ // Ignore nobuiltin of the CallBase, so that we can cover nobuiltin libcalls
+ // if requested via isInstrumentableLibFunc(). Note that isAllocationFn() is
+ // returning false for nobuiltin calls.
+ LibFunc Func;
+ if (TLI.getLibFunc(*Callee, Func)) {
+ if (isInstrumentableLibFunc(Func, CB, TLI))
+ return Func;
+ } else if (Options.Extended && getAllocTokenMetadata(CB)) {
+ return NotLibFunc;
+ }
+
+ return std::nullopt;
+}
+
+bool AllocToken::isInstrumentableLibFunc(LibFunc Func, const CallBase &CB,
+ const TargetLibraryInfo &TLI) {
+ if (ignoreInstrumentableLibFunc(Func))
+ return false;
+
+ if (isAllocationFn(&CB, &TLI))
+ return true;
+
+ switch (Func) {
+ // These libfuncs don't return normal pointers, and are therefore not handled
+ // by isAllocationFn().
+ case LibFunc_posix_memalign:
+ case LibFunc_size_returning_new:
+ case LibFunc_size_returning_new_hot_cold:
+ case LibFunc_size_returning_new_aligned:
+ case LibFunc_size_returning_new_aligned_hot_cold:
+ return true;
+
+ // See comment above ClCoverReplaceableNew.
+ case LibFunc_Znwj:
+ case LibFunc_ZnwjRKSt9nothrow_t:
+ case LibFunc_ZnwjSt11align_val_t:
+ case LibFunc_ZnwjSt11align_val_tRKSt9nothrow_t:
+ case LibFunc_Znwm:
+ case LibFunc_Znwm12__hot_cold_t:
+ case LibFunc_ZnwmRKSt9nothrow_t:
+ case LibFunc_ZnwmRKSt9nothrow_t12__hot_cold_t:
+ case LibFunc_ZnwmSt11align_val_t:
+ case LibFunc_ZnwmSt11align_val_t12__hot_cold_t:
+ case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t:
+ case LibFunc_ZnwmSt11align_val_tRKSt9nothrow_t12__hot_cold_t:
+ case LibFunc_Znaj:
+ case LibFunc_ZnajRKSt9nothrow_t:
+ case LibFunc_ZnajSt11align_val_t:
+ case LibFunc_ZnajSt11align_val_tRKSt9nothrow_t:
+ case LibFunc_Znam:
+ case LibFunc_Znam12__hot_cold_t:
+ case LibFunc_ZnamRKSt9nothrow_t:
+ case LibFunc_ZnamRKSt9nothrow_t12__hot_cold_t:
+ case LibFunc_ZnamSt11align_val_t:
+ case LibFunc_ZnamSt11align_val_t12__hot_cold_t:
+ case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t:
+ case LibFunc_ZnamSt11align_val_tRKSt9nothrow_t12__hot_cold_t:
+ return ClCoverReplaceableNew;
+
+ default:
+ return false;
+ }
+}
+
+bool AllocToken::ignoreInstrumentableLibFunc(LibFunc Func) {
+ switch (Func) {
+ case LibFunc_strdup:
+ case LibFunc_dunder_strdup:
+ case LibFunc_strndup:
+ case LibFunc_dunder_strndup:
+ return true;
+ default:
+ return false;
+ }
+}
+
+bool AllocToken::replaceAllocationCall(CallBase *CB, LibFunc Func,
+ OptimizationRemarkEmitter &ORE,
+ const TargetLibraryInfo &TLI) {
+ uint64_t TokenID = getToken(*CB, ORE);
+
+ FunctionCallee TokenAlloc = getTokenAllocFunction(*CB, TokenID, Func);
+ if (!TokenAlloc)
+ return false;
+ NumAllocationsInstrumented++;
+
+ if (Options.FastABI) {
+ assert(TokenAlloc.getFunctionType()->getNumParams() == CB->arg_size());
+ CB->setCalledFunction(TokenAlloc);
+ return true;
+ }
+
+ IRBuilder<> IRB(CB);
+ // Original args.
+ SmallVector<Value *, 4> NewArgs{CB->args()};
+ // Add token ID, truncated to IntPtrTy width.
+ NewArgs.push_back(ConstantInt::get(IntPtrTy, TokenID));
+ assert(TokenAlloc.getFunctionType()->getNumParams() == NewArgs.size());
+
+ // Preserve invoke vs call semantics for exception handling.
+ CallBase *NewCall;
+ if (auto *II = dyn_cast<InvokeInst>(CB)) {
+ NewCall = IRB.CreateInvoke(TokenAlloc, II->getNormalDest(),
+ II->getUnwindDest(), NewArgs);
+ } else {
+ NewCall = IRB.CreateCall(TokenAlloc, NewArgs);
+ cast<CallInst>(NewCall)->setTailCall(CB->isTailCall());
+ }
+ NewCall->setCallingConv(CB->getCallingConv());
+ NewCall->copyMetadata(*CB);
+ NewCall->setAttributes(CB->getAttributes());
+
+ // Replace all uses and delete the old call.
+ CB->replaceAllUsesWith(NewCall);
+ CB->eraseFromParent();
+ return true;
+}
+
+FunctionCallee AllocToken::getTokenAllocFunction(const CallBase &CB,
+ uint64_t TokenID,
+ LibFunc OriginalFunc) {
+ std::optional<std::pair<LibFunc, uint64_t>> Key;
+ if (OriginalFunc != NotLibFunc) {
+ Key = std::make_pair(OriginalFunc, Options.FastABI ? TokenID : 0);
+ auto It = TokenAllocFunctions.find(*Key);
+ if (It != TokenAllocFunctions.end())
+ return It->second;
+ }
+
+ const Function *Callee = CB.getCalledFunction();
+ if (!Callee)
+ return FunctionCallee();
+ const FunctionType *OldFTy = Callee->getFunctionType();
+ if (OldFTy->isVarArg())
+ return FunctionCallee();
+ // Copy params, and append token ID type.
+ Type *RetTy = OldFTy->getReturnType();
+ SmallVector<Type *, 4> NewParams{OldFTy->params()};
+ std::string TokenAllocName = ClFuncPrefix;
+ if (Options.FastABI)
+ TokenAllocName += utostr(TokenID) + "_";
+ else
+ NewParams.push_back(IntPtrTy); // token ID
+ TokenAllocName += Callee->getName();
+ FunctionType *NewFTy = FunctionType::get(RetTy, NewParams, false);
+ FunctionCallee TokenAlloc = Mod.getOrInsertFunction(TokenAllocName, NewFTy);
+ if (Function *F = dyn_cast<Function>(TokenAlloc.getCallee()))
+ F->copyAttributesFrom(Callee); // preserve attrs
+
+ if (Key.has_value())
+ TokenAllocFunctions[*Key] = TokenAlloc;
+ return TokenAlloc;
+}
+
+} // namespace
+
+AllocTokenPass::AllocTokenPass(AllocTokenOptions Opts)
+ : Options(std::move(Opts)) {}
+
+PreservedAnalyses AllocTokenPass::run(Module &M, ModuleAnalysisManager &MAM) {
+ AllocToken Pass(Options, M, MAM);
+ bool Modified = false;
+
+ for (Function &F : M) {
+ if (F.empty())
+ continue; // declaration
+ Modified |= Pass.instrumentFunction(F);
+ }
+
+ return Modified ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()
+ : PreservedAnalyses::all();
+}
diff --git a/llvm/lib/Transforms/Instrumentation/CMakeLists.txt b/llvm/lib/Transforms/Instrumentation/CMakeLists.txt
index 15fd421a41b0f..80576c61fd80c 100644
--- a/llvm/lib/Transforms/Instrumentation/CMakeLists.txt
+++ b/llvm/lib/Transforms/Instrumentation/CMakeLists.txt
@@ -1,5 +1,6 @@
add_llvm_component_library(LLVMInstrumentation
AddressSanitizer.cpp
+ AllocToken.cpp
BoundsChecking.cpp
CGProfile.cpp
ControlHeightReduction.cpp
diff --git a/llvm/lib/Transforms/Utils/CodeExtractor.cpp b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
index bbd1ed6a3ab2d..5ba6f95f5fae8 100644
--- a/llvm/lib/Transforms/Utils/CodeExtractor.cpp
+++ b/llvm/lib/Transforms/Utils/CodeExtractor.cpp
@@ -970,6 +970,7 @@ Function *CodeExtractor::constructFunctionDeclaration(
case Attribute::SanitizeMemTag:
case Attribute::SanitizeRealtime:
case Attribute::SanitizeRealtimeBlocking:
+ case Attribute::SanitizeAllocToken:
case Attribute::SpeculativeLoadHardening:
case Attribute::StackProtect:
case Attribute::StackProtectReq:
diff --git a/llvm/lib/Transforms/Utils/Local.cpp b/llvm/lib/Transforms/Utils/Local.cpp
index 123881e276584..5662c596108e5 100644
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@@ -3025,6 +3025,13 @@ static void combineMetadata(Instruction *K, const Instruction *J,
// Preserve !nosanitize if both K and J have it.
K->setMetadata(Kind, JMD);
break;
+ case LLVMContext::MD_alloc_token:
+ // Preserve !alloc_token if both K and J have it, and they are equal.
+ if (KMD == JMD)
+ K->setMetadata(Kind, JMD);
+ else
+ K->setMetadata(Kind, nullptr);
+ break;
}
}
// Set !invariant.group from J if J has it. If both instructions have it
diff --git a/llvm/test/Bitcode/attributes.ll b/llvm/test/Bitcode/attributes.ll
index 8c1a76365e1b4..aef7810fe2c3b 100644
--- a/llvm/test/Bitcode/attributes.ll
+++ b/llvm/test/Bitcode/attributes.ll
@@ -516,6 +516,11 @@ define void @f93() sanitize_realtime_blocking {
ret void;
}
+; CHECK: define void @f_sanitize_alloc_token() #55
+define void @f_sanitize_alloc_token() sanitize_alloc_token {
+ ret void;
+}
+
; CHECK: define void @f87() [[FNRETTHUNKEXTERN:#[0-9]+]]
define void @f87() fn_ret_thunk_extern { ret void }
@@ -627,6 +632,7 @@ define void @dead_on_return(ptr dead_on_return %p) {
; CHECK: attributes #52 = { nosanitize_bounds }
; CHECK: attributes #53 = { sanitize_realtime }
; CHECK: attributes #54 = { sanitize_realtime_blocking }
+; CHECK: attributes #55 = { sanitize_alloc_token }
; CHECK: attributes [[FNRETTHUNKEXTERN]] = { fn_ret_thunk_extern }
; CHECK: attributes [[SKIPPROFILE]] = { skipprofile }
; CHECK: attributes [[OPTDEBUG]] = { optdebug }
diff --git a/llvm/test/Bitcode/compatibility.ll b/llvm/test/Bitcode/compatibility.ll
index 0b5ce08c00a23..e21786e5ee330 100644
--- a/llvm/test/Bitcode/compatibility.ll
+++ b/llvm/test/Bitcode/compatibility.ll
@@ -1718,7 +1718,7 @@ exit:
; CHECK: select <2 x i1> <i1 true, i1 false>, <2 x i8> <i8 2, i8 3>, <2 x i8> <i8 3, i8 2>
call void @f.nobuiltin() builtin
- ; CHECK: call void @f.nobuiltin() #54
+ ; CHECK: call void @f.nobuiltin() #55
call fastcc noalias ptr @f.noalias() noinline
; CHECK: call fastcc noalias ptr @f.noalias() #12
@@ -2151,6 +2151,9 @@ declare void @f.sanitize_realtime() sanitize_realtime
declare void @f.sanitize_realtime_blocking() sanitize_realtime_blocking
; CHECK: declare void @f.sanitize_realtime_blocking() #53
+declare void @f.sanitize_alloc_token() sanitize_alloc_token
+; CHECK: declare void @f.sanitize_alloc_token() #54
+
; CHECK: declare nofpclass(snan) float @nofpclass_snan(float nofpclass(snan))
declare nofpclass(snan) float @nofpclass_snan(float nofpclass(snan))
@@ -2284,7 +2287,8 @@ define float @nofpclass_callsites(float %arg, { float } %arg1) {
; CHECK: attributes #51 = { sanitize_numerical_stability }
; CHECK: attributes #52 = { sanitize_realtime }
; CHECK: attributes #53 = { sanitize_realtime_blocking }
-; CHECK: attributes #54 = { builtin }
+; CHECK: attributes #54 = { sanitize_alloc_token }
+; CHECK: attributes #55 = { builtin }
;; Metadata
diff --git a/llvm/test/Instrumentation/AllocToken/basic.ll b/llvm/test/Instrumentation/AllocToken/basic.ll
new file mode 100644
index 0000000000000..099d37df264d6
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/basic.ll
@@ -0,0 +1,98 @@
+; 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
+
+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 @malloc(i64)
+declare ptr @calloc(i64, i64)
+declare ptr @realloc(ptr, i64)
+declare ptr @_Znwm(i64)
+declare ptr @_Znam(i64)
+declare void @free(ptr)
+declare void @_ZdlPv(ptr)
+declare i32 @foobar(i64)
+
+; Test basic allocation call rewriting
+define ptr @test_basic_rewriting() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_basic_rewriting(
+; CHECK-SAME: ) #[[ATTR5:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 64, i64 0)
+; CHECK-NEXT: [[TMP1:%.*]] = call ptr @__alloc_token_calloc(i64 8, i64 8, i64 1)
+; CHECK-NEXT: [[TMP2:%.*]] = call ptr @__alloc_token_realloc(ptr [[TMP0]], i64 128, i64 2)
+; CHECK-NEXT: ret ptr [[TMP2]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 64)
+ %ptr2 = call ptr @calloc(i64 8, i64 8)
+ %ptr3 = call ptr @realloc(ptr %ptr1, i64 128)
+ ret ptr %ptr3
+}
+
+; Test C++ operator rewriting
+define ptr @test_cpp_operators() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_cpp_operators(
+; CHECK-SAME: ) #[[ATTR5]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token__Znwm(i64 32, i64 3)
+; CHECK-NEXT: [[TMP1:%.*]] = call ptr @__alloc_token__Znam(i64 64, i64 4)
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @_Znwm(i64 32)
+ %ptr2 = call ptr @_Znam(i64 64)
+ ret ptr %ptr1
+}
+
+; Functions without sanitize_alloc_token do not get instrumented
+define ptr @without_attribute() {
+; CHECK-LABEL: define ptr @without_attribute() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[PTR:%.*]] = call ptr @malloc(i64 16)
+; CHECK-NEXT: ret ptr [[PTR]]
+;
+entry:
+ %ptr = call ptr @malloc(i64 16)
+ ret ptr %ptr
+}
+
+; Test that free/delete are untouched
+define void @test_free_untouched(ptr %ptr) sanitize_alloc_token {
+; CHECK-LABEL: define void @test_free_untouched(
+; CHECK-SAME: ptr [[PTR:%.*]]) #[[ATTR5]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: call void @free(ptr [[PTR]])
+; CHECK-NEXT: call void @_ZdlPv(ptr [[PTR]])
+; CHECK-NEXT: ret void
+;
+entry:
+ call void @free(ptr %ptr)
+ call void @_ZdlPv(ptr %ptr)
+ ret void
+}
+
+; Non-allocation functions are untouched
+define i32 @no_allocations(i32 %x) sanitize_alloc_token {
+; CHECK-LABEL: define i32 @no_allocations(
+; CHECK-SAME: i32 [[X:%.*]]) #[[ATTR5]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[RESULT:%.*]] = call i32 @foobar(i64 42)
+; CHECK-NEXT: ret i32 [[RESULT]]
+;
+entry:
+ %result = call i32 @foobar(i64 42)
+ ret i32 %result
+}
+
+; Test that tail calls are preserved
+define ptr @test_tail_call_preserved() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_tail_call_preserved(
+; CHECK-SAME: ) #[[ATTR5]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = tail call ptr @__alloc_token_malloc(i64 42, i64 5)
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %result = tail call ptr @malloc(i64 42)
+ ret ptr %result
+}
diff --git a/llvm/test/Instrumentation/AllocToken/basic32.ll b/llvm/test/Instrumentation/AllocToken/basic32.ll
new file mode 100644
index 0000000000000..944a452f4b4d7
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/basic32.ll
@@ -0,0 +1,31 @@
+; 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
+
+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"
+
+declare ptr @malloc(i32)
+declare ptr @_Znwm(i32)
+
+define ptr @test_basic_rewriting() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_basic_rewriting(
+; CHECK-SAME: ) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i32 64, i32 0)
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i32 64)
+ ret ptr %ptr1
+}
+
+define ptr @test_cpp_operators() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_cpp_operators(
+; CHECK-SAME: ) #[[ATTR2]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token__Znwm(i32 32, i32 1)
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @_Znwm(i32 32)
+ ret ptr %ptr1
+}
diff --git a/llvm/test/Instrumentation/AllocToken/extralibfuncs.ll b/llvm/test/Instrumentation/AllocToken/extralibfuncs.ll
new file mode 100644
index 0000000000000..5f08552c789f3
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/extralibfuncs.ll
@@ -0,0 +1,44 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Test for special libfuncs not automatically considered allocation functions.
+;
+; 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, i64} @__size_returning_new(i64)
+
+define ptr @test_extra_libfuncs() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_extra_libfuncs(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call { ptr, i64 } @__alloc_token___size_returning_new(i64 10, i64 2689373973731826898), !alloc_token [[META0:![0-9]+]]
+; CHECK-NEXT: [[PTR1:%.*]] = extractvalue { ptr, i64 } [[TMP0]], 0
+; CHECK-NEXT: ret ptr [[PTR1]]
+;
+entry:
+ %srn = call {ptr, i64} @__size_returning_new(i64 10), !alloc_token !0
+ %ptr1 = extractvalue {ptr, i64} %srn, 0
+ ret ptr %ptr1
+}
+
+declare ptr @_Znwm(i64) nobuiltin allocsize(0)
+declare ptr @_Znam(i64) nobuiltin allocsize(0)
+
+define ptr @test_replaceable_new() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_replaceable_new(
+; CHECK-SAME: ) #[[ATTR1]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token__Znwm(i64 32, i64 2689373973731826898), !alloc_token [[META0]]
+; CHECK-NEXT: [[TMP1:%.*]] = call ptr @__alloc_token__Znam(i64 64, i64 2689373973731826898), !alloc_token [[META0]]
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @_Znwm(i64 32), !alloc_token !0
+ %ptr2 = call ptr @_Znam(i64 64), !alloc_token !0
+ ret ptr %ptr1
+}
+
+!0 = !{!"int"}
+;.
+; CHECK: [[META0]] = !{!"int"}
+;.
diff --git a/llvm/test/Instrumentation/AllocToken/fast.ll b/llvm/test/Instrumentation/AllocToken/fast.ll
new file mode 100644
index 0000000000000..19a3ef6bb9ede
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/fast.ll
@@ -0,0 +1,42 @@
+; 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
+
+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 @malloc(i64)
+declare ptr @calloc(i64, i64)
+declare ptr @realloc(ptr, i64)
+declare ptr @_Znwm(i64)
+declare ptr @_Znam(i64)
+
+; Test basic allocation call rewriting
+define ptr @test_basic_rewriting() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_basic_rewriting(
+; CHECK-SAME: ) #[[ATTR4:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[PTR1:%.*]] = call ptr @__alloc_token_0_malloc(i64 64)
+; CHECK-NEXT: [[PTR2:%.*]] = call ptr @__alloc_token_1_calloc(i64 8, i64 8)
+; CHECK-NEXT: [[PTR3:%.*]] = call ptr @__alloc_token_2_realloc(ptr [[PTR1]], i64 128)
+; CHECK-NEXT: ret ptr [[PTR3]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 64)
+ %ptr2 = call ptr @calloc(i64 8, i64 8)
+ %ptr3 = call ptr @realloc(ptr %ptr1, i64 128)
+ ret ptr %ptr3
+}
+
+; Test C++ operator rewriting
+define ptr @test_cpp_operators() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_cpp_operators(
+; CHECK-SAME: ) #[[ATTR4]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[PTR1:%.*]] = call ptr @__alloc_token_0__Znwm(i64 32)
+; CHECK-NEXT: [[PTR2:%.*]] = call ptr @__alloc_token_1__Znam(i64 64)
+; CHECK-NEXT: ret ptr [[PTR1]]
+;
+entry:
+ %ptr1 = call ptr @_Znwm(i64 32)
+ %ptr2 = call ptr @_Znam(i64 64)
+ ret ptr %ptr1
+}
diff --git a/llvm/test/Instrumentation/AllocToken/ignore.ll b/llvm/test/Instrumentation/AllocToken/ignore.ll
new file mode 100644
index 0000000000000..b92a920ed11ee
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/ignore.ll
@@ -0,0 +1,29 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; Test for all allocation functions that should be ignored by default.
+;
+; 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 @strdup(ptr)
+declare ptr @__strdup(ptr)
+declare ptr @strndup(ptr, i64)
+declare ptr @__strndup(ptr, i64)
+
+define ptr @test_ignored_allocation_functions(ptr %ptr) sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_ignored_allocation_functions(
+; CHECK-SAME: ptr [[PTR:%.*]]) #[[ATTR2:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[PTR1:%.*]] = call ptr @strdup(ptr [[PTR]])
+; CHECK-NEXT: [[PTR2:%.*]] = call ptr @__strdup(ptr [[PTR]])
+; CHECK-NEXT: [[PTR3:%.*]] = call ptr @strndup(ptr [[PTR]], i64 42)
+; CHECK-NEXT: [[PTR4:%.*]] = call ptr @__strndup(ptr [[PTR]], i64 42)
+; CHECK-NEXT: ret ptr [[PTR1]]
+;
+entry:
+ %ptr1 = call ptr @strdup(ptr %ptr)
+ %ptr2 = call ptr @__strdup(ptr %ptr)
+ %ptr3 = call ptr @strndup(ptr %ptr, i64 42)
+ %ptr4 = call ptr @__strndup(ptr %ptr, i64 42)
+ ret ptr %ptr1
+}
diff --git a/llvm/test/Instrumentation/AllocToken/invoke.ll b/llvm/test/Instrumentation/AllocToken/invoke.ll
new file mode 100644
index 0000000000000..347c99a2e8f8d
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/invoke.ll
@@ -0,0 +1,123 @@
+; 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
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128"
+
+define ptr @test_invoke_malloc() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_invoke_malloc(
+; CHECK-SAME: ) #[[ATTR0:[0-9]+]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = invoke ptr @__alloc_token_malloc(i64 64, i64 0)
+; CHECK-NEXT: to label %[[NORMAL:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK: [[NORMAL]]:
+; CHECK-NEXT: ret ptr [[TMP0]]
+; CHECK: [[CLEANUP]]:
+; CHECK-NEXT: [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT: cleanup
+; CHECK-NEXT: ret ptr null
+;
+entry:
+ %ptr = invoke ptr @malloc(i64 64) to label %normal unwind label %cleanup
+
+normal:
+ ret ptr %ptr
+
+cleanup:
+ %lp = landingpad { ptr, i32 } cleanup
+ ret ptr null
+}
+
+define ptr @test_invoke_operator_new() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_invoke_operator_new(
+; CHECK-SAME: ) #[[ATTR0]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = invoke ptr @__alloc_token__Znwm(i64 32, i64 1)
+; CHECK-NEXT: to label %[[NORMAL:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK: [[NORMAL]]:
+; CHECK-NEXT: ret ptr [[TMP0]]
+; CHECK: [[CLEANUP]]:
+; CHECK-NEXT: [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT: cleanup
+; CHECK-NEXT: ret ptr null
+;
+entry:
+ %ptr = invoke ptr @_Znwm(i64 32) to label %normal unwind label %cleanup
+
+normal:
+ ret ptr %ptr
+
+cleanup:
+ %lp = landingpad { ptr, i32 } cleanup
+ ret ptr null
+}
+
+; Test complex exception flow with multiple invoke allocations
+define ptr @test_complex_invoke_flow() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_complex_invoke_flow(
+; CHECK-SAME: ) #[[ATTR0]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = invoke ptr @__alloc_token_malloc(i64 16, i64 2)
+; CHECK-NEXT: to label %[[FIRST_OK:.*]] unwind label %[[CLEANUP1:.*]]
+; CHECK: [[FIRST_OK]]:
+; CHECK-NEXT: [[TMP1:%.*]] = invoke ptr @__alloc_token__Znwm(i64 32, i64 3)
+; CHECK-NEXT: to label %[[SECOND_OK:.*]] unwind label %[[CLEANUP2:.*]]
+; CHECK: [[SECOND_OK]]:
+; CHECK-NEXT: ret ptr [[TMP0]]
+; CHECK: [[CLEANUP1]]:
+; CHECK-NEXT: [[LP1:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT: cleanup
+; CHECK-NEXT: ret ptr null
+; CHECK: [[CLEANUP2]]:
+; CHECK-NEXT: [[LP2:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT: cleanup
+; CHECK-NEXT: ret ptr null
+;
+entry:
+ %ptr1 = invoke ptr @malloc(i64 16) to label %first_ok unwind label %cleanup1
+
+first_ok:
+ %ptr2 = invoke ptr @_Znwm(i64 32) to label %second_ok unwind label %cleanup2
+
+second_ok:
+ ret ptr %ptr1
+
+cleanup1:
+ %lp1 = landingpad { ptr, i32 } cleanup
+ ret ptr null
+
+cleanup2:
+ %lp2 = landingpad { ptr, i32 } cleanup
+ ret ptr null
+}
+
+; Test mixed call/invoke
+define ptr @test_mixed_call_invoke() sanitize_alloc_token personality ptr @__gxx_personality_v0 {
+; CHECK-LABEL: define ptr @test_mixed_call_invoke(
+; CHECK-SAME: ) #[[ATTR0]] personality ptr @__gxx_personality_v0 {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 8, i64 4)
+; CHECK-NEXT: [[TMP1:%.*]] = invoke ptr @__alloc_token_malloc(i64 16, i64 5)
+; CHECK-NEXT: to label %[[NORMAL:.*]] unwind label %[[CLEANUP:.*]]
+; CHECK: [[NORMAL]]:
+; CHECK-NEXT: ret ptr [[TMP0]]
+; CHECK: [[CLEANUP]]:
+; CHECK-NEXT: [[LP:%.*]] = landingpad { ptr, i32 }
+; CHECK-NEXT: cleanup
+; CHECK-NEXT: ret ptr null
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 8)
+
+ %ptr2 = invoke ptr @malloc(i64 16) to label %normal unwind label %cleanup
+
+normal:
+ ret ptr %ptr1
+
+cleanup:
+ %lp = landingpad { ptr, i32 } cleanup
+ ret ptr null
+}
+
+declare ptr @malloc(i64)
+declare ptr @_Znwm(i64)
+declare i32 @__gxx_personality_v0(...)
diff --git a/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll b/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
new file mode 100644
index 0000000000000..e023ab6b11b1e
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
@@ -0,0 +1,85 @@
+; 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
+
+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 @malloc(i64)
+declare ptr @custom_malloc(i64)
+declare ptr @kmalloc(i64, i64)
+
+define ptr @test_libcall() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_libcall(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 64, i64 0)
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 64)
+ ret ptr %ptr1
+}
+
+define ptr @test_libcall_hint() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_libcall_hint(
+; CHECK-SAME: ) #[[ATTR1]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 64, i64 1), !alloc_token [[META0:![0-9]+]]
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 64), !alloc_token !0
+ ret ptr %ptr1
+}
+
+define ptr @test_nonlibcall_nohint() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_nonlibcall_nohint(
+; CHECK-SAME: ) #[[ATTR1]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[PTR1:%.*]] = call ptr @custom_malloc(i64 8)
+; CHECK-NEXT: [[PTR2:%.*]] = call ptr @kmalloc(i64 32, i64 0)
+; CHECK-NEXT: ret ptr [[PTR1]]
+;
+entry:
+ %ptr1 = call ptr @custom_malloc(i64 8)
+ %ptr2 = call ptr @kmalloc(i64 32, i64 0)
+ ret ptr %ptr1
+}
+
+define ptr @test_nonlibcall_hint() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_nonlibcall_hint(
+; CHECK-SAME: ) #[[ATTR1]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_custom_malloc(i64 8, i64 2), !alloc_token [[META0]]
+; CHECK-NEXT: [[TMP1:%.*]] = call ptr @__alloc_token_kmalloc(i64 32, i64 0, i64 3), !alloc_token [[META0]]
+; CHECK-NEXT: [[TMP2:%.*]] = call ptr @__alloc_token_custom_malloc(i64 64, i64 4), !alloc_token [[META0]]
+; CHECK-NEXT: [[TMP3:%.*]] = call ptr @__alloc_token_kmalloc(i64 128, i64 2, i64 5), !alloc_token [[META0]]
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @custom_malloc(i64 8), !alloc_token !0
+ %ptr2 = call ptr @kmalloc(i64 32, i64 0), !alloc_token !0
+ %ptr3 = call ptr @custom_malloc(i64 64), !alloc_token !0
+ %ptr4 = call ptr @kmalloc(i64 128, i64 2), !alloc_token !0
+ ret ptr %ptr1
+}
+
+; Functions without sanitize_alloc_token do not get instrumented
+define ptr @without_attribute() {
+; CHECK-LABEL: define ptr @without_attribute() {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[PTR1:%.*]] = call ptr @malloc(i64 64), !alloc_token [[META0]]
+; CHECK-NEXT: [[PTR2:%.*]] = call ptr @custom_malloc(i64 8), !alloc_token [[META0]]
+; CHECK-NEXT: [[PTR3:%.*]] = call ptr @kmalloc(i64 32, i64 0), !alloc_token [[META0]]
+; CHECK-NEXT: ret ptr [[PTR1]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 64), !alloc_token !0
+ %ptr2 = call ptr @custom_malloc(i64 8), !alloc_token !0
+ %ptr3 = call ptr @kmalloc(i64 32, i64 0), !alloc_token !0
+ ret ptr %ptr1
+}
+
+!0 = !{!"int"}
+;.
+; CHECK: [[META0]] = !{!"int"}
+;.
diff --git a/llvm/test/Instrumentation/AllocToken/remark.ll b/llvm/test/Instrumentation/AllocToken/remark.ll
new file mode 100644
index 0000000000000..a2404526ea53f
--- /dev/null
+++ b/llvm/test/Instrumentation/AllocToken/remark.ll
@@ -0,0 +1,38 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=inferattrs,alloc-token -pass-remarks=alloc-token -S 2>&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"
+
+declare ptr @malloc(i64)
+
+; CHECK-NOT: remark: <unknown>:0:0: Call to 'malloc' in 'test_has_metadata' without source-level type token
+; CHECK: remark: <unknown>:0:0: Call to 'malloc' in 'test_no_metadata' without source-level type token
+
+define ptr @test_has_metadata() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_has_metadata(
+; CHECK-SAME: ) #[[ATTR1:[0-9]+]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 64, i64 2689373973731826898), !alloc_token [[META0:![0-9]+]]
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 64), !alloc_token !0
+ ret ptr %ptr1
+}
+
+define ptr @test_no_metadata() sanitize_alloc_token {
+; CHECK-LABEL: define ptr @test_no_metadata(
+; CHECK-SAME: ) #[[ATTR1]] {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[TMP0:%.*]] = call ptr @__alloc_token_malloc(i64 32, i64 0)
+; CHECK-NEXT: ret ptr [[TMP0]]
+;
+entry:
+ %ptr1 = call ptr @malloc(i64 32)
+ ret ptr %ptr1
+}
+
+!0 = !{!"int"}
+;.
+; CHECK: [[META0]] = !{!"int"}
+;.
diff --git a/llvm/test/Transforms/Inline/attributes.ll b/llvm/test/Transforms/Inline/attributes.ll
index 42b1a3a29aec4..55ab430f201d6 100644
--- a/llvm/test/Transforms/Inline/attributes.ll
+++ b/llvm/test/Transforms/Inline/attributes.ll
@@ -26,6 +26,10 @@ define i32 @sanitize_memtag_callee(i32 %i) sanitize_memtag {
ret i32 %i
}
+define i32 @sanitize_alloc_token_callee(i32 %i) sanitize_alloc_token {
+ ret i32 %i
+}
+
define i32 @safestack_callee(i32 %i) safestack {
ret i32 %i
}
@@ -58,6 +62,10 @@ define i32 @alwaysinline_sanitize_memtag_callee(i32 %i) alwaysinline sanitize_me
ret i32 %i
}
+define i32 @alwaysinline_sanitize_alloc_token_callee(i32 %i) alwaysinline sanitize_alloc_token {
+ ret i32 %i
+}
+
define i32 @alwaysinline_safestack_callee(i32 %i) alwaysinline safestack {
ret i32 %i
}
@@ -184,6 +192,39 @@ define i32 @test_sanitize_memtag(i32 %arg) sanitize_memtag {
; CHECK-NEXT: ret i32
}
+; ---------------------------------------------------------------------------- ;
+
+; Can inline sanitize_alloc_token functions into a noattr function. The
+; attribute is *not* viral, otherwise may break code.
+define i32 @test_no_sanitize_alloc_token(i32 %arg) {
+; CHECK-LABEL: @test_no_sanitize_alloc_token(
+; CHECK-SAME: ) {
+; CHECK-NOT: call
+; CHECK: ret i32
+entry:
+ %x1 = call i32 @noattr_callee(i32 %arg)
+ %x2 = call i32 @sanitize_alloc_token_callee(i32 %x1)
+ %x3 = call i32 @alwaysinline_callee(i32 %x2)
+ %x4 = call i32 @alwaysinline_sanitize_alloc_token_callee(i32 %x3)
+ ret i32 %x4
+}
+
+; Can inline noattr functions into a sanitize_alloc_token function. If
+; inlinable noattr functions cannot be instrumented, they should be marked with
+; explicit noinline.
+define i32 @test_sanitize_alloc_token(i32 %arg) sanitize_alloc_token {
+; CHECK-LABEL: @test_sanitize_alloc_token(
+; CHECK-SAME: ) [[SANITIZE_ALLOC_TOKEN:.*]] {
+; CHECK-NOT: call
+; CHECK: ret i32
+entry:
+ %x1 = call i32 @noattr_callee(i32 %arg)
+ %x2 = call i32 @sanitize_alloc_token_callee(i32 %x1)
+ %x3 = call i32 @alwaysinline_callee(i32 %x2)
+ %x4 = call i32 @alwaysinline_sanitize_alloc_token_callee(i32 %x3)
+ ret i32 %x4
+}
+
define i32 @test_safestack(i32 %arg) safestack {
%x1 = call i32 @noattr_callee(i32 %arg)
%x2 = call i32 @safestack_callee(i32 %x1)
@@ -639,6 +680,7 @@ define i32 @loader_replaceable_caller() {
ret i32 %1
}
+; CHECK: attributes [[SANITIZE_ALLOC_TOKEN]] = { sanitize_alloc_token }
; CHECK: attributes [[SLH]] = { speculative_load_hardening }
; CHECK: attributes [[FPMAD_FALSE]] = { "less-precise-fpmad"="false" }
; CHECK: attributes [[FPMAD_TRUE]] = { "less-precise-fpmad"="true" }
diff --git a/llvm/test/Transforms/SimplifyCFG/merge-calls-alloc-token.ll b/llvm/test/Transforms/SimplifyCFG/merge-calls-alloc-token.ll
new file mode 100644
index 0000000000000..9bbe3eb371673
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/merge-calls-alloc-token.ll
@@ -0,0 +1,104 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=simplifycfg -S | FileCheck %s
+
+target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-i128:128-f80:128-n8:16:32:64-S128"
+
+declare ptr @_Znwm(i64)
+
+define ptr @test_merge_alloc_token_same(i1 %b) {
+; CHECK-LABEL: define ptr @test_merge_alloc_token_same(
+; CHECK-SAME: i1 [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[CALL:%.*]] = call ptr @_Znwm(i64 4), !alloc_token [[META0:![0-9]+]]
+; CHECK-NEXT: ret ptr [[CALL]]
+;
+entry:
+ br i1 %b, label %if.then, label %if.else
+
+if.then:
+ %call = call ptr @_Znwm(i64 4), !alloc_token !0
+ br label %if.end
+
+if.else:
+ %call1 = call ptr @_Znwm(i64 4), !alloc_token !0
+ br label %if.end
+
+if.end:
+ %x.0 = phi ptr [ %call, %if.then ], [ %call1, %if.else ]
+ ret ptr %x.0
+}
+
+define ptr @test_merge_alloc_token_different(i1 %b) {
+; CHECK-LABEL: define ptr @test_merge_alloc_token_different(
+; CHECK-SAME: i1 [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[CALL:%.*]] = call ptr @_Znwm(i64 4)
+; CHECK-NEXT: ret ptr [[CALL]]
+;
+entry:
+ br i1 %b, label %if.then, label %if.else
+
+if.then:
+ %call = call ptr @_Znwm(i64 4), !alloc_token !0
+ br label %if.end
+
+if.else:
+ %call1 = call ptr @_Znwm(i64 4), !alloc_token !1
+ br label %if.end
+
+if.end:
+ %x.0 = phi ptr [ %call, %if.then ], [ %call1, %if.else ]
+ ret ptr %x.0
+}
+
+define ptr @test_merge_alloc_token_some1(i1 %b) {
+; CHECK-LABEL: define ptr @test_merge_alloc_token_some1(
+; CHECK-SAME: i1 [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[CALL:%.*]] = call ptr @_Znwm(i64 4)
+; CHECK-NEXT: ret ptr [[CALL]]
+;
+entry:
+ br i1 %b, label %if.then, label %if.else
+
+if.then:
+ %call = call ptr @_Znwm(i64 4), !alloc_token !0
+ br label %if.end
+
+if.else:
+ %call1 = call ptr @_Znwm(i64 4)
+ br label %if.end
+
+if.end:
+ %x.0 = phi ptr [ %call, %if.then ], [ %call1, %if.else ]
+ ret ptr %x.0
+}
+
+define ptr @test_merge_alloc_token_some2(i1 %b) {
+; CHECK-LABEL: define ptr @test_merge_alloc_token_some2(
+; CHECK-SAME: i1 [[B:%.*]]) {
+; CHECK-NEXT: [[ENTRY:.*:]]
+; CHECK-NEXT: [[CALL:%.*]] = call ptr @_Znwm(i64 4)
+; CHECK-NEXT: ret ptr [[CALL]]
+;
+entry:
+ br i1 %b, label %if.then, label %if.else
+
+if.then:
+ %call = call ptr @_Znwm(i64 4)
+ br label %if.end
+
+if.else:
+ %call1 = call ptr @_Znwm(i64 4), !alloc_token !0
+ br label %if.end
+
+if.end:
+ %x.0 = phi ptr [ %call, %if.then ], [ %call1, %if.else ]
+ ret ptr %x.0
+}
+
+!0 = !{!"int"}
+!1 = !{!"char[4]"}
+;.
+; CHECK: [[META0]] = !{!"int"}
+;.
diff --git a/llvm/utils/emacs/llvm-mode.el b/llvm/utils/emacs/llvm-mode.el
index 660d0718f098c..240c13319f634 100644
--- a/llvm/utils/emacs/llvm-mode.el
+++ b/llvm/utils/emacs/llvm-mode.el
@@ -34,7 +34,7 @@
"inaccessiblemem_or_argmemonly" "inalloca" "inlinehint" "jumptable" "minsize" "mustprogress" "naked" "nobuiltin" "nonnull" "nocapture"
"nocallback" "nocf_check" "noduplicate" "noext" "nofree" "noimplicitfloat" "noinline" "nomerge" "nonlazybind" "noprofile" "noredzone" "noreturn"
"norecurse" "nosync" "noundef" "nounwind" "nosanitize_bounds" "nosanitize_coverage" "null_pointer_is_valid" "optdebug" "optforfuzzing" "optnone" "optsize" "preallocated" "readnone" "readonly" "returned" "returns_twice"
- "shadowcallstack" "signext" "speculatable" "speculative_load_hardening" "ssp" "sspreq" "sspstrong" "safestack" "sanitize_address" "sanitize_hwaddress" "sanitize_memtag"
+ "shadowcallstack" "signext" "speculatable" "speculative_load_hardening" "ssp" "sspreq" "sspstrong" "safestack" "sanitize_address" "sanitize_alloc_token" "sanitize_hwaddress" "sanitize_memtag"
"sanitize_thread" "sanitize_memory" "strictfp" "swifterror" "uwtable" "vscale_range" "willreturn" "writeonly" "zeroext") 'symbols) . font-lock-constant-face)
;; Variables
'("%[-a-zA-Z$._][-a-zA-Z$._0-9]*" . font-lock-variable-name-face)
diff --git a/llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn b/llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
index a8eb834c1da23..2c6204e758559 100644
--- a/llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
+++ b/llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
@@ -11,6 +11,7 @@ static_library("Instrumentation") {
]
sources = [
"AddressSanitizer.cpp",
+ "AllocToken.cpp",
"BlockCoverageInference.cpp",
"BoundsChecking.cpp",
"CGProfile.cpp",
diff --git a/llvm/utils/llvm.grm b/llvm/utils/llvm.grm
index 411323178bde1..dddfe3c301b65 100644
--- a/llvm/utils/llvm.grm
+++ b/llvm/utils/llvm.grm
@@ -173,6 +173,7 @@ FuncAttr ::= noreturn
| returns_twice
| nonlazybind
| sanitize_address
+ | sanitize_alloc_token
| sanitize_thread
| sanitize_memory
| mustprogress
diff --git a/llvm/utils/vim/syntax/llvm.vim b/llvm/utils/vim/syntax/llvm.vim
index e3b8ff8629559..e048caa20406a 100644
--- a/llvm/utils/vim/syntax/llvm.vim
+++ b/llvm/utils/vim/syntax/llvm.vim
@@ -163,6 +163,7 @@ syn keyword llvmKeyword
\ returns_twice
\ safestack
\ sanitize_address
+ \ sanitize_alloc_token
\ sanitize_hwaddress
\ sanitize_memory
\ sanitize_memtag
diff --git a/llvm/utils/vscode/llvm/syntaxes/ll.tmLanguage.yaml b/llvm/utils/vscode/llvm/syntaxes/ll.tmLanguage.yaml
index b64482336f404..1faaf6b26f301 100644
--- a/llvm/utils/vscode/llvm/syntaxes/ll.tmLanguage.yaml
+++ b/llvm/utils/vscode/llvm/syntaxes/ll.tmLanguage.yaml
@@ -258,6 +258,7 @@ patterns:
\\breturns_twice\\b|\
\\bsafestack\\b|\
\\bsanitize_address\\b|\
+ \\bsanitize_alloc_token\\b|\
\\bsanitize_hwaddress\\b|\
\\bsanitize_memory\\b|\
\\bsanitize_memtag\\b|\
More information about the llvm-commits
mailing list