Initial support for deallocation callbacks.

Adds hsa_amd_register_deallocation_callback and hsa_amd_deregister_deallocation_callback
to notify when HSA memory has been released.

Change-Id: I1f33cee250ca890e5c2e7fddfa4479aa5874651d

runtime/hsa-runtime/core/common/hsa_table_interface.cpp

@@ -1141,3 +1141,16 @@ hsa_status_t HSA_API hsa_amd_queue_set_priority(hsa_queue_t* queue,
                                                 hsa_amd_queue_priority_t priority) {
   return amdExtTable->hsa_amd_queue_set_priority_fn(queue, priority);
 }
+
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_register_deallocation_callback(void* ptr,
+                                                    hsa_amd_deallocation_callback_t callback,
+                                                    void* user_data) {
+  return amdExtTable->hsa_amd_register_deallocation_callback_fn(ptr, callback, user_data);
+}
+
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_deregister_deallocation_callback(void* ptr,
+                                                      hsa_amd_deallocation_callback_t callback) {
+  return amdExtTable->hsa_amd_deregister_deallocation_callback_fn(ptr, callback);
+}

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

@@ -75,6 +75,9 @@ template <class R, class... Args> class callback_t<R (*)(Args...)> {
   callback_t(func_t function_ptr) : function(function_ptr) {}
   callback_t& operator=(func_t function_ptr) { function = function_ptr; return *this; }
 
+  bool operator==(func_t function_ptr) { return function == function_ptr; }
+  bool operator!=(func_t function_ptr) { return function != function_ptr; }
+
   // Allows common function pointer idioms, such as if( func != nullptr )...
   // without allowing silent reversion to the original function pointer type.
   operator void*() { return reinterpret_cast<void*>(function); }

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

@@ -199,37 +199,48 @@ hsa_status_t HSA_API hsa_amd_interop_map_buffer(uint32_t num_agents,
 hsa_status_t HSA_API hsa_amd_interop_unmap_buffer(void* ptr);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_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 HSA_API hsa_amd_pointer_info(void* ptr, hsa_amd_pointer_info_t* info,
+                                          void* (*alloc)(size_t), uint32_t* num_agents_accessible,
+                                          hsa_agent_t** accessible);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_pointer_info_set_userdata(void* ptr, void* userdata);
+hsa_status_t HSA_API hsa_amd_pointer_info_set_userdata(void* ptr, void* userdata);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_ipc_memory_create(void* ptr, size_t len, hsa_amd_ipc_memory_t* handle);
+hsa_status_t HSA_API hsa_amd_ipc_memory_create(void* ptr, size_t len, hsa_amd_ipc_memory_t* handle);
 
 // Mirrors Amd Extension Apis
-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);
+hsa_status_t HSA_API 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);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_ipc_memory_detach(void* mapped_ptr);
+hsa_status_t HSA_API hsa_amd_ipc_memory_detach(void* mapped_ptr);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_ipc_signal_create(hsa_signal_t signal, hsa_amd_ipc_signal_t* handle);
+hsa_status_t HSA_API hsa_amd_ipc_signal_create(hsa_signal_t signal, hsa_amd_ipc_signal_t* handle);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_ipc_signal_attach(const hsa_amd_ipc_signal_t* handle, hsa_signal_t* signal);
+hsa_status_t HSA_API hsa_amd_ipc_signal_attach(const hsa_amd_ipc_signal_t* handle,
+                                               hsa_signal_t* signal);
 
 // Mirrors Amd Extension Apis
-hsa_status_t hsa_amd_register_system_event_handler(
-    hsa_amd_system_event_callback_t callback, void* data);
+hsa_status_t HSA_API hsa_amd_register_system_event_handler(hsa_amd_system_event_callback_t callback,
+                                                           void* data);
 
 // Mirrors Amd Extension Apis
 hsa_status_t HSA_API hsa_amd_queue_set_priority(hsa_queue_t* queue,
                                                 hsa_amd_queue_priority_t priority);
 
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_register_deallocation_callback(
+    void* ptr, hsa_amd_deallocation_callback_t callback, void* user_data);
+
+// Mirrors Amd Extension Apis
+hsa_status_t HSA_API hsa_amd_deregister_deallocation_callback(
+    void* ptr, hsa_amd_deallocation_callback_t callback);
+
 }  // end of AMD namespace
 
 #endif  // header guard

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

