fixup! address reviewer comments

Created using spr 1.3.8-beta.1

Author: melver

llvm/docs/LangRef.rst

@@ -2428,7 +2428,7 @@ For example:
     attributed with ``sanitize_realtime``.
     This attribute is incompatible with the ``sanitize_realtime`` attribute.
 ``sanitize_alloc_token``
-    This attributes indicates that implicit allocation token instrumentation
+    This attribute indicates that implicit allocation token instrumentation
     is enabled for this function.
 ``speculative_load_hardening``
     This attribute indicates that
@@ -8392,6 +8392,13 @@ 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.
 
 Module Flags Metadata
 =====================

llvm/docs/ReleaseNotes.md

@@ -166,9 +166,9 @@ Changes to Sanitizers
 Other Changes
 -------------
 
-* Introduces the `AllocToken` pass, an instrumentation pass designed to provide
-  tokens to memory allocators enabling various heap organization strategies,
-  such as heap partitioning.
+* 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}}
 ===============================================================

llvm/include/llvm/IR/FixedMetadataKinds.def

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

llvm/lib/Transforms/Instrumentation/AllocToken.cpp

@@ -69,27 +69,19 @@ enum class TokenMode : unsigned {
 
   /// Token ID based on allocated type hash.
   TypeHash = 2,
-
-  // Mode count - keep last
-  ModeCount
 };
 
 //===--- Command-line options ---------------------------------------------===//
 
-struct ModeParser : public cl::parser<unsigned> {
-  ModeParser(cl::Option &O) : cl::parser<unsigned>(O) {}
-  bool parse(cl::Option &O, StringRef ArgName, StringRef Arg, unsigned &Value) {
-    if (cl::parser<unsigned>::parse(O, ArgName, Arg, Value))
-      return true;
-    if (Value >= static_cast<unsigned>(TokenMode::ModeCount))
-      return O.error("'" + Arg + "' value invalid");
-    return false;
-  }
-};
-
-cl::opt<unsigned, false, ModeParser>
-    ClMode("alloc-token-mode", cl::desc("Token assignment mode"), cl::Hidden,
-           cl::init(static_cast<unsigned>(TokenMode::TypeHash)));
+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"),
@@ -106,7 +98,7 @@ cl::opt<bool>
 
 // 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_hint
