[ASan] Intercept rocr api's 'hsa_amd_vmem_address_reserve_align' and … (llvm#4111)

Author: Salinas, David

compiler-rt/lib/asan/asan_allocator.cpp

@@ -1304,6 +1304,9 @@ DECLARE_REAL(hsa_status_t, hsa_amd_ipc_memory_attach,
   const hsa_amd_ipc_memory_t *handle, size_t len, uint32_t num_agents,
   const hsa_agent_t *mapping_agents, void **mapped_ptr)
 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);
 
 namespace __asan {
 
@@ -1332,21 +1335,17 @@ hsa_status_t asan_hsa_amd_memory_pool_free(
   if (p) {
     instance.Deallocate(ptr, 0, 0, stack, FROM_MALLOC);
     return HSA_STATUS_SUCCESS;
-  } else {
-    return REAL(hsa_amd_memory_pool_free)(ptr);
   }
+  return REAL(hsa_amd_memory_pool_free)(ptr);
 }
 
 hsa_status_t asan_hsa_amd_agents_allow_access(
   uint32_t num_agents, const hsa_agent_t *agents, const uint32_t *flags,
   const void *ptr,
   BufferedStackTrace *stack) {
   void *p = get_allocator().GetBlockBegin(ptr);
-  if (p) {
-    return REAL(hsa_amd_agents_allow_access)(num_agents, agents, flags, p);
-  } else {
-    return REAL(hsa_amd_agents_allow_access)(num_agents, agents, flags, ptr);
-  }
+  return REAL(hsa_amd_agents_allow_access)(num_agents, agents, flags,
+                                           p ? p : ptr);
 }
 
 // For asan allocator, kMetadataSize is 0 and maximum redzone size is 2048. This
@@ -1394,5 +1393,59 @@ hsa_status_t asan_hsa_amd_ipc_memory_detach(void *mapped_ptr) {
       reinterpret_cast<void *>(reinterpret_cast<uptr>(mapped_ptr) - kPageSize_);
   return REAL(hsa_amd_ipc_memory_detach)(mapped_ptr_);
 }
+
+hsa_status_t asan_hsa_amd_vmem_address_reserve_align(
+    void** ptr, size_t size, uint64_t address, uint64_t alignment,
+    uint64_t flags, BufferedStackTrace* stack) {
+  // Bypass the tracking for a fixed address since it cannot be supported.
+  // Reasons:
+  //  1. Address may not meet the alignment/page-size requirement.
+  //  2. Requested range overlaps an existing reserved/mapped range.
+  //  3. Insufficient VA space to honor that exact placement.
+  if (address)
+    return REAL(hsa_amd_vmem_address_reserve_align)(ptr, size, address,
+                                                    alignment, flags);
+
+  if (alignment < kPageSize_)
+    alignment = kPageSize_;
+
+  if (UNLIKELY(!IsPowerOfTwo(alignment))) {
+    errno = errno_EINVAL;
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  AmdgpuAllocationInfo aa_info;
+  aa_info.alloc_func =
+      reinterpret_cast<void*>(asan_hsa_amd_vmem_address_reserve_align);
+  aa_info.memory_pool = {0};
+  aa_info.size = size;
+  aa_info.flags64 = flags;
+  aa_info.address = 0;
+  aa_info.alignment = alignment;
+  aa_info.ptr = nullptr;
+  SetErrnoOnNull(*ptr = instance.Allocate(size, alignment, stack, FROM_MALLOC,
+                                          false, &aa_info));
+
+  return aa_info.status;
+}
+
+hsa_status_t asan_hsa_amd_vmem_address_free(void* ptr, size_t size,
+                                            BufferedStackTrace* stack) {
+  if (UNLIKELY(!IsAligned(reinterpret_cast<uptr>(ptr), kPageSize_))) {
+    errno = errno_EINVAL;
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+  if (size == 0) {
+    errno = errno_EINVAL;
+    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  }
+
+  void* p = get_allocator().GetBlockBegin(ptr);
+  if (p) {
+    instance.Deallocate(ptr, 0, 0, stack, FROM_MALLOC);
+    return HSA_STATUS_SUCCESS;
+  }
+  return REAL(hsa_amd_vmem_address_free)(ptr, size);
+}
 }  // namespace __asan
 #endif

compiler-rt/lib/asan/asan_allocator.h

@@ -320,6 +320,13 @@ hsa_status_t asan_hsa_amd_ipc_memory_attach(
   const hsa_agent_t* mapping_agents, void** mapped_ptr);
 hsa_status_t asan_hsa_amd_ipc_memory_detach(
   void* mapped_ptr);
+hsa_status_t asan_hsa_amd_vmem_address_reserve_align(void** ptr, size_t size,
+                                                     uint64_t address,
+                                                     uint64_t alignment,
+                                                     uint64_t flags,
+                                                     BufferedStackTrace* stack);
+hsa_status_t asan_hsa_amd_vmem_address_free(void* ptr, size_t size,
+                                            BufferedStackTrace* stack);
 } // namespace __asan
 #endif
 

compiler-rt/lib/asan/asan_interceptors.cpp

@@ -897,6 +897,22 @@ INTERCEPTOR(hsa_status_t, hsa_amd_ipc_memory_detach, void* mapped_ptr) {
   return asan_hsa_amd_ipc_memory_detach(mapped_ptr);
 }
 
+INTERCEPTOR(hsa_status_t, hsa_amd_vmem_address_reserve_align, void** ptr,
+            size_t size, uint64_t address, uint64_t alignment, uint64_t flags) {
+  AsanInitFromRtl();
+  ENSURE_HSA_INITED();
+  GET_STACK_TRACE_MALLOC;
+  return asan_hsa_amd_vmem_address_reserve_align(ptr, size, address, alignment,
+                                                 flags, &stack);
+}
+
+INTERCEPTOR(hsa_status_t, hsa_amd_vmem_address_free, void* ptr, size_t size) {
+  AsanInitFromRtl();
+  ENSURE_HSA_INITED();
+  GET_STACK_TRACE_FREE;
+  return asan_hsa_amd_vmem_address_free(ptr, size, &stack);
+}
+
 void InitializeAmdgpuInterceptors() {
   ASAN_INTERCEPT_FUNC(hsa_memory_copy);
   ASAN_INTERCEPT_FUNC(hsa_amd_memory_pool_allocate);
@@ -909,6 +925,8 @@ void InitializeAmdgpuInterceptors() {
   ASAN_INTERCEPT_FUNC(hsa_amd_ipc_memory_create);
   ASAN_INTERCEPT_FUNC(hsa_amd_ipc_memory_attach);
   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);
 }
 
 void ENSURE_HSA_INITED() {

compiler-rt/lib/sanitizer_common/sanitizer_allocator_amdgpu.cpp

@@ -22,6 +22,11 @@ struct HsaMemoryFunctions {
                                void *(*alloc)(size_t),
                                uint32_t *num_agents_accessible,
                                hsa_agent_t **accessible);
+  hsa_status_t (*vmem_address_reserve_align)(void** ptr, size_t size,
+                                             uint64_t address,
+                                             uint64_t alignment,
+                                             uint64_t flags);
+  hsa_status_t (*vmem_address_free)(void* ptr, size_t size);
 };
 
 static HsaMemoryFunctions hsa_amd;
@@ -37,31 +42,49 @@ bool AmdgpuMemFuncs::Init() {
       RTLD_NEXT, "hsa_amd_memory_pool_free");
   hsa_amd.pointer_info = (decltype(hsa_amd.pointer_info))dlsym(
       RTLD_NEXT, "hsa_amd_pointer_info");
+  hsa_amd.vmem_address_reserve_align =
+      (decltype(hsa_amd.vmem_address_reserve_align))dlsym(
+          RTLD_NEXT, "hsa_amd_vmem_address_reserve_align");
+  hsa_amd.vmem_address_free = (decltype(hsa_amd.vmem_address_free))dlsym(
+      RTLD_NEXT, "hsa_amd_vmem_address_free");
   if (!hsa_amd.memory_pool_allocate || !hsa_amd.memory_pool_free ||
-      !hsa_amd.pointer_info)
+      !hsa_amd.pointer_info || !hsa_amd.vmem_address_reserve_align ||
+      !hsa_amd.vmem_address_free)
     return false;
