From a2d8199402fcddf60674bd070fca1fa7db7dda1c Mon Sep 17 00:00:00 2001
From: Rahul Manocha <rmanocha@amd.com>
Date: Fri, 1 Nov 2024 22:04:07 +0000
Subject: [PATCH] SWDEV-494409 - Fix destructor path for Static and Dynamic CO

1) For Dynamic CO variables, free the device pointer in
   DynCO destructor instead of DeviceVar destructor.
2) For Static CO Remove Fatbinary,
   only call hipFree for valid device Vars instead of all devices.

Change-Id: I84291f5371b2c05d1d0bcdb4f9c6bd122e7c9b21
---
 hipamd/src/hip_code_object.cpp | 23 +++++++++++++++++++----
 hipamd/src/hip_global.cpp      | 10 ++++++++++
 hipamd/src/hip_global.hpp      |  3 +++
 3 files changed, 32 insertions(+), 4 deletions(-)

diff --git a/hipamd/src/hip_code_object.cpp b/hipamd/src/hip_code_object.cpp
index 8ad95b7aa5..69d4258e79 100644
--- a/hipamd/src/hip_code_object.cpp
+++ b/hipamd/src/hip_code_object.cpp
@@ -965,6 +965,19 @@ DynCO::~DynCO() {
       hipError_t err = ihipFree(elem.second->getManagedVarPtr());
       assert(err == hipSuccess);
     }
+
+    if (elem.second->getVarKind() == Var::DVK_Variable) {
+      for (auto dev : g_devices) {
+        DeviceVar* dvar = nullptr;
+        hipError_t err = elem.second->getDeviceVarPtr(&dvar, dev->deviceId());
+        assert(err == hipSuccess);
+        if (dvar != nullptr) {
+          // free also deletes the device ptr
+          err = ihipFree(dvar->device_ptr());
+          assert(err == hipSuccess);
+        }
+      }
+    }
     delete elem.second;
   }
   vars_.clear();
@@ -1179,10 +1192,12 @@ hipError_t StatCO::removeFatBinary(FatBinaryInfo** module) {
       hipError_t err;
       for (auto dev : g_devices) {
         DeviceVar* dvar = nullptr;
-        IHIP_RETURN_ONFAIL((*it)->getStatDeviceVar(&dvar, dev->deviceId()));
-        // free also deletes the device ptr
-        err = ihipFree(dvar->device_ptr());
-        assert(err == hipSuccess);
+        IHIP_RETURN_ONFAIL((*it)->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**>((*it)->getManagedVarPtr())));
       assert(err == hipSuccess);
diff --git a/hipamd/src/hip_global.cpp b/hipamd/src/hip_global.cpp
index c63831b498..124b91090b 100644
--- a/hipamd/src/hip_global.cpp
+++ b/hipamd/src/hip_global.cpp
@@ -229,6 +229,16 @@ Var::~Var() {
   modules_ = nullptr;
 }
 
+hipError_t Var::getDeviceVarPtr(DeviceVar** dvar, int deviceId) {
+  guarantee((deviceId >= 0), "Invalid DeviceId, less than zero");
+  guarantee((static_cast<size_t>(deviceId) < g_devices.size()),
+            "Invalid DeviceId, greater than no of code objects");
+  guarantee((dVar_.size() == g_devices.size()),
+             "Device Var not initialized to size");
+  *dvar = dVar_[deviceId];
+  return hipSuccess;
+}
+
 hipError_t Var::getDeviceVar(DeviceVar** dvar, int deviceId, hipModule_t hmod) {
   guarantee((deviceId >= 0), "Invalid DeviceId, less than zero");
   guarantee((static_cast<size_t>(deviceId) < g_devices.size()),
diff --git a/hipamd/src/hip_global.hpp b/hipamd/src/hip_global.hpp
index 5d28da21be..22ba881809 100644
--- a/hipamd/src/hip_global.hpp
+++ b/hipamd/src/hip_global.hpp
@@ -121,6 +121,9 @@ public:
 
   //Return DeviceVar for module Generate/build if not already done so.
   hipError_t getStatDeviceVar(DeviceVar** dvar, int deviceId);
+
+  hipError_t getDeviceVarPtr(DeviceVar** dvar, int deviceId);
+
   void resize_dVar(size_t size) { dVar_.resize(size); }
 
   FatBinaryInfo** moduleInfo() { return modules_; };