From 8bdda9007104ddd90f8a77955821d2c765979ad3 Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Thu, 30 Nov 2023 16:02:25 +0000
Subject: [PATCH] SWDEV-434702 - get the nullstream in constant time in
 iHipWaitActiveStreams

Change-Id: Ib92cc4d2ca1b725cf30b24cdd711c873fd513d95
---
 hipamd/src/hip_internal.hpp |  2 +-
 hipamd/src/hip_stream.cpp   | 61 +++++++++++++++++++------------------
 2 files changed, 33 insertions(+), 30 deletions(-)

diff --git a/hipamd/src/hip_internal.hpp b/hipamd/src/hip_internal.hpp
index ac63466c17..7afe92d6e0 100644
--- a/hipamd/src/hip_internal.hpp
+++ b/hipamd/src/hip_internal.hpp
@@ -488,7 +488,7 @@ public:
     void Reset();
 
     hip::Stream* NullStream();
-    Stream* GetNullStream();
+    Stream* GetNullStream() const {return null_stream_;};
 
     void SetActiveStatus() {
       isActive_ = true;
diff --git a/hipamd/src/hip_stream.cpp b/hipamd/src/hip_stream.cpp
index 8babc5a917..44ecfd310c 100644
--- a/hipamd/src/hip_stream.cpp
+++ b/hipamd/src/hip_stream.cpp
@@ -198,41 +198,44 @@ bool Stream::existsActiveStreamForDevice(hip::Device* device) {
 void iHipWaitActiveStreams(hip::Stream* blocking_stream, bool wait_null_stream) {
   amd::Command::EventWaitList eventWaitList(0);
   bool submitMarker = 0;
-  {
+
+  auto waitForStream = [&submitMarker,
+                         &eventWaitList](hip::Stream* stream) {
+    if (amd::Command *command = stream->getLastQueuedCommand(true)) {
+      amd::Event &event = command->event();
+      // Check HW status of the ROCcrl event.
+      // Note: not all ROCclr modes support HW status
+      bool ready = stream->device().IsHwEventReady(event);
+      if (!ready) {
+        ready = (command->status() == CL_COMPLETE);
+      }
+      submitMarker |= stream->vdev()->isFenceDirty();
+      // Check the current active status
+      if (!ready) {
+        command->notifyCmdQueue();
+        eventWaitList.push_back(command);
+      } else {
+        command->release();
+      }
+    }
+  };
+
+  if (wait_null_stream) {
+    if (hip::Stream* null_stream = blocking_stream->GetDevice()->GetNullStream()) {
+      waitForStream(null_stream);
+    }
+  } else {
     amd::ScopedLock lock(streamSetLock);
 
     for (const auto& active_stream : streamSet) {
       // If it's the current device
       if ((&active_stream->device() == &blocking_stream->device()) &&
-          // Make sure it's a default stream
-          ((active_stream->Flags() & hipStreamNonBlocking) == 0) &&
-          // and it's not the current stream
-          (active_stream != blocking_stream) &&
-          // check for a wait on the null stream
-          (active_stream->Null() == wait_null_stream)) {
+        // Make sure it's a default stream
+        ((active_stream->Flags() & hipStreamNonBlocking) == 0) &&
+        // and it's not the current stream
+        (active_stream != blocking_stream)) {
         // Get the last valid command
-        amd::Command* command = active_stream->getLastQueuedCommand(true);
-        if (command != nullptr) {
-          amd::Event& event = command->event();
-          // Check HW status of the ROCcrl event.
-          // Note: not all ROCclr modes support HW status
-          bool ready = active_stream->device().IsHwEventReady(event);
-          if (!ready) {
-            ready = (command->status() == CL_COMPLETE);
-          }
-          submitMarker |= active_stream->vdev()->isFenceDirty();
-          // Check the current active status
-          if (!ready) {
-            command->notifyCmdQueue();
-            eventWaitList.push_back(command);
-          } else {
-            command->release();
-          }
-        }
-        // Nullstream, hence there is nothing else to wait
-        if (wait_null_stream) {
-          break;
-        }
+        waitForStream(active_stream);
       }
     }
   }