SWDEV-311271 - Fix Windows mempool management

Windows path still uses multi threading implementation. Hence, in graphs
all nodes are executed in a queue thread and that requires to manage
mempool in the queue thread. However, the spec allows to destory memory,
allocated in a graph, outside of the graph's execution. That may cause
mempool management to go out of sync.

Change-Id: I0ffb2244b3cb720455ed44d1b3e2487fa8959a77

hipamd/src/hip_device.cpp

@@ -91,11 +91,11 @@ void Device::RemoveMemoryPool(MemoryPool* pool) {
 }
 
 // ================================================================================================
-bool Device::FreeMemory(amd::Memory* memory, Stream* stream) {
+bool Device::FreeMemory(amd::Memory* memory, Stream* stream, Event* event) {
   amd::ScopedLock lock(lock_);
   // Search for memory in the entire list of pools
   for (auto it : mem_pools_) {
-    if (it->FreeMemory(memory, stream)) {
+    if (it->FreeMemory(memory, stream, event)) {
       return true;
     }
   }

hipamd/src/hip_graph_internal.hpp

@@ -2138,6 +2138,7 @@ class GraphMemAllocNode final : public GraphNode {
       auto graph = GetParentGraph();
       if (graph != nullptr) {
         assert(va_ != nullptr && "Runtime can't create a command for an invalid node!");
+        stream->GetDevice()->GetGraphMemoryPool()->SetGraphInUse();
         // Create command for memory mapping
         auto cmd = new VirtualMemAllocNode(*stream, amd::Event::EventWaitList{},
             va_, node_params_.bytesize, nullptr, graph);

hipamd/src/hip_internal.hpp

@@ -241,6 +241,7 @@ public:
 
   class Device;
   class MemoryPool;
+  class Event;
   class Stream : public amd::HostQueue {
   public:
     enum Priority : int { High = -1, Normal = 0, Low = 1 };
@@ -526,7 +527,7 @@ public:
     void RemoveMemoryPool(MemoryPool* pool);
 
     /// Free memory from the device
-    bool FreeMemory(amd::Memory* memory, Stream* stream);
+    bool FreeMemory(amd::Memory* memory, Stream* stream, Event* event = nullptr);
 
     /// Release freed memory from all pools on the current device
     void ReleaseFreedMemory();

hipamd/src/hip_mempool.cpp

@@ -95,22 +95,19 @@ hipError_t hipMallocAsync(void** dev_ptr, size_t size, hipStream_t stream) {
 // memory allocatiom in graph, which occurs in a worker thread, and host execution of hipFreeAsync
 class FreeAsyncCommand : public amd::Command {
  private:
-  void* ptr_;     //!< Virtual address for asynchronious free
+  void*   ptr_;         //!< Virtual address for asynchronious free
+  hip::Event* event_;   //!< HIP event, associated with this memory release
 
  public:
-  FreeAsyncCommand(amd::HostQueue& queue, void* ptr)
-      : amd::Command(queue, 1, amd::Event::nullWaitList), ptr_(ptr) {}
+  FreeAsyncCommand(amd::HostQueue& queue, void* ptr, hip::Event* event)
+      : amd::Command(queue, 1, amd::Event::nullWaitList), ptr_(ptr), event_(event) {}
 
   virtual void submit(device::VirtualDevice& device) final {
     size_t offset = 0;
     auto memory = getMemoryObject(ptr_, offset);
     if (memory != nullptr) {
       auto id = memory->getUserData().deviceId;
-      if (!AMD_DIRECT_DISPATCH) {
-        // Required for HIP events
-        hip::setCurrentDevice(id);
-      }
-      if (!g_devices[id]->FreeMemory(memory, static_cast<hip::Stream*>(queue()))) {
+      if (!g_devices[id]->FreeMemory(memory, static_cast<hip::Stream*>(queue()), event_)) {
         // @note It's not the most optimal logic.
         // The current implementation has unconditional waits
         if (ihipFree(ptr_) != hipSuccess) {
@@ -131,12 +128,55 @@ hipError_t hipFreeAsync(void* dev_ptr, hipStream_t stream) {
 
   auto hip_stream = (stream == nullptr) ? hip::getCurrentDevice()->NullStream()
                     : reinterpret_cast<hip::Stream*>(stream);
-  auto cmd = new FreeAsyncCommand(*hip_stream, dev_ptr);
-  if (cmd == nullptr) {
-    HIP_RETURN(hipErrorUnknown);
+
+  hip::Event* event = nullptr;
+  bool graph_in_use = false;
+
+  if (!AMD_DIRECT_DISPATCH) {
+    // Note: This logic is required for multithreading execution only and
+    // can reduce mempool reserved memory efficiency
+    for (auto dev : g_devices) {
+      graph_in_use |= dev->GetGraphMemoryPool()->GraphInUse();
+    }
+  }
+
+  if (!graph_in_use) {
+    size_t offset = 0;
+    auto memory = getMemoryObject(dev_ptr, offset);
+    if (memory != nullptr) {
+      auto id = memory->getUserData().deviceId;
+      if (!g_devices[id]->FreeMemory(memory, hip_stream, event)) {
+        // @note It's not the most optimal logic.
+        // The current implementation has unconditional waits
+        return ihipFree(dev_ptr);
+      }
+    }
+  } else {
+    if (!AMD_DIRECT_DISPATCH) {
+      // Add a marker to the stream to trace availability of this memory
+      // Note: MT path requires the marker command to be created in the host thread,
+      // so the queue thread could process it, because creating a command from the queue thread
+      // may block the execution
+      event = new hip::Event(0);
+      if (event != nullptr) {
+        if (hipSuccess !=
+            event->addMarker(reinterpret_cast<hipStream_t>(hip_stream), nullptr, true)) {
+          delete event;
+          event = nullptr;
+        } else {
+          // Make sure runtime sends a notification to the worker thread
+          auto result = event->ready(Query);
+        }
+      }
+    }
+    
+    auto cmd = new FreeAsyncCommand(*hip_stream, dev_ptr, event);
+    if (cmd == nullptr) {
+      HIP_RETURN(hipErrorUnknown);
+    }
+    cmd->enqueue();
+    cmd->release();
   }
-  cmd->enqueue();
-  cmd->release();
 
   HIP_RETURN(hipSuccess);
 }

hipamd/src/hip_mempool_impl.cpp

@@ -25,9 +25,9 @@
 namespace hip {
 
 // ================================================================================================
-void Heap::AddMemory(amd::Memory* memory, hip::Stream* stream) {
+void Heap::AddMemory(amd::Memory* memory, Stream* stream) {
   auto mem_size = memory->getSize();
-  allocations_.insert({{mem_size, memory}, {stream, nullptr}});
+  allocations_.insert({{mem_size, memory}, {stream}});
   total_size_ += mem_size;
   max_total_size_ = std::max(max_total_size_, total_size_);
 }
@@ -41,7 +41,7 @@ void Heap::AddMemory(amd::Memory* memory, const MemoryTimestamp& ts) {
 }
 
 // ================================================================================================
-amd::Memory* Heap::FindMemory(size_t size, hip::Stream* stream, bool opportunistic, void* dptr) {
+amd::Memory* Heap::FindMemory(size_t size, Stream* stream, bool opportunistic, void* dptr) {
   amd::Memory* memory = nullptr;
   auto start = allocations_.lower_bound({size, nullptr});
   for (auto it = start; it != allocations_.end();) {
@@ -139,7 +139,7 @@ bool Heap::ReleaseAllMemory() {
 }
 
 // ================================================================================================
-void Heap::RemoveStream(hip::Stream* stream) {
+void Heap::RemoveStream(Stream* stream) {
   for (auto it : allocations_) {
     it.second.safe_streams_.erase(stream);
   }
@@ -165,7 +165,7 @@ void Heap::SetAccess(hip::Device* device, bool enable) {
 }
 
 // ================================================================================================
-void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream, void* dptr) {
+void* MemoryPool::AllocateMemory(size_t size, Stream* stream, void* dptr) {
   amd::ScopedLock lock(lock_pool_ops_);
 
   void* dev_ptr = nullptr;
@@ -221,68 +221,78 @@ void* MemoryPool::AllocateMemory(size_t size, hip::Stream* stream, void* dptr) {
 }
 
 // ================================================================================================
-bool MemoryPool::FreeMemory(amd::Memory* memory, hip::Stream* stream) {
-  amd::ScopedLock lock(lock_pool_ops_);
+bool MemoryPool::FreeMemory(amd::Memory* memory, Stream* stream, Event* event) {
+  {
+    amd::ScopedLock lock(lock_pool_ops_);
 
-  // If the free heap grows over the busy heap, then force release
-  if (free_heap_.GetTotalSize() > busy_heap_.GetTotalSize()) {
-    // Use event base release to reduce memory pressure
-    constexpr size_t kBytesToHold = 0;
-    free_heap_.ReleaseAllMemory(kBytesToHold);
+    // If the free heap grows over the busy heap, then force release
+    if (AMD_DIRECT_DISPATCH && (free_heap_.GetTotalSize() > busy_heap_.GetTotalSize())) {
+      // Use event base release to reduce memory pressure
+      constexpr size_t kBytesToHold = 0;
+      free_heap_.ReleaseAllMemory(kBytesToHold);
 
-    // If free mmeory is less than 12.5% of total, then force wait release
-    size_t free = 0;
-    size_t total = 0;
-    hipError_t err = hipMemGetInfo(&free, &total);
-    if ((err == hipSuccess) && (free < (total >> 3))) {
-      constexpr bool kSafeRelease = true;
-      free_heap_.ReleaseAllMemory(free_heap_.GetTotalSize() >> 1, kSafeRelease);
-    }
-  }
-
-  MemoryTimestamp ts;
-  // Remove memory object from the busy pool
-  if (!busy_heap_.RemoveMemory(memory, &ts)) {
-    // This pool doesn't contain memory
-    return false;
-  }
-  ClPrint(amd::LOG_INFO, amd::LOG_MEM_POOL, "Pool FreeMem: %p, %p", memory->getSvmPtr(), memory);
-
-  auto ga = reinterpret_cast<hip::MemMapAllocUserData*>(memory->getUserData().data);
-  if (ga != nullptr) {
-    if (stream == nullptr) {
-      stream = g_devices[memory->getUserData().deviceId]->NullStream();
-    }
-    // Unmap virtual address from memory
-    auto cmd = new amd::VirtualMapCommand(*stream, amd::Command::EventWaitList{},
-                                          memory->getSvmPtr(), ga->size_, nullptr);
-    cmd->enqueue();
-    cmd->release();
-    memory->setSvmPtr(ga->ptr_);
-    // Free virtual address and destroy generic allocation object
-    ga->va_->release();
-    delete ga;
-    memory->getUserData().data = nullptr;
-  }
-
-  if (stream != nullptr) {
-    // 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);
+      // If free mmeory is less than 12.5% of total, then force wait release
+      size_t free = 0;
+      size_t total = 0;
+      hipError_t err = hipMemGetInfo(&free, &total);
+      if ((err == hipSuccess) && (free < (total >> 3))) {
+        constexpr bool kSafeRelease = true;
+        free_heap_.ReleaseAllMemory(free_heap_.GetTotalSize() >> 1, kSafeRelease);
       }
     }
-  } else {
-    // Assume a safe release from hipFree() if stream is nullptr
-    ts.SetEvent(nullptr);
-  }
-  free_heap_.AddMemory(memory, ts);
 
-  // Decrement the reference counter on the pool
+    MemoryTimestamp ts;
+    // Remove memory object from the busy pool
+    if (!busy_heap_.RemoveMemory(memory, &ts)) {
+      // This pool doesn't contain memory
+      return false;
+    }
+    ClPrint(amd::LOG_INFO, amd::LOG_MEM_POOL, "Pool FreeMem: %p, %p", memory->getSvmPtr(), memory);
+
+    auto ga = reinterpret_cast<hip::MemMapAllocUserData*>(memory->getUserData().data);
+    if (ga != nullptr) {
+      if (stream == nullptr) {
+        stream = g_devices[memory->getUserData().deviceId]->NullStream();
+      }
+      // Unmap virtual address from memory
+      auto cmd = new amd::VirtualMapCommand(*stream, amd::Command::EventWaitList{},
+                                            memory->getSvmPtr(), ga->size_, nullptr);
+      cmd->enqueue();
+      cmd->release();
+      memory->setSvmPtr(ga->ptr_);
+      // Free virtual address and destroy generic allocation object
+      ga->va_->release();
+      delete ga;
+      memory->getUserData().data = nullptr;
+    }
+
+    if (stream != nullptr) {
+      // The stream of destruction is a safe stream, because the app must handle sync
+      ts.AddSafeStream(stream);
+
+      if (event == nullptr) {
+        // 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);
+            // Make sure runtime sends a notification
+            auto result = e->ready(Query);
+          }
+        }
+      } else {
+        ts.SetEvent(event);
+      }
+    } else {
+      // Assume a safe release from hipFree() if stream is nullptr
+      ts.SetEvent(nullptr);
+    }
+    free_heap_.AddMemory(memory, ts);
+  }
+
+  // Decrement the reference counter on the pool.
+  // Note: It may delete memory pool for the last allocation. Thus, the scope lock can't include
+  // this call. 
   release();
 
   return true;
@@ -303,7 +313,7 @@ void MemoryPool::ReleaseFreedMemory() {
 }
 
 // ================================================================================================
-void MemoryPool::RemoveStream(hip::Stream* stream) {
+void MemoryPool::RemoveStream(Stream* stream) {
   amd::ScopedLock lock(lock_pool_ops_);
 
   free_heap_.RemoveStream(stream);
@@ -453,7 +463,7 @@ void MemoryPool::GetAccess(hip::Device* device, hipMemAccessFlags* flags) {
 }
 
 // ================================================================================================
-void MemoryPool::FreeAllMemory(hip::Stream* stream) {
+void MemoryPool::FreeAllMemory(Stream* stream) {
   while (!busy_heap_.Allocations().empty()) {
     FreeMemory(busy_heap_.Allocations().begin()->first.second, stream);
   }

hipamd/src/hip_mempool_impl.hpp

@@ -38,7 +38,7 @@ struct SharedMemPointer {
 };
 
 struct MemoryTimestamp {
-  MemoryTimestamp(hip::Stream* stream, hip::Event* event = nullptr): event_(event) {
+  MemoryTimestamp(hip::Stream* stream): event_(nullptr) {
     if (stream != nullptr) {
       safe_streams_.insert(stream);
     }
@@ -100,13 +100,13 @@ public:
   ~Heap() {}
 
   /// Adds allocation into the heap on a specific stream
-  void AddMemory(amd::Memory* memory, hip::Stream* stream);
+  void AddMemory(amd::Memory* memory, 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, void* dptr = nullptr);
+  amd::Memory* FindMemory(size_t size, Stream* stream, bool opportunistic, void* dptr = nullptr);
 
   /// Removes allocation from the map
   bool RemoveMemory(amd::Memory* memory, MemoryTimestamp* ts = nullptr);
@@ -118,7 +118,7 @@ public:
   bool ReleaseAllMemory();
 
   /// Remove the provided stream from the safe list
-  void RemoveStream(hip::Stream* stream);
+  void RemoveStream(Stream* stream);
 
   /// Enables P2P access to the provided device
   void SetAccess(hip::Device* device, bool enable);
@@ -211,10 +211,10 @@ class MemoryPool : public amd::ReferenceCountedObject {
   }
 
   /// The same stream can reuse memory without HIP event validation
-  void* AllocateMemory(size_t size, hip::Stream* stream, void* dptr = nullptr);
+  void* AllocateMemory(size_t size, Stream* stream, void* dptr = nullptr);
 
   /// Frees memory by placing memory object with HIP event into free_heap_
-  bool FreeMemory(amd::Memory* memory, hip::Stream* stream);
+  bool FreeMemory(amd::Memory* memory, Stream* stream, Event* event = nullptr);
 
   /// Check if memory is active and belongs to the busy heap
   bool IsBusyMemory(amd::Memory* memory) const { return busy_heap_.IsActiveMemory(memory); }
@@ -264,8 +264,10 @@ class MemoryPool : public amd::ReferenceCountedObject {
   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; }
+  bool GraphInUse() const { return (state_.graph_in_use_) ? true : false; }
+  void SetGraphInUse() { state_.graph_in_use_ = true; }
 
-private:
+ private:
   MemoryPool() = delete;
   MemoryPool(const MemoryPool&) = delete;
   MemoryPool& operator=(const MemoryPool&) = delete;
@@ -278,7 +280,8 @@ private:
       uint32_t opportunistic_ : 1;          //!< HIP event check is enabled
       uint32_t internal_dependencies_ : 1;  //!< Runtime adds internal events to handle memory
                                             //!< dependencies
-      uint32_t interprocess_ : 1;  //!< Memory pool can be used in interprocess communications
+      uint32_t interprocess_ : 1;   //!< Memory pool can be used in interprocess communications
+      uint32_t graph_in_use_ : 1;   //!< Memory pool was used in a graph execution
     };
     uint32_t value_;
   } state_;