From 3d862af4dffa9d516d24b0e113fcebee3288c7dd Mon Sep 17 00:00:00 2001
From: "Arandjelovic, Marko" <Marko.Arandjelovic@amd.com>
Date: Mon, 12 May 2025 16:56:10 +0200
Subject: [PATCH] SWDEV-512344 - Unmap all subbuffers (#214)

[ROCm/clr commit: c5ced8c3a2bd193f740790622058761b25adc336]
---
 projects/clr/hipamd/src/hip_vm.cpp | 42 +++++++++++++++++-------------
 1 file changed, 24 insertions(+), 18 deletions(-)

diff --git a/projects/clr/hipamd/src/hip_vm.cpp b/projects/clr/hipamd/src/hip_vm.cpp
index fc2533d8a3..184908ad3b 100644
--- a/projects/clr/hipamd/src/hip_vm.cpp
+++ b/projects/clr/hipamd/src/hip_vm.cpp
@@ -354,26 +354,32 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
     HIP_RETURN(hipErrorInvalidValue);
   }
 
-  amd::Memory* phys_mem_obj = vaddr_sub_obj->getUserData().phys_mem_obj;
-  if (phys_mem_obj == nullptr) {
-    HIP_RETURN(hipErrorInvalidValue);
+  address end_address = reinterpret_cast<address>(vaddr_sub_obj->getSvmPtr()) + size;
+  while (vaddr_sub_obj &&
+         reinterpret_cast<address>(vaddr_sub_obj->getSvmPtr()) + vaddr_sub_obj->getSize() <=
+             end_address) {
+    amd::Memory* phys_mem_obj = vaddr_sub_obj->getUserData().phys_mem_obj;
+    if (phys_mem_obj == nullptr) {
+      HIP_RETURN(hipErrorInvalidValue);
+    }
+
+    amd::Command* cmd = new amd::VirtualMapCommand(
+        *hip::getCurrentDevice()->NullStream(), amd::Command::EventWaitList{},
+        vaddr_sub_obj->getSvmPtr(), vaddr_sub_obj->getSize(), nullptr);
+
+    cmd->enqueue();
+    cmd->awaitCompletion();
+    cmd->release();
+    // restore the original pa of the generic allocation
+    hip::GenericAllocation* ga =
+        reinterpret_cast<hip::GenericAllocation*>(phys_mem_obj->getUserData().data);
+    ga->release();
+    address next_subbuffer_ptr =
+        reinterpret_cast<address>(vaddr_sub_obj->getSvmPtr()) + vaddr_sub_obj->getSize();
+    vaddr_sub_obj->release();
+    vaddr_sub_obj = amd::MemObjMap::FindMemObj(next_subbuffer_ptr);
   }
 
-  auto& queue = *g_devices[phys_mem_obj->getUserData().deviceId]->NullStream();
-
-  amd::Command* cmd = new amd::VirtualMapCommand(queue, amd::Command::EventWaitList{}, ptr, size,
-                                                 nullptr);
-  cmd->enqueue();
-  cmd->awaitCompletion();
-  cmd->release();
-  vaddr_sub_obj->release();
-
-  // restore the original pa of the generic allocation
-  hip::GenericAllocation* ga
-    = reinterpret_cast<hip::GenericAllocation*>(phys_mem_obj->getUserData().data);
-  ga->release();
-
   HIP_RETURN(hipSuccess);
 }
 } //namespace hip
-