[AllocToken] Introduce sanitize_alloc_token attribute and alloc_token metadata (#160131)

In preparation of adding the "AllocToken" pass, add the pre-requisite `sanitize_alloc_token` function attribute and `alloc_token` metadata. --- 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
2025-10-07 12:51:42 +02:00 · 2025-10-07 12:51:42 +02:00 · 224873d7ac
commit 224873d7ac
parent ed113e7904
17 changed files with 207 additions and 3 deletions
--- a/llvm/docs/LangRef.rst
+++ b/llvm/docs/LangRef.rst
@ -2529,6 +2529,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>`_
@ -8577,6 +8580,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
 =====================
--- 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 {
--- 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
--- a/llvm/include/llvm/IR/FixedMetadataKinds.def
+++ b/llvm/include/llvm/IR/FixedMetadataKinds.def
@ -56,3 +56,4 @@ 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_captures, "captures", 44)
+LLVM_FIXED_MD_KIND(MD_alloc_token, "alloc_token", 45)
--- 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:
--- 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:
--- a/llvm/lib/IR/Verifier.cpp
+++ b/llvm/lib/IR/Verifier.cpp
@ -543,6 +543,7 @@ private:
  void visitAliasScopeListMetadata(const MDNode *MD);
  void visitAccessGroupMetadata(const MDNode *MD);
  void visitCapturesMetadata(Instruction &I, const MDNode *Captures);
+  void visitAllocTokenMetadata(Instruction &I, MDNode *MD);

  template <class Ty> bool isValidMetadataArray(const MDTuple &N);
 #define HANDLE_SPECIALIZED_MDNODE_LEAF(CLASS) void visit##CLASS(const CLASS &N);
@ -5395,6 +5396,12 @@ void Verifier::visitCapturesMetadata(Instruction &I, const MDNode *Captures) {
  }
 }

+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) {
@ -5625,6 +5632,9 @@ void Verifier::visitInstruction(Instruction &I) {
  if (MDNode *Captures = I.getMetadata(LLVMContext::MD_captures))
    visitCapturesMetadata(I, Captures);

+  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);
--- 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:
--- a/llvm/lib/Transforms/Utils/Local.cpp
+++ b/llvm/lib/Transforms/Utils/Local.cpp
@ -3031,6 +3031,13 @@ static void combineMetadata(Instruction *K, const Instruction *J,
                      K->getContext(), MDNode::toCaptureComponents(JMD) |
                                           MDNode::toCaptureComponents(KMD)));
        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
--- 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 }
--- 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

--- 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" }
--- a/llvm/test/Transforms/SimplifyCFG/merge-calls-alloc-token.ll
+++ 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"}
+;.
--- 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)
--- 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
--- 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
--- 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|\