From b457d0ba8228906dc19ba186cd6c555a7ac2bf66 Mon Sep 17 00:00:00 2001
From: kjayapra-amd <karthik.jayaprakash@amd.com>
Date: Wed, 27 Dec 2023 15:59:42 -0500
Subject: [PATCH] SWDEV-413997 - Physical address size could be more than
 mapped size, only virtual address needs to be match the size of unmap
 request.

Change-Id: If0e973149d8ce1b6f01000ab951cf6e9f4b38e97


[ROCm/clr commit: 0f9c624f4ca41b2024dd013c63a865aeb53ef147]
---
 projects/clr/hipamd/src/hip_vm.cpp | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_vm.cpp b/projects/clr/hipamd/src/hip_vm.cpp
index bf76ffa490..2335c426bc 100644
--- a/projects/clr/hipamd/src/hip_vm.cpp
+++ b/projects/clr/hipamd/src/hip_vm.cpp
@@ -306,11 +306,17 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
     HIP_RETURN(hipErrorInvalidValue);
   }
 
-  amd::Memory* va = amd::MemObjMap::FindMemObj(ptr);
-  if (va && va->getSize() != size) {
+  amd::Memory* pa = amd::MemObjMap::FindMemObj(ptr);
+  if (pa == nullptr) {
     HIP_RETURN(hipErrorInvalidValue);
   }
-  auto& queue = *g_devices[va->getUserData().deviceId]->NullStream();
+
+  amd::Memory* va = amd::MemObjMap::FindVirtualMemObj(ptr);
+  if (va == nullptr && va->getSize() != size) {
+    HIP_RETURN(hipErrorInvalidValue);
+  }
+
+  auto& queue = *g_devices[pa->getUserData().deviceId]->NullStream();
 
   amd::Command* cmd = new amd::VirtualMapCommand(queue, amd::Command::EventWaitList{}, ptr, size,
                                                  nullptr);
@@ -318,9 +324,9 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
   cmd->awaitCompletion();
   cmd->release();
 
-  // restore the original va of the generic allocation
-  hip::GenericAllocation* ga = reinterpret_cast<hip::GenericAllocation*>(va->getUserData().data);
-  va->setSvmPtr(ga->genericAddress());
+  // restore the original pa of the generic allocation
+  hip::GenericAllocation* ga = reinterpret_cast<hip::GenericAllocation*>(pa->getUserData().data);
+  pa->setSvmPtr(ga->genericAddress());
 
   ga->release();