@@ -47,6 +47,8 @@
 
 #include <vector>
 #include <map>
+#include <memory>
+#include <tuple>
 #include <utility>
 
 #include "core/inc/hsa_ext_interface.h"
@@ -178,6 +180,11 @@ class Runtime {
   /// @retval ::HSA_STATUS_SUCCESS if @p ptr is successfully released.
   hsa_status_t FreeMemory(void* ptr);
 
+  hsa_status_t RegisterReleaseNotifier(void* ptr, hsa_amd_deallocation_callback_t callback,
+                                       void* user_data);
+
+  hsa_status_t DeregisterReleaseNotifier(void* ptr, hsa_amd_deallocation_callback_t callback);
+
   /// @brief Blocking memory copy from src to dst.
   ///
   /// @param [in] dst Memory address of the destination.
@@ -341,9 +348,16 @@ class Runtime {
     AllocationRegion(const MemoryRegion* region_arg, size_t size_arg)
         : region(region_arg), size(size_arg), user_ptr(nullptr) {}
 
+    struct notifier_t {
+      void* ptr;
+      AMD::callback_t<hsa_amd_deallocation_callback_t> callback;
+      void* user_data;
+    };
+
     const MemoryRegion* region;
     size_t size;
     void* user_ptr;
+    std::unique_ptr<std::vector<notifier_t>> notifiers;
   };
 
   struct AsyncEventsControl {

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

@@ -387,6 +387,8 @@ void HsaApiTable::UpdateAmdExts() {
   amd_ext_api.hsa_amd_memory_async_copy_rect_fn = AMD::hsa_amd_memory_async_copy_rect;
   amd_ext_api.hsa_amd_runtime_queue_create_register_fn = AMD::hsa_amd_runtime_queue_create_register;
   amd_ext_api.hsa_amd_memory_lock_to_pool_fn = AMD::hsa_amd_memory_lock_to_pool;
+  amd_ext_api.hsa_amd_register_deallocation_callback_fn = AMD::hsa_amd_register_deallocation_callback;
+  amd_ext_api.hsa_amd_deregister_deallocation_callback_fn = AMD::hsa_amd_deregister_deallocation_callback;
 }
 
 class Init {

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

@@ -961,6 +961,31 @@ hsa_status_t HSA_API hsa_amd_queue_set_priority(hsa_queue_t* queue,
   CATCH;
 }
 
+hsa_status_t hsa_amd_register_deallocation_callback(void* ptr,
+                                                    hsa_amd_deallocation_callback_t callback,
+                                                    void* user_data) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(ptr);
+  IS_BAD_PTR(callback);
+
+  return core::Runtime::runtime_singleton_->RegisterReleaseNotifier(ptr, callback, user_data);
+
+  CATCH;
+}
+
+hsa_status_t hsa_amd_deregister_deallocation_callback(void* ptr,
+                                                      hsa_amd_deallocation_callback_t callback) {
+  TRY;
+  IS_OPEN();
+  IS_BAD_PTR(ptr);
+  IS_BAD_PTR(callback);
+
+  return core::Runtime::runtime_singleton_->DeregisterReleaseNotifier(ptr, callback);
+
+  CATCH;
+}
+
 // For use by tools only - not in library export table.
 hsa_status_t hsa_amd_runtime_queue_create_register(hsa_amd_runtime_queue_notifier callback,
                                                    void* user_data) {

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

@@ -293,28 +293,93 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {
 
   const MemoryRegion* region = nullptr;
   size_t size = 0;
+  std::unique_ptr<std::vector<AllocationRegion::notifier_t>> notifiers;
+
+  {
+    ScopedAcquire<KernelMutex> lock(&memory_lock_);
+
+    std::map<const void*, AllocationRegion>::iterator it = allocation_map_.find(ptr);
+
+    if (it == allocation_map_.end()) {
+      debug_warning(false && "Can't find address in allocation map");
+      return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+    }
+    region = it->second.region;
+    size = it->second.size;
+
+    // Imported fragments can't be released with FreeMemory.
+    if (region == nullptr) {
+      assert(false && "Can't release imported memory with free.");
+      return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+    }
+
+    notifiers = std::move(it->second.notifiers);
+
+    allocation_map_.erase(it);
+
+    // Fast path to avoid doubling lock ops in the common case (no notifiers).
+    if (!notifiers) return region->Free(ptr, size);
+  }
+
+  // Notifiers can't run while holding the lock or the callback won't be able to manage memory.
+  // The memory triggering the notification has already been removed from the memory map so can't
+  // be double released during the callback.
+  for (auto& notifier : *notifiers) {
+    notifier.callback(notifier.ptr, notifier.user_data);
+  }
+
+  // Fragment allocator requires protection.
   ScopedAcquire<KernelMutex> lock(&memory_lock_);
-
-  std::map<const void*, AllocationRegion>::const_iterator it = allocation_map_.find(ptr);
-
-  if (it == allocation_map_.end()) {
-    assert(false && "Can't find address in allocation map");
-    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
-  }
-  region = it->second.region;
-  size = it->second.size;
-
-  // Imported fragments can't be released with FreeMemory.
-  if (region == nullptr) {
-    assert(false && "Can't release imported memory with free.");
-    return HSA_STATUS_ERROR_INVALID_ARGUMENT;
-  }
-
-  allocation_map_.erase(it);
-
   return region->Free(ptr, size);
 }
 
+hsa_status_t Runtime::RegisterReleaseNotifier(void* ptr, hsa_amd_deallocation_callback_t callback,
+                                              void* user_data) {
+  ScopedAcquire<KernelMutex> lock(&memory_lock_);
+  auto mem = allocation_map_.upper_bound(ptr);
+  if (mem != allocation_map_.begin()) {
+    mem--;
+
+    // No support for imported fragments yet.
+    if (mem->second.region == nullptr) return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+
+    if ((mem->first <= ptr) &&
+        (ptr < reinterpret_cast<const uint8_t*>(mem->first) + mem->second.size)) {
+      auto& notifiers = mem->second.notifiers;
+      if (!notifiers) notifiers.reset(new std::vector<AllocationRegion::notifier_t>);
+      AllocationRegion::notifier_t notifier = {
+          ptr, AMD::callback_t<hsa_amd_deallocation_callback_t>(callback), user_data};
+      notifiers->push_back(notifier);
+      return HSA_STATUS_SUCCESS;
+    }
+  }
+  return HSA_STATUS_ERROR_INVALID_ALLOCATION;
+}
+
+hsa_status_t Runtime::DeregisterReleaseNotifier(void* ptr,
+                                                hsa_amd_deallocation_callback_t callback) {
+  hsa_status_t ret = HSA_STATUS_ERROR_INVALID_ARGUMENT;
+  ScopedAcquire<KernelMutex> lock(&memory_lock_);
+  auto mem = allocation_map_.upper_bound(ptr);
+  if (mem != allocation_map_.begin()) {
+    mem--;
+    if ((mem->first <= ptr) &&
+        (ptr < reinterpret_cast<const uint8_t*>(mem->first) + mem->second.size)) {
+      auto& notifiers = mem->second.notifiers;
+      if (!notifiers) return HSA_STATUS_ERROR_INVALID_ARGUMENT;
+      for (size_t i = 0; i < notifiers->size(); i++) {
+        if (((*notifiers)[i].ptr == ptr) && ((*notifiers)[i].callback) == callback) {
+          (*notifiers)[i] = std::move((*notifiers)[notifiers->size() - 1]);
+          notifiers->pop_back();
+          i--;
+          ret = HSA_STATUS_SUCCESS;
+        }
+      }
+    }
+  }
+  return ret;
+}
+
 hsa_status_t Runtime::CopyMemory(void* dst, const void* src, size_t size) {
   // Choose agents from pointer info
   bool is_src_system = false;