https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/214754
>From d524bab72499646d9d45aa6c28ca32ee0999a212 Mon Sep 17 00:00:00 2001 From: Robert Imschweiler <[email protected]> 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: _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