-  else
-    return true;
+  return true;
 }
 
 void *AmdgpuMemFuncs::Allocate(uptr size, uptr alignment,
                                DeviceAllocationInfo *da_info) {
   AmdgpuAllocationInfo *aa_info =
       reinterpret_cast<AmdgpuAllocationInfo *>(da_info);
-
-  aa_info->status = hsa_amd.memory_pool_allocate(aa_info->memory_pool, size,
-                                                 aa_info->flags, &aa_info->ptr);
+  if (!aa_info->memory_pool.handle) {
+    aa_info->status = hsa_amd.vmem_address_reserve_align(
+        &aa_info->ptr, size, aa_info->address, aa_info->alignment,
+        aa_info->flags64);
+  } else {
+    aa_info->status = hsa_amd.memory_pool_allocate(
+        aa_info->memory_pool, size, aa_info->flags, &aa_info->ptr);
+  }
   if (aa_info->status != HSA_STATUS_SUCCESS)
     return nullptr;
 
   return aa_info->ptr;
 }
 
 void AmdgpuMemFuncs::Deallocate(void *p) {
-  UNUSED hsa_status_t status = hsa_amd.memory_pool_free(p);
+  DevicePointerInfo DevPtrInfo;
+  if (AmdgpuMemFuncs::GetPointerInfo(reinterpret_cast<uptr>(p), &DevPtrInfo)) {
+    if (DevPtrInfo.type == HSA_EXT_POINTER_TYPE_HSA) {
+      UNUSED hsa_status_t status = hsa_amd.memory_pool_free(p);
+    } else if (DevPtrInfo.type == HSA_EXT_POINTER_TYPE_RESERVED_ADDR) {
+      UNUSED hsa_status_t status =
+          hsa_amd.vmem_address_free(p, DevPtrInfo.map_size);
+    }
+  }
 }
 
