From 18edb0d97c4bb36536265692414cc882dfbc4b58 Mon Sep 17 00:00:00 2001
From: Anusha GodavarthySurya <anusha.godavarthysurya@amd.com>
Date: Thu, 1 Dec 2022 13:22:06 +0000
Subject: [PATCH] SWDEV-366636 - Fix performance drop in TF-RCCL models

Change-Id: Ie04219daacc29654d0f47e1c15ed5cd78d88d8c8


[ROCm/clr commit: a2785c68b2b8a71a8d4f9318dcd6da053976cc5d]
---
 projects/clr/hipamd/src/hip_event.cpp | 16 ++++++++--------
 projects/clr/hipamd/src/hip_event.hpp | 15 +++++++++++++--
 2 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_event.cpp b/projects/clr/hipamd/src/hip_event.cpp
index 747b69ce2f..d2d94602ec 100644
--- a/projects/clr/hipamd/src/hip_event.cpp
+++ b/projects/clr/hipamd/src/hip_event.cpp
@@ -30,21 +30,21 @@ namespace hip {
 static amd::Monitor eventSetLock{"Guards global event set"};
 static std::unordered_set<hipEvent_t> eventSet;
 
-bool Event::ready() {
+bool Event::ready(eventType type) {
   if (event_->status() != CL_COMPLETE) {
     event_->notifyCmdQueue();
   }
   // Check HW status of the ROCcrl event. Note: not all ROCclr modes support HW status
-  bool ready = g_devices[deviceId()]->devices()[0]->IsHwEventReady(*event_);
+  bool ready = CheckHwEvent(type);
   if (!ready) {
     ready = (event_->status() == CL_COMPLETE);
   }
   return ready;
 }
 
-bool EventDD::ready() {
+bool EventDD::ready(eventType type) {
   // Check HW status of the ROCcrl event. Note: not all ROCclr modes support HW status
-  bool ready = g_devices[deviceId()]->devices()[0]->IsHwEventReady(*event_);
+  bool ready = CheckHwEvent(type);
   // FIXME: Remove status check entirely
   if (!ready) {
     ready = (event_->status() == CL_COMPLETE);
@@ -60,7 +60,7 @@ hipError_t Event::query() {
     return hipSuccess;
   }
 
-  return ready() ? hipSuccess : hipErrorNotReady;
+  return ready(Query) ? hipSuccess : hipErrorNotReady;
 }
 
 hipError_t Event::synchronize() {
@@ -108,7 +108,7 @@ hipError_t Event::elapsedTime(Event& eStop, float& ms) {
       return hipErrorInvalidHandle;
     }
 
-    if (!ready()) {
+    if (!ready(ElapsedTime)) {
       return hipErrorNotReady;
     }
 
@@ -124,7 +124,7 @@ hipError_t Event::elapsedTime(Event& eStop, float& ms) {
     return hipErrorInvalidHandle;
   }
 
-  if (!ready() || !eStop.ready()) {
+  if (!ready(ElapsedTime) || !eStop.ready(ElapsedTime)) {
     return hipErrorNotReady;
   }
 
@@ -199,7 +199,7 @@ hipError_t Event::streamWait(hipStream_t stream, uint flags) {
   hip::Stream* hip_stream = hip::getStream(stream);
   // Access to event_ object must be lock protected
   amd::ScopedLock lock(lock_);
-  if ((event_ == nullptr) || (event_->command().queue() == hip_stream) || ready()) {
+  if ((event_ == nullptr) || (event_->command().queue() == hip_stream) || ready(StreamWait)) {
     return hipSuccess;
   }
   if (!event_->notifyCmdQueue()) {
diff --git a/projects/clr/hipamd/src/hip_event.hpp b/projects/clr/hipamd/src/hip_event.hpp
index 91a8193d48..74f57ca163 100644
--- a/projects/clr/hipamd/src/hip_event.hpp
+++ b/projects/clr/hipamd/src/hip_event.hpp
@@ -89,6 +89,7 @@ class EventMarker : public amd::Marker {
   }
 };
 
+enum eventType { Query, StreamWait, ElapsedTime };
 class Event {
   /// event recorded on stream where capture is active
   bool onCapture_;
@@ -96,6 +97,16 @@ class Event {
   hipStream_t captureStream_ = nullptr;
   /// Previous captured nodes before event record
   std::vector<hipGraphNode_t> nodesPrevToRecorded_;
+ protected:
+  bool CheckHwEvent(eventType type) {
+    bool ready;
+    if (type == Query) {
+      ready = g_devices[deviceId()]->devices()[0]->IsHwEventReadyForcedWait(*event_);
+    } else {
+      ready = g_devices[deviceId()]->devices()[0]->IsHwEventReady(*event_);
+    }
+    return ready;
+  }
 
  public:
   Event(unsigned int flags) : flags(flags), lock_("hipEvent_t", true),
@@ -170,7 +181,7 @@ class Event {
     return hipErrorInvalidConfiguration;
   }
   virtual bool awaitEventCompletion();
-  virtual bool ready();
+  virtual bool ready(eventType type);
   virtual int64_t time(bool getStartTs) const;
 
  protected:
@@ -190,7 +201,7 @@ class EventDD : public Event {
   virtual ~EventDD() {}
 
   virtual bool awaitEventCompletion();
-  virtual bool ready();
+  virtual bool ready(eventType type);
   virtual int64_t time(bool getStartTs) const;
 };