[Offload] Provide proper memory management for Images on host device (#146066)

The `unloadBinaryImpl` method on the host plugin is now implemented
properly (rather than just being a stub). When an image is unloaded,
it is deallocated and the library associated with it is closed.

2 files changed, 9 insertions, 3 deletions

diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index b5addc13d664..162b149ab483 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1198,6 +1198,8 @@ struct GenericPluginTy {
     return reinterpret_cast<Ty *>(Allocator.Allocate(sizeof(Ty), alignof(Ty)));
   }
 
+  template <typename Ty> void free(Ty *Mem) { Allocator.Deallocate(Mem); }
+
   /// Get the reference to the global handler of this plugin.
   GenericGlobalHandlerTy &getGlobalHandler() {
     assert(GlobalHandler && "Global handler not initialized");
diff --git a/offload/plugins-nextgen/host/src/rtl.cpp b/offload/plugins-nextgen/host/src/rtl.cpp
index a35910aece98..d950572265b4 100644
--- a/offload/plugins-nextgen/host/src/rtl.cpp
+++ b/offload/plugins-nextgen/host/src/rtl.cpp
@@ -151,7 +151,12 @@ struct GenELF64DeviceTy : public GenericDeviceTy {
   ///
   /// TODO: This currently does nothing, and should be implemented as part of
   /// broader memory handling logic for this plugin
-  Error unloadBinaryImpl(DeviceImageTy *) override { return Plugin::success(); }
+  Error unloadBinaryImpl(DeviceImageTy *Image) override {
+    auto Elf = reinterpret_cast<GenELF64DeviceImageTy *>(Image);
+    DynamicLibrary::closeLibrary(Elf->getDynamicLibrary());
+    Plugin.free(Elf);
+    return Plugin::success();
+  }
 
   /// Deinitialize the device, which is a no-op
   Error deinitImpl() override { return Plugin::success(); }
@@ -212,8 +217,7 @@ struct GenELF64DeviceTy : public GenericDeviceTy {
 
     // Load the temporary file as a dynamic library.
     std::string ErrMsg;
-    DynamicLibrary DynLib =
-        DynamicLibrary::getPermanentLibrary(TmpFileName, &ErrMsg);
+    DynamicLibrary DynLib = DynamicLibrary::getLibrary(TmpFileName, &ErrMsg);
 
     // Check if the loaded library is valid.
     if (!DynLib.isValid())