SWDEV-311271 - Make sure memory pool can accept default stream

Add lock protection for access to the pool list.
Remove destroyed stream from the list of the safe streams

Change-Id: I1863b89bd3f5e188c161227cc790c3adaf72cc58


[ROCm/clr commit: 5957ff9f7b]

projects/clr/hipamd/src/hip_device.cpp

@@ -36,6 +36,17 @@ amd::HostQueue* Device::NullStream(bool skip_alloc) {
   return null_queue;
 }
 
+// ================================================================================================
+Stream* Device::GetNullStream() {
+  amd::HostQueue* null_queue = null_stream_.asHostQueue();
+  if (null_queue == nullptr) {
+    return nullptr;
+  }
+  // Wait for all active streams before executing commands on the default
+  iHipWaitActiveStreams(null_queue);
+  return &null_stream_;
+}
+
 // ================================================================================================
 bool Device::Create() {
   // Create default memory pool
@@ -50,6 +61,7 @@ bool Device::Create() {
 
 // ================================================================================================
 void Device::AddMemoryPool(MemoryPool* pool) {
+  amd::ScopedLock lock(lock_);
   if (auto it = mem_pools_.find(pool); it == mem_pools_.end()) {
     mem_pools_.insert(pool);
   }
@@ -57,6 +69,7 @@ void Device::AddMemoryPool(MemoryPool* pool) {
 
 // ================================================================================================
 void Device::RemoveMemoryPool(MemoryPool* pool) {
+  amd::ScopedLock lock(lock_);
   if (auto it = mem_pools_.find(pool); it != mem_pools_.end()) {
     mem_pools_.erase(it);
   }
@@ -64,6 +77,7 @@ void Device::RemoveMemoryPool(MemoryPool* pool) {
 
 // ================================================================================================
 bool Device::FreeMemory(amd::Memory* memory, Stream* stream) {
+  amd::ScopedLock lock(lock_);
   // Search for memory in the entire list of pools
   for (auto& it : mem_pools_) {
     if (it->FreeMemory(memory, stream)) {
@@ -75,12 +89,22 @@ bool Device::FreeMemory(amd::Memory* memory, Stream* stream) {
 
 // ================================================================================================
 void Device::ReleaseFreedMemory(Stream* stream) {
+  amd::ScopedLock lock(lock_);
   // Search for memory in the entire list of pools
   for (auto& it : mem_pools_) {
     it->ReleaseFreedMemory(stream);
   }
 }
 
+// ================================================================================================
+void Device::RemoveStreamFromPools(Stream* stream) {
+  amd::ScopedLock lock(lock_);
+  // Update all pools with the destroyed stream
+  for (auto& it : mem_pools_) {
+    it->RemoveStream(stream);
+  }
+}
+
 // ================================================================================================
 Device::~Device() {
   if (default_mem_pool_ != nullptr) {

projects/clr/hipamd/src/hip_internal.hpp

@@ -391,6 +391,7 @@ namespace hip {
     unsigned int getFlags() const { return flags_; }
     void setFlags(unsigned int flags) { flags_ = flags; }
     amd::HostQueue* NullStream(bool skip_alloc = false);
+    Stream* GetNullStream();
 
     void SaveQueue(amd::HostQueue* queue) {
       amd::ScopedLock lock(lock_);
@@ -431,6 +432,9 @@ namespace hip {
 
     /// Release freed memory from all pools on the current device
     void ReleaseFreedMemory(Stream* stream);
+
+    /// Removes a destroyed stream from the safe list of memory pools
+    void RemoveStreamFromPools(Stream* stream);
   };
 
   /// Current thread's device

projects/clr/hipamd/src/hip_mempool.cpp

@@ -58,26 +58,29 @@ hipError_t hipDeviceGetMemPool(hipMemPool_t* mem_pool, int device) {
 // ================================================================================================
 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)) {
+  if ((dev_ptr == nullptr) || (size == 0) || (!hip::isValid(stream))) {
     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));
+  auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->GetNullStream() :
+    reinterpret_cast<hip::Stream*>(stream);
+  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(size, hip_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)) {
+  if ((dev_ptr == nullptr) || (!hip::isValid(stream))) {
     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))) {
+  auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->GetNullStream() :
+    reinterpret_cast<hip::Stream*>(stream);
+  if (!g_devices[id]->FreeMemory(memory, hip_stream)) {
     //! @todo It's not the most optimal logic. The current implementation has unconditional waits
     HIP_RETURN(ihipFree(dev_ptr));
    }
@@ -185,11 +188,12 @@ hipError_t hipMallocFromPoolAsync(
     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)) {
+  if ((dev_ptr == nullptr) || (size == 0) || (mem_pool == nullptr) || (!hip::isValid(stream))) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(
-      size, reinterpret_cast<hip::Stream*>(stream));
+  auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->GetNullStream() :
+    reinterpret_cast<hip::Stream*>(stream);
+  *dev_ptr = reinterpret_cast<hip::MemoryPool*>(mem_pool)->AllocateMemory(size, hip_stream);
   HIP_RETURN(hipSuccess);
 }
 

projects/clr/hipamd/src/hip_mempool_impl.cpp

@@ -120,6 +120,13 @@ bool Heap::ReleaseAllMemory(hip::Stream* stream) {
   return true;
 }
 
+// ================================================================================================
+void Heap::RemoveStream(hip::Stream* stream) {
+  for (auto it = allocations_.begin(); it != allocations_.end();) {
+    it->second.safe_streams_.erase(stream);
+  }
+}
+
 // ================================================================================================
 void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream) {
   amd::ScopedLock lock(lock_pool_ops_);
@@ -197,6 +204,13 @@ void MemoryPool::ReleaseFreedMemory(hip::Stream* stream) {
   free_heap_.ReleaseAllMemory(stream);
 }
 
+// ================================================================================================
+void MemoryPool::RemoveStream(hip::Stream* stream) {
+  amd::ScopedLock lock(lock_pool_ops_);
+
+  free_heap_.RemoveStream(stream);
+}
+
 // ================================================================================================
 void MemoryPool::TrimTo(size_t min_bytes_to_hold) {
   amd::ScopedLock lock(lock_pool_ops_);

projects/clr/hipamd/src/hip_mempool_impl.hpp

@@ -104,6 +104,9 @@ public:
   /// Releases all memory, safe to the provided stream, until the threshold value is met
   bool ReleaseAllMemory(hip::Stream* stream);
 
+  /// Remove the provided stream from the safe list
+  void RemoveStream(hip::Stream* stream);
+
   /// Heap doesn't have any allocations
   bool IsEmpty() const { return (allocations_.size() == 0) ? true : false; }
 
@@ -122,6 +125,7 @@ public:
   /// Set maximum total, allocated by the heap
   void SetMaxTotalSize(uint64_t value) { max_total_size_ = value; }
 
+  /// Erases single allocation form the heap's map
   std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator EraseAllocaton(
     std::unordered_map<amd::Memory*, MemoryTimestamp>::iterator& it);
 
@@ -171,6 +175,9 @@ public:
   /// @note The caller must make sure it's safe to release memory
   void ReleaseFreedMemory(hip::Stream* stream = nullptr);
 
+  /// Removes a stream from tracking
+  void RemoveStream(hip::Stream* stream);
+
   /// Releases all allocations in MemoryPool
   void ReleaseAllMemory();
 

projects/clr/hipamd/src/hip_stream.cpp

@@ -436,6 +436,8 @@ hipError_t hipStreamDestroy(hipStream_t stream) {
   }
   hip::Stream* s = reinterpret_cast<hip::Stream*>(stream);
 
+  s->GetDevice()->RemoveStreamFromPools(s);
+
   amd::ScopedLock lock(g_captureStreamsLock);
   const auto& g_it = std::find(g_captureStreams.begin(), g_captureStreams.end(), s);
   if (g_it != g_captureStreams.end()) {