+// instrument non-libcall allocation function calls with !alloc_token
 // metadata.
 cl::opt<bool>
     ClExtended("alloc-token-extended",
@@ -140,14 +132,14 @@ STATISTIC(NumAllocations, "Allocations found");
 
 //===----------------------------------------------------------------------===//
 
-/// Returns the !alloc_token_hint metadata if available.
+/// Returns the !alloc_token metadata if available.
 ///
 /// Expected format is: !{<type-name>}
-MDNode *getAllocTokenHintMetadata(const CallBase &CB) {
-  MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token_hint);
+MDNode *getAllocTokenMetadata(const CallBase &CB) {
+  MDNode *Ret = CB.getMetadata(LLVMContext::MD_alloc_token);
   if (!Ret)
     return nullptr;
-  assert(Ret->getNumOperands() == 1 && "bad !alloc_token_hint");
+  assert(Ret->getNumOperands() == 1 && "bad !alloc_token");
   assert(isa<MDString>(Ret->getOperand(0)));
   return Ret;
 }
@@ -198,7 +190,7 @@ class TypeHashMode : public ModeBase {
   using ModeBase::ModeBase;
 
   uint64_t operator()(const CallBase &CB, OptimizationRemarkEmitter &ORE) {
-    if (MDNode *N = getAllocTokenHintMetadata(CB)) {
+    if (MDNode *N = getAllocTokenMetadata(CB)) {
       MDString *S = cast<MDString>(N->getOperand(0));
       return boundedToken(xxHash64(S->getString()));
     }
@@ -212,7 +204,7 @@ class TypeHashMode : public ModeBase {
       ore::NV FuncNV("Function", CB.getParent()->getParent());
       const Function *Callee = CB.getCalledFunction();
       ore::NV CalleeNV("Callee", Callee ? Callee->getName() : "<unknown>");
-      return OptimizationRemark(DEBUG_TYPE, "NoAllocTokenHint", &CB)
+      return OptimizationRemark(DEBUG_TYPE, "NoAllocToken", &CB)
              << "Call to '" << CalleeNV << "' in '" << FuncNV
              << "' without source-level type token";
     });
@@ -235,7 +227,7 @@ class AllocToken {
       : Options(transformOptionsFromCl(std::move(Opts))), Mod(M),
         FAM(MAM.getResult<FunctionAnalysisManagerModuleProxy>(M).getManager()),
         Mode(IncrementMode(*Options.MaxTokens)) {
-    switch (static_cast<TokenMode>(ClMode.getValue())) {
+    switch (ClMode.getValue()) {
     case TokenMode::Increment:
       break;
     case TokenMode::Random:
@@ -244,9 +236,6 @@ class AllocToken {
     case TokenMode::TypeHash:
       Mode.emplace<TypeHashMode>(*Options.MaxTokens);
       break;
-    case TokenMode::ModeCount:
-      llvm_unreachable("");
-      break;
     }
   }
 
@@ -255,10 +244,11 @@ class AllocToken {
 private:
   /// Returns true for !isAllocationFn() functions that are also eligible for
   /// instrumentation.
-  bool isInstrumentableLibFunc(LibFunc Func) const;
+  static bool isInstrumentableLibFunc(LibFunc Func, const Value *V,
+                                      const TargetLibraryInfo *TLI);
 
   /// Returns true for isAllocationFn() functions that we should ignore.
-  bool ignoreInstrumentableLibFunc(LibFunc Func) const;
+  static bool ignoreInstrumentableLibFunc(LibFunc Func);
 
   /// Replace a call/invoke with a call/invoke to the allocation function
   /// with token ID.
@@ -277,6 +267,7 @@ class AllocToken {
 
   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;
@@ -316,28 +307,29 @@ bool AllocToken::instrumentFunction(Function &F) {
     if (TLI.getLibFunc(*Callee, Func)) {
       if (ignoreInstrumentableLibFunc(Func))
         continue;
-      if (isInstrumentableLibFunc(Func) || isAllocationFn(CB, &TLI))
+      if (isInstrumentableLibFunc(Func, CB, &TLI))
         AllocCalls.emplace_back(CB, Func);
-    } else if (Options.Extended && getAllocTokenHintMetadata(*CB)) {
+    } else if (Options.Extended && getAllocTokenMetadata(*CB)) {
       AllocCalls.emplace_back(CB, NotLibFunc);
     }
   }
 
-  bool Modified = false;
+  if (AllocCalls.empty())
+    return false;
 
-  if (!AllocCalls.empty()) {
-    for (auto &[CB, Func] : AllocCalls) {
-      replaceAllocationCall(CB, Func, ORE, TLI);
-    }
-    NumAllocations += AllocCalls.size();
-    NumFunctionsInstrumented++;
-    Modified = true;
-  }
+  for (auto &[CB, Func] : AllocCalls)
+    replaceAllocationCall(CB, Func, ORE, TLI);
+  NumAllocations += AllocCalls.size();
+  NumFunctionsInstrumented++;
 
-  return Modified;
+  return true;
 }
 
-bool AllocToken::isInstrumentableLibFunc(LibFunc Func) const {
+bool AllocToken::isInstrumentableLibFunc(LibFunc Func, const Value *V,
+                                         const TargetLibraryInfo *TLI) {
+  if (isAllocationFn(V, TLI))
+    return true;
+
   switch (Func) {
   case LibFunc_posix_memalign:
   case LibFunc_size_returning_new:
@@ -375,7 +367,7 @@ bool AllocToken::isInstrumentableLibFunc(LibFunc Func) const {
   }
 }
 
-bool AllocToken::ignoreInstrumentableLibFunc(LibFunc Func) const {
+bool AllocToken::ignoreInstrumentableLibFunc(LibFunc Func) {
   switch (Func) {
   case LibFunc_strdup:
   case LibFunc_dunder_strdup:
@@ -395,16 +387,17 @@ void AllocToken::replaceAllocationCall(CallBase *CB, LibFunc Func,
   FunctionCallee TokenAlloc = getTokenAllocFunction(*CB, TokenID, Func);
   if (!TokenAlloc)
     return;
+  if (Options.FastABI) {
+    assert(TokenAlloc.getFunctionType()->getNumParams() == CB->arg_size());
+    CB->setCalledFunction(TokenAlloc);
+    return;
+  }
 
   IRBuilder<> IRB(CB);
-
   // Original args.
   SmallVector<Value *, 4> NewArgs{CB->args()};
-  if (!Options.FastABI) {
-    // Add token ID.
-    NewArgs.push_back(
-        ConstantInt::get(Type::getInt64Ty(Mod.getContext()), TokenID));
-  }
+  // Add token ID.
+  NewArgs.push_back(ConstantInt::get(IntPtrTy, TokenID));
   assert(TokenAlloc.getFunctionType()->getNumParams() == NewArgs.size());
 
   // Preserve invoke vs call semantics for exception handling.
@@ -443,15 +436,13 @@ FunctionCallee AllocToken::getTokenAllocFunction(const CallBase &CB,
   if (OldFTy->isVarArg())
     return FunctionCallee();
   // Copy params, and append token ID type.
-  LLVMContext &C = Mod.getContext();
   Type *RetTy = OldFTy->getReturnType();
   SmallVector<Type *, 4> NewParams{OldFTy->params()};
   std::string TokenAllocName = ClFuncPrefix;
-  if (Options.FastABI) {
+  if (Options.FastABI)
     TokenAllocName += utostr(TokenID) + "_";
-  } else {
-    NewParams.push_back(Type::getInt64Ty(C)); // token ID
-  }
+  else
+    NewParams.push_back(IntPtrTy); // token ID
   FunctionType *NewFTy = FunctionType::get(RetTy, NewParams, false);
   // Remove leading '_' - we add our own.
   StringRef No_ = Callee->getName().drop_while([](char C) { return C == '_'; });
@@ -480,5 +471,6 @@ PreservedAnalyses AllocTokenPass::run(Module &M, ModuleAnalysisManager &MAM) {
     Modified |= Pass.instrumentFunction(F);
   }
 
-  return Modified ? PreservedAnalyses::none() : PreservedAnalyses::all();
+  return Modified ? PreservedAnalyses::none().preserveSet<CFGAnalyses>()
+                  : PreservedAnalyses::all();
 }

llvm/lib/Transforms/Utils/Local.cpp

@@ -3025,8 +3025,8 @@ 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_hint:
-        // Preserve !alloc_token_hint if both K and J have it.
+      case LLVMContext::MD_alloc_token:
+        // Preserve !alloc_token if both K and J have it.
         K->setMetadata(Kind, JMD);
         break;
       }

llvm/test/Instrumentation/AllocToken/basic.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=0 -S | FileCheck %s
+; 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"
 target triple = "x86_64-unknown-linux-gnu"

llvm/test/Instrumentation/AllocToken/basic32.ll

@@ -0,0 +1,25 @@
+; 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"
+target triple = "i386-pc-linux-gnu"
+
+declare ptr @malloc(i32)
+declare ptr @_Znwm(i32)
+
+; CHECK-LABEL: @test_basic_rewriting
+define ptr @test_basic_rewriting() sanitize_alloc_token {
+entry:
+  ; CHECK: [[PTR1:%[0-9]]] = call ptr @__alloc_token_malloc(i32 64, i32 0)
+  ; CHECK-NOT: call ptr @malloc(
+  %ptr1 = call ptr @malloc(i32 64)
+  ret ptr %ptr1
+}
+
+; CHECK-LABEL: @test_cpp_operators
+define ptr @test_cpp_operators() sanitize_alloc_token {
+entry:
+  ; CHECK: call ptr @__alloc_token_Znwm(i32 32, i32 1)
+  ; CHECK-NOT: call ptr @_Znwm(
+  %ptr1 = call ptr @_Znwm(i32 32)
+  ret ptr %ptr1
+}

llvm/test/Instrumentation/AllocToken/extralibfuncs.ll

@@ -11,7 +11,7 @@ declare {ptr, i64} @__size_returning_new(i64)
 define ptr @test_extra_libfuncs() sanitize_alloc_token {
 entry:
   ; CHECK: call {{.*}} @__alloc_token_size_returning_new(
-  %srn = call {ptr, i64} @__size_returning_new(i64 10), !alloc_token_hint !0
+  %srn = call {ptr, i64} @__size_returning_new(i64 10), !alloc_token !0
   %ptr1  = extractvalue {ptr, i64} %srn, 0
   ret ptr %ptr1
 }
@@ -23,9 +23,9 @@ declare ptr @_Znam(i64) nobuiltin allocsize(0)
 define ptr @test_replaceable_new() sanitize_alloc_token {
 entry:
   ; CHECK: call ptr @__alloc_token_Znwm(
-  %ptr1 = call ptr @_Znwm(i64 32), !alloc_token_hint !0
+  %ptr1 = call ptr @_Znwm(i64 32), !alloc_token !0
   ; CHECK: call ptr @__alloc_token_Znam(
-  %ptr2 = call ptr @_Znam(i64 64), !alloc_token_hint !0
+  %ptr2 = call ptr @_Znam(i64 64), !alloc_token !0
   ret ptr %ptr1
 }
 

llvm/test/Instrumentation/AllocToken/fast.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=0 -alloc-token-fast-abi -alloc-token-max=3 -S | FileCheck %s
+; 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"
 target triple = "x86_64-unknown-linux-gnu"
@@ -13,7 +13,7 @@ declare ptr @_Znam(i64)
 ; CHECK-LABEL: @test_basic_rewriting
 define ptr @test_basic_rewriting() sanitize_alloc_token {
 entry:
-  ; CHECK: [[PTR1:%[0-9]]] = call ptr @__alloc_token_0_malloc(i64 64)
+  ; CHECK: [[PTR1:%ptr[0-9]]] = call ptr @__alloc_token_0_malloc(i64 64)
   ; CHECK: call ptr @__alloc_token_1_calloc(i64 8, i64 8)
   ; CHECK: call ptr @__alloc_token_2_realloc(ptr [[PTR1]], i64 128)
   ; CHECK-NOT: call ptr @malloc(

llvm/test/Instrumentation/AllocToken/invoke.ll

@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes=inferattrs,alloc-token -alloc-token-mode=0 -S | FileCheck %s
+; 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"
 target triple = "x86_64-unknown-linux-gnu"