[llvm] [openmp] [offload] Pool host and shared allocations (PR #214754)
Robert Imschweiler via llvm-commits
llvm-commits at lists.llvm.org
Fri Aug 7 08:35:13 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/214754
>From 9ecc43d40d73fb6f614a1d1fba6b46bdfab2438f Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 5 Aug 2026 10:24:27 -0500
Subject: [PATCH] [offload] Pool host and shared allocations
Route them through a memory manager, like the device allocations. Also,
move the registration as pinned memory to the plugin site since only the
plugin knows if the corresponding host/shared memory is pinned.
Improves performance of affected allocations by ballpark 1,000x.
---
offload/plugins-nextgen/amdgpu/src/rtl.cpp | 12 ++++
.../common/include/MemoryManager.h | 20 ++++--
.../common/include/PluginInterface.h | 20 +++++-
.../common/src/PluginInterface.cpp | 69 ++++++++-----------
openmp/docs/design/Runtimes.rst | 11 +--
5 files changed, 80 insertions(+), 52 deletions(-)
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 0f0de67b20c7c..e74cf03071fa7 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -2728,6 +2728,12 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
"no memory pool for the specified allocation kind");
+ // See allocate() for the registration of host / shared memory as pinned
+ // memory.
+ if (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)
+ if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
+ return Err;
+
if (auto Err = MemoryPool->deallocate(TgtPtr))
return Err;
@@ -4529,6 +4535,12 @@ Expected<void *> AMDGPUDeviceTy::allocate(size_t Size, void *,
// Enable all valid kernel agents to access the buffer.
if (auto Err = MemoryPool->enableAccess(Alloc, Size, Agents))
return std::move(Err);
+
+ // Register host / shared memory as pinned memory, so that transfers reading
+ // from it can take a device-accessible path.
+ if (Kind == TARGET_ALLOC_HOST || Kind == TARGET_ALLOC_SHARED)
+ if (auto Err = PinnedAllocs.registerHostBuffer(Alloc, Alloc, Size))
+ return std::move(Err);
}
return Alloc;
diff --git a/offload/plugins-nextgen/common/include/MemoryManager.h b/offload/plugins-nextgen/common/include/MemoryManager.h
index 883ac40c269ea..9712a461ad974 100644
--- a/offload/plugins-nextgen/common/include/MemoryManager.h
+++ b/offload/plugins-nextgen/common/include/MemoryManager.h
@@ -138,21 +138,24 @@ class MemoryManagerTy {
/// The reference to a device allocator
DeviceAllocatorTy &DeviceAllocator;
+ /// The kind of device for which the memory manager is allocated
+ TargetAllocTy DeviceKind;
/// The threshold to manage memory using memory manager. If the request size
/// is larger than \p SizeThreshold, the allocation will not be managed by the
/// memory manager.
- size_t SizeThreshold = 1U << 13;
+ size_t SizeThreshold = DefaultSizeThreshold;
/// Request memory from target device
Expected<void *> allocateOnDevice(size_t Size, void *HstPtr,
size_t Alignment) const {
- return DeviceAllocator.allocate(Size, HstPtr, TARGET_ALLOC_DEVICE,
- Alignment);
+ return DeviceAllocator.allocate(Size, HstPtr, DeviceKind, Alignment);
}
/// Deallocate data on device
- Error deleteOnDevice(void *Ptr) const { return DeviceAllocator.free(Ptr); }
+ Error deleteOnDevice(void *Ptr) const {
+ return DeviceAllocator.free(Ptr, DeviceKind);
+ }
/// This function is called when it tries to allocate memory on device but the
/// device returns out of memory. It will first free all memory in the
@@ -216,11 +219,14 @@ class MemoryManagerTy {
}
public:
+ static constexpr size_t DefaultSizeThreshold = 1U << 13;
+
/// Constructor. If \p Threshold is non-zero, then the default threshold will
/// be overwritten by \p Threshold.
- MemoryManagerTy(DeviceAllocatorTy &DeviceAllocator, size_t Threshold = 0)
+ MemoryManagerTy(DeviceAllocatorTy &DeviceAllocator, size_t Threshold = 0,
+ TargetAllocTy DeviceKind = TARGET_ALLOC_DEVICE)
: FreeLists(NumBuckets), FreeListLocks(NumBuckets),
- DeviceAllocator(DeviceAllocator) {
+ DeviceAllocator(DeviceAllocator), DeviceKind(DeviceKind) {
if (Threshold)
SizeThreshold = Threshold;
}
@@ -359,6 +365,8 @@ class MemoryManagerTy {
/// threshold and the second element represents whether user disables memory
/// manager explicitly by setting the var to 0. If user doesn't specify
/// anything, returns <0, true>.
+ /// Note that this only affects the device memory manager, not the manager for
+ /// host or shared memory.
static std::pair<size_t, bool> getSizeThresholdFromEnv() {
static UInt64Envar MemoryManagerThreshold(
"LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD", 0);
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index b3675e5a8700f..50be67886c527 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1368,8 +1368,12 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
/// only necessary for unhosted targets like the GPU.
virtual bool shouldSetupRPCServer() const { return false; }
- /// Pointer to the memory manager or nullptr if not available.
+ /// Pointer to the device memory manager or nullptr if not available.
MemoryManagerTy *MemoryManager;
+ /// Memory managers for the host and shared allocation kinds or nullptr if not
+ /// available.
+ MemoryManagerTy *HostMemoryManager;
+ MemoryManagerTy *SharedMemoryManager;
/// Per device setting of MemoryManager's Threshold
virtual size_t getMemoryManagerSizeThreshold() { return 0; }
@@ -1406,6 +1410,20 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
/// Record and replay manager.
RecordReplayTy *RecordReplay = nullptr;
+ /// Return the memory manager for the given allocation kind.
+ MemoryManagerTy *getMemoryManagerFor(TargetAllocTy Kind) {
+ switch (Kind) {
+ case TARGET_ALLOC_DEFAULT:
+ case TARGET_ALLOC_DEVICE:
+ return MemoryManager;
+ case TARGET_ALLOC_HOST:
+ return HostMemoryManager;
+ case TARGET_ALLOC_SHARED:
+ return SharedMemoryManager;
+ }
+ return nullptr;
+ }
+
protected:
/// Environment variables defined by the LLVM OpenMP implementation
/// regarding the initial number of streams and events.
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 70e289060a446..0d6aede7a1426 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -490,7 +490,8 @@ uint32_t GenericKernelTy::getEffectiveNumBlocks(
GenericDeviceTy::GenericDeviceTy(GenericPluginTy &Plugin, int32_t DeviceId,
int32_t NumDevices,
const llvm::omp::GV &OMPGridValues)
- : Plugin(Plugin), MemoryManager(nullptr), OMP_TeamLimit("OMP_TEAM_LIMIT"),
+ : Plugin(Plugin), MemoryManager(nullptr), HostMemoryManager(nullptr),
+ SharedMemoryManager(nullptr), OMP_TeamLimit("OMP_TEAM_LIMIT"),
OMP_NumTeams("OMP_NUM_TEAMS"),
OMP_TeamsThreadLimit("OMP_TEAMS_THREAD_LIMIT"),
OMPX_DebugKind("LIBOMPTARGET_DEVICE_RTL_DEBUG"),
@@ -607,6 +608,14 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
ThresholdMM = getMemoryManagerSizeThreshold();
MemoryManager = new MemoryManagerTy(*this, ThresholdMM);
}
+ if (!OMPX_TrackAllocationTraces) {
+ // Keep the threshold for pooling sizes conservative since we're dealing
+ // with pinned memory for the host.
+ HostMemoryManager = new MemoryManagerTy(
+ *this, MemoryManagerTy::DefaultSizeThreshold, TARGET_ALLOC_HOST);
+ SharedMemoryManager = new MemoryManagerTy(
+ *this, MemoryManagerTy::DefaultSizeThreshold, TARGET_ALLOC_SHARED);
+ }
return Plugin::success();
}
@@ -654,6 +663,12 @@ Error GenericDeviceTy::deinit(GenericPluginTy &Plugin) {
if (MemoryManager)
delete MemoryManager;
MemoryManager = nullptr;
+ if (HostMemoryManager)
+ delete HostMemoryManager;
+ HostMemoryManager = nullptr;
+ if (SharedMemoryManager)
+ delete SharedMemoryManager;
+ SharedMemoryManager = nullptr;
if (RecordReplay) {
if (auto Err = RecordReplay->deinit())
@@ -1014,22 +1029,15 @@ Expected<void *> GenericDeviceTy::dataAlloc(int64_t Size, void *HostPtr,
if (RecordReplay && RecordReplay->isRecordingOrReplaying())
return RecordReplay->allocate(Size);
- switch (Kind) {
- case TARGET_ALLOC_DEFAULT:
- case TARGET_ALLOC_DEVICE:
- if (MemoryManager) {
- auto AllocOrErr = MemoryManager->allocate(Size, HostPtr, Alignment);
- if (!AllocOrErr)
- return AllocOrErr.takeError();
- Alloc = *AllocOrErr;
- if (!Alloc)
- return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
- "failed to allocate from memory manager");
- break;
- }
- [[fallthrough]];
- case TARGET_ALLOC_HOST:
- case TARGET_ALLOC_SHARED: {
+ if (MemoryManagerTy *MM = getMemoryManagerFor(Kind)) {
+ auto AllocOrErr = MM->allocate(Size, HostPtr, Alignment);
+ if (!AllocOrErr)
+ return AllocOrErr.takeError();
+ Alloc = *AllocOrErr;
+ if (!Alloc)
+ return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
+ "failed to allocate from memory manager");
+ } else {
auto AllocOrErr = allocate(Size, HostPtr, Kind, Alignment);
if (!AllocOrErr)
return AllocOrErr.takeError();
@@ -1038,7 +1046,6 @@ Expected<void *> GenericDeviceTy::dataAlloc(int64_t Size, void *HostPtr,
return Plugin::error(ErrorCode::OUT_OF_RESOURCES,
"failed to allocate from device allocator");
}
- }
// Report error if the memory manager or the device allocator did not return
// any memory buffer.
@@ -1047,11 +1054,6 @@ Expected<void *> GenericDeviceTy::dataAlloc(int64_t Size, void *HostPtr,
"invalid target data allocation kind or requested "
"allocator not implemented yet");
- // Register allocated buffer as pinned memory if the type is host memory.
- if (Kind == TARGET_ALLOC_HOST)
- if (auto Err = PinnedAllocs.registerHostBuffer(Alloc, Alloc, Size))
- return std::move(Err);
-
// Keep track of the allocation stack if we track allocation traces.
if (OMPX_TrackAllocationTraces) {
std::string StackTrace;
@@ -1108,26 +1110,13 @@ Error GenericDeviceTy::dataDelete(void *TgtPtr, TargetAllocTy Kind) {
ATI->DeallocationTrace = StackTrace;
}
- switch (Kind) {
- case TARGET_ALLOC_DEFAULT:
- case TARGET_ALLOC_DEVICE:
- if (MemoryManager) {
- if (auto Err = MemoryManager->free(TgtPtr))
- return Err;
- break;
- }
- [[fallthrough]];
- case TARGET_ALLOC_HOST:
- case TARGET_ALLOC_SHARED:
- if (auto Err = free(TgtPtr, Kind))
+ if (MemoryManagerTy *MM = getMemoryManagerFor(Kind)) {
+ if (auto Err = MM->free(TgtPtr))
return Err;
+ } else if (auto Err = free(TgtPtr, Kind)) {
+ return Err;
}
- // Unregister deallocated pinned memory buffer if the type is host memory.
- if (Kind == TARGET_ALLOC_HOST)
- if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
- return Err;
-
return Plugin::success();
}
diff --git a/openmp/docs/design/Runtimes.rst b/openmp/docs/design/Runtimes.rst
index 14e300a0f531f..0e75925c78742 100644
--- a/openmp/docs/design/Runtimes.rst
+++ b/openmp/docs/design/Runtimes.rst
@@ -784,11 +784,12 @@ LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD
"""""""""""""""""""""""""""""""""""""
``LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD`` sets the threshold size for which the
-``libomptarget`` memory manager will handle the allocation. Any allocations
-larger than this threshold will not use the memory manager and be freed after
-the device kernel exits. The default threshold value is ``8KB``. If
-``LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD`` is set to ``0`` the memory manager
-will be completely disabled.
+``libomptarget`` device memory manager will handle the allocation. Any
+allocations larger than this threshold will not use the memory manager and be
+freed after the device kernel exits. The default threshold value is ``8KB``. If
+``LIBOMPTARGET_MEMORY_MANAGER_THRESHOLD`` is set to ``0`` the device memory
+manager will be completely disabled.
+This has no effect on the host or shared memory managers.
.. _libomptarget_info:
More information about the llvm-commits
mailing list