diff --git a/projects/clr/hipamd/src/hip_code_object.cpp b/projects/clr/hipamd/src/hip_code_object.cpp index e2277c6026..8c2eef23f2 100644 --- a/projects/clr/hipamd/src/hip_code_object.cpp +++ b/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(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_); diff --git a/projects/clr/hipamd/src/hip_code_object.hpp b/projects/clr/hipamd/src/hip_code_object.hpp index b625af364e..20664686dd 100644 --- a/projects/clr/hipamd/src/hip_code_object.hpp +++ b/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); diff --git a/projects/clr/hipamd/src/hip_platform.cpp b/projects/clr/hipamd/src/hip_platform.cpp index 117793e09a..3b7b56d0d2 100644 --- a/projects/clr/hipamd/src/hip_platform.cpp +++ b/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) { diff --git a/projects/clr/rocclr/platform/runtime.cpp b/projects/clr/rocclr/platform/runtime.cpp index c3a1aed75f..51f5fe2f85 100644 --- a/projects/clr/rocclr/platform/runtime.cpp +++ b/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 RuntimeTearDown::external_; +amd::Monitor listenerLock{}; +std::vector RuntimeTearDown::external_{}; +std::vector> + 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"); diff --git a/projects/clr/rocclr/platform/runtime.hpp b/projects/clr/rocclr/platform/runtime.hpp index c923142c40..5439080872 100644 --- a/projects/clr/rocclr/platform/runtime.hpp +++ b/projects/clr/rocclr/platform/runtime.hpp @@ -21,6 +21,7 @@ #ifndef RUNTIME_HPP_ #define RUNTIME_HPP_ +#include #include "top.hpp" #include "thread/thread.hpp" @@ -61,13 +62,18 @@ class Runtime : AllStatic { /*@}*/ class RuntimeTearDown : public HeapObject { - static std::vector external_; - public: + using TearDownCallback = std::function; + RuntimeTearDown() {} ~RuntimeTearDown(); static void RegisterObject(ReferenceCountedObject* obj); + static void RegisterTearDownCallback(const std::string& msg, TearDownCallback func); + + private: + static std::vector external_; + static std::vector> tear_down_funcs_; }; } // namespace amd