diff --git a/clang/docs/AllocToken.rst b/clang/docs/AllocToken.rst index bda84669456ce..be66ac461fef0 100644 --- a/clang/docs/AllocToken.rst +++ b/clang/docs/AllocToken.rst @@ -49,6 +49,39 @@ change or removal. These may (experimentally) be selected with ``-mllvm * ``increment``: This mode assigns a simple, incrementally increasing token ID to each allocation site. +The following command-line options affect generated token IDs: + +* ``-falloc-token-max=`` + Configures the maximum number of tokens. No max by default (tokens bounded + by ``SIZE_MAX``). + +Querying Token IDs with ``__builtin_infer_alloc_token`` +======================================================= + +For use cases where the token ID must be known at compile time, Clang provides +a builtin function: + +.. code-block:: c + + size_t __builtin_infer_alloc_token(, ...); + +This builtin returns the token ID inferred from its argument expressions, which +mirror arguments normally passed to any allocation function. The argument +expressions are **unevaluated**, so it can be used with expressions that would +have side effects without any runtime impact. + +For example, it can be used as follows: + +.. code-block:: c + + struct MyType { ... }; + void *__partition_alloc(size_t size, size_t partition); + #define partition_alloc(...) __partition_alloc(__VA_ARGS__, __builtin_infer_alloc_token(__VA_ARGS__)) + + void foo(void) { + MyType *x = partition_alloc(sizeof(*x)); + } + Allocation Token Instrumentation ================================ @@ -70,16 +103,6 @@ example: // Instrumented: ptr = __alloc_token_malloc(size, ); -The following command-line options affect generated token IDs: - -* ``-falloc-token-max=`` - Configures the maximum number of tokens. No max by default (tokens bounded - by ``SIZE_MAX``). - - .. code-block:: console - - % clang++ -fsanitize=alloc-token -falloc-token-max=512 example.cc - Runtime Interface ----------------- diff --git a/clang/docs/ReleaseNotes.rst b/clang/docs/ReleaseNotes.rst index 9a0d69c6c1b0e..09aa4e94bc673 100644 --- a/clang/docs/ReleaseNotes.rst +++ b/clang/docs/ReleaseNotes.rst @@ -261,6 +261,9 @@ Non-comprehensive list of changes in this release allocator-level heap organization strategies. A feature to instrument all allocation functions with a token ID can be enabled via the ``-fsanitize=alloc-token`` flag. +- A builtin ``__builtin_infer_alloc_token(, ...)`` is provided to allow + compile-time querying of allocation token IDs, where the builtin arguments + mirror those normally passed to an allocation function. New Compiler Flags ------------------ diff --git a/clang/include/clang/Basic/Builtins.td b/clang/include/clang/Basic/Builtins.td index 468121f7d20ab..3cda0aa50666e 100644 --- a/clang/include/clang/Basic/Builtins.td +++ b/clang/include/clang/Basic/Builtins.td @@ -1286,6 +1286,12 @@ def AllocaWithAlignUninitialized : Builtin { let Prototype = "void*(size_t, _Constant size_t)"; } +def InferAllocToken : Builtin { + let Spellings = ["__builtin_infer_alloc_token"]; + let Attributes = [NoThrow, Const, Pure, CustomTypeChecking, UnevaluatedArguments]; + let Prototype = "size_t(...)"; +} + def CallWithStaticChain : Builtin { let Spellings = ["__builtin_call_with_static_chain"]; let Attributes = [NoThrow, CustomTypeChecking]; diff --git a/clang/lib/CodeGen/BackendUtil.cpp b/clang/lib/CodeGen/BackendUtil.cpp index 2d959827d6972..7a7dfa3633eeb 100644 --- a/clang/lib/CodeGen/BackendUtil.cpp +++ b/clang/lib/CodeGen/BackendUtil.cpp @@ -799,16 +799,6 @@ static void addSanitizers(const Triple &TargetTriple, MPM.addPass(DataFlowSanitizerPass(LangOpts.NoSanitizeFiles, PB.getVirtualFileSystemPtr())); } - - if (LangOpts.Sanitize.has(SanitizerKind::AllocToken)) { - if (Level == OptimizationLevel::O0) { - // The default pass builder only infers libcall function attrs when - // optimizing, so we insert it here because we need it for accurate - // memory allocation function detection. - MPM.addPass(InferFunctionAttrsPass()); - } - MPM.addPass(AllocTokenPass(getAllocTokenOptions(CodeGenOpts))); - } }; if (ClSanitizeOnOptimizerEarlyEP) { PB.registerOptimizerEarlyEPCallback( @@ -851,6 +841,22 @@ static void addSanitizers(const Triple &TargetTriple, } } +static void addAllocTokenPass(const Triple &TargetTriple, + const CodeGenOptions &CodeGenOpts, + const LangOptions &LangOpts, PassBuilder &PB) { + PB.registerOptimizerLastEPCallback( + [&](ModulePassManager &MPM, OptimizationLevel Level, ThinOrFullLTOPhase) { + if (Level == OptimizationLevel::O0 && + LangOpts.Sanitize.has(SanitizerKind::AllocToken)) { + // The default pass builder only infers libcall function attrs when + // optimizing, so we insert it here because we need it for accurate + // memory allocation function detection with -fsanitize=alloc-token. + MPM.addPass(InferFunctionAttrsPass()); + } + MPM.addPass(AllocTokenPass(getAllocTokenOptions(CodeGenOpts))); + }); +} + void EmitAssemblyHelper::RunOptimizationPipeline( BackendAction Action, std::unique_ptr &OS, std::unique_ptr &ThinLinkOS, BackendConsumer *BC) { @@ -1105,6 +1111,7 @@ void EmitAssemblyHelper::RunOptimizationPipeline( if (!IsThinLTOPostLink) { addSanitizers(TargetTriple, CodeGenOpts, LangOpts, PB); addKCFIPass(TargetTriple, LangOpts, PB); + addAllocTokenPass(TargetTriple, CodeGenOpts, LangOpts, PB); } if (std::optional Options = diff --git a/clang/lib/CodeGen/CGBuiltin.cpp b/clang/lib/CodeGen/CGBuiltin.cpp index 9ee810c9d5775..8a85789a256d5 100644 --- a/clang/lib/CodeGen/CGBuiltin.cpp +++ b/clang/lib/CodeGen/CGBuiltin.cpp @@ -4525,6 +4525,15 @@ RValue CodeGenFunction::EmitBuiltinExpr(const GlobalDecl GD, unsigned BuiltinID, return RValue::get(AI); } + case Builtin::BI__builtin_infer_alloc_token: { + llvm::MDNode *MDN = buildAllocToken(E); + llvm::Value *MDV = MetadataAsValue::get(getLLVMContext(), MDN); + llvm::Function *F = + CGM.getIntrinsic(llvm::Intrinsic::alloc_token_id, {IntPtrTy}); + llvm::CallBase *TokenID = Builder.CreateCall(F, MDV); + return RValue::get(TokenID); + } + case Builtin::BIbzero: case Builtin::BI__builtin_bzero: { Address Dest = EmitPointerWithAlignment(E->getArg(0)); diff --git a/clang/lib/CodeGen/CGExpr.cpp b/clang/lib/CodeGen/CGExpr.cpp index e8255b0554da8..c6f87e4c78fe3 100644 --- a/clang/lib/CodeGen/CGExpr.cpp +++ b/clang/lib/CodeGen/CGExpr.cpp @@ -1322,10 +1322,7 @@ typeContainsPointer(QualType T, return false; } -void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) { - assert(SanOpts.has(SanitizerKind::AllocToken) && - "Only needed with -fsanitize=alloc-token"); - +llvm::MDNode *CodeGenFunction::buildAllocToken(QualType AllocType) { llvm::MDBuilder MDB(getLLVMContext()); // Get unique type name. @@ -1344,14 +1341,19 @@ void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) { const bool ContainsPtr = typeContainsPointer(AllocType, VisitedRD, IncompleteType); if (!ContainsPtr && IncompleteType) - return; + return nullptr; auto *ContainsPtrC = Builder.getInt1(ContainsPtr); auto *ContainsPtrMD = MDB.createConstant(ContainsPtrC); // Format: !{, } - auto *MDN = - llvm::MDNode::get(CGM.getLLVMContext(), {TypeNameMD, ContainsPtrMD}); - CB->setMetadata(llvm::LLVMContext::MD_alloc_token, MDN); + return llvm::MDNode::get(CGM.getLLVMContext(), {TypeNameMD, ContainsPtrMD}); +} + +void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, QualType AllocType) { + assert(SanOpts.has(SanitizerKind::AllocToken) && + "Only needed with -fsanitize=alloc-token"); + CB->setMetadata(llvm::LLVMContext::MD_alloc_token, + buildAllocToken(AllocType)); } namespace { @@ -1445,7 +1447,7 @@ QualType inferPossibleTypeFromCastExpr(const CallExpr *CallE, } } // end anonymous namespace -void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, const CallExpr *E) { +llvm::MDNode *CodeGenFunction::buildAllocToken(const CallExpr *E) { QualType AllocType; // First check arguments. for (const Expr *Arg : E->arguments()) { @@ -1460,7 +1462,15 @@ void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, const CallExpr *E) { AllocType = inferPossibleTypeFromCastExpr(E, CurCast); // Emit if we were able to infer the type. if (!AllocType.isNull()) - EmitAllocToken(CB, AllocType); + return buildAllocToken(AllocType); + return nullptr; +} + +void CodeGenFunction::EmitAllocToken(llvm::CallBase *CB, const CallExpr *E) { + assert(SanOpts.has(SanitizerKind::AllocToken) && + "Only needed with -fsanitize=alloc-token"); + if (llvm::MDNode *MDN = buildAllocToken(E)) + CB->setMetadata(llvm::LLVMContext::MD_alloc_token, MDN); } CodeGenFunction::ComplexPairTy CodeGenFunction:: diff --git a/clang/lib/CodeGen/CodeGenFunction.h b/clang/lib/CodeGen/CodeGenFunction.h index 1f0be2d8756de..8c4c1c8c2dc95 100644 --- a/clang/lib/CodeGen/CodeGenFunction.h +++ b/clang/lib/CodeGen/CodeGenFunction.h @@ -3352,9 +3352,14 @@ class CodeGenFunction : public CodeGenTypeCache { SanitizerAnnotateDebugInfo(ArrayRef Ordinals, SanitizerHandler Handler); - /// Emit additional metadata used by the AllocToken instrumentation. + /// Build metadata used by the AllocToken instrumentation. + llvm::MDNode *buildAllocToken(QualType AllocType); + /// Emit and set additional metadata used by the AllocToken instrumentation. void EmitAllocToken(llvm::CallBase *CB, QualType AllocType); - /// Emit additional metadata used by the AllocToken instrumentation, + /// Build additional metadata used by the AllocToken instrumentation, + /// inferring the type from an allocation call expression. + llvm::MDNode *buildAllocToken(const CallExpr *E); + /// Emit and set additional metadata used by the AllocToken instrumentation, /// inferring the type from an allocation call expression. void EmitAllocToken(llvm::CallBase *CB, const CallExpr *E); diff --git a/clang/lib/Sema/SemaChecking.cpp b/clang/lib/Sema/SemaChecking.cpp index 063db05665af1..953428e9e8f34 100644 --- a/clang/lib/Sema/SemaChecking.cpp +++ b/clang/lib/Sema/SemaChecking.cpp @@ -1498,6 +1498,24 @@ static void builtinAllocaAddrSpace(Sema &S, CallExpr *TheCall) { TheCall->setType(S.Context.getPointerType(RT)); } +static bool checkBuiltinInferAllocToken(Sema &S, CallExpr *TheCall) { + if (S.checkArgCountAtLeast(TheCall, 1)) + return true; + + for (Expr *Arg : TheCall->arguments()) { + // If argument is dependent on a template parameter, we can't resolve now. + if (Arg->isTypeDependent() || Arg->isValueDependent()) + continue; + // Reject void types. + QualType ArgTy = Arg->IgnoreParenImpCasts()->getType(); + if (ArgTy->isVoidType()) + return S.Diag(Arg->getBeginLoc(), diag::err_param_with_void_type); + } + + TheCall->setType(S.Context.UnsignedLongLongTy); + return false; +} + namespace { enum PointerAuthOpKind { PAO_Strip, @@ -2779,6 +2797,10 @@ Sema::CheckBuiltinFunctionCall(FunctionDecl *FDecl, unsigned BuiltinID, builtinAllocaAddrSpace(*this, TheCall); } break; + case Builtin::BI__builtin_infer_alloc_token: + if (checkBuiltinInferAllocToken(*this, TheCall)) + return ExprError(); + break; case Builtin::BI__arithmetic_fence: if (BuiltinArithmeticFence(TheCall)) return ExprError(); diff --git a/clang/test/CodeGen/lto-newpm-pipeline.c b/clang/test/CodeGen/lto-newpm-pipeline.c index ea9784a76f923..dceaaf136ebfc 100644 --- a/clang/test/CodeGen/lto-newpm-pipeline.c +++ b/clang/test/CodeGen/lto-newpm-pipeline.c @@ -32,10 +32,12 @@ // CHECK-FULL-O0-NEXT: Running pass: AlwaysInlinerPass // CHECK-FULL-O0-NEXT: Running analysis: ProfileSummaryAnalysis // CHECK-FULL-O0-NEXT: Running pass: CoroConditionalWrapper +// CHECK-FULL-O0-NEXT: Running pass: AllocTokenPass +// CHECK-FULL-O0-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis +// CHECK-FULL-O0-NEXT: Running analysis: TargetLibraryAnalysis // CHECK-FULL-O0-NEXT: Running pass: CanonicalizeAliasesPass // CHECK-FULL-O0-NEXT: Running pass: NameAnonGlobalPass // CHECK-FULL-O0-NEXT: Running pass: AnnotationRemarksPass -// CHECK-FULL-O0-NEXT: Running analysis: TargetLibraryAnalysis // CHECK-FULL-O0-NEXT: Running pass: VerifierPass // CHECK-FULL-O0-NEXT: Running pass: BitcodeWriterPass @@ -46,10 +48,12 @@ // CHECK-THIN-O0-NEXT: Running pass: AlwaysInlinerPass // CHECK-THIN-O0-NEXT: Running analysis: ProfileSummaryAnalysis // CHECK-THIN-O0-NEXT: Running pass: CoroConditionalWrapper +// CHECK-THIN-O0-NEXT: Running pass: AllocTokenPass +// CHECK-THIN-O0-NEXT: Running analysis: OptimizationRemarkEmitterAnalysis +// CHECK-THIN-O0-NEXT: Running analysis: TargetLibraryAnalysis // CHECK-THIN-O0-NEXT: Running pass: CanonicalizeAliasesPass // CHECK-THIN-O0-NEXT: Running pass: NameAnonGlobalPass // CHECK-THIN-O0-NEXT: Running pass: AnnotationRemarksPass -// CHECK-THIN-O0-NEXT: Running analysis: TargetLibraryAnalysis // CHECK-THIN-O0-NEXT: Running pass: VerifierPass // CHECK-THIN-O0-NEXT: Running pass: ThinLTOBitcodeWriterPass diff --git a/clang/test/CodeGenCXX/alloc-token-builtin.cpp b/clang/test/CodeGenCXX/alloc-token-builtin.cpp new file mode 100644 index 0000000000000..cf8d11d03c764 --- /dev/null +++ b/clang/test/CodeGenCXX/alloc-token-builtin.cpp @@ -0,0 +1,76 @@ +// Test IR generation of the builtin without evaluating the LLVM intrinsic. +// RUN: %clang_cc1 -triple x86_64-linux-gnu -Werror -std=c++20 -emit-llvm -disable-llvm-passes %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-CODEGEN +// RUN: %clang_cc1 -triple x86_64-linux-gnu -Werror -std=c++20 -emit-llvm -falloc-token-max=2 %s -o - | FileCheck %s --check-prefixes=CHECK,CHECK-LOWER + +extern "C" void *my_malloc(unsigned long, unsigned long); + +struct NoPtr { + int x; + long y; +}; + +struct WithPtr { + int a; + char *buf; +}; + +int unevaluated_fn(); + +// CHECK-LABEL: @_Z16test_builtin_intv( +unsigned long test_builtin_int() { + // CHECK-CODEGEN: call i64 @llvm.alloc.token.id.i64(metadata ![[MD_INT:[0-9]+]]) + // CHECK-LOWER: ret i64 0 + return __builtin_infer_alloc_token(sizeof(1)); +} + +// CHECK-LABEL: @_Z16test_builtin_ptrv( +unsigned long test_builtin_ptr() { + // CHECK-CODEGEN: call i64 @llvm.alloc.token.id.i64(metadata ![[MD_PTR:[0-9]+]]) + // CHECK-LOWER: ret i64 1 + return __builtin_infer_alloc_token(sizeof(int *)); +} + +// CHECK-LABEL: @_Z25test_builtin_struct_noptrv( +unsigned long test_builtin_struct_noptr() { + // CHECK-CODEGEN: call i64 @llvm.alloc.token.id.i64(metadata ![[MD_NOPTR:[0-9]+]]) + // CHECK-LOWER: ret i64 0 + return __builtin_infer_alloc_token(sizeof(NoPtr)); +} + +// CHECK-LABEL: @_Z25test_builtin_struct_w_ptrv( +unsigned long test_builtin_struct_w_ptr() { + // CHECK-CODEGEN: call i64 @llvm.alloc.token.id.i64(metadata ![[MD_WITHPTR:[0-9]+]]) + // CHECK-LOWER: ret i64 1 + return __builtin_infer_alloc_token(sizeof(WithPtr), 123); +} + +// CHECK-LABEL: @_Z24test_builtin_unevaluatedv( +unsigned long test_builtin_unevaluated() { + // CHECK-NOT: call{{.*}}unevaluated_fn + // CHECK-CODEGEN: call i64 @llvm.alloc.token.id.i64(metadata ![[MD_INT:[0-9]+]]) + // CHECK-LOWER: ret i64 0 + return __builtin_infer_alloc_token(sizeof(int) * unevaluated_fn()); +} + +// CHECK-LABEL: @_Z36test_builtin_unsequenced_unevaluatedi( +void test_builtin_unsequenced_unevaluated(int x) { + // CHECK: add nsw + // CHECK-NOT: add nsw + // CHECK-CODEGEN: %[[REG:[0-9]+]] = call i64 @llvm.alloc.token.id.i64(metadata ![[MD_UNKNOWN:[0-9]+]]) + // CHECK-CODEGEN: call{{.*}}@my_malloc({{.*}}, i64 noundef %[[REG]]) + // CHECK-LOWER: call{{.*}}@my_malloc({{.*}}, i64 noundef 0) + my_malloc(++x, __builtin_infer_alloc_token(++x)); +} + +// CHECK-LABEL: @_Z20test_builtin_unknownv( +unsigned long test_builtin_unknown() { + // CHECK-CODEGEN: call i64 @llvm.alloc.token.id.i64(metadata ![[MD_UNKNOWN:[0-9]+]]) + // CHECK-LOWER: ret i64 0 + return __builtin_infer_alloc_token(4096); +} + +// CHECK-CODEGEN: ![[MD_INT]] = !{!"int", i1 false} +// CHECK-CODEGEN: ![[MD_PTR]] = !{!"int *", i1 true} +// CHECK-CODEGEN: ![[MD_NOPTR]] = !{!"NoPtr", i1 false} +// CHECK-CODEGEN: ![[MD_WITHPTR]] = !{!"WithPtr", i1 true} +// CHECK-CODEGEN: ![[MD_UNKNOWN]] = !{} diff --git a/llvm/include/llvm/IR/Intrinsics.td b/llvm/include/llvm/IR/Intrinsics.td index 96da698538314..fe66b3f437e96 100644 --- a/llvm/include/llvm/IR/Intrinsics.td +++ b/llvm/include/llvm/IR/Intrinsics.td @@ -2841,7 +2841,15 @@ def int_ptrauth_blend : def int_ptrauth_sign_generic : DefaultAttrsIntrinsic<[llvm_i64_ty], [llvm_i64_ty, llvm_i64_ty], [IntrNoMem]>; +//===----------------- AllocToken Intrinsics ------------------------------===// + +// Return the token ID for the given !alloc_token metadata. +def int_alloc_token_id : + DefaultAttrsIntrinsic<[llvm_anyint_ty], [llvm_metadata_ty], + [IntrNoMem, NoUndef]>; + //===----------------------------------------------------------------------===// + //===------- Convergence Intrinsics ---------------------------------------===// def int_experimental_convergence_entry diff --git a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp index 40720ae4b39ae..e9c136f2711e9 100644 --- a/llvm/lib/Transforms/Instrumentation/AllocToken.cpp +++ b/llvm/lib/Transforms/Instrumentation/AllocToken.cpp @@ -31,6 +31,7 @@ #include "llvm/IR/InstIterator.h" #include "llvm/IR/InstrTypes.h" #include "llvm/IR/Instructions.h" +#include "llvm/IR/IntrinsicInst.h" #include "llvm/IR/Metadata.h" #include "llvm/IR/Module.h" #include "llvm/IR/PassManager.h" @@ -140,9 +141,19 @@ STATISTIC(NumAllocationsInstrumented, "Allocations instrumented"); /// /// Expected format is: !{, } MDNode *getAllocTokenMetadata(const CallBase &CB) { - MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token); - if (!Ret) - return nullptr; + MDNode *Ret = nullptr; + if (auto *II = dyn_cast(&CB); + II && II->getIntrinsicID() == Intrinsic::alloc_token_id) { + auto *MDV = cast(II->getArgOperand(0)); + Ret = cast(MDV->getMetadata()); + // If the intrinsic has an empty MDNode, type inference failed. + if (Ret->getNumOperands() == 0) + return nullptr; + } else { + Ret = CB.getMetadata(LLVMContext::MD_alloc_token); + if (!Ret) + return nullptr; + } assert(Ret->getNumOperands() == 2 && "bad !alloc_token"); assert(isa(Ret->getOperand(0))); assert(isa(Ret->getOperand(1))); @@ -315,6 +326,9 @@ class AllocToken { FunctionCallee getTokenAllocFunction(const CallBase &CB, uint64_t TokenID, LibFunc OriginalFunc); + /// Lower alloc_token_* intrinsics. + void replaceIntrinsicInst(IntrinsicInst *II, OptimizationRemarkEmitter &ORE); + /// Return the token ID from metadata in the call. uint64_t getToken(const CallBase &CB, OptimizationRemarkEmitter &ORE) { return std::visit([&](auto &&Mode) { return Mode(CB, ORE); }, Mode); @@ -336,21 +350,30 @@ 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(F); auto &TLI = FAM.getResult(F); SmallVector, 4> AllocCalls; + SmallVector IntrinsicInsts; // Collect all allocation calls to avoid iterator invalidation. for (Instruction &I : instructions(F)) { + // Collect all alloc_token_* intrinsics. + if (auto *II = dyn_cast(&I); + II && II->getIntrinsicID() == Intrinsic::alloc_token_id) { + IntrinsicInsts.emplace_back(II); + continue; + } + + // Only instrument functions that have the sanitize_alloc_token attribute. + if (!F.hasFnAttribute(Attribute::SanitizeAllocToken)) + continue; + if (F.hasFnAttribute(Attribute::DisableSanitizerInstrumentation)) + continue; + auto *CB = dyn_cast(&I); if (!CB) continue; @@ -359,11 +382,19 @@ bool AllocToken::instrumentFunction(Function &F) { } bool Modified = false; - for (auto &[CB, Func] : AllocCalls) - Modified |= replaceAllocationCall(CB, Func, ORE, TLI); - if (Modified) - NumFunctionsInstrumented++; + if (!AllocCalls.empty()) { + for (auto &[CB, Func] : AllocCalls) + Modified |= replaceAllocationCall(CB, Func, ORE, TLI); + if (Modified) + NumFunctionsInstrumented++; + } + + for (auto *II : IntrinsicInsts) { + replaceIntrinsicInst(II, ORE); + Modified = true; + } + return Modified; } @@ -528,6 +559,16 @@ FunctionCallee AllocToken::getTokenAllocFunction(const CallBase &CB, return TokenAlloc; } +void AllocToken::replaceIntrinsicInst(IntrinsicInst *II, + OptimizationRemarkEmitter &ORE) { + assert(II->getIntrinsicID() == Intrinsic::alloc_token_id); + + uint64_t TokenID = getToken(*II, ORE); + Value *V = ConstantInt::get(IntPtrTy, TokenID); + II->replaceAllUsesWith(V); + II->eraseFromParent(); +} + } // namespace AllocTokenPass::AllocTokenPass(AllocTokenOptions Opts) diff --git a/llvm/test/Instrumentation/AllocToken/intrinsic.ll b/llvm/test/Instrumentation/AllocToken/intrinsic.ll new file mode 100644 index 0000000000000..13aaa90008a7c --- /dev/null +++ b/llvm/test/Instrumentation/AllocToken/intrinsic.ll @@ -0,0 +1,32 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5 +; Test that the alloc-token pass lowers the intrinsic to a constant token ID. +; +; RUN: opt < %s -passes=alloc-token -alloc-token-mode=typehashpointersplit -alloc-token-max=2 -S | FileCheck %s + +target datalayout = "e-m:e-p270:32:32-p271:32:32-p272:64:64-i64:64-f80:128-n8:16:32:64-S128" +target triple = "x86_64-unknown-linux-gnu" + +declare i64 @llvm.alloc.token.id.i64(metadata) + +define i64 @test_intrinsic_lowering() { +; CHECK-LABEL: define i64 @test_intrinsic_lowering() { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: ret i64 0 +; +entry: + %token_no_ptr = call i64 @llvm.alloc.token.id.i64(metadata !0) + ret i64 %token_no_ptr +} + +define i64 @test_intrinsic_lowering_ptr() { +; CHECK-LABEL: define i64 @test_intrinsic_lowering_ptr() { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: ret i64 1 +; +entry: + %token_with_ptr = call i64 @llvm.alloc.token.id.i64(metadata !1) + ret i64 %token_with_ptr +} + +!0 = !{!"NoPointerType", i1 false} +!1 = !{!"PointerType", i1 true} diff --git a/llvm/test/Instrumentation/AllocToken/intrinsic32.ll b/llvm/test/Instrumentation/AllocToken/intrinsic32.ll new file mode 100644 index 0000000000000..eb5dbbe91a83e --- /dev/null +++ b/llvm/test/Instrumentation/AllocToken/intrinsic32.ll @@ -0,0 +1,32 @@ +; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5 +; Test that the alloc-token pass lowers the intrinsic to a constant token ID. +; +; RUN: opt < %s -passes=alloc-token -alloc-token-mode=typehashpointersplit -alloc-token-max=2 -S | FileCheck %s + +target datalayout = "e-p:32:32:32-i1:8:8-i8:8:8-i16:16:16-i32:32:32-i64:32:64-f32:32:32-f64:32:64-v64:64:64-v128:128:128-a0:0:64-f80:32:32-n8:16:32-S128" +target triple = "i386-pc-linux-gnu" + +declare i32 @llvm.alloc.token.id.i32(metadata) + +define i32 @test_intrinsic_lowering() { +; CHECK-LABEL: define i32 @test_intrinsic_lowering() { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: ret i32 0 +; +entry: + %token_no_ptr = call i32 @llvm.alloc.token.id.i32(metadata !0) + ret i32 %token_no_ptr +} + +define i32 @test_intrinsic_lowering_ptr() { +; CHECK-LABEL: define i32 @test_intrinsic_lowering_ptr() { +; CHECK-NEXT: [[ENTRY:.*:]] +; CHECK-NEXT: ret i32 1 +; +entry: + %token_with_ptr = call i32 @llvm.alloc.token.id.i32(metadata !1) + ret i32 %token_with_ptr +} + +!0 = !{!"NoPointerType", i1 false} +!1 = !{!"PointerType", i1 true}