[Openmp-commits] [llvm] [openmp] [offload] Use pinned memory for KLE (PR #213767)
Robert Imschweiler via Openmp-commits
openmp-commits at lists.llvm.org
Fri Aug 7 07:36:22 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/213767
>From 022265c76bd74cc16714960cdb91961fd6fb42a7 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 5 Aug 2026 10:24:00 -0500
Subject: [PATCH 1/5] [offload] Do not pool memory while allocation traces are
requested
Would otherwise hide use-after-free because memory stays valid if it's
in the pool.
---
offload/plugins-nextgen/common/src/PluginInterface.cpp | 8 ++++++--
1 file changed, 6 insertions(+), 2 deletions(-)
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index e539cbfc55324..d4f3d8b62a272 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -598,9 +598,13 @@ Error GenericDeviceTy::init(GenericPluginTy &Plugin) {
GridValues.GV_Max_WG_Size =
std::min(GridValues.GV_Max_WG_Size, uint32_t(OMP_TeamsThreadLimit));
- // Enable the memory manager if required.
+ // Enable the memory manager if required. Memory that the manager pools is
+ // not returned to the driver when it is freed, so a later access to it
+ // neither faults nor is reported by the allocation tracker. Leave the pool
+ // disabled while allocation traces are requested, so that we don't mask
+ // use-after-free (since they don't fault if the memory is still in the pool).
auto [ThresholdMM, EnableMM] = MemoryManagerTy::getSizeThresholdFromEnv();
- if (EnableMM) {
+ if (EnableMM && !OMPX_TrackAllocationTraces) {
if (ThresholdMM == 0)
ThresholdMM = getMemoryManagerSizeThreshold();
MemoryManager = new MemoryManagerTy(*this, ThresholdMM);
>From 35c52f4f5d6ba85ffdf5ec1438780bc7caf883e2 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Thu, 6 Aug 2026 11:51:52 -0500
Subject: [PATCH 2/5] [offload] Respect alignment in memory manager
Also fix alignment comparison in AMDGPU rtl.cpp.
---
offload/plugins-nextgen/amdgpu/src/rtl.cpp | 2 +-
offload/plugins-nextgen/common/include/MemoryManager.h | 9 +++++++--
2 files changed, 8 insertions(+), 3 deletions(-)
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 29b5006f0d7b6..0f0de67b20c7c 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -339,7 +339,7 @@ struct AMDGPUMemoryPoolTy {
// compared with the alignment of the memory allocated using the given pool.
// If the default alignment is greater than or equal to the alignment
// requested by the user, it would still meet the user's requirements.
- if (Alignment > 0 && Alignment >= PoolAllocationAlignment) {
+ if (Alignment > 0 && Alignment > PoolAllocationAlignment) {
return Plugin::error(ErrorCode::UNSUPPORTED,
"requested alignment (%lu) larger than maximum "
"supported pool alignment (%lu)",
diff --git a/offload/plugins-nextgen/common/include/MemoryManager.h b/offload/plugins-nextgen/common/include/MemoryManager.h
index 4b57be45e7551..883ac40c269ea 100644
--- a/offload/plugins-nextgen/common/include/MemoryManager.h
+++ b/offload/plugins-nextgen/common/include/MemoryManager.h
@@ -13,6 +13,7 @@
#ifndef LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H
#define LLVM_OPENMP_LIBOMPTARGET_PLUGINS_COMMON_MEMORYMANAGER_H
+#include <algorithm>
#include <cassert>
#include <functional>
#include <list>
@@ -25,6 +26,7 @@
#include "Shared/Utils.h"
#include "omptarget.h"
+#include "llvm/Support/Alignment.h"
#include "llvm/Support/Error.h"
using namespace llvm::offload::debug;
@@ -268,9 +270,12 @@ class MemoryManagerTy {
NodeTy TempNode(Size, nullptr);
std::lock_guard<std::mutex> LG(FreeListLocks[B]);
- const auto Itr = List.find(TempNode);
+ auto [First, Last] = List.equal_range(TempNode);
- if (Itr != List.end()) {
+ auto Itr = std::find_if(First, Last, [Alignment](const NodeTy &N) {
+ return Alignment == 0 || isAddrAligned(Align(Alignment), N.Ptr);
+ });
+ if (Itr != Last) {
NodePtr = &Itr->get();
List.erase(Itr);
}
>From 73db29cd0403c6f01c0b838e24310653cfe6321a 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 3/5] [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 d4f3d8b62a272..389d2ffbf4395 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"),
@@ -609,6 +610,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();
}
@@ -656,6 +665,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())
@@ -1016,22 +1031,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();
@@ -1040,7 +1048,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.
@@ -1049,11 +1056,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;
@@ -1110,26 +1112,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:
>From d4ce6a6138bc4af6cb85156ab8945493cb1228ef Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Fri, 7 Aug 2026 02:43:12 -0500
Subject: [PATCH 4/5] [offload] Thread allocation kind through async info
Claude assisted with this patch.
---
offload/include/Shared/APITypes.h | 7 ++-
.../common/include/PluginInterface.h | 4 +-
.../common/src/PluginInterface.cpp | 44 ++++++++++++++++---
3 files changed, 44 insertions(+), 11 deletions(-)
diff --git a/offload/include/Shared/APITypes.h b/offload/include/Shared/APITypes.h
index 47d8c49bf7ef2..67edc7c9f60b7 100644
--- a/offload/include/Shared/APITypes.h
+++ b/offload/include/Shared/APITypes.h
@@ -23,6 +23,9 @@
#include <cstddef>
#include <cstdint>
#include <mutex>
+#include <utility>
+
+enum TargetAllocTy : int32_t;
extern "C" {
@@ -75,8 +78,8 @@ struct __tgt_async_info {
void *Queue = nullptr;
/// A collection of allocations that are associated with this stream and that
- /// should be freed after finalization.
- llvm::SmallVector<void *, 2> AssociatedAllocations;
+ /// should be freed after finalization. Also store the type of the allocation.
+ llvm::SmallVector<std::pair<void *, TargetAllocTy>, 2> AssociatedAllocations;
/// Mutex to guard access to AssociatedAllocations and the Queue.
std::mutex Mutex;
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 50be67886c527..0b6a2223a75ce 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -168,9 +168,9 @@ struct AsyncInfoWrapperTy {
/// Register \p Ptr as an associated allocation that is freed after
/// finalization.
- void freeAllocationAfterSynchronization(void *Ptr) {
+ void freeAllocationAfterSynchronization(void *Ptr, TargetAllocTy Kind) {
std::lock_guard<std::mutex> AllocationGuard(AsyncInfoPtr->Mutex);
- AsyncInfoPtr->AssociatedAllocations.push_back(Ptr);
+ AsyncInfoPtr->AssociatedAllocations.push_back({Ptr, Kind});
}
private:
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 389d2ffbf4395..f4dcff3879ccd 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -137,7 +137,8 @@ GenericKernelTy::getKernelLaunchEnvironment(
return AllocOrErr.takeError();
// Remember to free the memory later.
- AsyncInfoWrapper.freeAllocationAfterSynchronization(*AllocOrErr);
+ AsyncInfoWrapper.freeAllocationAfterSynchronization(
+ *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE);
/// Use the KLE in the __tgt_async_info to ensure a stable address for the
/// async data transfer.
@@ -159,7 +160,8 @@ GenericKernelTy::getKernelLaunchEnvironment(
return AllocOrErr.takeError();
LocalKLE.ReductionBuffer = *AllocOrErr;
// Remember to free the memory later.
- AsyncInfoWrapper.freeAllocationAfterSynchronization(*AllocOrErr);
+ AsyncInfoWrapper.freeAllocationAfterSynchronization(
+ *AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE);
}
INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(),
@@ -289,7 +291,7 @@ Error GenericKernelTy::launch(GenericDeviceTy &GenericDevice, void **ArgPtrs,
DynBlockMemConfTy &DynBlockMemConf = *DynBlockMemConfOrErr;
if (DynBlockMemConf.FallbackPtr)
AsyncInfoWrapper.freeAllocationAfterSynchronization(
- DynBlockMemConf.FallbackPtr);
+ DynBlockMemConf.FallbackPtr, TargetAllocTy::TARGET_ALLOC_DEVICE);
auto KernelLaunchEnvOrErr =
getKernelLaunchEnvironment(GenericDevice, KernelArgs, DynBlockMemConf,
@@ -976,7 +978,7 @@ Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo,
return Plugin::error(ErrorCode::INVALID_ARGUMENT,
"invalid async info queue");
- SmallVector<void *> AllocsToDelete{};
+ SmallVector<std::pair<void *, TargetAllocTy>> AllocsToDelete{};
{
std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex};
@@ -989,8 +991,8 @@ Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo,
std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations);
}
- for (auto *Ptr : AllocsToDelete)
- if (auto Err = dataDelete(Ptr, TargetAllocTy::TARGET_ALLOC_DEVICE))
+ for (auto [Ptr, Kind] : AllocsToDelete)
+ if (auto Err = dataDelete(Ptr, Kind))
return Err;
return Plugin::success();
@@ -1003,7 +1005,35 @@ Error GenericDeviceTy::queryAsync(__tgt_async_info *AsyncInfo,
return Plugin::error(ErrorCode::INVALID_ARGUMENT,
"invalid async info queue");
- return queryAsyncImpl(*AsyncInfo, ReleaseQueue, IsQueueWorkCompleted);
+ bool WorkCompleted = false;
+ SmallVector<std::pair<void *, TargetAllocTy>> AllocsToDelete{};
+
+ {
+ // Query and collect under the mutex, as synchronize does. Querying outside
+ // it would let an operation issued in between have its allocations freed
+ // here while it is still using them.
+ std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex};
+ if (auto Err = queryAsyncImpl(*AsyncInfo, ReleaseQueue, &WorkCompleted)) {
+ if (IsQueueWorkCompleted)
+ *IsQueueWorkCompleted = WorkCompleted;
+ return Err;
+ }
+
+ // An async info belonging to a nowait task is never synchronized, so this
+ // is the only completion notification it ever gets. Without releasing its
+ // allocations here they are never freed at all.
+ if (WorkCompleted)
+ std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations);
+ }
+
+ for (auto [Ptr, Kind] : AllocsToDelete)
+ if (auto Err = dataDelete(Ptr, Kind))
+ return Err;
+
+ if (IsQueueWorkCompleted)
+ *IsQueueWorkCompleted = WorkCompleted;
+
+ return Plugin::success();
}
Error GenericDeviceTy::memoryVAMap(void **Addr, void *VAddr, size_t *RSize) {
>From b5bd0cd9b2f25c2b3659cae942ff2fcdb6f826f5 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Mon, 3 Aug 2026 16:38:02 -0500
Subject: [PATCH 5/5] [offload] Use pinned memory for KLE
Reduce kernel launch latency by using the fast path "pinned host memory
-> device memory" for submitting the kernel launch environment to the
device.
Claude assisted with this patch.
---
offload/plugins-nextgen/amdgpu/src/rtl.cpp | 2 +
.../common/include/PluginInterface.h | 14 ++++
.../common/src/PluginInterface.cpp | 49 ++++++++++++-
.../offloading/kernel_launch_environment.c | 71 +++++++++++++++++++
4 files changed, 134 insertions(+), 2 deletions(-)
create mode 100644 offload/test/offloading/kernel_launch_environment.c
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index e74cf03071fa7..327764102c43d 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -2839,6 +2839,8 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
return true;
}
+ bool hasFastSubmitFromPinnedMemory() const override { return true; }
+
/// Submit data to the device (host to device transfer).
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
AsyncInfoWrapperTy &AsyncInfoWrapper) override {
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 0b6a2223a75ce..1092a6ced6a6d 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -982,6 +982,20 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
virtual Error queryAsyncImpl(__tgt_async_info &AsyncInfo, bool ReleaseQueue,
bool *IsQueueWorkCompleted) = 0;
+ /// Indicate whether dataSubmitImpl has a faster path for host buffers that
+ /// are registered as pinned memory. If a plugin returns true, the kernel
+ /// launch environment is copied into a pinned host buffer before it is
+ /// submitted, so that its transfer takes that path.
+ virtual bool hasFastSubmitFromPinnedMemory() const { return false; }
+
+ /// Allocate a pinned host buffer to stage a kernel launch environment. The
+ /// caller owns it until it registers it with
+ /// AsyncInfoWrapperTy::freeAllocationAfterSynchronization, which releases it
+ /// once the transfer reading it has completed. Returns nullptr if staging is
+ /// unavailable, in which case the caller must submit the launch environment
+ /// from ordinary host memory.
+ KernelLaunchEnvironmentTy *getPinnedLaunchEnvBuffer();
+
/// Check whether the architecture supports VA management
virtual bool supportVAManagement() const { return false; }
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index f4dcff3879ccd..32357ac678b2d 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -164,17 +164,36 @@ GenericKernelTy::getKernelLaunchEnvironment(
*AllocOrErr, TargetAllocTy::TARGET_ALLOC_DEVICE);
}
+ // Copy into a pinned buffer if the plugin transfers those faster: dataAlloc
+ // registers TARGET_ALLOC_HOST memory in PinnedAllocs, so the dataSubmit
+ // below can take the one-step copy path.
+ const void *LaunchEnvSrc = &LocalKLE;
+ auto *PinnedKLE = GenericDevice.getPinnedLaunchEnvBuffer();
+ if (PinnedKLE) {
+ *PinnedKLE = LocalKLE;
+ LaunchEnvSrc = PinnedKLE;
+ }
+
INFO(OMP_INFOTYPE_DATA_TRANSFER, GenericDevice.getDeviceId(),
"Copying data from host to device, HstPtr=" DPxMOD ", TgtPtr=" DPxMOD
", Size=%" PRId64 ", Name=KernelLaunchEnv\n",
- DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr),
+ DPxPTR(LaunchEnvSrc), DPxPTR(*AllocOrErr),
sizeof(KernelLaunchEnvironmentTy));
- auto Err = GenericDevice.dataSubmit(*AllocOrErr, &LocalKLE,
+ auto Err = GenericDevice.dataSubmit(*AllocOrErr, LaunchEnvSrc,
sizeof(KernelLaunchEnvironmentTy),
AsyncInfoWrapper);
if (Err)
return Err;
+
+ // Register the staging buffer only now that the transfer reading it has
+ // been issued. Registering it earlier would let a concurrent finalization
+ // that observes the queue as complete release it while the transfer is
+ // still in flight.
+ if (PinnedKLE)
+ AsyncInfoWrapper.freeAllocationAfterSynchronization(
+ PinnedKLE, TargetAllocTy::TARGET_ALLOC_HOST);
+
return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
}
@@ -1036,6 +1055,32 @@ Error GenericDeviceTy::queryAsync(__tgt_async_info *AsyncInfo,
return Plugin::success();
}
+KernelLaunchEnvironmentTy *GenericDeviceTy::getPinnedLaunchEnvBuffer() {
+ if (!hasFastSubmitFromPinnedMemory())
+ return nullptr;
+
+ // While recording or replaying, dataAlloc serves every allocation kind from
+ // the record-replay device memory pool, so it cannot give us host memory.
+ if (RecordReplay && RecordReplay->isRecordingOrReplaying())
+ return nullptr;
+
+ auto AllocOrErr =
+ dataAlloc(sizeof(KernelLaunchEnvironmentTy), /*HostPtr=*/nullptr,
+ TargetAllocTy::TARGET_ALLOC_HOST, /*Alignment=*/0);
+ if (!AllocOrErr) {
+ // Staging is optional, so fall back to unpinned memory. Consume the
+ // error unconditionally: ODBG does not evaluate its operands unless
+ // debugging is enabled.
+ std::string ErrStr = toString(AllocOrErr.takeError());
+ ODBG(OLDT_Alloc) << "Failed to allocate a pinned buffer for the kernel "
+ "launch environment, submitting it unstaged: "
+ << ErrStr;
+ return nullptr;
+ }
+
+ return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
+}
+
Error GenericDeviceTy::memoryVAMap(void **Addr, void *VAddr, size_t *RSize) {
return Plugin::error(ErrorCode::UNSUPPORTED,
"device does not support VA Management");
diff --git a/offload/test/offloading/kernel_launch_environment.c b/offload/test/offloading/kernel_launch_environment.c
new file mode 100644
index 0000000000000..71a31e9d1bc6d
--- /dev/null
+++ b/offload/test/offloading/kernel_launch_environment.c
@@ -0,0 +1,71 @@
+// Stress the kernel launch environment (KLE) transfer.
+//
+// A cross-team reduction gives the launch a KLE, which the plugin stages in a
+// host buffer and copies to the device asynchronously. The KLE carries that
+// launch's own reduction buffer, so overlapping launches must neither share a
+// staging buffer nor release one while its transfer is still in flight: either
+// makes a launch reduce into another launch's buffer. Every launch here
+// accumulates a distinct value, so that shows up as a wrong sum.
+//
+// RUN: %libomptarget-compile-generic -fopenmp-offload-mandatory
+// RUN: %libomptarget-run-generic
+// RUN: %libomptarget-compileopt-generic -fopenmp-offload-mandatory
+// RUN: %libomptarget-run-generic
+//
+// REQUIRES: gpu
+
+#include <omp.h>
+#include <stdio.h>
+
+#define NUM_LAUNCHES 32
+#define N 4096
+
+// Launch k accumulates k + 1 per element.
+static long expected(int k) { return (long)(k + 1) * N; }
+
+static int check(const char *Phase, long *Results) {
+ int Errors = 0;
+ for (int k = 0; k < NUM_LAUNCHES; k++)
+ if (Results[k] != expected(k)) {
+ fprintf(stderr, "%s: launch %d reduced to %ld, expected %ld\n", Phase, k,
+ Results[k], expected(k));
+ Errors++;
+ }
+ return Errors;
+}
+
+int main(void) {
+ static long Results[NUM_LAUNCHES];
+ int Errors = 0;
+
+ // Launches issued back to back without an intervening synchronization, so
+ // that several KLE transfers are outstanding at once.
+ for (int k = 0; k < NUM_LAUNCHES; k++) {
+ Results[k] = 0;
+#pragma omp target teams distribute parallel for map(tofrom : Results[k : 1]) \
+ reduction(+ : Results[k]) firstprivate(k) nowait
+ for (int i = 0; i < N; i++)
+ Results[k] += k + 1;
+ }
+#pragma omp taskwait
+ Errors += check("nowait", Results);
+
+ // Same, but with the launches and the synchronizations spread over several
+ // host threads: a thread finalizing its queue must not release a staging
+ // buffer that another thread's launch is still using.
+#pragma omp parallel for num_threads(8)
+ for (int k = 0; k < NUM_LAUNCHES; k++) {
+ long Sum = 0;
+#pragma omp target teams distribute parallel for map(tofrom : Sum) \
+ reduction(+ : Sum) firstprivate(k)
+ for (int i = 0; i < N; i++)
+ Sum += k + 1;
+ Results[k] = Sum;
+ }
+ Errors += check("threaded", Results);
+
+ if (Errors)
+ return 1;
+ printf("PASS\n");
+ return 0;
+}
More information about the Openmp-commits
mailing list