[llvm-branch-commits] [llvm] [offload] Use pinned memory for KLE (PR #213767)
Robert Imschweiler via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Fri Aug 7 07:56:41 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/213767
>From 6a3f9560ef0f0644f82c4e55e7154c39b623c487 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] [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 4ea3eb7a46151..21b2d7304e531 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);
}
@@ -1034,6 +1053,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 llvm-branch-commits
mailing list