From 52f5b0fa1676803094951aa27f33bf53bead451b Mon Sep 17 00:00:00 2001
From: German Andryeyev <German.Andryeyev@amd.com>
Date: Thu, 24 Jun 2021 15:02:58 -0400
Subject: [PATCH] SWDEV-290160 - Add blocking command for callbacks in IPC and
 graph

Change-Id: I1474e19db6ebf2559e1b2c8dd7d06139bd699835
---
 hipamd/src/hip_event.cpp          | 15 ++++++++++++++-
 hipamd/src/hip_graph_internal.cpp | 13 +++++++++++++
 2 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/hipamd/src/hip_event.cpp b/hipamd/src/hip_event.cpp
index 898b2406e7..3f37b39c51 100755
--- a/hipamd/src/hip_event.cpp
+++ b/hipamd/src/hip_event.cpp
@@ -289,7 +289,8 @@ bool createIpcEventShmemIfNeeded(hip::Event::ihipIpcEvent_t& ipc_evt) {
   ipc_evt.ipc_name_ = name_template;
   ipc_evt.ipc_name_.replace(0, 5, "/hip_");
   if (!amd::Os::MemoryMapFileTruncated(ipc_evt.ipc_name_.c_str(),
-      const_cast<const void**> (reinterpret_cast<void**>(&(ipc_evt.ipc_shmem_))), sizeof(hip::ihipIpcEventShmem_t))) {
+      const_cast<const void**> (reinterpret_cast<void**>(&(ipc_evt.ipc_shmem_))),
+      sizeof(hip::ihipIpcEventShmem_t))) {
     return false;
   }
   ipc_evt.ipc_shmem_->owners = 1;
@@ -306,6 +307,7 @@ bool createIpcEventShmemIfNeeded(hip::Event::ihipIpcEvent_t& ipc_evt) {
 #endif
 }
 
+// ================================================================================================
 hipError_t hipEventRecord(hipEvent_t event, hipStream_t stream) {
   HIP_INIT_API(hipEventRecord, event, stream);
 
@@ -346,6 +348,17 @@ hipError_t hipEventRecord(hipEvent_t event, hipStream_t stream) {
     }
     command->enqueue();
     tEvent.notifyCmdQueue();
+
+    // Add the new barrier to stall the stream, until the callback is done
+    amd::Command::EventWaitList eventWaitList;;
+    eventWaitList.push_back(command);
+    amd::Command* block_command = new amd::Marker(*queue, !kMarkerDisableFlush, eventWaitList);
+    if (block_command == nullptr) {
+      return hipErrorInvalidValue;
+    }
+    block_command->enqueue();
+    block_command->release();
+
     // Update read index to indicate new signal.
     int expected = write_index - 1;
     while (!e->ipc_evt_.ipc_shmem_->read_index.compare_exchange_weak(expected, write_index)) {
diff --git a/hipamd/src/hip_graph_internal.cpp b/hipamd/src/hip_graph_internal.cpp
index 75eb6f6726..4a8f00ceb5 100644
--- a/hipamd/src/hip_graph_internal.cpp
+++ b/hipamd/src/hip_graph_internal.cpp
@@ -355,10 +355,23 @@ hipError_t hipGraphExec::Run(hipStream_t stream) {
   if (!event.setCallback(CL_COMPLETE, hipGraphExec::ResetGraph, command)) {
     return hipErrorInvalidHandle;
   }
+
   hipGraphExec::activeGraphExec_[command] = this;
   lastEnqueuedGraphCmd_ = command;
   bExecPending_.store(true);
   command->enqueue();
+
+  // Add the new barrier to stall the stream, until the callback is done
+  amd::Command::EventWaitList eventWaitList;;
+  eventWaitList.push_back(command);
+  amd::Command* block_command = new amd::Marker(*queue, !kMarkerDisableFlush, eventWaitList);
+  if (block_command == nullptr) {
+    return hipErrorInvalidValue;
+  }
+  block_command->enqueue();
+  block_command->release();
+
   command->release();
+
   return hipSuccess;
 }
\ No newline at end of file