From 3b7e1b928e227f9e3cd36149672c0eebf6b287f8 Mon Sep 17 00:00:00 2001
From: jujiang <julia.jiang@amd.com>
Date: Mon, 13 Jul 2020 15:57:15 -0400
Subject: [PATCH] SWDEV-242337-Fix hipHostRegister test failure for MGPU

Change-Id: I748dac0c5a47fa00184f065d9ef716573aa70b12
---
 rocclr/hip_memory.cpp | 11 ++++++++---
 1 file changed, 8 insertions(+), 3 deletions(-)

diff --git a/rocclr/hip_memory.cpp b/rocclr/hip_memory.cpp
index 553b59a77c..a7b23bcc4c 100755
--- a/rocclr/hip_memory.cpp
+++ b/rocclr/hip_memory.cpp
@@ -725,9 +725,14 @@ hipError_t hipHostUnregister(void* hostPtr) {
     amd::Memory* mem = getMemoryObject(hostPtr, offset);
 
     if(mem) {
-      for (const auto& device: hip::getCurrentDevice()->devices()) {
-        const device::Memory* devMem = mem->getDeviceMemory(*device);
-        amd::MemObjMap::RemoveMemObj(reinterpret_cast<void*>(devMem->virtualAddress()));
+      for (const auto& device: g_devices) {
+        const device::Memory* devMem = mem->getDeviceMemory(*device->devices()[0]);
+        if (devMem != nullptr) {
+          void* vAddr = reinterpret_cast<void*>(devMem->virtualAddress());
+          if (amd::MemObjMap::FindMemObj(vAddr)) {
+            amd::MemObjMap::RemoveMemObj(vAddr);
+          }
+        }
       }
       amd::MemObjMap::RemoveMemObj(hostPtr);
       mem->release();