[llvm] [offload] Use pinned memory for KLE (PR #213767)
via llvm-commits
llvm-commits at lists.llvm.org
Tue Aug 4 01:54:25 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-offload
Author: Robert Imschweiler (ro-i)
<details>
<summary>Changes</summary>
Reduce kernel launch latency by using the fast path "pinned host memory -> device memory" for submitting the kernel launch environment.
Claude assisted with this patch.
---
Full diff: https://github.com/llvm/llvm-project/pull/213767.diff
5 Files Affected:
- (modified) offload/include/Shared/APITypes.h (+6)
- (modified) offload/plugins-nextgen/amdgpu/src/rtl.cpp (+4)
- (modified) offload/plugins-nextgen/common/include/PluginInterface.h (+38)
- (modified) offload/plugins-nextgen/common/src/PluginInterface.cpp (+105-3)
- (modified) offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp (+74)
``````````diff
diff --git a/offload/include/Shared/APITypes.h b/offload/include/Shared/APITypes.h
index 47d8c49bf7ef2..9a5783560318f 100644
--- a/offload/include/Shared/APITypes.h
+++ b/offload/include/Shared/APITypes.h
@@ -85,6 +85,12 @@ struct __tgt_async_info {
/// ensure it is a valid location while the transfer to the device is
/// happening.
KernelLaunchEnvironmentTy KernelLaunchEnvironment;
+
+ /// Pinned host copies of the launch environment, one per kernel issued on
+ /// this async info; empty if the plugin does not stage. The device reads
+ /// them asynchronously, so each launch needs its own. The device that handed
+ /// them out reclaims them once this async info's operations have completed.
+ llvm::SmallVector<void *, 2> PinnedKernelLaunchEnvironments;
};
/// This struct contains all of the arguments to a target kernel region launch.
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 65ef83c394f6a..c9e48f7f82125 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -2833,6 +2833,10 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
return true;
}
+ /// Submitting from pinned host memory avoids the intermediate buffer and the
+ /// host-to-host copy of the two-step path in dataSubmitImpl.
+ 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 67ebfbc943fdc..3afbf4ea034c8 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -173,6 +173,16 @@ struct AsyncInfoWrapperTy {
AsyncInfoPtr->AssociatedAllocations.push_back(Ptr);
}
+ /// Hand the pinned launch environment buffer \p Ptr to this async info,
+ /// which returns it to the device's pool once its operations have completed.
+ /// Call only after the transfer reading \p Ptr has been issued: a buffer
+ /// registered before that could be recycled by a concurrent synchronization
+ /// that observes the queue as complete.
+ void recycleLaunchEnvAfterSynchronization(void *Ptr) {
+ std::lock_guard<std::mutex> AllocationGuard(AsyncInfoPtr->Mutex);
+ AsyncInfoPtr->PinnedKernelLaunchEnvironments.push_back(Ptr);
+ }
+
private:
GenericDeviceTy &Device;
__tgt_async_info LocalAsyncInfo;
@@ -982,6 +992,19 @@ 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; }
+
+ /// Take a pinned buffer from this device's pool. The caller owns it until it
+ /// passes it to AsyncInfoWrapperTy::recycleLaunchEnvAfterSynchronization;
+ /// one that is never passed there is only freed at device deinit. 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; }
@@ -1400,6 +1423,21 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
/// Record and replay manager.
RecordReplayTy *RecordReplay = nullptr;
+ /// Return \p AsyncInfo's pinned buffers to the pool. Call only once its
+ /// operations have completed, and with its mutex held.
+ void recyclePinnedLaunchEnvBuffers(__tgt_async_info &AsyncInfo);
+
+ /// Free every pinned buffer this device handed out, including those still
+ /// owned by an async info. This is the only place they are released, so one
+ /// whose async info never reported completion is simply not reused.
+ Error destroyPinnedLaunchEnvBuffers();
+
+ /// All pinned buffers allocated for staging launch environments, and the
+ /// subset of them not currently owned by an async info.
+ llvm::SmallVector<void *> PinnedLaunchEnvBuffers;
+ llvm::SmallVector<void *> FreePinnedLaunchEnvBuffers;
+ std::mutex PinnedLaunchEnvMutex;
+
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 7b821e77df179..6e914541621b9 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -162,17 +162,32 @@ GenericKernelTy::getKernelLaunchEnvironment(
AsyncInfoWrapper.freeAllocationAfterSynchronization(*AllocOrErr);
}
+ // 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;
+
+ // The transfer is issued, so the async info can now take the buffer.
+ if (PinnedKLE)
+ AsyncInfoWrapper.recycleLaunchEnvAfterSynchronization(PinnedKLE);
+
return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
}
@@ -647,6 +662,10 @@ Error GenericDeviceTy::deinit(GenericPluginTy &Plugin) {
return Err;
LoadedImages.clear();
+ // Release the pinned staging buffers while the device is still alive.
+ if (auto Err = destroyPinnedLaunchEnvBuffers())
+ return Err;
+
// Delete the memory manager before deinitializing the device. Otherwise,
// we may delete device allocations after the device is deinitialized.
if (MemoryManager)
@@ -968,6 +987,9 @@ Error GenericDeviceTy::synchronize(__tgt_async_info *AsyncInfo,
return Err;
std::swap(AllocsToDelete, AsyncInfo->AssociatedAllocations);
+
+ // The device no longer reads from this async info's pinned buffers.
+ recyclePinnedLaunchEnvBuffers(*AsyncInfo);
}
for (auto *Ptr : AllocsToDelete)
@@ -984,7 +1006,87 @@ 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;
+
+ // Query and recycle under the mutex, as synchronize does. Querying outside
+ // it would let a launch issued in between register a buffer that this call
+ // then recycles, although the transfer reading it is still in flight.
+ std::lock_guard<std::mutex> AllocationGuard{AsyncInfo->Mutex};
+ auto Err = queryAsyncImpl(*AsyncInfo, ReleaseQueue, &WorkCompleted);
+
+ // An async info belonging to a nowait task is never synchronized, so this is
+ // its only completion notification; without recycling here it never reuses.
+ if (!Err && WorkCompleted)
+ recyclePinnedLaunchEnvBuffers(*AsyncInfo);
+
+ if (IsQueueWorkCompleted)
+ *IsQueueWorkCompleted = WorkCompleted;
+
+ return Err;
+}
+
+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;
+
+ void *Buffer = nullptr;
+ {
+ std::lock_guard<std::mutex> Guard{PinnedLaunchEnvMutex};
+ if (!FreePinnedLaunchEnvBuffers.empty())
+ Buffer = FreePinnedLaunchEnvBuffers.pop_back_val();
+ }
+
+ if (!Buffer) {
+ 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;
+ }
+ Buffer = *AllocOrErr;
+
+ std::lock_guard<std::mutex> Guard{PinnedLaunchEnvMutex};
+ PinnedLaunchEnvBuffers.push_back(Buffer);
+ }
+
+ return static_cast<KernelLaunchEnvironmentTy *>(Buffer);
+}
+
+void GenericDeviceTy::recyclePinnedLaunchEnvBuffers(
+ __tgt_async_info &AsyncInfo) {
+ if (AsyncInfo.PinnedKernelLaunchEnvironments.empty())
+ return;
+
+ std::lock_guard<std::mutex> Guard{PinnedLaunchEnvMutex};
+ FreePinnedLaunchEnvBuffers.append(AsyncInfo.PinnedKernelLaunchEnvironments);
+ AsyncInfo.PinnedKernelLaunchEnvironments.clear();
+}
+
+Error GenericDeviceTy::destroyPinnedLaunchEnvBuffers() {
+ llvm::SmallVector<void *> Buffers;
+ {
+ std::lock_guard<std::mutex> Guard{PinnedLaunchEnvMutex};
+ std::swap(Buffers, PinnedLaunchEnvBuffers);
+ FreePinnedLaunchEnvBuffers.clear();
+ }
+
+ for (void *Buffer : Buffers)
+ if (auto Err = dataDelete(Buffer, TargetAllocTy::TARGET_ALLOC_HOST))
+ return Err;
+
+ return Plugin::success();
}
Error GenericDeviceTy::memoryVAMap(void **Addr, void *VAddr, size_t *RSize) {
diff --git a/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp b/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
index c159829e7e40e..d739ff5bb0073 100644
--- a/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
+++ b/offload/unittests/OffloadAPI/kernel/olLaunchKernel.cpp
@@ -226,6 +226,80 @@ TEST_P(olLaunchKernelLocalMemTest, Success) {
ASSERT_SUCCESS(olMemFree(Mem));
}
+// A kernel requesting dynamic shared memory needs a launch environment, which
+// the plugin may transfer from a host buffer it reads asynchronously. Issuing
+// several such kernels on one queue without synchronizing must not let one
+// clobber a preceding launch's environment. Two rounds, so that the second
+// reuses the buffers the first released.
+TEST_P(olLaunchKernelLocalMemTest, MultipleLaunchesWithoutSync) {
+ SKIP_KNOWN_FAILURE(LevelZero{"unsupported DynSharedMemory"});
+
+ constexpr uint32_t NumLaunches = 8;
+ constexpr uint32_t NumRounds = 2;
+
+ LaunchArgs.NumGroups.x = 4;
+ LaunchArgs.DynSharedMemory = 64 * sizeof(uint32_t);
+
+ const uint32_t NumElements = LaunchArgs.GroupSize.x * LaunchArgs.NumGroups.x;
+
+ for (uint32_t Round = 0; Round < NumRounds; Round++) {
+ // Each launch writes its own buffer, so a clobbered environment cannot be
+ // masked by another launch producing the expected values.
+ std::vector<void *> Mems(NumLaunches);
+ for (auto &Mem : Mems) {
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED,
+ NumElements * sizeof(uint32_t), &Mem));
+
+ void *ArgPtrs[] = {&Mem};
+ size_t ArgSizes[] = {sizeof(Mem)};
+ ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
+ std::size(ArgPtrs), ArgPtrs, ArgSizes));
+ }
+
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ for (auto *Mem : Mems) {
+ uint32_t *Data = (uint32_t *)Mem;
+ for (uint32_t I = 0; I < NumElements; I++)
+ ASSERT_EQ(Data[I], (I % 64) * 2);
+ ASSERT_SUCCESS(olMemFree(Mem));
+ }
+ }
+}
+
+// Same, but with the launches and synchronizations of one queue spread over
+// several threads: a synchronization must not reclaim a launch environment
+// that another thread's launch is still using.
+TEST_P(olLaunchKernelLocalMemTest, MultipleLaunchesThreaded) {
+ SKIP_KNOWN_FAILURE(LevelZero{"unsupported DynSharedMemory"});
+
+ LaunchArgs.NumGroups.x = 4;
+ LaunchArgs.DynSharedMemory = 64 * sizeof(uint32_t);
+
+ const uint32_t NumElements = LaunchArgs.GroupSize.x * LaunchArgs.NumGroups.x;
+
+ threadify([&](size_t) {
+ for (uint32_t Round = 0; Round < 8; Round++) {
+ void *Mem;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED,
+ NumElements * sizeof(uint32_t), &Mem));
+
+ void *ArgPtrs[] = {&Mem};
+ size_t ArgSizes[] = {sizeof(Mem)};
+ ASSERT_SUCCESS(olLaunchKernel(Queue, Device, Kernel, &LaunchArgs, nullptr,
+ std::size(ArgPtrs), ArgPtrs, ArgSizes));
+
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ uint32_t *Data = (uint32_t *)Mem;
+ for (uint32_t I = 0; I < NumElements; I++)
+ ASSERT_EQ(Data[I], (I % 64) * 2);
+
+ ASSERT_SUCCESS(olMemFree(Mem));
+ }
+ });
+}
+
TEST_P(olLaunchKernelLocalMemReductionTest, Success) {
SKIP_KNOWN_FAILURE(LevelZero{"unsupported DynSharedMemory"});
``````````
</details>
https://github.com/llvm/llvm-project/pull/213767
More information about the llvm-commits
mailing list