From 544318fffefe852f87f2969cb466033fc4e214a3 Mon Sep 17 00:00:00 2001
From: sdashmiz <shadi.dashmiz@amd.com>
Date: Wed, 22 Feb 2023 16:16:13 -0500
Subject: [PATCH] SWDEV-360031 - Correct APIs behaviour

- correct error for hipStreamWaitEvent when event recorded before
  capture
- correct hipEventSync when event is synced during capture

Signed-off-by: sdashmiz <shadi.dashmiz@amd.com>
Change-Id: I7ecbed5621eaf323846d4ccb20ec112aaa8a5757
---
 hipamd/src/hip_event.cpp  | 19 ++++++++++++++++---
 hipamd/src/hip_stream.cpp |  8 ++++++++
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/hipamd/src/hip_event.cpp b/hipamd/src/hip_event.cpp
index d2d94602ec..601a1f5bca 100644
--- a/hipamd/src/hip_event.cpp
+++ b/hipamd/src/hip_event.cpp
@@ -350,8 +350,12 @@ hipError_t hipEventDestroy(hipEvent_t event) {
   }
 
   hip::Event* e = reinterpret_cast<hip::Event*>(event);
-  if (e->GetCaptureStream() != nullptr) {
-    reinterpret_cast<hip::Stream*>(e->GetCaptureStream())->EraseCaptureEvent(event);
+  // There is a possibility that stream destroy be called first
+  hipStream_t s = e->GetCaptureStream();
+  if (hip::isValid(s)) {
+    if (e->GetCaptureStream() != nullptr) {
+      reinterpret_cast<hip::Stream*>(e->GetCaptureStream())->EraseCaptureEvent(event);
+    }
   }
   delete e;
   HIP_RETURN(hipSuccess);
@@ -386,6 +390,7 @@ hipError_t hipEventRecord_common(hipEvent_t event, hipStream_t stream) {
   }
   hip::Event* e = reinterpret_cast<hip::Event*>(event);
   hip::Stream* hip_stream = hip::getStream(stream);
+  e->SetCaptureStream(stream);
   if (g_devices[e->deviceId()]->devices()[0] != &hip_stream->device()) {
     return hipErrorInvalidHandle;
   }
@@ -409,8 +414,16 @@ hipError_t hipEventSynchronize(hipEvent_t event) {
   if (event == nullptr) {
     HIP_RETURN(hipErrorInvalidHandle);
   }
-
   hip::Event* e = reinterpret_cast<hip::Event*>(event);
+  hip::Stream* s = reinterpret_cast<hip::Stream*>(e->GetCaptureStream());
+  if ((s != nullptr) && (s->GetCaptureStatus() == hipStreamCaptureStatusActive)) {
+    if (e->GetCaptureStatus() == false) {
+      return HIP_RETURN(hipErrorStreamCaptureUnsupported);
+    }
+  }
+  if (hip::Stream::StreamCaptureOngoing(e->GetCaptureStream()) == true) {
+    HIP_RETURN(hipErrorStreamCaptureUnsupported);
+  }
   HIP_RETURN(e->synchronize());
 }
 
diff --git a/hipamd/src/hip_stream.cpp b/hipamd/src/hip_stream.cpp
index d9bc512181..7f832af8aa 100644
--- a/hipamd/src/hip_stream.cpp
+++ b/hipamd/src/hip_stream.cpp
@@ -522,6 +522,14 @@ hipError_t hipStreamWaitEvent_common(hipStream_t stream, hipEvent_t event, unsig
   }
 
   hip::Event* e = reinterpret_cast<hip::Event*>(event);
+  if ((e->GetCaptureStream() != nullptr) &&
+      (reinterpret_cast<hip::Stream*>(e->GetCaptureStream())->GetCaptureStatus()
+      == hipStreamCaptureStatusActive)) {
+    // If stream is capturing but event is not recorded on event's stream.
+    if (e->GetCaptureStatus() == false) {
+      return hipErrorStreamCaptureIsolation;
+    }
+  }
   return e->streamWait(stream, flags);
 }