[llvm] c7274fc - [AllocToken] Introduce AllocToken instrumentation pass (#156838)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Oct 7 04:30:11 PDT 2025
Author: Marco Elver
Date: 2025-10-07T13:30:07+02:00
New Revision: c7274fce2f16d9d7343390217d46c345fd42dc7d
URL: https://github.com/llvm/llvm-project/commit/c7274fce2f16d9d7343390217d46c345fd42dc7d
DIFF: https://github.com/llvm/llvm-project/commit/c7274fce2f16d9d7343390217d46c345fd42dc7d.diff
LOG: [AllocToken] Introduce AllocToken instrumentation pass (#156838)
Introduce `AllocToken`, an instrumentation pass designed to provide
tokens to memory allocators enabling various heap organization
strategies, such as heap partitioning.
Initially, the pass instruments functions marked with a new attribute
`sanitize_alloc_token` by rewriting allocation calls to include a token
ID, appended as a function argument with the default ABI.
The design aims to provide a flexible framework for implementing
different token generation schemes. It currently supports the following
token modes:
- TypeHash (default): token IDs based on a hash of the allocated type
- Random: statically-assigned pseudo-random token IDs
- Increment: incrementing token IDs per TU
For the `TypeHash` mode introduce support for `!alloc_token` metadata:
the metadata can be attached to allocation calls to provide richer
semantic
information to be consumed by the AllocToken pass. Optimization remarks
can be enabled to show where no metadata was available.
An alternative "fast ABI" is provided, where instead of passing the
token ID as an argument (e.g., `__alloc_token_malloc(size, id)`), the
token ID is directly encoded into the name of the called function (e.g.,
`__alloc_token_0_malloc(size)`). Where the maximum tokens is small, this
offers more efficient instrumentation by avoiding the overhead of
passing an additional argument at each allocation site.
Link: https://discourse.llvm.org/t/rfc-a-framework-for-allocator-partitioning-hints/87434 [1]
---
This change is part of the following series:
1. https://github.com/llvm/llvm-project/pull/160131
2. https://github.com/llvm/llvm-project/pull/156838
3. https://github.com/llvm/llvm-project/pull/162098
4. https://github.com/llvm/llvm-project/pull/162099
5. https://github.com/llvm/llvm-project/pull/156839
6. https://github.com/llvm/llvm-project/pull/156840
7. https://github.com/llvm/llvm-project/pull/156841
8. https://github.com/llvm/llvm-project/pull/156842
Added:
llvm/include/llvm/Transforms/Instrumentation/AllocToken.h
llvm/lib/Transforms/Instrumentation/AllocToken.cpp
llvm/test/Instrumentation/AllocToken/basic.ll
llvm/test/Instrumentation/AllocToken/basic32.ll
llvm/test/Instrumentation/AllocToken/extralibfuncs.ll
llvm/test/Instrumentation/AllocToken/fast.ll
llvm/test/Instrumentation/AllocToken/ignore.ll
llvm/test/Instrumentation/AllocToken/invoke.ll
llvm/test/Instrumentation/AllocToken/nonlibcalls.ll
llvm/test/Instrumentation/AllocToken/remark.ll
Modified:
llvm/docs/ReleaseNotes.md
llvm/lib/Passes/PassBuilder.cpp
llvm/lib/Passes/PassRegistry.def
llvm/lib/Transforms/Instrumentation/CMakeLists.txt
llvm/utils/gn/secondary/llvm/lib/Transforms/Instrumentation/BUILD.gn
Removed:
################################################################################
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/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/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 88550ea055012..c5c0d643d1498 100644
--- a/llvm/lib/Passes/PassRegistry.def
+++ b/llvm/lib/Passes/PassRegistry.def
@@ -125,6 +125,7 @@ MODULE_PASS("openmp-opt", OpenMPOptPass())
MODULE_PASS("openmp-opt-postlink",
OpenMPOptPass(ThinOrFullLTOPhase::FullLTOPostLink))
MODULE_PASS("partial-inliner", PartialInlinerPass())
+MODULE_PASS("alloc-token", AllocTokenPass())
MODULE_PASS("pgo-icall-prom", PGOIndirectCallPromotion())
MODULE_PASS("pgo-instr-gen", PGOInstrumentationGen())
MODULE_PASS("pgo-instr-use", PGOInstrumentationUse())
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
diff erent 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/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/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",
More information about the llvm-commits
mailing list