From a8edb8d467ebb5678d2f4506bd01efd3aaeddcab Mon Sep 17 00:00:00 2001
From: Ioannis Assiouras <Ioannis.Assiouras@amd.com>
Date: Thu, 6 Feb 2025 14:59:57 +0000
Subject: [PATCH] SWDEV-508435 - Use the stream of the src/dst image memory
 object in A2H and H2A commands

Change-Id: I9b776a54760a4633d5f84cf7b467d2d3ba8cbdde
---
 hipamd/src/hip_memory.cpp | 28 ++++++++++++++++++++++++++--
 1 file changed, 26 insertions(+), 2 deletions(-)

diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index 66f4675fcb..2f07a5227f 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -1973,8 +1973,20 @@ hipError_t ihipMemcpyHtoACommand(amd::Command*& command, amd::Image* dstImage,
     }
     command = copyCommand;
   } else {
+
+    hip::Stream* pStream = stream;
+    amd::Device* queueDevice = &stream->device();
+    amd::Command::EventWaitList waitList;
+    if (queueDevice != dstImage->GetDeviceById()) {
+      pStream = hip::getNullStream(dstImage->GetDeviceById()->context());
+      amd::Command* cmd = stream->getLastQueuedCommand(true);
+      if (cmd != nullptr) {
+        waitList.push_back(cmd);
+      }
+    }
+
     amd::WriteMemoryCommand* writeMemCmd = new amd::WriteMemoryCommand(
-      *stream, CL_COMMAND_WRITE_IMAGE, amd::Command::EventWaitList{}, *dstImage, dstOrigin,
+      *pStream, CL_COMMAND_WRITE_IMAGE, waitList, *dstImage, dstOrigin,
       copyRegion, static_cast<const char*>(srcHost) + start, srcRowPitch, srcSlicePitch,
       copyMetadata);
     if (writeMemCmd == nullptr) {
@@ -2010,8 +2022,20 @@ hipError_t ihipMemcpyAtoHCommand(amd::Command*& command, void* dstHost, amd::Coo
     }
     command = copyCommand;
   } else {
+
+    hip::Stream* pStream = stream;
+    amd::Device* queueDevice = &stream->device();
+    amd::Command::EventWaitList waitList;
+    if (queueDevice != srcImage->GetDeviceById()) {
+      pStream = hip::getNullStream(srcImage->GetDeviceById()->context());
+      amd::Command* cmd = stream->getLastQueuedCommand(true);
+      if (cmd != nullptr) {
+        waitList.push_back(cmd);
+      }
+    }
+
     amd::ReadMemoryCommand* readMemCmd = new amd::ReadMemoryCommand(
-      *stream, CL_COMMAND_READ_IMAGE, amd::Command::EventWaitList{}, *srcImage, srcOrigin,
+      *pStream, CL_COMMAND_READ_IMAGE, waitList, *srcImage, srcOrigin,
       copyRegion, static_cast<char*>(dstHost) + start, dstRowPitch, dstSlicePitch,
       copyMetadata);