SWDEV-512344 - Unmap all subbuffers

Since hipMemMap can be called for multiple device handles on the same virtual memory, the same is true for hipMemUnmap, meaning that virtual memory can be "partially unmapped".

This means that the unmap function can be called for a specific part of the reserved address, meaning that only the designated subbuffer should be released. If unmap is called on the entire reserved memory, then all subbuffers should be released.

The main point is that for every hsa_amd_vmem_map, there should be a corresponding hsa_amd_vmem_unmap. Otherwise, if entire memory is unmapped by a single unmap call, then HSA will report the memory as "in use" if an attempt is made to delete it.

Change-Id: I039308eafb820decfb1c09f60347f26cdad1a362

hipamd/src/hip_vm.cpp

@@ -366,12 +366,6 @@ hipError_t hipMemUnmap(void* ptr, size_t size) {
   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);
 }

rocclr/device/rocm/rocvirtual.cpp

@@ -2880,21 +2880,29 @@ void VirtualGPU::submitVirtualMap(amd::VirtualMapCommand& vcmd) {
     dispatchBarrierPacket(kBarrierPacketHeader, false);
     Barriers().WaitCurrent();
 
-    amd::Memory* vaddr_sub_obj = amd::MemObjMap::FindMemObj(vcmd.ptr());
-    assert(vaddr_sub_obj != nullptr);
-
-    // Unmap the object, since the physical addr is set.
-    if ((hsa_status = hsa_amd_vmem_unmap(vaddr_sub_obj->getSvmPtr(), vcmd.size()))
-                        == HSA_STATUS_SUCCESS) {
-      // 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());
-      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;
+    size_t total_unmapped_buffers_size = 0;
+    auto sub_buffers = vaddr_base_obj->subBuffers();
+    for (auto buffer : sub_buffers) {
+      if (total_unmapped_buffers_size + buffer->getSize() <= vcmd.size()) {
+        // Unmap the object, since the physical addr is set.
+        if ((hsa_status = hsa_amd_vmem_unmap(buffer->getSvmPtr(), buffer->getSize())) ==
+            HSA_STATUS_SUCCESS) {
+          buffer->getContext().devices()[0]->DestroyVirtualBuffer(buffer);
+          amd::MemObjMap::RemoveMemObj(buffer->getSvmPtr());
+          if (buffer->getUserData().phys_mem_obj != nullptr) {
+            auto& phys_mem_user_data = buffer->getUserData().phys_mem_obj->getUserData();
+            phys_mem_user_data.vaddr_mem_obj = nullptr;
+            if (phys_mem_user_data.data != nullptr) {
+              reinterpret_cast<amd::RuntimeObject*>(phys_mem_user_data.data)->release();
+            }
+            buffer->getUserData().phys_mem_obj = nullptr;
+            total_unmapped_buffers_size += buffer->getSize();
+          }
+          buffer->release();
+        } else {
+          LogError("HSA Command: hsa_amd_vmem_unmap failed");
+        }
       }
-    } else {
-      LogError("HSA Command: hsa_amd_vmem_unmap failed");
     }
   }