[Offload] Provide proper memory management for Images on host device

[RossBrunton]

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.

[llvmbot]

@llvm/pr-subscribers-offload

Author: Ross Brunton (RossBrunton)

Changes

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.

---

Full diff: https://github.com/llvm/llvm-project/pull/146066.diff

2 Files Affected:

• (modified) offload/plugins-nextgen/common/include/PluginInterface.h (+2)
• (modified) offload/plugins-nextgen/host/src/rtl.cpp (+7-3)

diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index fbc798faec24b..3c03116550636 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 a35910aece986..8747540ced9f0 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(Image);
+    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())