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/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;