-bool AmdgpuMemFuncs::GetPointerInfo(uptr ptr, DevivePointerInfo *ptr_info) {
+bool AmdgpuMemFuncs::GetPointerInfo(uptr ptr, DevicePointerInfo* ptr_info) {
   hsa_amd_pointer_info_t info;
   info.size = sizeof(hsa_amd_pointer_info_t);
   hsa_status_t status =
@@ -70,8 +93,12 @@ bool AmdgpuMemFuncs::GetPointerInfo(uptr ptr, DevivePointerInfo *ptr_info) {
   if (status != HSA_STATUS_SUCCESS)
     return false;
 
-  ptr_info->map_beg = reinterpret_cast<uptr>(info.agentBaseAddress);
+  if (info.type == HSA_EXT_POINTER_TYPE_RESERVED_ADDR)
+    ptr_info->map_beg = reinterpret_cast<uptr>(info.hostBaseAddress);
+  else if (info.type == HSA_EXT_POINTER_TYPE_HSA)
+    ptr_info->map_beg = reinterpret_cast<uptr>(info.agentBaseAddress);
   ptr_info->map_size = info.sizeInBytes;
+  ptr_info->type = reinterpret_cast<hsa_amd_pointer_type_t>(info.type);
 
   return true;
 }

compiler-rt/lib/sanitizer_common/sanitizer_allocator_amdgpu.h

@@ -20,7 +20,7 @@ class AmdgpuMemFuncs {
   static void *Allocate(uptr size, uptr alignment,
                         DeviceAllocationInfo *da_info);
   static void Deallocate(void *p);
-  static bool GetPointerInfo(uptr ptr, DevivePointerInfo *ptr_info);
+  static bool GetPointerInfo(uptr ptr, DevicePointerInfo* ptr_info);
   static uptr GetPageSize();
 };
 
@@ -32,8 +32,11 @@ struct AmdgpuAllocationInfo : public DeviceAllocationInfo {
   hsa_status_t status;
   void *alloc_func;
   hsa_amd_memory_pool_t memory_pool;
-  size_t size;
-  uint32_t flags;
+  u64 alignment;
+  u64 address;
+  u64 flags64;
+  usize size;
+  u32 flags;
   void *ptr;
 };
 #endif  // SANITIZER_AMDGPU

compiler-rt/lib/sanitizer_common/sanitizer_allocator_device.h

@@ -31,7 +31,8 @@ struct DeviceAllocationInfo {
   DeviceAllocationType type_;
 };
 
-struct DevivePointerInfo {
+struct DevicePointerInfo {
+  u64 type;
   uptr map_beg;
   uptr map_size;
 };
@@ -165,7 +166,7 @@ class DeviceAllocatorT {
              : nullptr;
   }
 
-  void *GetBlockBegin(const void *ptr) const {
+  void* GetBlockBegin(const void* ptr) const {
     Header header;
     if (!mem_funcs_inited_) return nullptr;
     uptr p = reinterpret_cast<uptr>(ptr);
@@ -182,7 +183,7 @@ class DeviceAllocatorT {
     if (!nearest_chunk)
       return nullptr;
     if (p != nearest_chunk) {
-      Header *h = GetHeader(nearest_chunk, &header);
+      Header* h = GetHeader(nearest_chunk, &header);
       CHECK_GE(nearest_chunk, h->map_beg);
       CHECK_LT(nearest_chunk, h->map_beg + h->map_size);
       CHECK_LE(nearest_chunk, p);
@@ -297,7 +298,7 @@ class DeviceAllocatorT {
     return mem_funcs_inited_;
   }
 
-  typedef DevivePointerInfo Header;
+  typedef DevicePointerInfo Header;
 
   Header *GetHeaderAnyPointer(uptr p, Header* h) const {
     CHECK(IsAligned(p, page_size_));