From 64799f183f6abccd24573092496b97e65c7ff211 Mon Sep 17 00:00:00 2001
From: anusha GodavarthySurya <Anusha.GodavarthySurya@amd.com>
Date: Wed, 25 Aug 2021 08:16:52 -0700
Subject: [PATCH] SWDEV-300471 - updated hipFree to wait for all streams to
 finish

Change-Id: I4983ef48a9e720ed03128d83a8d23ad11024fd8f


[ROCm/clr commit: c9d9fd992a6c6494e591779f0dde03254446f1ac]
---
 projects/clr/hipamd/src/hip_context.cpp        | 10 ++++++++++
 projects/clr/hipamd/src/hip_device_runtime.cpp |  2 +-
 projects/clr/hipamd/src/hip_internal.hpp       |  4 +++-
 projects/clr/hipamd/src/hip_memory.cpp         |  4 +++-
 projects/clr/hipamd/src/hip_stream.cpp         |  6 ++++--
 5 files changed, 21 insertions(+), 5 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_context.cpp b/projects/clr/hipamd/src/hip_context.cpp
index 3779c317d2..caf6b195c2 100644
--- a/projects/clr/hipamd/src/hip_context.cpp
+++ b/projects/clr/hipamd/src/hip_context.cpp
@@ -115,6 +115,16 @@ amd::HostQueue* getNullStream(amd::Context& ctx) {
   return nullptr;
 }
 
+// ================================================================================================
+int getDeviceID(amd::Context& ctx) {
+  for (auto& it : g_devices) {
+    if (it->asContext() == &ctx) {
+      return it->deviceId();
+    }
+  }
+  return -1;
+}
+
 // ================================================================================================
 amd::HostQueue* getNullStream() {
   Device* device = getCurrentDevice();
diff --git a/projects/clr/hipamd/src/hip_device_runtime.cpp b/projects/clr/hipamd/src/hip_device_runtime.cpp
index d2c9d30bfa..7ba0e5d0d7 100644
--- a/projects/clr/hipamd/src/hip_device_runtime.cpp
+++ b/projects/clr/hipamd/src/hip_device_runtime.cpp
@@ -467,7 +467,7 @@ hipError_t hipDeviceSynchronize ( void ) {
 
   queue->finish();
 
-  hip::Stream::syncNonBlockingStreams();
+  hip::Stream::syncNonBlockingStreams(hip::getCurrentDevice()->deviceId());
 
   HIP_RETURN(hipSuccess);
 }
diff --git a/projects/clr/hipamd/src/hip_internal.hpp b/projects/clr/hipamd/src/hip_internal.hpp
index 6fdded7ff5..bcc6954435 100644
--- a/projects/clr/hipamd/src/hip_internal.hpp
+++ b/projects/clr/hipamd/src/hip_internal.hpp
@@ -210,7 +210,7 @@ namespace hip {
     const std::vector<uint32_t> GetCUMask() const { return cuMask_; }
 
     /// Sync all non-blocking streams
-    static void syncNonBlockingStreams();
+    static void syncNonBlockingStreams(int deviceId = -1);
 
     /// Destroy all streams on a given device
     static void destroyAllStreams(int deviceId);
@@ -325,6 +325,8 @@ namespace hip {
   extern amd::HostQueue* getNullStream(amd::Context&);
   /// Get default stream of the thread
   extern amd::HostQueue* getNullStream();
+  /// Get device ID associated with the ROCclr context
+  int getDeviceID(amd::Context& ctx);
   /// Check if stream is valid
   extern bool isValid(hipStream_t& stream);
 };
diff --git a/projects/clr/hipamd/src/hip_memory.cpp b/projects/clr/hipamd/src/hip_memory.cpp
index 1f1d3e16b8..0aed8e2318 100644
--- a/projects/clr/hipamd/src/hip_memory.cpp
+++ b/projects/clr/hipamd/src/hip_memory.cpp
@@ -86,7 +86,7 @@ hipError_t ihipFree(void *ptr)
 
   size_t offset = 0;
   amd::Memory* memory_object = getMemoryObject(ptr, offset);
-
+  int deviceID = -1;
   if (memory_object != nullptr) {
     // Check if it's an allocation in system memory and can be shared across all devices
     if (memory_object->getMemFlags() & CL_MEM_SVM_FINE_GRAIN_BUFFER) {
@@ -102,7 +102,9 @@ hipError_t ihipFree(void *ptr)
     } else {
       // Wait on the device, associated with the current memory object
       hip::getNullStream(memory_object->getContext())->finish();
+      deviceID = hip::getDeviceID(memory_object->getContext());
     }
+    hip::Stream::syncNonBlockingStreams(deviceID);
     amd::SvmBuffer::free(memory_object->getContext(), ptr);
     return hipSuccess;
   }
diff --git a/projects/clr/hipamd/src/hip_stream.cpp b/projects/clr/hipamd/src/hip_stream.cpp
index 36b0060bc1..99f15e6277 100644
--- a/projects/clr/hipamd/src/hip_stream.cpp
+++ b/projects/clr/hipamd/src/hip_stream.cpp
@@ -148,11 +148,13 @@ int Stream::DeviceId(const hipStream_t hStream) {
   return deviceId;
 }
 
-void Stream::syncNonBlockingStreams() {
+void Stream::syncNonBlockingStreams(int deviceId) {
   amd::ScopedLock lock(streamSetLock);
   for (auto& it : streamSet) {
     if (it->Flags() & hipStreamNonBlocking) {
-      it->asHostQueue()->finish();
+      if (deviceId == -1 || it->DeviceId() == deviceId) {
+        it->asHostQueue()->finish();
+      }
     }
   }
 }