From ae3b053ddfd85494c08f7f563c7c37d3a26a6261 Mon Sep 17 00:00:00 2001
From: Pengda Xie <pengda.xie@amd.com>
Date: Fri, 28 Feb 2025 11:07:43 -0800
Subject: [PATCH] SWDEV-518317 - Don't attempt to remove managedVars when map
 is empty

Change-Id: I25c33487dc08f96c087b6acc1abe42a4a666a609
---
 hipamd/src/hip_code_object.cpp | 75 +++++++++++++++++++---------------
 1 file changed, 41 insertions(+), 34 deletions(-)

diff --git a/hipamd/src/hip_code_object.cpp b/hipamd/src/hip_code_object.cpp
index 8700531484..bb18f45e16 100644
--- a/hipamd/src/hip_code_object.cpp
+++ b/hipamd/src/hip_code_object.cpp
@@ -1296,43 +1296,50 @@ FatBinaryInfo** StatCO::addFatBinary(const void* data, bool initialized, bool& s
 hipError_t StatCO::removeFatBinary(FatBinaryInfo** module) {
   amd::ScopedLock lock(sclock_);
 
-  auto hostVarsIter = module_to_hostVars_.find(module);
-  if (hostVarsIter != module_to_hostVars_.end()) {
-    for (auto& hostVar : hostVarsIter->second) {
-      auto varIter = vars_.find(hostVar);
-      if (varIter == vars_.end()) {
-        LogPrintfError("removeFatBinary: Unable to find module 0x%x hostVar 0x%x",
-                        module, hostVar);
-      } else {
-        delete varIter->second;
-        vars_.erase(varIter);
-      }
-    }
-    module_to_hostVars_.erase(hostVarsIter);
-  } else {
-    LogPrintfError("removeFatBinary: Unable to find module 0x%x hostVars", module);
-  }
-
-  auto managedVarsIter = managedVars_.find(module);
-  if (managedVarsIter != managedVars_.end()) {
-    for (auto& managedVar : managedVarsIter->second) {
-      hipError_t err;
-      for (auto dev : g_devices) {
-        DeviceVar* dvar = nullptr;
-        IHIP_RETURN_ONFAIL(managedVar->getDeviceVarPtr(&dvar, dev->deviceId()));
-        if (dvar != nullptr) {
-          // free also deletes the device ptr
-          err = ihipFree(dvar->device_ptr());
-          assert(err == hipSuccess);
+  if (!module_to_hostVars_.empty()) {
+    auto hostVarsIter = module_to_hostVars_.find(module);
+    if (hostVarsIter != module_to_hostVars_.end()) {
+      for (auto& hostVar : hostVarsIter->second) {
+        auto varIter = vars_.find(hostVar);
+        if (varIter == vars_.end()) {
+          LogPrintfError(
+            "removeFatBinary: Unable to find module 0x%x hostVar 0x%x",
+            module, hostVar);
+        } else {
+          delete varIter->second;
+          vars_.erase(varIter);
         }
       }
-      err = ihipFree(*(static_cast<void**>(managedVar->getManagedVarPtr())));
-      assert(err == hipSuccess);
-      delete managedVar;
+      module_to_hostVars_.erase(hostVarsIter);
+    } else {
+      LogPrintfError(
+        "removeFatBinary: Unable to find module 0x%x hostVars", module);
+    }
+  }
+
+  if (!managedVars_.empty()) {
+    auto managedVarsIter = managedVars_.find(module);
+    if (managedVarsIter != managedVars_.end()) {
+      for (auto& managedVar : managedVarsIter->second) {
+        hipError_t err;
+        for (auto dev : g_devices) {
+          DeviceVar* dvar = nullptr;
+          IHIP_RETURN_ONFAIL(managedVar->getDeviceVarPtr(&dvar, dev->deviceId()));
+          if (dvar != nullptr) {
+            // free also deletes the device ptr
+            err = ihipFree(dvar->device_ptr());
+            assert(err == hipSuccess);
+          }
+        }
+        err = ihipFree(*(static_cast<void**>(managedVar->getManagedVarPtr())));
+        assert(err == hipSuccess);
+        delete managedVar;
+      }
+      managedVars_.erase(managedVarsIter);
+    } else {
+      LogPrintfError("removeFatBinary: Unable to find module 0x%x managedVars",
+                     module);
     }
-    managedVars_.erase(managedVarsIter);
-  } else {
-    LogPrintfError("removeFatBinary: Unable to find module 0x%x managedVars", module);
   }
 
   auto hostFuncsIter = module_to_hostFunctions_.find(module);