[llvm] [Offload] Generalize memory manager's deallocation handling (PR #228566)

[Hansang Bae via llvm-commits]

https://github.com/hansangbae created https://github.com/llvm/llvm-project/pull/228566

This change adds flexibility in the memory deallocation to support memory deallocation without explicit memory kind input.

Assisted by Claude Code.

>From b6759374a4740ba53493b36503ad4adfdac40c1e Mon Sep 17 00:00:00 2001
From: Hansang Bae <hansang.bae at intel.com>
Date: Fri, 2 Oct 2026 12:40:14 -0700
Subject: [PATCH] [Offload] Generalize memory manager's deallocation handling

This change adds flexibility in the memory deallocation to support
memory deallocation path without explicit memory kind input.

Assisted by Claude Code.
---
 .../common/include/MemoryManager.h            |  6 ++++++
 .../common/include/PluginInterface.h          |  4 ++++
 .../common/src/PluginInterface.cpp            | 21 +++++++++++++++++++
 3 files changed, 31 insertions(+)

diff --git a/offload/plugins-nextgen/common/include/MemoryManager.h b/offload/plugins-nextgen/common/include/MemoryManager.h
index e01f721cd82d0..386c0b17beb1b 100644
--- a/offload/plugins-nextgen/common/include/MemoryManager.h
+++ b/offload/plugins-nextgen/common/include/MemoryManager.h
@@ -253,6 +253,12 @@ class MemoryManagerTy {
     }
   }
 
+  /// Return true if \p TgtPtr was allocated through this memory manager.
+  bool isManaged(void *TgtPtr) {
+    std::lock_guard<std::mutex> G(MapTableLock);
+    return PtrToNodeTable.count(TgtPtr);
+  }
+
   /// Allocate memory of size \p Size from target device. \p HstPtr is used to
   /// assist the allocation.
   Expected<void *> allocate(size_t Size, void *HstPtr, size_t Alignment) {
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 33740c3af6e16..026eba0ea3c4d 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -902,6 +902,10 @@ struct PluginContextTy {
                                              TargetAllocTy Kind);
   MemoryManagerTy *getHostMemoryManager();
 
+  /// Return the kind of the existing host or shared memory manager that owns
+  /// \p Ptr, or TARGET_ALLOC_DEFAULT if neither does. Never creates a manager.
+  TargetAllocTy getPooledKind(GenericDeviceTy &Device, void *Ptr);
+
   llvm::DenseMap<std::pair<GenericDeviceTy *, int>,
                  std::unique_ptr<MemoryManagerTy>>
       DeviceMemoryManagers;
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index d68a4a7ec1b26..27c6042b24c20 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -994,6 +994,20 @@ MemoryManagerTy *PluginContextTy::getHostMemoryManager() {
   return HostMemoryManager.get();
 }
 
+TargetAllocTy PluginContextTy::getPooledKind(GenericDeviceTy &Device,
+                                             void *Ptr) {
+  std::lock_guard<std::mutex> Lock(MemoryManagersMutex);
+  if (HostMemoryManager && HostMemoryManager->isManaged(Ptr))
+    return TARGET_ALLOC_HOST;
+
+  auto It = DeviceMemoryManagers.find(
+      {&Device, static_cast<int>(TARGET_ALLOC_SHARED)});
+  if (It != DeviceMemoryManagers.end() && It->second->isManaged(Ptr))
+    return TARGET_ALLOC_SHARED;
+
+  return TARGET_ALLOC_DEFAULT;
+}
+
 Expected<void *> PluginContextTy::allocate(GenericDeviceTy &Device,
                                            int64_t Size, void *HostPtr,
                                            TargetAllocTy Kind,
@@ -1030,6 +1044,13 @@ Error PluginContextTy::deallocate(GenericDeviceTy &Device, void *Ptr,
   if (auto *RR = Device.getRecordReplay(); RR && RR->isRecordingOrReplaying())
     return Device.dataDelete(Ptr, Kind);
 
+  // omp_target_free passes TARGET_ALLOC_DEFAULT regardless of how the memory
+  // was allocated. Return pooled host and shared allocations to the manager
+  // that owns them; otherwise they are freed on the device behind the pool's
+  // back and the stale node can be handed out or freed again later.
+  if (Kind == TARGET_ALLOC_DEFAULT)
+    Kind = getPooledKind(Device, Ptr);
+
   MemoryManagerTy *MM = (Kind == TARGET_ALLOC_HOST)
                             ? getHostMemoryManager()
                             : getDeviceMemoryManagerFor(Device, Kind);