Improve ASan implementations of hsa api calls based on @bing-ma

suggestions.

This commit again adds the interception of 'hsa_amd_pointer_info' call.

Author: ampandey-AMD

compiler-rt/lib/asan/asan_allocator.cpp

@@ -1400,6 +1400,9 @@ 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_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 {
 
@@ -1454,21 +1457,18 @@ static struct AP32<LocalAddressSpaceView> AP_;
 
 hsa_status_t asan_hsa_amd_ipc_memory_create(void* ptr, size_t len,
                                             hsa_amd_ipc_memory_t* handle) {
-  static_assert(AP_.kMetadataSize == 0, "Expression below requires this");
   void* ptr_ = get_allocator().GetBlockBegin(ptr);
-  size_t len_ = get_allocator().GetActuallyAllocatedSize(ptr);
-
-  uptr p = reinterpret_cast<uptr>(ptr);
-  uptr p_ = reinterpret_cast<uptr>(ptr_);
-
-  if (p == p_)
+  AsanChunk* m = ptr_
+                     ? instance.GetAsanChunkByAddr(reinterpret_cast<uptr>(ptr_))
+                     : nullptr;
+  if (ptr_ && m) {
+    static_assert(AP_.kMetadataSize == 0, "Expression below requires this");
+    uptr p = reinterpret_cast<uptr>(ptr);
+    uptr p_ = reinterpret_cast<uptr>(ptr_);
+    size_t len_ = len;
+    if (p == p_ + kPageSize_ && len == m->UsedSize())
+      len_ = get_allocator().GetActuallyAllocatedSize(ptr);
     return REAL(hsa_amd_ipc_memory_create)(ptr_, len_, handle);
-
-  if (p == p_ + kPageSize_) {
-    AsanChunk* m = instance.GetAsanChunkByAddr(p_);
-    if (m && len == m->UsedSize())
-      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);
 }
@@ -1547,5 +1547,28 @@ 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* p = get_allocator().GetBlockBegin(ptr);
+  AsanChunk* m = instance.GetAsanChunkByAddr(reinterpret_cast<uptr>(p));
+  hsa_status_t status;
+  if (p && m)
+    status = REAL(hsa_amd_pointer_info)(ptr, info, alloc, num_agents_accessible,
+                                        accessible);
+  if (status == HSA_STATUS_SUCCESS && info && p && m) {
+    static_assert(AP_.kMetadataSize == 0, "Expression below requires this");
+    info->agentBaseAddress = reinterpret_cast<void*>(
+        reinterpret_cast<uptr>(info->agentBaseAddress) + kPageSize_);
+    info->hostBaseAddress = reinterpret_cast<void*>(
+        reinterpret_cast<uptr>(info->hostBaseAddress) + kPageSize_);
+    info->sizeInBytes = m->UsedSize();
+  }
+  return status;
+}
+
 }  // namespace __asan
 #endif

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
 

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() {