rocr/driver: add ASAN header page management to Driver class

Add ASAN header page management to Driver

- Add ReplaceAsanHeaderPage and ReturnAsanHeaderPage to Driver interface
- Implement ASAN functions in KfdDriver using hsaKmt calls

Signed-off-by: Honglei Huang <Honglei1.Huang@amd.com>

runtime/hsa-runtime/core/driver/kfd/amd_kfd_driver.cpp

@@ -716,5 +716,19 @@ hsa_status_t KfdDriver::RegisterSharedHandle(const HsaSharedMemoryHandle* share_
   return HSA_STATUS_SUCCESS;
 }
 
+hsa_status_t KfdDriver::ReplaceAsanHeaderPage(void* mem) const {
+  if (HSAKMT_CALL(hsaKmtReplaceAsanHeaderPage(mem)) != HSAKMT_STATUS_SUCCESS) {
+    return HSA_STATUS_ERROR;
+  }
+  return HSA_STATUS_SUCCESS;
+}
+
+hsa_status_t KfdDriver::ReturnAsanHeaderPage(void* mem) const {
+  if (HSAKMT_CALL(hsaKmtReturnAsanHeaderPage(mem)) != HSAKMT_STATUS_SUCCESS) {
+    return HSA_STATUS_ERROR;
+  }
+  return HSA_STATUS_SUCCESS;
+}
+
 } // namespace AMD
 } // namespace rocr

runtime/hsa-runtime/core/inc/amd_kfd_driver.h

@@ -139,6 +139,8 @@ public:
   hsa_status_t ShareMemory(void* mem, size_t size, HsaSharedMemoryHandle* share_mem) const override;
   hsa_status_t RegisterSharedHandle(const HsaSharedMemoryHandle* share_mem, void** mem,
                                     uint64_t* size) const override;
+  hsa_status_t ReplaceAsanHeaderPage(void* mem) const override;
+  hsa_status_t ReturnAsanHeaderPage(void* mem) const override;
 
   hsa_status_t OpenSMI(uint32_t node_id, int* fd) const override;
 

runtime/hsa-runtime/core/inc/driver.h

@@ -370,6 +370,22 @@ public:
     return HSA_STATUS_ERROR_INVALID_AGENT;
   }
 
+  /// @brief Replaces the ASAN header page with a valid one.
+  /// @param[in] mem Pointer to the memory to be replaced.
+  /// @return HSA_STATUS_SUCCESS if the ASAN header page was successfully replaced, or an error
+  /// code.
+  virtual hsa_status_t ReplaceAsanHeaderPage(void* mem) const {
+    return HSA_STATUS_ERROR_INVALID_AGENT;
+  }
+
+  /// @brief Returns the ASAN header page to its original state.
+  /// @param[in] mem Pointer to the memory to be returned.
+  /// @return HSA_STATUS_SUCCESS if the ASAN header page was successfully returned, or an error
+  /// code.
+  virtual hsa_status_t ReturnAsanHeaderPage(void* mem) const {
+    return HSA_STATUS_ERROR_INVALID_AGENT;
+  }
+
   /// Unique identifier for supported kernel-mode drivers.
   const DriverType kernel_driver_type_;
 

runtime/hsa-runtime/core/runtime/runtime.cpp

@@ -368,7 +368,7 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {
   }
 
   if (alloc_flags & core::MemoryRegion::AllocateAsan)
-    assert(HSAKMT_CALL(hsaKmtReturnAsanHeaderPage(ptr)) == HSAKMT_STATUS_SUCCESS);
+    assert(region->owner()->driver().ReturnAsanHeaderPage(ptr) == HSA_STATUS_SUCCESS);
 
   const hsa_status_t err = region->Free(ptr, size);
   if (err != HSA_STATUS_SUCCESS) {