[compiler-rt] [compiler-rt][asan] Add AMDGPU ASan support via HSA API interceptors. (PR #192240)

[Amit Kumar Pandey via llvm-commits]

================
@@ -0,0 +1,209 @@
+//===-- sanitizer_allocator_amdgpu.cpp --------------------------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// Part of the Sanitizer Allocator.
+//
+//===----------------------------------------------------------------------===//
+#if SANITIZER_AMDGPU
+#  include <dlfcn.h>  // For dlsym
+
+#  include "sanitizer_allocator.h"
+#  include "sanitizer_atomic.h"
+
+namespace __sanitizer {
+struct HsaFunctions {
+  // -------------- Memory Allocate/Deallocate Functions ----------------
+  hsa_status_t (*memory_pool_allocate)(hsa_amd_memory_pool_t memory_pool,
+                                       size_t size, uint32_t flags, void** ptr);
+  hsa_status_t (*memory_pool_free)(void* ptr);
+  hsa_status_t (*pointer_info)(void* ptr, hsa_amd_pointer_info_t* info,
+                               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);
+
+  // ----------------- System Event Register Function -------------------
+  hsa_status_t (*register_system_event_handler)(
+      hsa_amd_system_event_callback_t callback, void* data);
+};
+
+static HsaFunctions hsa_amd;
+
+// Always align to page boundary to match current ROCr behavior
+static const size_t kPageSize_ = 4096;
+
+static atomic_uint8_t amdgpu_runtime_shutdown{0};
+static atomic_uint8_t amdgpu_event_registered{0};
+
+#  define LOAD_HSA_FUNC_WITH_ERROR_CHECK(func, name, success)         \
+    func = (decltype(func))dlsym(RTLD_NEXT, name);                    \
+    if (!func) {                                                      \
+      VReport(2, "Amdgpu Init: Failed to load " #name " function\n"); \
+      success = false;                                                \
+    }
+
+// Check AMDGPU runtime shutdown state
+bool AmdgpuMemFuncs::IsAmdgpuRuntimeShutdown() {
+  return static_cast<bool>(
+      atomic_load(&amdgpu_runtime_shutdown, memory_order_acquire));
+}
+
+// Notify AMDGPU runtime shutdown to allocator
+void AmdgpuMemFuncs::NotifyAmdgpuRuntimeShutdown() {
+  uint8_t shutdown = 0;
+  if (atomic_compare_exchange_strong(&amdgpu_runtime_shutdown, &shutdown, 1,
+                                     memory_order_acq_rel)) {
+    VReport(2, "Amdgpu Allocator: AMDGPU runtime shutdown detected\n");
+  }
+}
+
+// Clear shutdown state when hsa_init() succeeds again (re-init after shutdown).
+// Resets amdgpu_runtime_shutdown so allocator operations are enabled, and
+// amdgpu_event_registered so RegisterSystemEventHandlers() will register the
+// shutdown callback for the new runtime instance.
+void AmdgpuMemFuncs::ClearAmdgpuRuntimeShutdownState() {
+  atomic_store(&amdgpu_runtime_shutdown, 0, memory_order_release);
+  atomic_store(&amdgpu_event_registered, 0, memory_order_release);
+}
+
+bool AmdgpuMemFuncs::Init() {
+  bool success = true;
+  LOAD_HSA_FUNC_WITH_ERROR_CHECK(hsa_amd.memory_pool_allocate,
+                                 "hsa_amd_memory_pool_allocate", success);
+  LOAD_HSA_FUNC_WITH_ERROR_CHECK(hsa_amd.memory_pool_free,
+                                 "hsa_amd_memory_pool_free", success);
+  LOAD_HSA_FUNC_WITH_ERROR_CHECK(hsa_amd.pointer_info, "hsa_amd_pointer_info",
+                                 success);
+  LOAD_HSA_FUNC_WITH_ERROR_CHECK(hsa_amd.vmem_address_reserve_align,
+                                 "hsa_amd_vmem_address_reserve_align", success);
+  LOAD_HSA_FUNC_WITH_ERROR_CHECK(hsa_amd.vmem_address_free,
+                                 "hsa_amd_vmem_address_free", success);
+  LOAD_HSA_FUNC_WITH_ERROR_CHECK(hsa_amd.register_system_event_handler,
+                                 "hsa_amd_register_system_event_handler",
+                                 success);
+  if (!success) {
+    VReport(1, "Amdgpu Init: Failed to load AMDGPU runtime functions\n");
+    return false;
+  }
+  return true;
+}
+
+void* AmdgpuMemFuncs::Allocate(uptr size, uptr alignment,
+                               DeviceAllocationInfo* da_info) {
+  // Do not allocate if AMDGPU runtime is shutdown
+  if (UNLIKELY(IsAmdgpuRuntimeShutdown())) {
+    VReport(1,
+            "Amdgpu Allocate: Runtime shutdown, skipping allocation for size "
+            "%zu alignment %zu\n",
+            size, alignment);
+    return nullptr;
----------------
ampandey-1995 wrote:

Resolved in commit https://github.com/llvm/llvm-project/commit/f7830b34a0ba295dabbed03700e222028b5ac8c8

https://github.com/llvm/llvm-project/pull/192240