SWDEV-562761 - Cleanup static fatbin on runtime teardown (#1873)

projects/clr/hipamd/src/hip_code_object.cpp

@@ -361,6 +361,62 @@ hipError_t StatCO::removeFatBinary(FatBinaryInfo** module) {
   return hipSuccess;
 }
 
+// =================================================================================================
+void StatCO::RemoveAllFatBinaries() {
+  amd::ScopedLock lock(sclock_);
+
+  // Clear mapping tables that associate modules with host-side constructs
+  module_to_hostModule_.clear();
+  module_to_hostFunctions_.clear();
+  module_to_hostVars_.clear();
+
+  // Delete all registered variables and clear the container
+  for (auto const& [_, var] : vars_) {
+    delete var;
+  }
+  vars_.clear();
+
+  // Clean up managed variables - these require special handling for memory on each device
+  for (auto& [_, managed_vars] : managedVars_) {
+    for (auto& managed_var : managed_vars) {
+      // Free device-specific allocations across all devices
+      for (auto dev : g_devices) {
+        DeviceVar* dvar = nullptr;
+        if (managed_var->getDeviceVarPtr(&dvar, dev->deviceId()) == hipSuccess && dvar) {
+          // Free device memory (also deletes the device ptr)
+          [[maybe_unused]] hipError_t err = ihipFree(dvar->device_ptr());
+          assert(err == hipSuccess);
+        }
+      }
+
+      // Free the managed memory allocation itself
+      void** managed_ptr = static_cast<void**>(managed_var->getManagedVarPtr());
+      if (managed_var->getAllocFlag()) {
+        // Memory was allocated with ihipMallocManaged - use ihipFree
+        [[maybe_unused]] hipError_t err = ihipFree(*managed_ptr);
+        assert(err == hipSuccess);
+      } else {
+        // Memory was allocated with OS-level allocator - use OS release
+        amd::Os::releaseMemory(*managed_ptr, managed_var->getSize());
+      }
+      delete managed_var;
+    }
+  }
+  managedVars_.clear();
+
+  // Delete all registered functions and clear the container
+  for (auto const& [_, func] : functions_) {
+    delete func;
+  }
+  functions_.clear();
+
+  // Delete all fat binary info objects and clear the modules container
+  for (auto const& [_, fb_info] : modules_) {
+    delete fb_info;
+  }
+  modules_.clear();
+}
+
 hipError_t StatCO::registerStatFunction(const void* hostFunction, Function* func) {
   amd::ScopedLock lock(sclock_);
 

projects/clr/hipamd/src/hip_code_object.hpp

@@ -156,6 +156,7 @@ class StatCO : public CodeObject {
   FatBinaryInfo** addFatBinary(const void* data, bool initialized, bool& success);
   hipError_t removeFatBinary(FatBinaryInfo** module);
   hipError_t digestFatBinary(const void* data, FatBinaryInfo*& programs);
+  void RemoveAllFatBinaries();
 
   // Register vars/funcs given to use from __hipRegister[Var/Func/ManagedVar]
   hipError_t registerStatFunction(const void* hostFunction, Function* func);

projects/clr/hipamd/src/hip_platform.cpp

@@ -813,6 +813,9 @@ void PlatformState::init() {
   for (auto& it : statCO_.functions_) {
     it.second->resize_dFunc(g_devices.size());
   }
+  amd::RuntimeTearDown::RegisterTearDownCallback("PlatformState static fatbin cleanup", [this]() {
+    statCO_.RemoveAllFatBinaries();
+  });
 }
 
 hipError_t PlatformState::loadModule(hipModule_t* module, const char* fname, const void* image) {

projects/clr/rocclr/platform/runtime.cpp

@@ -105,17 +105,29 @@ void Runtime::tearDown() {
 }
 
 // ~RuntimeTearDown() will reference listenerLock.
-// listenerLock will be constructed ealier and destructed later than
+// listenerLock will be constructed earlier and destructed later than
 // runtime_tear_down.
-amd::Monitor listenerLock("Hostcall listener lock");
-std::vector<ReferenceCountedObject*> RuntimeTearDown::external_;
+amd::Monitor listenerLock{};
+std::vector<ReferenceCountedObject*> RuntimeTearDown::external_{};
+std::vector<std::pair<std::string, RuntimeTearDown::TearDownCallback>> 
+  RuntimeTearDown::tear_down_funcs_{};
+class RuntimeTearDown runtime_tear_down{};
 
+// =================================================================================================
 RuntimeTearDown::~RuntimeTearDown() {
+  ClPrint(amd::LOG_INFO, amd::LOG_INIT, "Begin runtime teardown");
 #if !defined(_WIN32) && !defined(BUILD_STATIC_LIBS)
   // Only perform destruction if process matches the initialization,
-  // to avoid a call with the child process after fork()
+  // to avoid a call with the child process after fork().
   if (amd::IS_HIP && amd::Os::getProcessId() == Runtime::pid()) {
+    // Execute teardown funcs in reverse order of registration.
+    for (auto it = tear_down_funcs_.rbegin(); it != tear_down_funcs_.rend(); ++it) {
+      ClPrint(amd::LOG_DEBUG, amd::LOG_INIT, "~RuntimeTearDown invoke callback: %s",
+              it->first.c_str());
+      it->second();
+    }
     for (auto it : external_) {
+      ClPrint(amd::LOG_DEBUG, amd::LOG_INIT, "~RuntimeTearDown release external object: %p", it);
       it->release();
     }
     Runtime::tearDown();
@@ -123,10 +135,19 @@ RuntimeTearDown::~RuntimeTearDown() {
 #endif
 }
 
-void RuntimeTearDown::RegisterObject(ReferenceCountedObject* obj) { external_.push_back(obj); }
+// =================================================================================================
+void RuntimeTearDown::RegisterObject(ReferenceCountedObject* obj) {
+  external_.push_back(obj);
+  ClPrint(amd::LOG_DEBUG, amd::LOG_INIT, "RuntimeTearDown registered external object: %p", obj);
+}
 
-class RuntimeTearDown runtime_tear_down;
+// =================================================================================================
+void RuntimeTearDown::RegisterTearDownCallback(const std::string& msg, TearDownCallback func) {
+  tear_down_funcs_.emplace_back(msg, std::move(func));
+  ClPrint(amd::LOG_DEBUG, amd::LOG_INIT, "RuntimeTearDown registered callback: %s", msg.c_str());
+}
 
+// =================================================================================================
 uint ReferenceCountedObject::retain() {
   uint prev = referenceCount_.fetch_add(1, std::memory_order_relaxed);
   assert(prev != 0 && "An object with count==0 is invalid");

projects/clr/rocclr/platform/runtime.hpp

@@ -21,6 +21,7 @@
 #ifndef RUNTIME_HPP_
 #define RUNTIME_HPP_
 
+#include <functional>
 #include "top.hpp"
 #include "thread/thread.hpp"
 
@@ -61,13 +62,18 @@ class Runtime : AllStatic {
 /*@}*/
 
 class RuntimeTearDown : public HeapObject {
-  static std::vector<ReferenceCountedObject*> external_;
-
  public:
+  using TearDownCallback = std::function<void()>;
+
   RuntimeTearDown() {}
   ~RuntimeTearDown();
 
   static void RegisterObject(ReferenceCountedObject* obj);
+  static void RegisterTearDownCallback(const std::string& msg, TearDownCallback func);
+
+ private:
+  static std::vector<ReferenceCountedObject*> external_;
+  static std::vector<std::pair<std::string, TearDownCallback>> tear_down_funcs_;
 };
 
 }  // namespace amd