[llvm-branch-commits] [llvm] [offload] Write small unordered host-to-device transfers directly on AMDGPU (PR #231569)
Robert Imschweiler via llvm-branch-commits
llvm-branch-commits at lists.llvm.org
Tue Oct 13 10:04:01 PDT 2026
https://github.com/ro-i updated https://github.com/llvm/llvm-project/pull/231569
>From 089a75d3f90a16908edf183b9b29eb34bfd12929 Mon Sep 17 00:00:00 2001
From: Robert Imschweiler <robert.imschweiler at amd.com>
Date: Fri, 9 Oct 2026 10:17:38 -0500
Subject: [PATCH] [offload] Write small unordered host-to-device transfers
directly on AMDGPU
---
offload/libompaccsupport/device.cpp | 6 +-
.../amdgpu/dynamic_hsa/hsa_ext_amd.h | 22 +++
offload/plugins-nextgen/amdgpu/src/rtl.cpp | 137 +++++++++++++++++-
.../common/include/PluginInterface.h | 9 +-
.../common/src/PluginInterface.cpp | 5 +-
offload/plugins-nextgen/cuda/src/rtl.cpp | 3 +-
offload/plugins-nextgen/host/src/rtl.cpp | 3 +-
.../level_zero/include/L0Device.h | 3 +-
.../level_zero/src/L0Device.cpp | 3 +-
.../offloading/kernel_launch_environment.c | 73 ++++++++++
10 files changed, 251 insertions(+), 13 deletions(-)
create mode 100644 offload/test/offloading/kernel_launch_environment.c
diff --git a/offload/libompaccsupport/device.cpp b/offload/libompaccsupport/device.cpp
index 48d490025f8385..70a9b1ff50a1dc 100644
--- a/offload/libompaccsupport/device.cpp
+++ b/offload/libompaccsupport/device.cpp
@@ -605,8 +605,10 @@ static llvm::Expected<KernelLaunchEnvironmentTy *> getKernelLaunchEnvironment(
DPxPTR(&LocalKLE), DPxPTR(*AllocOrErr),
sizeof(KernelLaunchEnvironmentTy));
- if (auto Err = GenericDevice.dataSubmit(
- *AllocOrErr, &LocalKLE, sizeof(KernelLaunchEnvironmentTy), AsyncInfo))
+ // The allocation is fresh, so nothing pending on the queue accesses it.
+ if (auto Err = GenericDevice.dataSubmit(*AllocOrErr, &LocalKLE,
+ sizeof(KernelLaunchEnvironmentTy),
+ AsyncInfo, /*Unordered=*/true))
return Err;
return static_cast<KernelLaunchEnvironmentTy *>(*AllocOrErr);
}
diff --git a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h
index c736ec0759841d..908327fd22517e 100644
--- a/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h
+++ b/offload/plugins-nextgen/amdgpu/dynamic_hsa/hsa_ext_amd.h
@@ -62,8 +62,30 @@ typedef enum {
typedef enum {
HSA_AMD_AGENT_MEMORY_POOL_INFO_ACCESS = 0,
+ HSA_AMD_AGENT_MEMORY_POOL_INFO_NUM_LINK_HOPS = 1,
+ HSA_AMD_AGENT_MEMORY_POOL_INFO_LINK_INFO = 2,
} hsa_amd_agent_memory_pool_info_t;
+typedef enum {
+ HSA_AMD_LINK_INFO_TYPE_HYPERTRANSPORT = 0,
+ HSA_AMD_LINK_INFO_TYPE_QPI = 1,
+ HSA_AMD_LINK_INFO_TYPE_PCIE = 2,
+ HSA_AMD_LINK_INFO_TYPE_INFINBAND = 3,
+ HSA_AMD_LINK_INFO_TYPE_XGMI = 4,
+} hsa_amd_link_info_type_t;
+
+typedef struct hsa_amd_memory_pool_link_info_s {
+ uint32_t min_latency;
+ uint32_t max_latency;
+ uint32_t min_bandwidth;
+ uint32_t max_bandwidth;
+ bool atomic_support_32bit;
+ bool atomic_support_64bit;
+ bool coherent_support;
+ hsa_amd_link_info_type_t link_type;
+ uint32_t numa_distance;
+} hsa_amd_memory_pool_link_info_t;
+
typedef enum {
HSA_AMD_MEMORY_POOL_ACCESS_NEVER_ALLOWED = 0,
} hsa_amd_memory_pool_access_t;
diff --git a/offload/plugins-nextgen/amdgpu/src/rtl.cpp b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
index 98ee066bf4b69f..af2925a6462ba4 100644
--- a/offload/plugins-nextgen/amdgpu/src/rtl.cpp
+++ b/offload/plugins-nextgen/amdgpu/src/rtl.cpp
@@ -14,8 +14,10 @@
#include <cassert>
#include <cstddef>
#include <cstdint>
+#include <cstring>
#include <deque>
#include <functional>
+#include <map>
#include <mutex>
#include <string>
#include <system_error>
@@ -378,6 +380,25 @@ struct AMDGPUMemoryPoolTy {
return Access != HSA_AMD_MEMORY_POOL_ACCESS_NEVER_ALLOWED;
}
+ /// Returns if the \p Agent reaches the memory pool over XGMI only, i.e.,
+ /// without any PCIe (or other) hop along the way.
+ bool isLinkedOverXGMI(hsa_agent_t Agent) {
+ uint32_t NumHops = 0;
+ if (hsa_amd_agent_memory_pool_get_info(
+ Agent, MemoryPool, HSA_AMD_AGENT_MEMORY_POOL_INFO_NUM_LINK_HOPS,
+ &NumHops) ||
+ NumHops == 0)
+ return false;
+ llvm::SmallVector<hsa_amd_memory_pool_link_info_t> Links(NumHops);
+ if (hsa_amd_agent_memory_pool_get_info(
+ Agent, MemoryPool, HSA_AMD_AGENT_MEMORY_POOL_INFO_LINK_INFO,
+ Links.data()))
+ return false;
+ return llvm::all_of(Links, [](const hsa_amd_memory_pool_link_info_t &L) {
+ return L.link_type == HSA_AMD_LINK_INFO_TYPE_XGMI;
+ });
+ }
+
/// Allow the device to access a specific allocation.
Error enableAccess(void *Ptr, int64_t Size,
const llvm::SmallVector<hsa_agent_t> &Agents) const {
@@ -525,6 +546,31 @@ struct AMDGPUMemoryManagerTy : public DeviceAllocatorTy {
MemoryManagerTy *MemoryManager;
};
+/// Largest size for direct host to device transfer.
+static constexpr size_t MaxDirectWriteSize = 2 * 4096;
+
+/// Copy \p Size bytes from \p Src to the device memory at \p Dst.
+/// \p OverXGMI tells whether the host reaches that memory over XGMI.
+///
+/// The host writes device memory through a write-combining mapping, which the
+/// release store on the doorbell does not order, so the fence drains it. Over
+/// PCIe the stores are posted and may still be in flight after that; reading
+/// back the last byte, written again after the fence, waits for them to land.
+/// XGMI needs no read-back. HSA exposes no call for this: it is the sequence of
+/// ROCr's internal GpuAgent::PcieWcFlush and of HIP for kernel arguments in
+/// device memory.
+static void publishToDeviceMemory(void *Dst, const void *Src, size_t Size,
+ bool OverXGMI) {
+ auto *Last = reinterpret_cast<volatile char *>(Dst) + Size - 1;
+ std::memcpy(Dst, Src, Size);
+ std::atomic_thread_fence(std::memory_order_seq_cst);
+ if (OverXGMI)
+ return;
+ *Last = reinterpret_cast<const char *>(Src)[Size - 1];
+ std::atomic_thread_fence(std::memory_order_seq_cst);
+ (void)*Last;
+}
+
/// Class implementing the AMDGPU device images' properties.
struct AMDGPUDeviceImageTy : public DeviceImageTy {
/// Create the AMDGPU image with the id and the target image pointer.
@@ -2220,6 +2266,9 @@ struct AMDHostDeviceTy : public AMDGenericDeviceTy {
/// Get one of the host agents. Return always the first agent.
hsa_agent_t getAgent() const override { return Agents[0]; }
+ /// Get all the host agents.
+ llvm::ArrayRef<hsa_agent_t> getAgents() const { return Agents; }
+
/// Get a memory pool for fine-grained allocations.
AMDGPUMemoryPoolTy &getFineGrainedMemoryPool() {
assert(!FineGrainedMemoryPools.empty() && "No fine-grained mempool");
@@ -2287,6 +2336,17 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
if (auto Err = initMemoryPools())
return Err;
+ HostWritesDeviceMemory =
+ !CoarseGrainedMemoryPools.empty() &&
+ CoarseGrainedMemoryPools[0]->canAccess(HostDevice.getAgent());
+ // The writing thread may run on any host agent, so all of them must reach
+ // the device over XGMI.
+ HostLinkIsXGMI =
+ HostWritesDeviceMemory &&
+ llvm::all_of(HostDevice.getAgents(), [&](hsa_agent_t HostAgent) {
+ return CoarseGrainedMemoryPools[0]->isLinkedOverXGMI(HostAgent);
+ });
+
char GPUName[64];
if (auto Err = getDeviceAttr(HSA_AGENT_INFO_NAME, GPUName))
return Err;
@@ -2735,6 +2795,11 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
if (auto Err = PinnedAllocs.unregisterHostBuffer(TgtPtr))
return Err;
+ if (MemoryPool == CoarseGrainedMemoryPools[0] && HostWritesDeviceMemory) {
+ std::lock_guard<std::mutex> Lock(HostMappedAllocsMutex);
+ HostMappedAllocs.erase(TgtPtr);
+ }
+
if (auto Err = MemoryPool->deallocate(TgtPtr))
return Err;
@@ -2840,12 +2905,39 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
return true;
}
+ /// Whether a transfer of \p Size bytes should be done via direct H2D
+ /// transfer.
+ bool allowsDirectWrite(int64_t Size) const {
+ return HostWritesDeviceMemory && Size > 0 &&
+ Size <= int64_t(MaxDirectWriteSize);
+ }
+
+ /// Whether the host can write \p Size bytes of device memory at \p Ptr,
+ /// i.e., whether they lie within an allocation that allocate() mapped for
+ /// the host.
+ bool isHostWritable(const void *Ptr, int64_t Size) const {
+ std::lock_guard<std::mutex> Lock(HostMappedAllocsMutex);
+ auto It = HostMappedAllocs.upper_bound(Ptr);
+ if (It == HostMappedAllocs.begin())
+ return false;
+ --It;
+ return utils::getPtrDiff(Ptr, It->first) + Size <= int64_t(It->second);
+ }
+
/// Submit data to the device (host to device transfer).
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
- AsyncInfoWrapperTy &AsyncInfoWrapper) override {
+ AsyncInfoWrapperTy &AsyncInfoWrapper,
+ bool Unordered) override {
AMDGPUStreamTy *Stream = nullptr;
void *PinnedPtr = nullptr;
+ // The direct H2D write is independent of any queueing, so Unordered is a
+ // requirement.
+ if (Unordered && allowsDirectWrite(Size) && isHostWritable(TgtPtr, Size)) {
+ publishToDeviceMemory(TgtPtr, HstPtr, Size, HostLinkIsXGMI);
+ return Plugin::success();
+ }
+
// Use one-step asynchronous operation when host memory is already pinned.
if (void *PinnedPtr =
PinnedAllocs.getDeviceAccessiblePtrFromPinnedBuffer(HstPtr)) {
@@ -3721,6 +3813,18 @@ struct AMDGPUDeviceTy : public GenericDeviceTy, AMDGenericDeviceTy {
/// Reference to the host device.
AMDHostDeviceTy &HostDevice;
+ /// Whether the host can map this device's memory (large BAR) and write small
+ /// transfers into it directly.
+ bool HostWritesDeviceMemory = false;
+
+ /// Whether the host reaches this device's memory over XGMI rather than PCIe.
+ bool HostLinkIsXGMI = false;
+
+ /// The allocations mapped for the host, by base address, with their sizes.
+ /// Ordered, so that a pointer into one of them finds it.
+ std::map<const void *, size_t> HostMappedAllocs;
+ mutable std::mutex HostMappedAllocsMutex;
+
/// The current size of the stack that will be used in cases where it could
/// not be statically determined.
uint64_t StackSize = 16 * 1024 /* 16 KB */;
@@ -4608,9 +4712,38 @@ Expected<void *> AMDGPUDeviceTy::allocate(size_t Size, void *,
return MemoryPool->canAccess(Agent);
});
+ // Map small device allocations for the host too, so that it can write
+ // transfers into them directly (see dataSubmitImpl).
+ bool MapForHost =
+ MemoryPool == CoarseGrainedMemoryPools[0] && allowsDirectWrite(Size);
+ if (MapForHost)
+ Agents.push_back(HostDevice.getAgent());
+
// Enable all valid kernel agents to access the buffer.
- if (auto Err = MemoryPool->enableAccess(Alloc, Size, Agents))
+ Error Err = MemoryPool->enableAccess(Alloc, Size, Agents);
+
+ // The host mapping is only an optimization. If it fails, e.g., because the
+ // host-visible window of device memory is exhausted, fall back to the
+ // device agents alone.
+ if (Err && MapForHost) {
+ std::string Msg = toString(std::move(Err));
+ ODBG(OLDT_Alloc) << "Failed to map device allocation " << Alloc
+ << " for the host: " << Msg;
+ MapForHost = false;
+ Agents.pop_back();
+ Err = MemoryPool->enableAccess(Alloc, Size, Agents);
+ }
+
+ if (Err) {
+ if (auto FreeErr = MemoryPool->deallocate(Alloc))
+ return joinErrors(std::move(Err), std::move(FreeErr));
return std::move(Err);
+ }
+
+ if (MapForHost) {
+ std::lock_guard<std::mutex> Lock(HostMappedAllocsMutex);
+ HostMappedAllocs.insert_or_assign(Alloc, Size);
+ }
// Register host / shared memory as pinned memory, so that transfers reading
// from it can take a device-accessible path.
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index 9c02166259e229..80c238761fbe0d 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1008,11 +1008,14 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
void *&BaseDevAccessiblePtr,
size_t &BaseSize) const = 0;
- /// Submit data to the device (host to device transfer).
+ /// Submit data to the device (host to device transfer). Set \p Unordered if
+ /// no operation pending on the queue accesses \p TgtPtr, so that the
+ /// transfer need not be ordered after them.
Error dataSubmit(void *TgtPtr, const void *HstPtr, int64_t Size,
- __tgt_async_info *AsyncInfo);
+ __tgt_async_info *AsyncInfo, bool Unordered = false);
virtual Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
- AsyncInfoWrapperTy &AsyncInfoWrapper) = 0;
+ AsyncInfoWrapperTy &AsyncInfoWrapper,
+ bool Unordered) = 0;
/// Retrieve data from the device (device to host transfer).
Error dataRetrieve(void *HstPtr, const void *TgtPtr, int64_t Size,
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 9355d3e6d27ec4..1bd33ae5cf04c2 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -723,10 +723,11 @@ Error GenericDeviceTy::dataDelete(void *TgtPtr, TargetAllocTy Kind,
}
Error GenericDeviceTy::dataSubmit(void *TgtPtr, const void *HstPtr,
- int64_t Size, __tgt_async_info *AsyncInfo) {
+ int64_t Size, __tgt_async_info *AsyncInfo,
+ bool Unordered) {
AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
- auto Err = dataSubmitImpl(TgtPtr, HstPtr, Size, AsyncInfoWrapper);
+ auto Err = dataSubmitImpl(TgtPtr, HstPtr, Size, AsyncInfoWrapper, Unordered);
AsyncInfoWrapper.finalize(Err);
return Err;
}
diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index 4215ec1e4cd5e8..d8ad1fa8885131 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -823,7 +823,8 @@ struct CUDADeviceTy : public GenericDeviceTy {
/// Submit data to the device (host to device transfer).
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
- AsyncInfoWrapperTy &AsyncInfoWrapper) override {
+ AsyncInfoWrapperTy &AsyncInfoWrapper,
+ bool Unordered) override {
if (auto Err = setContext())
return Err;
diff --git a/offload/plugins-nextgen/host/src/rtl.cpp b/offload/plugins-nextgen/host/src/rtl.cpp
index faff6d9f89e356..90d087e46368c8 100644
--- a/offload/plugins-nextgen/host/src/rtl.cpp
+++ b/offload/plugins-nextgen/host/src/rtl.cpp
@@ -264,7 +264,8 @@ struct GenELF64DeviceTy : public GenericDeviceTy {
/// Submit data to the device (host to device transfer).
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
- AsyncInfoWrapperTy &AsyncInfoWrapper) override {
+ AsyncInfoWrapperTy &AsyncInfoWrapper,
+ bool Unordered) override {
std::memcpy(TgtPtr, HstPtr, Size);
return Plugin::success();
}
diff --git a/offload/plugins-nextgen/level_zero/include/L0Device.h b/offload/plugins-nextgen/level_zero/include/L0Device.h
index ba84bed58b28af..413069e0e4daf6 100644
--- a/offload/plugins-nextgen/level_zero/include/L0Device.h
+++ b/offload/plugins-nextgen/level_zero/include/L0Device.h
@@ -545,7 +545,8 @@ class L0DeviceTy final : public GenericDeviceTy {
Error queryAsyncImpl(__tgt_async_info &AsyncInfo, bool ReleaseQueue,
bool *IsQueueWorkCompleted) override;
Error dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
- AsyncInfoWrapperTy &AsyncInfoWrapper) override;
+ AsyncInfoWrapperTy &AsyncInfoWrapper,
+ bool Unordered) override;
Error dataRetrieveImpl(void *HstPtr, const void *TgtPtr, int64_t Size,
AsyncInfoWrapperTy &AsyncInfoWrapper) override;
Error dataMemcpyImpl(void *DstPtr, const void *SrcPtr, int64_t Size,
diff --git a/offload/plugins-nextgen/level_zero/src/L0Device.cpp b/offload/plugins-nextgen/level_zero/src/L0Device.cpp
index 89f5f6b2d2620c..9816b6dceed299 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Device.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Device.cpp
@@ -344,7 +344,8 @@ Error L0DeviceTy::free(void *TgtPtr, TargetAllocTy Kind) {
}
Error L0DeviceTy::dataSubmitImpl(void *TgtPtr, const void *HstPtr, int64_t Size,
- AsyncInfoWrapperTy &AsyncInfoWrapper) {
+ AsyncInfoWrapperTy &AsyncInfoWrapper,
+ bool Unordered) {
if (Size == 0)
return Plugin::success();
diff --git a/offload/test/offloading/kernel_launch_environment.c b/offload/test/offloading/kernel_launch_environment.c
new file mode 100644
index 00000000000000..95cea9c3e7f5b2
--- /dev/null
+++ b/offload/test/offloading/kernel_launch_environment.c
@@ -0,0 +1,73 @@
+// 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
+//
+// 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