diff --git a/clang/lib/Headers/__clang_hip_runtime_wrapper.h b/clang/lib/Headers/__clang_hip_runtime_wrapper.h --- a/clang/lib/Headers/__clang_hip_runtime_wrapper.h +++ b/clang/lib/Headers/__clang_hip_runtime_wrapper.h @@ -80,12 +80,25 @@ #if HIP_VERSION_MAJOR * 100 + HIP_VERSION_MINOR >= 405 extern "C" __device__ unsigned long long __ockl_dm_alloc(unsigned long long __size); extern "C" __device__ void __ockl_dm_dealloc(unsigned long long __addr); +#if __has_feature(address_sanitizer) +extern "C" __device__ unsigned long long __asan_malloc_impl(unsigned long long __size, unsigned long long __pc); +extern "C" __device__ void __asan_free_impl(unsigned long long __addr, unsigned long long __pc); +__attribute__((noinline, weak)) __device__ void *malloc(__hip_size_t __size) { + unsigned long long __pc = (unsigned long long)__builtin_return_address(0); + return (void *)__asan_malloc_impl(__size, __pc); +} +__attribute__((noinline, weak)) __device__ void free(void *__ptr) { + unsigned long long __pc = (unsigned long long)__builtin_return_address(0); + __asan_free_impl((unsigned long long)__ptr, __pc); +} +#else __attribute__((weak)) inline __device__ void *malloc(__hip_size_t __size) { return (void *) __ockl_dm_alloc(__size); } __attribute__((weak)) inline __device__ void free(void *__ptr) { __ockl_dm_dealloc((unsigned long long)__ptr); } +#endif // __has_feature(address_sanitizer) #else // HIP version check #if __HIP_ENABLE_DEVICE_MALLOC__ __device__ void *__hip_malloc(__hip_size_t __size); diff --git a/clang/test/Headers/hip-header.hip b/clang/test/Headers/hip-header.hip --- a/clang/test/Headers/hip-header.hip +++ b/clang/test/Headers/hip-header.hip @@ -32,6 +32,14 @@ // RUN: -target-cpu gfx906 -emit-llvm %s -fcuda-is-device -o - \ // RUN: -DHIP_VERSION_MAJOR=4 -DHIP_VERSION_MINOR=5 \ // RUN: -D__HIPCC_RTC__ | FileCheck -check-prefixes=CHECK,MALLOC %s +// RUN: %clang_cc1 -no-opaque-pointers -include __clang_hip_runtime_wrapper.h \ +// RUN: -internal-isystem %S/../../lib/Headers/cuda_wrappers \ +// RUN: -internal-isystem %S/Inputs/include \ +// RUN: -triple amdgcn-amd-amdhsa -aux-triple x86_64-unknown-unknown \ +// RUN: -target-cpu gfx906 -emit-llvm %s -fcuda-is-device -o - \ +// RUN: -DHIP_VERSION_MAJOR=4 -DHIP_VERSION_MINOR=5 \ +// RUN: -fsanitize=address -disable-llvm-passes -D__HIPCC_RTC__ \ +// RUN: | FileCheck -check-prefixes=MALLOC-ASAN %s // expected-no-diagnostics @@ -130,6 +138,9 @@ // CHECK-LABEL: define weak {{.*}}i8* @malloc(i64 // MALLOC: call i64 @__ockl_dm_alloc // NOMALLOC: call void @llvm.trap +// MALLOC-ASAN-LABEL: define weak {{.*}}i8* @malloc(i64 +// MALLOC-ASAN: call i8* @llvm.returnaddress(i32 0) +// MALLOC-ASAN: call i64 @__asan_malloc_impl(i64 {{.*}}, i64 {{.*}}) __device__ void test_malloc(void *a) { a = malloc(42); } @@ -139,6 +150,9 @@ // CHECK-LABEL: define weak {{.*}}void @free(i8* // MALLOC: call void @__ockl_dm_dealloc // NOMALLOC: call void @llvm.trap +// MALLOC-ASAN-LABEL: define weak {{.*}}void @free(i8* +// MALLOC-ASAN: call i8* @llvm.returnaddress(i32 0) +// MALLOC-ASAN: call void @__asan_free_impl(i64 {{.*}}, i64 {{.*}}) __device__ void test_free(void *a) { free(a); }