From 7eecbdd9db5884f118d722140ff4ce1ddb9bdfc0 Mon Sep 17 00:00:00 2001
From: Jaydeep Patel <jaydeepkumar.patel@amd.com>
Date: Thu, 12 Jan 2023 07:08:30 +0000
Subject: [PATCH] SWDEV-374360 - Handle free for external memory.

Change-Id: Ibfaf30eef371499a3c1f931a73cce76c83d8f7f3
---
 hipamd/src/hip_memory.cpp | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/hipamd/src/hip_memory.cpp b/hipamd/src/hip_memory.cpp
index 37a53e269e..b72087907e 100644
--- a/hipamd/src/hip_memory.cpp
+++ b/hipamd/src/hip_memory.cpp
@@ -110,7 +110,13 @@ hipError_t ihipFree(void *ptr) {
     hip::Stream::syncNonBlockingStreams(device_id);
     // Find out if memory belongs to any memory pool
     if (!g_devices[device_id]->FreeMemory(memory_object, nullptr)) {
-      amd::SvmBuffer::free(memory_object->getContext(), ptr);
+      // External mem is not svm.
+      if (memory_object->isInterop()) {
+        amd::MemObjMap::RemoveMemObj(ptr);
+        memory_object->release();
+      } else {
+        amd::SvmBuffer::free(memory_object->getContext(), ptr);
+      }
     }
     return hipSuccess;
   }
@@ -163,7 +169,8 @@ hipError_t hipExternalMemoryGetMappedBuffer(
     HIP_RETURN(hipErrorInvalidValue);
   }
   *devPtr = reinterpret_cast<void*>(devMem->virtualAddress() + bufferDesc->offset);
-
+  amd::MemObjMap::AddMemObj(*devPtr, buf);
+  buf->retain();
   HIP_RETURN(hipSuccess);
 }