[HIP] Support ASAN with malloc/free

[yaxunl]

Summary

Device side malloc/free needs special implementation for ASAN.

[arsenm]

clang/test/Headers/hip-header.hip

142

Should convert this test to opaque pointers before or after

[yaxunl]

clang/test/Headers/hip-header.hip

142

will do it after

clang/lib/Headers/__clang_hip_runtime_wrapper.h

 
 #ifdef __cplusplus
 extern "C" {
 #endif //__cplusplus
 
 #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);
 __device__ void *__hip_free(void *__ptr);
 __attribute__((weak)) inline __device__ void *malloc(__hip_size_t __size) {
   return __hip_malloc(__size);
 }
 __attribute__((weak)) inline __device__ void free(void *__ptr) {

clang/test/Headers/hip-header.hip

 // RUN:   -D__HIPCC_RTC__ -std=c++14 | FileCheck -check-prefixes=CHECK,CXX14 %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:   -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
 
 // Check support for pure and deleted virtual functions
 struct base {
   __host__
   __device__
   virtual void pv() = 0;

 
 // Check that device malloc and free do not conflict with std headers.
 #include <cstdlib>
 // CHECK-LABEL: define{{.*}}@_Z11test_malloc
 // CHECK: call {{.*}}i8* @malloc(i64
 // 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);
 }
 
 // CHECK-LABEL: define{{.*}}@_Z9test_free
 // CHECK: call {{.*}}void @free(i8*
 // 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);
 }