From 1b9c17779bbd913f30b06ee3051bd8b8dd60cd40 Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Mon, 9 Sep 2024 15:30:50 +0100
Subject: [PATCH] SWDEV-501921 - Introduce an activeQueues set

The new set tracks only the queues that have a command
submitted to them. This allows for fast iteration
in waitActiveStreams.

Change-Id: I2c832eefa01280d9a87a5f57874d36d2e9441de7
(cherry picked from commit bcc545e6b893b09ea3b10badc92be6289bacafa9)
---
 hipamd/src/hip_device.cpp        |  7 +++----
 rocclr/device/device.hpp         | 20 ++++++++++++++++++++
 rocclr/platform/commandqueue.cpp |  5 +++++
 rocclr/platform/commandqueue.hpp |  3 +++
 4 files changed, 31 insertions(+), 4 deletions(-)

diff --git a/hipamd/src/hip_device.cpp b/hipamd/src/hip_device.cpp
index a45b1ea55..951fe02ab 100644
--- a/hipamd/src/hip_device.cpp
+++ b/hipamd/src/hip_device.cpp
@@ -185,10 +185,9 @@ void Device::WaitActiveStreams(hip::Stream* blocking_stream, bool wait_null_stre
       waitForStream(null_stream_);
     }
   } else {
-    amd::ScopedLock lock(streamSetLock);
-
-    for (const auto& active_stream : streamSet) {
-      // If it's the current device
+    auto activeQueues = blocking_stream->device().getActiveQueues();
+    for (const auto& command : activeQueues) {
+      hip::Stream* active_stream = static_cast<hip::Stream*>(command);
       if (// Make sure it's a default stream
         ((active_stream->Flags() & hipStreamNonBlocking) == 0) &&
         // and it's not the current stream
diff --git a/rocclr/device/device.hpp b/rocclr/device/device.hpp
index 32ae204e2..52c6a18e0 100644
--- a/rocclr/device/device.hpp
+++ b/rocclr/device/device.hpp
@@ -2082,6 +2082,24 @@ class Device : public RuntimeObject {
     return false;
   }
 
+  //! Returns the queues that have at least one submitted command
+  std::vector<amd::CommandQueue*> getActiveQueues() {
+     amd::ScopedLock lock(activeQueuesLock_);
+     return std::vector<amd::CommandQueue*>(activeQueues.begin(), activeQueues.end());
+  }
+
+  //! Adds the queue to the set of active command queues
+  void addToActiveQueues(amd::CommandQueue* commandQueue) {
+     amd::ScopedLock lock(activeQueuesLock_);
+     activeQueues.insert(commandQueue);
+  }
+
+  //! Removes the queue from the set of active command queues
+  void removeFromActiveQueues(amd::CommandQueue* commandQueue) {
+    amd::ScopedLock lock(activeQueuesLock_);
+    activeQueues.erase(commandQueue);
+  }
+
   // Notifies device about context destroy
   virtual void ContextDestroy() {}
 
@@ -2131,6 +2149,8 @@ class Device : public RuntimeObject {
   uint64_t stack_size_{1024};       //!< Device stack size
   device::Memory* initial_heap_buffer_;   //!< Initial heap buffer
   uint64_t initial_heap_size_{HIP_INITIAL_DM_SIZE};  //!< Initial device heap size
+  amd::Monitor activeQueuesLock_ {"Guards access to the activeQueues set"};
+  std::unordered_set<amd::CommandQueue*> activeQueues; //!< The set of active queues
  private:
   const Isa *isa_;                //!< Device isa
   bool IsTypeMatching(cl_device_type type, bool offlineDevices);
diff --git a/rocclr/platform/commandqueue.cpp b/rocclr/platform/commandqueue.cpp
index b9e3056ff..310f41e74 100644
--- a/rocclr/platform/commandqueue.cpp
+++ b/rocclr/platform/commandqueue.cpp
@@ -69,6 +69,7 @@ bool HostQueue::terminate() {
         // Note that if lastCommand isn't a marker, it may not be lastEnqueueCommand_ now
         // after lastCommand->awaitCompletion() is called.
         if (lastEnqueueCommand_ != nullptr) {
+          device_.removeFromActiveQueues(this);
           lastEnqueueCommand_ ->release(); // lastEnqueueCommand_ should be a marker
           lastEnqueueCommand_ = nullptr;
         }
@@ -162,6 +163,7 @@ void HostQueue::finish(bool cpu_wait) {
       // Runtime can clear the last command only if no other submissions occured
       // during finish()
       if (command == lastEnqueueCommand_) {
+        device_.removeFromActiveQueues(this);
         lastEnqueueCommand_->release();
         lastEnqueueCommand_ = nullptr;
       }
@@ -278,6 +280,9 @@ void HostQueue::append(Command& command) {
 
   if (prevLastEnqueueCommand != nullptr) {
     prevLastEnqueueCommand->release();
+  } else {
+    // The queue becomes active. Add it to the set of activeQueues.
+    device_.addToActiveQueues(this);
   }
 }
 
diff --git a/rocclr/platform/commandqueue.hpp b/rocclr/platform/commandqueue.hpp
index c149f171f..9500fc883 100644
--- a/rocclr/platform/commandqueue.hpp
+++ b/rocclr/platform/commandqueue.hpp
@@ -267,6 +267,9 @@ class HostQueue : public CommandQueue {
     // Release the last command in the batch
     if (lastEnqueueCommand_ != nullptr) {
       lastEnqueueCommand_->release();
+    } else {
+      // The queue becomes active. Add it to the set of activeQueues.
+      device_.addToActiveQueues(this);
     }
 
     // Extra retain for the last command