https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/213767
>From 22ac1fd23b14c4c835a4db71b24a231a57f39bcd Mon Sep 17 00:00:00 2001 From: Robert Imschweiler <[email protected]> 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 2ba4d80978f4a..9b426fb5866f7 100644 --- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp +++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp @@ -2840,6 +2840,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 2f8d52c830a64..a789ed382ec00 100644 --- a/offload/plugins-nextgen/common/include/PluginInterface.h +++ b/offload/plugins-nextgen/common/include/PluginInterface.h @@ -1027,6 +1027,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 9a76b41647a70..764dabbf162dc 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); } @@ -990,6 +1009,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; +} _______________________________________________ llvm-branch-commits mailing list [email protected] https://lists.llvm.org/cgi-bin/mailman/listinfo/llvm-branch-commits
