SWDEV-311271 - Initial mempool implementation

HIP_MEM_POOL_SUPPORT controls memory pool support in runtime.
Currently it's disabled by default. The initial change doesn't
include: IPC, MGPU, virtual memory alloc, suballoc, defragmentation,
internal dependencies.

Change-Id: Ibed8528ebec698b045ebb247e49c0ecd6e587ed7

hipamd/src/CMakeLists.txt

@@ -112,6 +112,8 @@ target_sources(amdhip64 PRIVATE
   hip_hmm.cpp
   hip_intercept.cpp
   hip_memory.cpp
+  hip_mempool.cpp
+  hip_mempool_impl.cpp
   hip_module.cpp
   hip_peer.cpp
   hip_platform.cpp

hipamd/src/amdhip.def

@@ -83,6 +83,23 @@ hipMalloc
 hipMalloc3D
 hipMalloc3DArray
 hipMallocManaged
+hipDeviceGetDefaultMemPool
+hipDeviceSetMemPool
+hipDeviceGetMemPool
+hipMallocAsync
+hipFreeAsync
+hipMemPoolTrimTo
+hipMemPoolSetAttribute
+hipMemPoolGetAttribute
+hipMemPoolSetAccess
+hipMemPoolGetAccess
+hipMemPoolCreate
+hipMemPoolDestroy
+hipMallocFromPoolAsync
+hipMemPoolExportToShareableHandle
+hipMemPoolImportFromShareableHandle
+hipMemPoolExportPointer
+hipMemPoolImportPointer
 hipArrayCreate
 hipArray3DCreate
 hipArrayDestroy

hipamd/src/hip_context.cpp

@@ -63,7 +63,11 @@ bool init() {
     if (context && CL_SUCCESS != context->create(nullptr)) {
       context->release();
     } else {
-      g_devices.push_back(new Device(context, i));
+      auto device = new Device(context, i);
+      if ((device == nullptr) || !device->Create()) {
+        return false;
+      }
+      g_devices.push_back(device);
     }
   }
 

hipamd/src/hip_device.cpp

@@ -21,6 +21,7 @@
 #include <hip/hip_runtime.h>
 
 #include "hip_internal.hpp"
+#include "hip_mempool_impl.hpp"
 
 namespace hip {
 
@@ -35,6 +36,58 @@ amd::HostQueue* Device::NullStream(bool skip_alloc) {
   return null_queue;
 }
 
+// ================================================================================================
+bool Device::Create() {
+  // Create default memory pool
+  default_mem_pool_ = new MemoryPool(this);
+  if (default_mem_pool_ == nullptr) {
+    return false;
+  }
+  // Current is default pool after device creation
+  current_mem_pool_ = default_mem_pool_;
+  return true;
+}
+
+// ================================================================================================
+void Device::AddMemoryPool(MemoryPool* pool) {
+  if (auto it = mem_pools_.find(pool); it == mem_pools_.end()) {
+    mem_pools_.insert(pool);
+  }
+}
+
+// ================================================================================================
+void Device::RemoveMemoryPool(MemoryPool* pool) {
+  if (auto it = mem_pools_.find(pool); it != mem_pools_.end()) {
+    mem_pools_.erase(it);
+  }
+}
+
+// ================================================================================================
+bool Device::FreeMemory(amd::Memory* memory, Stream* stream) {
+  // Search for memory in the entire list of pools
+  for (auto& it : mem_pools_) {
+    if (it->FreeMemory(memory, stream)) {
+      return true;
+    }
+  }
+  return false;
+}
+
+// ================================================================================================
+void Device::ReleaseFreedMemory(Stream* stream) {
+  // Search for memory in the entire list of pools
+  for (auto& it : mem_pools_) {
+    it->ReleaseFreedMemory(stream);
+  }
+}
+
+// ================================================================================================
+Device::~Device() {
+  if (default_mem_pool_ != nullptr) {
+    default_mem_pool_->release();
+  }
+}
+
 }
 
 namespace amd {

hipamd/src/hip_device_runtime.cpp

@@ -329,6 +329,9 @@ hipError_t hipDeviceGetAttribute(int* pi, hipDeviceAttribute_t attr, int device)
   case hipDeviceAttributeFineGrainSupport:
     *pi = static_cast<int>(g_devices[device]->devices()[0]->isFineGrainSupported());
     break;
+  case hipDeviceAttributeMemoryPoolsSupported:
+    *pi = HIP_MEM_POOL_SUPPORT;
+    break;
   default:
     HIP_RETURN(hipErrorInvalidValue);
   }

hipamd/src/hip_hcc.def.in

@@ -83,6 +83,23 @@ hipMalloc
 hipMalloc3D
 hipMalloc3DArray
 hipMallocManaged
+hipDeviceGetDefaultMemPool
+hipDeviceSetMemPool
+hipDeviceGetMemPool
+hipMallocAsync
+hipFreeAsync
+hipMemPoolTrimTo
+hipMemPoolSetAttribute
+hipMemPoolGetAttribute
+hipMemPoolSetAccess
+hipMemPoolGetAccess
+hipMemPoolCreate
+hipMemPoolDestroy
+hipMallocFromPoolAsync
+hipMemPoolExportToShareableHandle
+hipMemPoolImportFromShareableHandle
+hipMemPoolExportPointer
+hipMemPoolImportPointer
 hipArrayCreate
 hipArray3DCreate
 hipArrayDestroy

hipamd/src/hip_hcc.map.in

@@ -400,6 +400,23 @@ local:
 hip_5.1 {
 global:
     hipDeviceGetUuid;
+    hipDeviceGetDefaultMemPool;
+    hipDeviceSetMemPool;
+    hipDeviceGetMemPool;
+    hipMallocAsync;
+    hipFreeAsync;
+    hipMemPoolTrimTo;
+    hipMemPoolSetAttribute;
+    hipMemPoolGetAttribute;
+    hipMemPoolSetAccess;
+    hipMemPoolGetAccess;
+    hipMemPoolCreate;
+    hipMemPoolDestroy;
+    hipMallocFromPoolAsync;
+    hipMemPoolExportToShareableHandle;
+    hipMemPoolImportFromShareableHandle;
+    hipMemPoolExportPointer;
+    hipMemPoolImportPointer;
 local:
     *;
 } hip_5.0;

hipamd/src/hip_internal.hpp

@@ -1,4 +1,4 @@
-/* Copyright (c) 2015 - 2021 Advanced Micro Devices, Inc.
+/* Copyright (c) 2015 - 2022 Advanced Micro Devices, Inc.
 
  Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
@@ -195,6 +195,7 @@ class accelerator_view;
 };
 namespace hip {
   class Device;
+  class MemoryPool;
   class Stream {
   public:
     enum Priority : int { High = -1, Normal = 0, Low = 1 };
@@ -244,6 +245,8 @@ namespace hip {
     void Finish() const;
     /// Get device ID associated with the current stream;
     int DeviceId() const;
+    /// Get HIP device associated with the stream
+    Device* GetDevice() const { return device_; }
     /// Get device ID associated with a stream;
     static int DeviceId(const hipStream_t hStream);
     /// Returns if stream is null stream
@@ -344,15 +347,23 @@ namespace hip {
 
     std::vector<amd::HostQueue*> queues_;
 
+    MemoryPool* default_mem_pool_;
+    MemoryPool* current_mem_pool_;
+
+    std::set<MemoryPool*> mem_pools_;
+
   public:
     Device(amd::Context* ctx, int devId): context_(ctx),
-         deviceId_(devId),
-         null_stream_(this, Stream::Priority::Normal, 0, true),
+        deviceId_(devId),
+        null_stream_(this, Stream::Priority::Normal, 0, true),
          flags_(hipDeviceScheduleSpin),
-         isActive_(false)
+        isActive_(false),
+        default_mem_pool_(nullptr),
+        current_mem_pool_(nullptr)
         { assert(ctx != nullptr); }
-    ~Device() {}
+    ~Device();
 
+    bool Create();
     amd::Context* asContext() const { return context_; }
     int deviceId() const { return deviceId_; }
     void retain() const { context_->retain(); }
@@ -397,6 +408,29 @@ namespace hip {
       }
       return false;
     }
+
+    /// Set the current memory pool on the device
+    void SetCurrentMemoryPool(MemoryPool* pool = nullptr) {
+      current_mem_pool_ = (pool == nullptr) ? default_mem_pool_ : pool;
+    }
+
+    /// Get the current memory pool on the device
+    MemoryPool* GetCurrentMemoryPool() const { return current_mem_pool_; }
+
+    /// Get the default memory pool on the device
+    MemoryPool* GetDefaultMemoryPool() const { return default_mem_pool_; }
+
+    /// Add memory pool to the device
+    void AddMemoryPool(MemoryPool* pool);
+
+    /// Remove memory pool from the device
+    void RemoveMemoryPool(MemoryPool* pool);
+
+    /// Free memory from the device
+    bool FreeMemory(amd::Memory* memory, Stream* stream);
+
+    /// Release freed memory from all pools on the current device
+    void ReleaseFreedMemory(Stream* stream);
   };
 
   /// Current thread's device

hipamd/src/hip_mempool.cpp

@@ -0,0 +1,229 @@
+/* Copyright (c) 2022 Advanced Micro Devices, Inc.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE. */
+
+#include "hip_mempool_impl.hpp"
+
+/**
+ * API interfaces
+ */
+extern hipError_t ihipFree(void* ptr);
+
+// ================================================================================================
+hipError_t hipDeviceGetDefaultMemPool(hipMemPool_t* mem_pool, int device) {
+  HIP_INIT_API(hipDeviceGetDefaultMemPool, mem_pool, device);
+  if ((mem_pool == nullptr) || (device >= g_devices.size())) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  *mem_pool = reinterpret_cast<hipMemPool_t>(g_devices[device]->GetDefaultMemoryPool());
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipDeviceSetMemPool(int device, hipMemPool_t mem_pool) {
+  HIP_INIT_API(hipDeviceSetMemPool, device, mem_pool);
+  if ((mem_pool == nullptr) || (device >= g_devices.size())) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  g_devices[device]->SetCurrentMemoryPool(reinterpret_cast<hip::MemoryPool*>(mem_pool));
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipDeviceGetMemPool(hipMemPool_t* mem_pool, int device) {
+  HIP_INIT_API(hipDeviceGetMemPool, mem_pool, device);
+    if ((mem_pool == nullptr) || (device >= g_devices.size())) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  *mem_pool = reinterpret_cast<hipMemPool_t>(g_devices[device]->GetCurrentMemoryPool());
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMallocAsync(void** dev_ptr, size_t size, hipStream_t stream) {
+  HIP_INIT_API(hipMallocAsync, dev_ptr, size, stream);
+  if ((dev_ptr == nullptr) || (size == 0) || (stream == nullptr)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto device = reinterpret_cast<hip::Stream*>(stream)->GetDevice();
+  auto mem_pool = device->GetCurrentMemoryPool();
+  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(
+      size, reinterpret_cast<hip::Stream*>(stream));
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipFreeAsync(void* dev_ptr, hipStream_t stream) {
+  HIP_INIT_API(hipFreeAsync, dev_ptr, stream);
+  if ((dev_ptr == nullptr) || (stream == nullptr)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  size_t offset = 0;
+  auto memory = getMemoryObject(dev_ptr, offset);
+  auto id = memory->getUserData().deviceId;
+  if (!g_devices[id]->FreeMemory(memory, reinterpret_cast<hip::Stream*>(stream))) {
+    //! @todo It's not the most optimal logic. The current implementation has unconditional waits
+    HIP_RETURN(ihipFree(dev_ptr));
+   }
+
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolTrimTo(hipMemPool_t mem_pool, size_t min_bytes_to_hold) {
+  HIP_INIT_API(hipMemPoolTrimTo, mem_pool, min_bytes_to_hold);
+  if (mem_pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  hip::MemoryPool* hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+  hip_mem_pool->TrimTo(min_bytes_to_hold);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolSetAttribute(hipMemPool_t mem_pool, hipMemPoolAttr attr, void* value) {
+  HIP_INIT_API(hipMemPoolSetAttribute, mem_pool, attr, value);
+  if (mem_pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+  HIP_RETURN(hip_mem_pool->SetAttribute(attr, value));
+}
+
+// ================================================================================================
+hipError_t hipMemPoolGetAttribute(hipMemPool_t mem_pool, hipMemPoolAttr attr, void* value) {
+  HIP_INIT_API(hipMemPoolGetAttribute, mem_pool, attr, value);
+  if (mem_pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  if (mem_pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+  HIP_RETURN(hip_mem_pool->GetAttribute(attr, value));
+}
+
+// ================================================================================================
+hipError_t hipMemPoolSetAccess(
+    hipMemPool_t mem_pool,
+    const hipMemAccessDesc* desc_list,
+    size_t count) {
+  HIP_INIT_API(hipMemPoolSetAccess, mem_pool, desc_list, count);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolGetAccess(
+    hipMemAccessFlags* flags,
+    hipMemPool_t mem_pool,
+    hipMemLocation* location) {
+  HIP_INIT_API(hipMemPoolGetAccess, flags, mem_pool, location);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolCreate(hipMemPool_t* mem_pool, const hipMemPoolProps* pool_props) {
+  HIP_INIT_API(hipMemPoolCreate, mem_pool, pool_props);
+  if (mem_pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  // Make sure the pool creation occurs on a valid device
+  if ((pool_props->location.type != hipMemLocationTypeDevice) ||
+      (pool_props->location.id >= g_devices.size())) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  auto device = g_devices[pool_props->location.id];
+  auto pool = new hip::MemoryPool(device);
+  if (pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  *mem_pool = reinterpret_cast<hipMemPool_t>(pool);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolDestroy(hipMemPool_t mem_pool) {
+  HIP_INIT_API(hipMemPoolDestroy, mem_pool);
+  if (mem_pool == nullptr) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  hip::MemoryPool* hip_mem_pool = reinterpret_cast<hip::MemoryPool*>(mem_pool);
+  hip_mem_pool->ReleaseFreedMemory();
+
+  auto device = hip_mem_pool->Device();
+
+  // Force default pool if the current one is destroyed
+  if (hip_mem_pool == device->GetCurrentMemoryPool()) {
+    device->SetCurrentMemoryPool(device->GetDefaultMemoryPool());
+  }
+
+  hip_mem_pool->release();
+
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMallocFromPoolAsync(
+    void** dev_ptr,
+    size_t size,
+    hipMemPool_t mem_pool,
+    hipStream_t stream) {
+  HIP_INIT_API(hipMallocFromPoolAsync, dev_ptr, size, mem_pool, stream);
+  if ((dev_ptr == nullptr) || (size == 0) || (mem_pool == nullptr) || (stream == nullptr)) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(
+      size, reinterpret_cast<hip::Stream*>(stream));
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolExportToShareableHandle(
+    void*                      shared_handle,
+    hipMemPool_t               mem_pool,
+    hipMemAllocationHandleType handle_type,
+    unsigned int               flags) {
+  HIP_INIT_API(hipMemPoolExportToShareableHandle, shared_handle, mem_pool, handle_type, flags);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolImportFromShareableHandle(
+    hipMemPool_t*              mem_pool,
+    void*                      shared_handle,
+    hipMemAllocationHandleType handle_type,
+    unsigned int               flags) {
+  HIP_INIT_API(hipMemPoolImportFromShareableHandle, mem_pool, shared_handle, handle_type, flags);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolExportPointer(hipMemPoolPtrExportData* export_data, void* ptr) {
+  HIP_INIT_API(hipMemPoolExportPointer, export_data, ptr);
+  HIP_RETURN(hipSuccess);
+}
+
+// ================================================================================================
+hipError_t hipMemPoolImportPointer(
+    void**                   ptr,
+    hipMemPool_t             mem_pool,
+    hipMemPoolPtrExportData* export_data) {
+  HIP_INIT_API(hipMemPoolImportPointer, ptr, mem_pool, export_data);
+  HIP_RETURN(hipSuccess);
+}

hipamd/src/hip_mempool_impl.cpp

@@ -0,0 +1,299 @@
+/* Copyright (c) 2022 Advanced Micro Devices, Inc.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE. */
+
+#include "hip_mempool_impl.hpp"
+
+namespace hip {
+
+// ================================================================================================
+void Heap::AddMemory(amd::Memory* memory, hip::Stream* stream) {
+  allocations_.insert({memory, {stream, nullptr}});
+  total_size_ += memory->getSize();
+  max_total_size_ = std::max(max_total_size_, total_size_);
+}
+
+// ================================================================================================
+void Heap::AddMemory(amd::Memory* memory, const MemoryTimestamp& ts) {
+  allocations_.insert({memory, ts});
+  total_size_ += memory->getSize();
+  max_total_size_ = std::max(max_total_size_, total_size_);
+}
+
+// ================================================================================================
+amd::Memory* Heap::FindMemory(size_t size, hip::Stream* stream, bool opportunistic) {
+  amd::Memory* memory = nullptr;
+  for (auto it = allocations_.begin(); it != allocations_.end();) {
+    // Check if size can match and it's safe to use this resource
+    if ((it->first->getSize() >= size) && (it->second.IsSafeFind(stream, opportunistic))) {
+      memory = it->first;
+      total_size_ -= memory->getSize();
+      // Remove found allocation from the map
+      it = allocations_.erase(it);
+      break;
+    } else {
+      ++it;
+    }
+  }
+  return memory;
+}
+
+// ================================================================================================
+bool Heap::RemoveMemory(amd::Memory* memory, MemoryTimestamp* ts) {
+  if (auto it = allocations_.find(memory); it != allocations_.end()) {
+    if (ts != nullptr) {
+      // Preserve timestamp info for possible reuse later
+      *ts = it->second;
+    } else {
+      // Runtime will delete the timestamp object, hence make sure HIP event is released
+      it->second.Wait();
+      it->second.SetEvent(nullptr);
+    }
+    total_size_ -= memory->getSize();
+    allocations_.erase(it);
+    return true;
+  }
+  return false;
+}
+
+// ================================================================================================
+std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator
+Heap::EraseAllocaton(std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator& it) {
+  const device::Memory* dev_mem = it->first->getDeviceMemory(*device_->devices()[0]);
+  amd::SvmBuffer::free(it->first->getContext(), reinterpret_cast<void*>(dev_mem->virtualAddress()));
+  total_size_ -= it->first->getSize();
+  // Clear HIP event
+  it->second.SetEvent(nullptr);
+  // Remove the allocation from the map
+  return allocations_.erase(it);
+}
+
+// ================================================================================================
+bool Heap::ReleaseAllMemory(size_t min_bytes_to_hold, bool safe_release) {
+  for (auto it = allocations_.begin(); it != allocations_.end();) {
+    // Make sure the heap is smaller than the minimum value to hold
+    if (total_size_ <= min_bytes_to_hold) {
+      return true;
+    }
+    // Safe release forces unconditional wait for memory
+    if (safe_release) {
+      it->second.Wait();
+    }
+    if (it->second.IsSafeRelease()) {
+      it = EraseAllocaton(it);
+    } else {
+      ++it;
+    }
+  }
+  return true;
+}
+
+// ================================================================================================
+bool Heap::ReleaseAllMemory(hip::Stream* stream) {
+  for (auto it = allocations_.begin(); it != allocations_.end();) {
+    // Make sure the heap holds the minimum number of bytes
+    if (total_size_ <= release_threshold_) {
+      return true;
+    }
+    if (it->second.IsSafeRelease()) {
+      it = EraseAllocaton(it);
+    } else {
+      ++it;
+    }
+  }
+  return true;
+}
+
+// ================================================================================================
+void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  void* dev_ptr = nullptr;
+  amd::Memory* memory = free_heap_.FindMemory(size, stream, Opportunistic());
+  if (memory == nullptr) {
+    amd::Context* context = device_->asContext();
+    const auto& dev_info = context->devices()[0]->info();
+    if (dev_info.maxMemAllocSize_ < size) {
+      return nullptr;
+    }
+
+    dev_ptr = amd::SvmBuffer::malloc(*context, 0, size, dev_info.memBaseAddrAlign_, nullptr);
+    if (dev_ptr == nullptr) {
+      size_t free = 0, total =0;
+      hipError_t err = hipMemGetInfo(&free, &total);
+      if (err == hipSuccess) {
+        LogPrintfError("Allocation failed : Device memory : required :%zu | free :%zu | total :%zu \n",
+          size, free, total);
+      }
+      return nullptr;
+    }
+
+    size_t offset = 0;
+    memory = getMemoryObject(dev_ptr, offset);
+    // Saves the current device id so that it can be accessed later
+    memory->getUserData().deviceId = device_->deviceId();
+  } else {
+    free_heap_.RemoveMemory(memory);
+    const device::Memory* dev_mem = memory->getDeviceMemory(*device_->devices()[0]);
+    dev_ptr = reinterpret_cast<void*>(dev_mem->virtualAddress());
+  }
+  // Place the allocated memory into the busy heap
+  busy_heap_.AddMemory(memory, stream);
+
+  // Increment the reference counter on the pool
+  retain();
+
+  return dev_ptr;
+}
+
+// ================================================================================================
+bool MemoryPool::FreeMemory(amd::Memory* memory, hip::Stream* stream) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  MemoryTimestamp ts;
+  // Remove memory object fro the busy pool
+  if (!busy_heap_.RemoveMemory(memory, &ts)) {
+    // This pool doesn't contain memory
+    return false;
+  }
+  // The stream of destruction is a safe stream, because the app must handle sync
+  ts.AddSafeStream(stream);
+
+  // Add a marker to the stream to trace availability of this memory
+  Event* e = new hip::Event(0);
+  if (e != nullptr) {
+    if (hipSuccess == e->addMarker(reinterpret_cast<hipStream_t>(stream), nullptr, true)) {
+      ts.SetEvent(e);
+    }
+  }
+  free_heap_.AddMemory(memory, ts);
+
+  // Decrement the reference counter on the pool
+  release();
+
+  return true;
+}
+
+// ================================================================================================
+void MemoryPool::ReleaseFreedMemory(hip::Stream* stream) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  free_heap_.ReleaseAllMemory(stream);
+}
+
+// ================================================================================================
+void MemoryPool::TrimTo(size_t min_bytes_to_hold) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  free_heap_.ReleaseAllMemory(min_bytes_to_hold);
+}
+
+// ================================================================================================
+hipError_t MemoryPool::SetAttribute(hipMemPoolAttr attr, void* value) {
+  amd::ScopedLock lock(lock_pool_ops_);
+  uint64_t reset;
+
+  switch (attr) {
+    case hipMemPoolReuseFollowEventDependencies:
+      // Enable/disable HIP events tracking from the app's dependencies
+      state_.event_dependencies_ = *reinterpret_cast<int32_t*>(value);
+      break;
+    case hipMemPoolReuseAllowOpportunistic:
+      // Enable/disable HIP event check for freed memory
+      state_.opportunistic_ = *reinterpret_cast<int32_t*>(value);
+      break;
+    case  hipMemPoolReuseAllowInternalDependencies:
+      // Enable/disable internal extra dependencies introduced in runtime
+      state_.internal_dependencies_ = *reinterpret_cast<int32_t*>(value);
+      break;
+    case hipMemPoolAttrReleaseThreshold:
+      free_heap_.SetReleaseThreshold(*reinterpret_cast<uint64_t*>(value));
+      break;
+    case hipMemPoolAttrReservedMemCurrent:
+      // Should be GetAttribute only
+      return hipErrorInvalidValue;
+      break;
+    case hipMemPoolAttrReservedMemHigh:
+      reset = *reinterpret_cast<uint64_t*>(value);
+      // Only 0 is accepted
+      if (reset != 0) {
+        return hipErrorInvalidValue;
+      }
+      free_heap_.SetMaxTotalSize(reset);
+    case hipMemPoolAttrUsedMemCurrent:
+      // Should be GetAttribute only
+      return hipErrorInvalidValue;
+      break;
+    case hipMemPoolAttrUsedMemHigh:
+      reset = *reinterpret_cast<uint64_t*>(value);
+      // Only 0 is accepted
+      if (reset != 0) {
+        return hipErrorInvalidValue;
+      }
+      busy_heap_.SetMaxTotalSize(reset);
+      break;
+    default:
+      return hipErrorInvalidValue;
+  }
+  return hipSuccess;
+}
+
+// ================================================================================================
+hipError_t MemoryPool::GetAttribute(hipMemPoolAttr attr, void* value) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  switch (attr) {
+    case hipMemPoolReuseFollowEventDependencies:
+      // Enable/disable HIP events tracking from the app's dependencies
+      *reinterpret_cast<int32_t*>(value) = EventDependencies();
+      break;
+    case hipMemPoolReuseAllowOpportunistic:
+      // Enable/disable HIP event check for freed memory
+      *reinterpret_cast<int32_t*>(value) = Opportunistic();
+      break;
+    case  hipMemPoolReuseAllowInternalDependencies:
+      // Enable/disable internal extra dependencies introduced in runtime
+      *reinterpret_cast<int32_t*>(value) = InternalDependencies();
+      break;
+    case hipMemPoolAttrReleaseThreshold:
+      *reinterpret_cast<uint64_t*>(value) = free_heap_.GetReleaseThreshold();
+      break;
+    case hipMemPoolAttrReservedMemCurrent:
+      // All allocate memory by the pool in OS
+      *reinterpret_cast<uint64_t*>(value) = busy_heap_.GetTotalSize() + free_heap_.GetTotalSize();
+      break;
+    case hipMemPoolAttrReservedMemHigh:
+      // High watermark of all allocated memory in OS, since the last reset
+      *reinterpret_cast<uint64_t*>(value) = busy_heap_.GetTotalSize() + free_heap_.GetMaxTotalSize();
+      break;
+    case hipMemPoolAttrUsedMemCurrent:
+      // Total currently used memory by the pool
+      *reinterpret_cast<uint64_t*>(value) = busy_heap_.GetTotalSize();
+      break;
+    case hipMemPoolAttrUsedMemHigh:
+      // High watermark of all used memoryS, since the last reset
+      *reinterpret_cast<uint64_t*>(value) = busy_heap_.GetMaxTotalSize();
+      break;
+    default:
+      return hipErrorInvalidValue;
+  }
+  return hipSuccess;
+}
+
+}

hipamd/src/hip_mempool_impl.hpp

@@ -0,0 +1,212 @@
+/* Copyright (c) 2022 Advanced Micro Devices, Inc.
+
+ Permission is hereby granted, free of charge, to any person obtaining a copy
+ of this software and associated documentation files (the "Software"), to deal
+ in the Software without restriction, including without limitation the rights
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+ copies of the Software, and to permit persons to whom the Software is
+ furnished to do so, subject to the following conditions:
+
+ The above copyright notice and this permission notice shall be included in
+ all copies or substantial portions of the Software.
+
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+ THE SOFTWARE. */
+
+#pragma once
+
+#include <hip/hip_runtime.h>
+#include "hip_event.hpp"
+#include "hip_internal.hpp"
+#include <unordered_map>
+#include <unordered_set>
+
+namespace hip {
+
+class Device;
+class Stream;
+
+struct MemoryTimestamp {
+  MemoryTimestamp(hip::Stream* stream, hip::Event* event = nullptr): event_(event) {
+    safe_streams_.insert(stream);
+  }
+  MemoryTimestamp(): event_(nullptr) {}
+
+  /// Adds a safe stream to the list of stream for possible reuse
+  void AddSafeStream(hip::Stream* stream) {
+    if (safe_streams_.find(stream) != safe_streams_.end()) {
+      safe_streams_.insert(stream);
+    }
+  }
+  /// Changes last known valid event asociated with memory
+  void SetEvent(hip::Event* event) {
+    delete event_;
+    event_ = event;
+  }
+  /// Wait for memory to be available
+  void Wait() {
+    if (event_ != nullptr) {
+      auto hip_error = event_->synchronize();
+    }
+  }
+  /// Returns if memory object is safe for reuse
+  bool IsSafeFind(hip::Stream* stream = nullptr, bool opportunistic = true) {
+    bool result = false;
+    if (safe_streams_.find(stream) != safe_streams_.end()) {
+      // A safe stream doesn't require TS validation
+      result = true;
+    } else if (opportunistic && (event_ != nullptr)) {
+      // Check HIP event for a retired status
+      result = (event_->query() == hipSuccess) ? true : false;
+    }
+    return result;
+  }
+  /// Returns if memory object is safe for reuse
+  bool IsSafeRelease() {
+    bool result = true;
+    if (event_ != nullptr) {
+      // Check HIP event for a retired status
+      result = (event_->query() == hipSuccess) ? true : false;
+    }
+    return result;
+  }
+
+  std::unordered_set<hip::Stream*>  safe_streams_;  //!< Safe streams for memory reuse
+  hip::Event*   event_;   //!< Last known HIP event, associated with the memory object
+};
+
+class Heap : public amd::EmbeddedObject {
+public:
+  Heap(hip::Device* device):
+    total_size_(0), max_total_size_(0), release_threshold_(0), device_(device) {}
+  ~Heap() {}
+
+  /// Adds allocation into the heap on a specific stream
+  void AddMemory(amd::Memory* memory, hip::Stream* stream);
+
+  /// Adds allocation into the heap with specific TS
+  void AddMemory(amd::Memory* memory, const MemoryTimestamp& ts);
+
+  /// Finds memory object with the specified size
+  amd::Memory* FindMemory(size_t size, hip::Stream* stream, bool opportunistic);
+
+  /// Removes allocation from the map
+  bool RemoveMemory(amd::Memory* memory, MemoryTimestamp* ts = nullptr);
+
+  /// Releases all memory, until the threshold value is met
+  bool ReleaseAllMemory(size_t min_bytes_to_hold = std::numeric_limits<size_t>::max(), bool safe_release = false);
+
+  /// Releases all memory, safe to the provided stream, until the threshold value is met
+  bool ReleaseAllMemory(hip::Stream* stream);
+
+  /// Heap doesn't have any allocations
+  bool IsEmpty() const { return (allocations_.size() == 0) ? true : false; }
+
+  /// Set the memory release threshold
+  void SetReleaseThreshold(uint64_t value) { release_threshold_ = value; }
+
+  /// Set the memory release threshold
+  uint64_t GetReleaseThreshold() const { return release_threshold_; }
+
+  /// Get the size of all allocations in the heap
+  uint64_t GetTotalSize() const { return total_size_; }
+
+  /// Get the size of all allocations in the heap
+  uint64_t GetMaxTotalSize() const { return max_total_size_; }
+
+  /// Set maximum total, allocated by the heap
+  void SetMaxTotalSize(uint64_t value) { max_total_size_ = value; }
+
+  std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator EraseAllocaton(
+    std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator& it);
+
+private:
+  Heap() = delete;
+  Heap(const Heap&) = delete;
+  Heap& operator=(const Heap&) = delete;
+
+  std::unordered_map<amd::Memory*, MemoryTimestamp> allocations_;   //!< Map of allocations on a specific stream
+  uint64_t total_size_;         //!< Size of all allocations in the heap
+  uint64_t max_total_size_;     //!< Maximum heap allocation size
+  uint64_t release_threshold_;  //!< Threshold size in bytes for memory release from heap, default 0
+
+  hip::Device*  device_;    //!< Hip device the allocations will reside
+};
+
+/// Allocates memory in the pool on the specified stream and places the allocation into busy_heap_
+/// @note: the logic also will look in free_heap for possible reuse.
+/// hipMemPoolReuseAllowOpportunistic option will validate if HIP event,
+/// associated with memory is done, then reuse can be performed.
+class MemoryPool : public amd::ReferenceCountedObject {
+public:
+  MemoryPool(hip::Device* device):
+    busy_heap_(device),
+    free_heap_(device),
+    lock_pool_ops_("Pool operations", true), device_(device) {
+      device_->AddMemoryPool(this);
+      state_.event_dependencies_ = 1;
+      state_.opportunistic_ = 1;
+      state_.internal_dependencies_ = 1;
+    }
+  virtual ~MemoryPool() {
+    assert(busy_heap_.IsEmpty() && "Can't destroy pool with busy allocations!");
+    constexpr bool kSafeRelease = true;
+    free_heap_.ReleaseAllMemory(0, kSafeRelease);
+    // Remove memory pool from the list of all pool on the current device
+    device_->RemoveMemoryPool(this);
+  }
+
+  /// The same stream can reuse memory without HIP event validation
+  void* AllocateMemory(size_t size, hip::Stream* stream);
+
+  /// Frees memory by placing memory object with HIP event into free_heap_
+  bool FreeMemory(amd::Memory* memory, hip::Stream* stream);
+
+  /// Releases all allocations from free_heap_. It can be called on Stream or Device synchronization
+  /// @note The caller must make sure it's safe to release memory
+  void ReleaseFreedMemory(hip::Stream* stream = nullptr);
+
+  /// Releases all allocations in MemoryPool
+  void ReleaseAllMemory();
+
+  /// Trims the pool until it has only min_bytes_to_hold
+  void TrimTo(size_t min_bytes_to_hold);
+
+  /// Trims the pool until it has only min_bytes_to_hold
+  hip::Device* Device() const { return device_; }
+
+  /// Set memory pool control attributes
+  hipError_t SetAttribute(hipMemPoolAttr attr, void* value);
+
+  /// Get memory pool control attributes
+  hipError_t GetAttribute(hipMemPoolAttr attr, void* value);
+
+  /// Accessors for the pool state
+  bool EventDependencies() const { return (state_.event_dependencies_) ? true : false; }
+  bool Opportunistic() const { return (state_.opportunistic_) ? true : false; }
+  bool InternalDependencies() const { return (state_.internal_dependencies_) ? true : false; }
+
+private:
+  MemoryPool() = delete;
+  MemoryPool(const MemoryPool&) = delete;
+  MemoryPool& operator=(const MemoryPool&) = delete;
+
+
+  Heap busy_heap_;    //!< Heap of busy allocations
+  Heap free_heap_;    //!< Heap of freed allocations
+  struct {
+    uint32_t event_dependencies_ : 1;     //!< Event dependencies tracking is enabled
+    uint32_t opportunistic_ : 1;          //!< HIP event check is enabled
+    uint32_t internal_dependencies_ : 1;  //!< Runtime adds internal events to handle memory dependencies
+  } state_;
+
+  amd::Monitor  lock_pool_ops_;  //!< Access to the pool must be lock protected
+  hip::Device*  device_;    //!< Hip device the heap will reside
+};
+
+} // Mamespace hip

hipamd/src/hip_stream.cpp

@@ -407,6 +407,12 @@ hipError_t hipStreamSynchronize(hipStream_t stream) {
   // Wait for the current host queue
   hip::getQueue(stream)->finish();
 
+  // Make sure runtime releases memory for all memory pools on the device,
+  // associated with the queue
+  auto hip_stream = reinterpret_cast<hip::Stream*>(stream);
+  auto device = (hip_stream == nullptr) ? hip::getCurrentDevice() : hip_stream->GetDevice();
+  device->ReleaseFreedMemory(hip_stream);
+
   HIP_RETURN(hipSuccess);
 }