diff --git a/compiler-rt/lib/asan/asan_allocator.cpp b/compiler-rt/lib/asan/asan_allocator.cpp index 06c827c41eacc..5bcae291428f6 100644 --- a/compiler-rt/lib/asan/asan_allocator.cpp +++ b/compiler-rt/lib/asan/asan_allocator.cpp @@ -1399,7 +1399,10 @@ DECLARE_REAL(hsa_status_t, hsa_amd_ipc_memory_attach, DECLARE_REAL(hsa_status_t, hsa_amd_ipc_memory_detach, void *mapped_ptr) DECLARE_REAL(hsa_status_t, hsa_amd_vmem_address_reserve_align, void** ptr, size_t size, uint64_t address, uint64_t alignment, uint64_t flags) -DECLARE_REAL(hsa_status_t, hsa_amd_vmem_address_free, void* ptr, size_t size); +DECLARE_REAL(hsa_status_t, hsa_amd_vmem_address_free, void* ptr, size_t size) +DECLARE_REAL(hsa_status_t, hsa_amd_pointer_info, const void* ptr, + hsa_amd_pointer_info_t* info, void* (*alloc)(size_t), + uint32_t* num_agents_accessible, hsa_agent_t** accessible) namespace __asan { @@ -1452,18 +1455,22 @@ static struct AP64 AP_; static struct AP32 AP_; #endif -hsa_status_t asan_hsa_amd_ipc_memory_create(void *ptr, size_t len, - hsa_amd_ipc_memory_t * handle) { - void *ptr_; - size_t len_ = get_allocator().GetActuallyAllocatedSize(ptr); - if (len_) { +hsa_status_t asan_hsa_amd_ipc_memory_create(void* ptr, size_t len, + hsa_amd_ipc_memory_t* handle) { + void* ptr_ = get_allocator().GetBlockBegin(ptr); + AsanChunk* m = ptr_ + ? instance.GetAsanChunkByAddr(reinterpret_cast(ptr_)) + : nullptr; + if (ptr_ && m) { static_assert(AP_.kMetadataSize == 0, "Expression below requires this"); - ptr_ = reinterpret_cast(reinterpret_cast(ptr) - kPageSize_); - } else { - ptr_ = ptr; - len_ = len; + uptr p = reinterpret_cast(ptr); + uptr p_ = reinterpret_cast(ptr_); + if (p == p_ + kPageSize_ && len == m->UsedSize()) { + size_t len_ = get_allocator().GetActuallyAllocatedSize(ptr_); + return REAL(hsa_amd_ipc_memory_create)(ptr_, len_, handle); + } } - return REAL(hsa_amd_ipc_memory_create)(ptr_, len_, handle); + return REAL(hsa_amd_ipc_memory_create)(ptr, len, handle); } hsa_status_t asan_hsa_amd_ipc_memory_attach(const hsa_amd_ipc_memory_t *handle, @@ -1540,5 +1547,36 @@ hsa_status_t asan_hsa_amd_vmem_address_free(void* ptr, size_t size, } return REAL(hsa_amd_vmem_address_free)(ptr, size); } + +hsa_status_t asan_hsa_amd_pointer_info(const void* ptr, + hsa_amd_pointer_info_t* info, + void* (*alloc)(size_t), + uint32_t* num_agents_accessible, + hsa_agent_t** accessible) { + void* ptr_ = get_allocator().GetBlockBegin(ptr); + AsanChunk* m = ptr_ + ? instance.GetAsanChunkByAddr(reinterpret_cast(ptr_)) + : nullptr; + if (ptr_ && m) { + hsa_status_t status = REAL(hsa_amd_pointer_info)( + ptr_, info, alloc, num_agents_accessible, accessible); + if (status == HSA_STATUS_SUCCESS && info) { + static_assert(AP_.kMetadataSize == 0, "Expression below requires this"); + // Adjust base address of agent,host and sizeInBytes so as to return + // the actual pointer information of user allocation rather than asan + // allocation. Asan allocation pointer info can be acquired using internal + // 'GetPointerInfo' + info->agentBaseAddress = reinterpret_cast( + reinterpret_cast(info->agentBaseAddress) + kPageSize_); + info->hostBaseAddress = reinterpret_cast( + reinterpret_cast(info->hostBaseAddress) + kPageSize_); + info->sizeInBytes = m->UsedSize(); + } + return status; + } + return REAL(hsa_amd_pointer_info)(ptr, info, alloc, num_agents_accessible, + accessible); +} + } // namespace __asan #endif diff --git a/compiler-rt/lib/asan/asan_allocator.h b/compiler-rt/lib/asan/asan_allocator.h index ced10f62b7a58..f33e8d3b2819e 100644 --- a/compiler-rt/lib/asan/asan_allocator.h +++ b/compiler-rt/lib/asan/asan_allocator.h @@ -341,6 +341,11 @@ hsa_status_t asan_hsa_amd_vmem_address_reserve_align(void** ptr, size_t size, BufferedStackTrace* stack); hsa_status_t asan_hsa_amd_vmem_address_free(void* ptr, size_t size, BufferedStackTrace* stack); +hsa_status_t asan_hsa_amd_pointer_info(const void* ptr, + hsa_amd_pointer_info_t* info, + void* (*alloc)(size_t), + uint32_t* num_agents_accessible, + hsa_agent_t** accessible); } // namespace __asan #endif diff --git a/compiler-rt/lib/asan/asan_interceptors.cpp b/compiler-rt/lib/asan/asan_interceptors.cpp index 0951a77b1b93e..c04d532f909b1 100644 --- a/compiler-rt/lib/asan/asan_interceptors.cpp +++ b/compiler-rt/lib/asan/asan_interceptors.cpp @@ -948,6 +948,15 @@ INTERCEPTOR(hsa_status_t, hsa_amd_vmem_address_free, void* ptr, size_t size) { return asan_hsa_amd_vmem_address_free(ptr, size, &stack); } +INTERCEPTOR(hsa_status_t, hsa_amd_pointer_info, const void* ptr, + hsa_amd_pointer_info_t* info, void* (*alloc)(size_t), + uint32_t* num_agents_accessible, hsa_agent_t** accessible) { + AsanInitFromRtl(); + ENSURE_HSA_INITED(); + return asan_hsa_amd_pointer_info(ptr, info, alloc, num_agents_accessible, + accessible); +} + void InitializeAmdgpuInterceptors() { ASAN_INTERCEPT_FUNC(hsa_memory_copy); ASAN_INTERCEPT_FUNC(hsa_amd_memory_pool_allocate); @@ -962,6 +971,7 @@ void InitializeAmdgpuInterceptors() { ASAN_INTERCEPT_FUNC(hsa_amd_ipc_memory_detach); ASAN_INTERCEPT_FUNC(hsa_amd_vmem_address_reserve_align); ASAN_INTERCEPT_FUNC(hsa_amd_vmem_address_free); + ASAN_INTERCEPT_FUNC(hsa_amd_pointer_info); } void ENSURE_HSA_INITED() {