[llvm-branch-commits] [llvm] [offload] Copy the kernel launch environment with a direct-write memcpy (PR #231923)
Robert Imschweiler via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Sat Oct 17 05:27:56 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/231923
>From a138ea0258670ed6eefa1923cf0e59f633f6ab48 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Wed, 14 Oct 2026 07:12:35 -0500
Subject: [PATCH] [offload] Copy the kernel launch environment with a
direct-write memcpy
Claude assisted with this patch.
---
offload/include/PluginManager.h | 6 ++
offload/libompaccsupport/PluginManager.cpp | 13 ++++
offload/libompaccsupport/device.cpp | 12 ++-
.../offloading/kernel_launch_environment.c | 76 +++++++++++++++++++
4 files changed, 104 insertions(+), 3 deletions(-)
create mode 100644 offload/test/offloading/kernel_launch_environment.c
diff --git a/offload/include/PluginManager.h b/offload/include/PluginManager.h
index 3ca38f05ab743b..119b826d0eab8a 100644
--- a/offload/include/PluginManager.h
+++ b/offload/include/PluginManager.h
@@ -138,6 +138,9 @@ struct PluginManager {
/// Returns the number of plugins that are active.
int getNumActivePlugins() const;
+ /// Return liboffload's host device, which identifies host memory to it.
+ ol_device_handle_t getHostDevice() const { return HostDevice; }
+
private:
bool RTLsLoaded = false;
llvm::SmallVector<__tgt_bin_desc *> DelayedBinDesc;
@@ -145,6 +148,9 @@ struct PluginManager {
// List of all plugins, in use or not.
llvm::SmallVector<GenericPluginTy *> Plugins;
+ // liboffload's host device.
+ ol_device_handle_t HostDevice = nullptr;
+
// Mapping of device handles to the OpenMP device identifier.
llvm::DenseMap<ol_device_handle_t, int32_t> DeviceIds;
diff --git a/offload/libompaccsupport/PluginManager.cpp b/offload/libompaccsupport/PluginManager.cpp
index c85be9d62b31f3..6ddbdc4e83149d 100644
--- a/offload/libompaccsupport/PluginManager.cpp
+++ b/offload/libompaccsupport/PluginManager.cpp
@@ -60,6 +60,19 @@ void PluginManager::init() {
this))
REPORT() << "Failed to iterate platforms: " << toString(std::move(Err));
+ if (auto Err = iterateDevices([this](ol_device_handle_t Device) {
+ ol_platform_handle_t Platform;
+ ol_platform_backend_t Backend;
+ if (!HostDevice &&
+ !olGetDeviceInfo(Device, OL_DEVICE_INFO_PLATFORM, sizeof(Platform),
+ &Platform) &&
+ !olGetPlatformInfo(Platform, OL_PLATFORM_INFO_BACKEND,
+ sizeof(Backend), &Backend) &&
+ Backend == OL_PLATFORM_BACKEND_HOST)
+ HostDevice = Device;
+ }))
+ REPORT() << "Failed to iterate devices: " << toString(std::move(Err));
+
ODBG(ODT_Init) << "RTLs loaded!";
}
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 7d5d094bee0963..ea964c9e02710e 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -645,9 +645,15 @@ getKernelLaunchEnvironment(DeviceTy &Device, uint32_t OmpABIVersion,
DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr),
sizeof(KernelLaunchEnvironmentTy));
- if (auto Err = GenericDevice.dataSubmit(
- *AllocOrErr, &LocalKLE, sizeof(KernelLaunchEnvironmentTy), AsyncInfo))
- return Err;
+ // The allocation is fresh, so no work pending on the queue accesses it.
+ if (auto Res = olMemcpyWithFlags(
+ AsyncInfo.getQueue(), *AllocOrErr, Device.DeviceHandle, &LocalKLE,
+ Device.PM->getHostDevice(), sizeof(KernelLaunchEnvironmentTy),
+ OL_MEMCPY_FLAG_DIRECT_WRITE))
+ return createError(ErrorCode::BackendFailure,
+ "failed to copy the kernel launch environment to "
+ "device %d: %s",
+ Device.DeviceID, Res->Details);
return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
}
diff --git a/offload/test/offloading/kernel_launch_environment.c b/offload/test/offloading/kernel_launch_environment.c
new file mode 100644
index 00000000000000..3a1026b627a392
--- /dev/null
+++ b/offload/test/offloading/kernel_launch_environment.c
@@ -0,0 +1,76 @@
+// Stress the kernel launch environment (KLE) of concurrent launches.
+//
+// A cross-team reduction gives the launch a KLE, which carries that launch's
+// own reduction buffer. The plugin may either copy the KLE to the device
+// asynchronously or, where device memory is host-writable, store it in place.
+// Both have to hold up when launches overlap: every launch needs its own KLE,
+// and a launch's KLE has to be visible to the device before its kernel runs
+// and stay intact until it finishes. Otherwise a launch reduces into another
+// launch's buffer, or reads a KLE that has not landed yet. Every launch here
+// accumulates a distinct value, so either 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
+// RUN: rm -rf %t.rr && mkdir -p %t.rr
+// RUN: env LIBOMPTARGET_RECORD=1 LIBOMPTARGET_RECORD_DIR=%t.rr \
+// 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 launches are in flight 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 disturb the launch
+ // environment of a launch another thread still has in flight.
+#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