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;

runtime/hsa-runtime/hsacore.so.def

@@ -218,6 +218,8 @@ global:
 	hsa_amd_ipc_signal_attach;
 	hsa_amd_register_system_event_handler;
 	hsa_amd_queue_set_priority;
+	hsa_amd_register_deallocation_callback;
+	hsa_amd_deregister_deallocation_callback;
 
 local:
     *;

runtime/hsa-runtime/inc/hsa_api_trace.h

@@ -180,6 +180,8 @@ struct AmdExtTable {
   decltype(hsa_amd_memory_async_copy_rect)* hsa_amd_memory_async_copy_rect_fn;
   decltype(hsa_amd_runtime_queue_create_register)* hsa_amd_runtime_queue_create_register_fn;
   decltype(hsa_amd_memory_lock_to_pool)* hsa_amd_memory_lock_to_pool_fn;
+  decltype(hsa_amd_register_deallocation_callback)* hsa_amd_register_deallocation_callback_fn;
+  decltype(hsa_amd_deregister_deallocation_callback)* hsa_amd_deregister_deallocation_callback_fn;
 };
 
 // Table to export HSA Core Runtime Apis

runtime/hsa-runtime/inc/hsa_ext_amd.h

@@ -1814,7 +1814,7 @@ typedef hsa_status_t (*hsa_amd_system_event_callback_t)(const hsa_amd_event_t* e
  *
  * @retval HSA_STATUS_ERROR_INVALID_ARGUMENT @p event is invalid.
  */
-hsa_status_t hsa_amd_register_system_event_handler(hsa_amd_system_event_callback_t callback,
+hsa_status_t HSA_API hsa_amd_register_system_event_handler(hsa_amd_system_event_callback_t callback,
                                                    void* data);
 
 /**
@@ -1854,6 +1854,66 @@ typedef enum hsa_amd_queue_priority_s {
 hsa_status_t HSA_API hsa_amd_queue_set_priority(hsa_queue_t* queue,
                                                 hsa_amd_queue_priority_t priority);
 
+/**
+ * @brief Deallocation notifier function type.
+ */
+typedef void (*hsa_amd_deallocation_callback_t)(void* ptr, void* user_data);
+
+/**
+ * @brief Registers a deallocation notifier monitoring for release of agent
+ * accessible address @p ptr.  If successful, @p callback will be invoked when
+ * @p ptr is removed from accessibility from all agents.
+ *
+ * Notification callbacks are automatically deregistered when they are invoked.
+ *
+ * Note: The current version supports notifications of address release
+ * originating from ::hsa_amd_memory_pool_free.  Support for other address
+ * release APIs will follow.
+ *
+ * @param[in] ptr Agent accessible address to monitor for deallocation.  Passed
+ * to @p callback.
+ *
+ * @param[in] callback Notifier to be invoked when @p ptr is released from
+ * agent accessibility.
+ *
+ * @param[in] user_data User provided value passed to @p callback.  May be NULL.
+ *
+ * @retval ::HSA_STATUS_SUCCESS The notifier registered successfully
+ *
+ * @retval ::HSA_STATUS_ERROR_NOT_INITIALIZED The HSA runtime has not been
+ * initialized.
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ALLOCATION @p ptr does not refer to a valid agent accessible
+ * address.
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT @p callback is NULL or @p ptr is NULL.
+ *
+ * @retval ::HSA_STATUS_ERROR_OUT_OF_RESOURCES if there is a failure in allocating
+ * necessary resources
+ */
+hsa_status_t HSA_API hsa_amd_register_deallocation_callback(void* ptr,
+                                                    hsa_amd_deallocation_callback_t callback,
+                                                    void* user_data);
+
+/**
+ * @brief Removes a deallocation notifier previously registered with
+ * ::hsa_amd_register_deallocation_callback.  Arguments must be identical to
+ * those given in ::hsa_amd_register_deallocation_callback.
+ *
+ * @param[in] ptr Agent accessible address which was monitored for deallocation.
+ *
+ * @param[in] callback Notifier to be removed.
+ *
+ * @retval ::HSA_STATUS_SUCCESS The notifier has been removed successfully.
+ *
+ * @retval ::HSA_STATUS_ERROR_NOT_INITIALIZED The HSA runtime has not been
+ * initialized.
+ *
+ * @retval ::HSA_STATUS_ERROR_INVALID_ARGUMENT The given notifier was not registered.
+ */
+hsa_status_t HSA_API hsa_amd_deregister_deallocation_callback(void* ptr,
+                                                      hsa_amd_deallocation_callback_t callback);
+
 #ifdef __cplusplus
 }  // end extern "C" block
 #endif