[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:35:56 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 0692ee2164dc82135cafd645c0860a36257b93c9 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.
---
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 bfb4fbea1bd328fd51df913b7dc515f1fba969e8 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
Allocations registered for release after finalization were all freed as
TARGET_ALLOC_DEVICE, so only device memory could be registered. Store the
kind alongside the pointer and free with it.
Release them in queryAsync as well. An async info that belongs to a nowait
task is never synchronized, so queryAsync reporting the queue as complete is
the only notification it gets and its allocations were leaked. Collect them
under the async info's mutex, as synchronize does, so that an operation
issued between the query and the release does not have its allocations freed
while it is still using them.
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 fefb28de440ab9b98163665e7d366bd972402ba3 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