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

[Ross Brunton via llvm-commits]

https://github.com/RossBrunton updated https://github.com/llvm/llvm-project/pull/146066

>From db0ed69f9d2de7d75277f5cee49a2ab069aadf9f Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Fri, 27 Jun 2025 13:17:36 +0100
Subject: [PATCH 1/2] [Offload] Provide proper memory management for Images on
 host device

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.
---
 .../plugins-nextgen/common/include/PluginInterface.h   |  2 ++
 offload/plugins-nextgen/host/src/rtl.cpp               | 10 +++++++---
 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 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())

>From 2625ed5412608d70dd09d65299d1f89e21402d9f Mon Sep 17 00:00:00 2001
From: Ross Brunton <ross at codeplay.com>
Date: Fri, 27 Jun 2025 13:24:01 +0100
Subject: [PATCH 2/2] Free elf rather than image

---
 offload/plugins-nextgen/host/src/rtl.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/offload/plugins-nextgen/host/src/rtl.cpp b/offload/plugins-nextgen/host/src/rtl.cpp
index 8747540ced9f0..d950572265b4c 100644
--- a/offload/plugins-nextgen/host/src/rtl.cpp
+++ b/offload/plugins-nextgen/host/src/rtl.cpp
@@ -154,7 +154,7 @@ struct GenELF64DeviceTy : public GenericDeviceTy {
   Error unloadBinaryImpl(DeviceImageTy *Image) override {
     auto Elf = reinterpret_cast<GenELF64DeviceImageTy *>(Image);
     DynamicLibrary::closeLibrary(Elf->getDynamicLibrary());
-    Plugin.free(Image);
+    Plugin.free(Elf);
     return Plugin::success();
   }