8 files changed, 490 insertions, 0 deletions
diff --git a/llvm/test/Instrumentation/AllocToken/basic.ll b/llvm/test/Instrumentation/AllocToken/basic.ll
new file mode 100644
index 0000000..099d37d
--- /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 0000000..944a452
--- /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 0000000..5f08552
--- /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 0000000..19a3ef6
--- /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 0000000..b92a920
--- /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 0000000..347c99a
--- /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 0000000..e023ab6b
--- /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 0000000..a2404526
--- /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"}
+;.