From b1c0f73229eddb733f8e04934e6f4a053b172c7a Mon Sep 17 00:00:00 2001
From: German Andryeyev <German.Andryeyev@amd.com>
Date: Tue, 23 Apr 2024 19:15:14 -0400
Subject: [PATCH] SWDEV-353281 - Corret VA unmap

Make sure graph mempool unmaps VA on release

Change-Id: Id3f1bd8d0115b533ae60aa5ba3676b8bf7e5b961


[ROCm/clr commit: 5c1804aa14b29336cfab845b2165b3949dc9a7f9]
---
 projects/clr/hipamd/src/hip_mempool_impl.cpp   | 18 ++++++++++--------
 projects/clr/rocclr/device/pal/palvirtual.cpp  |  6 +++++-
 projects/clr/rocclr/device/rocm/rocvirtual.cpp |  6 +++++-
 projects/clr/rocclr/platform/memory.hpp        |  1 +
 4 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_mempool_impl.cpp b/projects/clr/hipamd/src/hip_mempool_impl.cpp
index 851836039f..dc091c71e3 100644
--- a/projects/clr/hipamd/src/hip_mempool_impl.cpp
+++ b/projects/clr/hipamd/src/hip_mempool_impl.cpp
@@ -209,8 +209,7 @@ void* MemoryPool::AllocateMemory(size_t size, Stream* stream, void* dptr) {
       }
     }
   } else {
-    const device::Memory* dev_mem = memory->getDeviceMemory(*device_->devices()[0]);
-    dev_ptr = reinterpret_cast<void*>(dev_mem->virtualAddress());
+    dev_ptr = memory->getSvmPtr();
   }
   // Place the allocated memory into the busy heap
   ts.AddSafeStream(stream);
@@ -259,14 +258,17 @@ bool MemoryPool::FreeMemory(amd::Memory* memory, Stream* stream, Event* event) {
     }
     ClPrint(amd::LOG_INFO, amd::LOG_MEM_POOL, "Pool FreeMem: %p, %p", memory->getSvmPtr(), memory);
 
-    if (stream == nullptr) {
+    if (memory->getUserData().vaddr_mem_obj != nullptr) {
+      auto va_mem = memory->getUserData().vaddr_mem_obj;
+      if (stream == nullptr) {
         stream = g_devices[memory->getUserData().deviceId]->NullStream();
+      }
+      // Unmap virtual address from memory
+      auto cmd = new amd::VirtualMapCommand(*stream, amd::Command::EventWaitList{},
+                                            va_mem->getSvmPtr(), va_mem->getSize(), nullptr);
+      cmd->enqueue();
+      cmd->release();
     }
-    // Unmap virtual address from memory
-    auto cmd = new amd::VirtualMapCommand(*stream, amd::Command::EventWaitList{},
-                                          memory->getSvmPtr(), memory->getSize(), nullptr);
-    cmd->enqueue();
-    cmd->release();
 
     if (stream != nullptr) {
       // The stream of destruction is a safe stream, because the app must handle sync
diff --git a/projects/clr/rocclr/device/pal/palvirtual.cpp b/projects/clr/rocclr/device/pal/palvirtual.cpp
index bb49c3af9f..a1aaa86ce5 100644
--- a/projects/clr/rocclr/device/pal/palvirtual.cpp
+++ b/projects/clr/rocclr/device/pal/palvirtual.cpp
@@ -2230,11 +2230,15 @@ void VirtualGPU::submitVirtualMap(amd::VirtualMapCommand& vcmd) {
       assert(amd::MemObjMap::FindMemObj(vcmd.ptr()) == nullptr);
       amd::MemObjMap::AddMemObj(vcmd.ptr(), vaddr_mem_obj);
       vaddr_mem_obj->getUserData().phys_mem_obj = vcmd.memory();
+      vcmd.memory()->getUserData().vaddr_mem_obj = vaddr_mem_obj;
     } else {
       // assert the vaddr_mem_obj is mapped and needs to be removed
       assert(amd::MemObjMap::FindMemObj(vcmd.ptr()) != nullptr);
       amd::MemObjMap::RemoveMemObj(vcmd.ptr());
-      vaddr_mem_obj->getUserData().phys_mem_obj = nullptr;
+      if (vaddr_mem_obj->getUserData().phys_mem_obj != nullptr) {
+        vaddr_mem_obj->getUserData().phys_mem_obj->getUserData().vaddr_mem_obj = nullptr;
+        vaddr_mem_obj->getUserData().phys_mem_obj = nullptr;
+      }
     }
   }
   profilingEnd(vcmd);
diff --git a/projects/clr/rocclr/device/rocm/rocvirtual.cpp b/projects/clr/rocclr/device/rocm/rocvirtual.cpp
index 4279f39bef..2ecfc53ef3 100644
--- a/projects/clr/rocclr/device/rocm/rocvirtual.cpp
+++ b/projects/clr/rocclr/device/rocm/rocvirtual.cpp
@@ -2598,6 +2598,7 @@ void VirtualGPU::submitVirtualMap(amd::VirtualMapCommand& vcmd) {
       assert(amd::MemObjMap::FindMemObj(vcmd.ptr()) == nullptr);
       amd::MemObjMap::AddMemObj(vcmd.ptr(), vaddr_sub_obj);
       vaddr_sub_obj->getUserData().phys_mem_obj = phys_mem_obj;
+      phys_mem_obj->getUserData().vaddr_mem_obj = vaddr_sub_obj;
     } else {
       LogError("HSA Command: hsa_amd_vmem_map failed!");
     }
@@ -2614,7 +2615,10 @@ void VirtualGPU::submitVirtualMap(amd::VirtualMapCommand& vcmd) {
       // assert the va is mapped and needs to be removed
       vaddr_sub_obj->getContext().devices()[0]->DestroyVirtualBuffer(vaddr_sub_obj);
       amd::MemObjMap::RemoveMemObj(vcmd.ptr());
-      vaddr_sub_obj->getUserData().phys_mem_obj = nullptr;
+      if (vaddr_sub_obj->getUserData().phys_mem_obj != nullptr) {
+        vaddr_sub_obj->getUserData().phys_mem_obj->getUserData().vaddr_mem_obj = nullptr;
+        vaddr_sub_obj->getUserData().phys_mem_obj = nullptr;
+      }
     } else {
       LogError("HSA Command: hsa_amd_vmem_unmap failed");
     }
diff --git a/projects/clr/rocclr/platform/memory.hpp b/projects/clr/rocclr/platform/memory.hpp
index 69f98111de..26da54ba0b 100644
--- a/projects/clr/rocclr/platform/memory.hpp
+++ b/projects/clr/rocclr/platform/memory.hpp
@@ -150,6 +150,7 @@ class Memory : public amd::RuntimeObject {
      int deviceId = 0;     //!< Device ID memory is allocated on
      void* data = nullptr; //!< Opaque user data from CL or HIP or etc.
      amd::Memory* phys_mem_obj = nullptr; //<! Physical mem obj, only set on virtual mem
+     amd::Memory* vaddr_mem_obj = nullptr; //<! Virtual address mem obj, only set on virtual mem
      uint64_t hsa_handle = 0; //!<Opaque hsa handle saved for Virtual memories
      unsigned int flags = 0; //!< HIP memory flags
      //! hipMallocPitch allocates buffer using width & height and returns pitch & device pointer.