[llvm] [offload] Add olMemPrefetch to liboffload (PR #206752)
Łukasz Plewa via llvm-commits
llvm-commits at lists.llvm.org
Tue Jun 30 08:51:00 PDT 2026
https://github.com/lplewa created https://github.com/llvm/llvm-project/pull/206752
Adds olMemPrefetch and the OL_USM_MIGRATION_FLAG_* enum, wired through a new dataPrefetchImpl on the plugin interface with Level Zero and CUDA implementations. Backends without prefetch support inherit the default no-op.
Prefetch is a hint that migrates a USM allocation between the host and a device so subsequent accesses on the target side are faster.
Assisted-by: Claude
>From 977bb8a1948cc1d38de7a847dbd1c14dcfa0e771 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Plewa?= <lukasz.plewa at intel.com>
Date: Mon, 29 Jun 2026 11:34:55 +0200
Subject: [PATCH] [offload] Add olMemPrefetch to liboffload
Adds olMemPrefetch and the OL_USM_MIGRATION_FLAG_* enum, wired through
a new dataPrefetchImpl on the plugin interface with Level Zero and
CUDA implementations. Backends without prefetch support inherit the
default no-op.
Prefetch is a hint that migrates a USM allocation between the host
and a device so subsequent accesses on the target side are faster.
---
offload/liboffload/API/Memory.td | 36 ++++++++
offload/liboffload/src/OffloadImpl.cpp | 12 +++
.../common/include/PluginInterface.h | 11 +++
.../common/src/PluginInterface.cpp | 8 ++
.../cuda/dynamic_cuda/cuda.cpp | 3 +
.../plugins-nextgen/cuda/dynamic_cuda/cuda.h | 9 ++
offload/plugins-nextgen/cuda/src/rtl.cpp | 35 ++++++++
.../level_zero/include/L0CmdListManager.h | 6 ++
.../level_zero/include/L0Device.h | 2 +
.../level_zero/include/L0Queue.h | 9 ++
.../level_zero/src/L0Device.cpp | 17 ++++
offload/unittests/OffloadAPI/CMakeLists.txt | 1 +
.../OffloadAPI/memory/olMemPrefetch.cpp | 88 +++++++++++++++++++
13 files changed, 237 insertions(+)
create mode 100644 offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
diff --git a/offload/liboffload/API/Memory.td b/offload/liboffload/API/Memory.td
index 78f136801885a..7eac24caf7ea1 100644
--- a/offload/liboffload/API/Memory.td
+++ b/offload/liboffload/API/Memory.td
@@ -214,3 +214,39 @@ def olMemUnregister : Function {
];
let returns = [];
}
+
+def ol_usm_migration_flags_t : Typedef {
+ let desc = "USM migration flags, indicating the direction data is migrated in.";
+ let value = "uint32_t";
+}
+
+def ol_usm_migration_flag_t : Enum {
+ let desc = "USM migration flags, indicating the direction data is migrated in.";
+ let is_bit_field = 1;
+ let etors = [
+ Etor<"HOST_TO_DEVICE", "Migrate data from the host to the device specified to olMemPrefetch">,
+ Etor<"DEVICE_TO_HOST", "Migrate data from the device specified to olMemPrefetch back to the host">,
+ ];
+}
+
+def olMemPrefetch : Function {
+ let desc = "Enqueue a command to migrate USM memory between the host and a device.";
+ let details = [
+ "Prefetch is a hint that migrates a USM allocation between the host and a device "
+ "so subsequent accesses on the target side are faster. The contents of the memory "
+ "are unchanged - only its physical location may change.",
+ "Prefetching may not be supported for all devices or allocation types. If it is not "
+ "supported the prefetch is ignored and the call still returns success.",
+ ];
+ let params = [
+ Param<"ol_queue_handle_t", "Queue", "handle of the queue", PARAM_IN>,
+ Param<"const void*", "Mem", "pointer into the USM allocation to be prefetched", PARAM_IN>,
+ Param<"size_t", "Size", "size in bytes of the region to be prefetched. Zero is allowed and is a no-op.", PARAM_IN>,
+ Param<"ol_usm_migration_flags_t", "Flags", "direction of the migration; exactly one of the OL_USM_MIGRATION_FLAG_* values", PARAM_IN>,
+ ];
+ let returns = [
+ Return<"OL_ERRC_INVALID_ENUMERATION", [
+ "Flags contains bits other than the defined OL_USM_MIGRATION_FLAG_* values"
+ ]>,
+ ];
+}
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index a36081f27b5ee..3754ae7d2b5b2 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -1030,6 +1030,18 @@ Error olMemFill_impl(ol_queue_handle_t Queue, void *Ptr, size_t PatternSize,
Queue->AsyncInfo);
}
+Error olMemPrefetch_impl(ol_queue_handle_t Queue, const void *Mem, size_t Size,
+ ol_usm_migration_flags_t Flags) {
+ if ((Flags & ~(OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE |
+ OL_USM_MIGRATION_FLAG_DEVICE_TO_HOST)) != 0)
+ return createOffloadError(ErrorCode::INVALID_ENUMERATION,
+ "olMemPrefetch flags '%i' are invalid", Flags);
+
+ bool ToHost = (Flags & OL_USM_MIGRATION_FLAG_DEVICE_TO_HOST) != 0;
+ return Queue->Device->Device->dataPrefetch(Mem, Size, ToHost,
+ Queue->AsyncInfo);
+}
+
Error olCreateProgram_impl(ol_device_handle_t Device, const void *ProgData,
size_t ProgDataSize, ol_program_handle_t *Program) {
StringRef Buffer(reinterpret_cast<const char *>(ProgData), ProgDataSize);
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index dc21abf1a334a..e299b710aaf89 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1043,6 +1043,17 @@ struct GenericDeviceTy : public DeviceAllocatorTy {
int64_t PatternSize, int64_t Size,
AsyncInfoWrapperTy &AsyncInfoWrapper) = 0;
+ /// Prefetch the memory range \p Mem of size \p Size to the device.
+ /// \p ToHost selects the direction of the migration: when true the data is
+ /// moved towards the host, otherwise towards the device. Backends that do
+ /// not natively support prefetching treat the call as a no-op.
+ Error dataPrefetch(const void *Mem, int64_t Size, bool ToHost,
+ __tgt_async_info *AsyncInfo);
+ virtual Error dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
+ AsyncInfoWrapperTy &AsyncInfoWrapper) {
+ return Plugin::success();
+ }
+
/// Run the kernel associated with \p EntryPtr
Error launchKernel(void *EntryPtr, void **ArgPtrs, ptrdiff_t *ArgOffsets,
KernelArgsTy &KernelArgs,
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index c77182cb03ec9..5f1f68653dce9 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -1155,6 +1155,14 @@ Error GenericDeviceTy::dataFill(void *TgtPtr, const void *PatternPtr,
return Err;
}
+Error GenericDeviceTy::dataPrefetch(const void *Mem, int64_t Size, bool ToHost,
+ __tgt_async_info *AsyncInfo) {
+ AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
+ auto Err = dataPrefetchImpl(Mem, Size, ToHost, AsyncInfoWrapper);
+ AsyncInfoWrapper.finalize(Err);
+ return Err;
+}
+
Error GenericDeviceTy::launchKernel(void *EntryPtr, void **ArgPtrs,
ptrdiff_t *ArgOffsets,
KernelArgsTy &KernelArgs,
diff --git a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp
index cfa68a122623b..14105fc4cc6bf 100644
--- a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp
+++ b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.cpp
@@ -67,6 +67,9 @@ DLWRAP(cuMemFree, 1)
DLWRAP(cuMemFreeHost, 1)
DLWRAP(cuMemFreeAsync, 2)
+DLWRAP(cuMemPrefetchAsync, 4)
+DLWRAP(cuPointerGetAttribute, 3)
+
DLWRAP(cuModuleGetFunction, 3)
DLWRAP(cuModuleGetGlobal, 4)
diff --git a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h
index c832157921cfb..fc455dd51d03e 100644
--- a/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h
+++ b/offload/plugins-nextgen/cuda/dynamic_cuda/cuda.h
@@ -38,6 +38,7 @@ typedef struct CUuuid_st {
} CUuuid;
#define CU_DEVICE_INVALID ((CUdevice)(-2))
+#define CU_DEVICE_CPU ((CUdevice)(-1))
typedef unsigned long long CUmemGenericAllocationHandle_v1;
typedef CUmemGenericAllocationHandle_v1 CUmemGenericAllocationHandle;
@@ -381,6 +382,14 @@ CUresult cuMemFree(CUdeviceptr);
CUresult cuMemFreeHost(void *);
CUresult cuMemFreeAsync(CUdeviceptr, CUstream);
+CUresult cuMemPrefetchAsync(CUdeviceptr, size_t, CUdevice, CUstream);
+
+typedef enum CUpointer_attribute_enum {
+ CU_POINTER_ATTRIBUTE_IS_MANAGED = 8
+} CUpointer_attribute;
+
+CUresult cuPointerGetAttribute(void *, CUpointer_attribute, CUdeviceptr);
+
CUresult cuModuleGetFunction(CUfunction *, CUmodule, const char *);
CUresult cuModuleGetGlobal(CUdeviceptr *, size_t *, CUmodule, const char *);
diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index fc7d9e083f85e..9d04e3d433296 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -927,6 +927,41 @@ struct CUDADeviceTy : public GenericDeviceTy {
return Plugin::check(Res, "error in cuMemset: %s");
}
+ /// Prefetch managed memory to the device or back to the host.
+ Error dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
+ AsyncInfoWrapperTy &AsyncInfoWrapper) override {
+ if (Size == 0)
+ return Plugin::success();
+ if (auto Err = setContext())
+ return Err;
+
+ CUstream Stream;
+ if (auto Err = getStream(AsyncInfoWrapper, Stream))
+ return Err;
+
+ // Certain cuda devices and Windows do not have support for some Unified
+ // Memory features. cuMemPrefetchAsync requires concurrent memory access
+ // for managed memory. Therefore, ignore prefetch hint if concurrent managed
+ // memory access is not available.
+ int ConcurrentManagedAccess = 0;
+ if (getDeviceAttrRaw(CU_DEVICE_ATTRIBUTE_CONCURRENT_MANAGED_ACCESS,
+ ConcurrentManagedAccess) != CUDA_SUCCESS ||
+ !ConcurrentManagedAccess)
+ return Plugin::success();
+
+ // Prefetch only works with USM (managed) memory; ignore the hint otherwise.
+ unsigned int IsManaged = 0;
+ if (cuPointerGetAttribute(&IsManaged, CU_POINTER_ATTRIBUTE_IS_MANAGED,
+ (CUdeviceptr)Mem) != CUDA_SUCCESS ||
+ !IsManaged)
+ return Plugin::success();
+
+ CUdevice Dst = ToHost ? CU_DEVICE_CPU : Device;
+ CUresult Res =
+ cuMemPrefetchAsync((CUdeviceptr)Mem, (size_t)Size, Dst, Stream);
+ return Plugin::check(Res, "error in cuMemPrefetchAsync: %s");
+ }
+
/// Initialize the async info for interoperability purposes.
Error initAsyncInfoImpl(AsyncInfoWrapperTy &AsyncInfoWrapper) override {
if (auto Err = setContext())
diff --git a/offload/plugins-nextgen/level_zero/include/L0CmdListManager.h b/offload/plugins-nextgen/level_zero/include/L0CmdListManager.h
index e93b1a8948053..94ddef231ff11 100644
--- a/offload/plugins-nextgen/level_zero/include/L0CmdListManager.h
+++ b/offload/plugins-nextgen/level_zero/include/L0CmdListManager.h
@@ -96,6 +96,12 @@ class L0CmdListManagerTy {
return Plugin::success();
}
+ Error appendMemoryPrefetch(const void *Ptr, size_t Size) {
+ std::lock_guard<std::mutex> Lock(Mtx);
+ CALL_ZE_RET_ERROR(zeCommandListAppendMemoryPrefetch, CmdList, Ptr, Size);
+ return Plugin::success();
+ }
+
Error appendLaunchKernel(ze_kernel_handle_t Kernel,
const ze_group_count_t *pLaunchFuncArgs,
ze_event_handle_t SignalEvent = nullptr,
diff --git a/offload/plugins-nextgen/level_zero/include/L0Device.h b/offload/plugins-nextgen/level_zero/include/L0Device.h
index 4aa0633339d25..9f55b1e55b342 100644
--- a/offload/plugins-nextgen/level_zero/include/L0Device.h
+++ b/offload/plugins-nextgen/level_zero/include/L0Device.h
@@ -507,6 +507,8 @@ class L0DeviceTy final : public GenericDeviceTy {
Error dataFillImpl(void *TgtPtr, const void *PatternPtr, int64_t PatternSize,
int64_t Size,
AsyncInfoWrapperTy &AsyncInfoWrapper) override;
+ Error dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
+ AsyncInfoWrapperTy &AsyncInfoWrapper) override;
Error synchronizeImpl(__tgt_async_info &AsyncInfo,
bool ReleaseQueue) override;
Error queryAsyncImpl(__tgt_async_info &AsyncInfo, bool ReleaseQueue,
diff --git a/offload/plugins-nextgen/level_zero/include/L0Queue.h b/offload/plugins-nextgen/level_zero/include/L0Queue.h
index 60558650b75b3..ca477376fc0a6 100644
--- a/offload/plugins-nextgen/level_zero/include/L0Queue.h
+++ b/offload/plugins-nextgen/level_zero/include/L0Queue.h
@@ -74,6 +74,12 @@ class L0QueueTy {
return memoryFillImpl(Ptr, Pattern, PatternSize, Size);
}
+ Error memoryPrefetch(const void *Ptr, size_t Size) {
+ if (Size == 0)
+ return Plugin::success();
+ return memoryPrefetchImpl(Ptr, Size);
+ }
+
Error dispatchLaunchKernel(ze_kernel_handle_t Kernel, L0LaunchEnvTy &KEnv,
ze_event_handle_t SignalEvent = nullptr,
uint32_t NumWaitEvents = 0,
@@ -136,6 +142,9 @@ class L0QueueTy {
size_t PatternSize, size_t Size) {
return CmdList->appendMemoryFill(Ptr, Pattern, PatternSize, Size);
}
+ virtual Error memoryPrefetchImpl(const void *Ptr, size_t Size) {
+ return CmdList->appendMemoryPrefetch(Ptr, Size);
+ }
virtual Error dataFenceImpl() = 0;
virtual Error appendSignalEventImpl(ze_event_handle_t Event) {
diff --git a/offload/plugins-nextgen/level_zero/src/L0Device.cpp b/offload/plugins-nextgen/level_zero/src/L0Device.cpp
index 9a7d534c56c40..d18741a19e75c 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Device.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Device.cpp
@@ -685,6 +685,23 @@ Error L0DeviceTy::dataFillImpl(void *TgtPtr, const void *PatternPtr,
AsyncInfoWrapper);
}
+Error L0DeviceTy::dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
+ AsyncInfoWrapperTy &AsyncInfoWrapper) {
+ if (Size == 0)
+ return Plugin::success();
+
+ // Level Zero only supports prefetching memory to the device. A prefetch
+ // request targeting the host is treated as a no-op.
+ if (ToHost)
+ return Plugin::success();
+
+ __tgt_async_info *AsyncInfo = AsyncInfoWrapper;
+ auto QueueOrErr = getOrCreateQueue(AsyncInfo);
+ if (!QueueOrErr)
+ return QueueOrErr.takeError();
+ return (*QueueOrErr)->memoryPrefetch(Mem, Size);
+}
+
Expected<void *> L0DeviceTy::dataAlloc(size_t Size, size_t Align, int32_t Kind,
intptr_t Offset, bool UserAlloc,
bool DevMalloc, uint32_t MemAdvice,
diff --git a/offload/unittests/OffloadAPI/CMakeLists.txt b/offload/unittests/OffloadAPI/CMakeLists.txt
index d960da4b98f82..16d6b633d885f 100644
--- a/offload/unittests/OffloadAPI/CMakeLists.txt
+++ b/offload/unittests/OffloadAPI/CMakeLists.txt
@@ -32,6 +32,7 @@ add_offload_unittest("memory"
memory/olMemFill.cpp
memory/olMemFree.cpp
memory/olMemcpy.cpp
+ memory/olMemPrefetch.cpp
memory/olGetMemInfo.cpp
memory/olGetMemInfoSize.cpp
memory/olMemRegister.cpp)
diff --git a/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp b/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
new file mode 100644
index 0000000000000..ef077de46dace
--- /dev/null
+++ b/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
@@ -0,0 +1,88 @@
+//===------- Offload API tests - olMemPrefetch ----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "../common/Fixtures.hpp"
+#include <OffloadAPI.h>
+#include <gtest/gtest.h>
+
+using olMemPrefetchTest = OffloadQueueTest;
+OFFLOAD_TESTS_INSTANTIATE_DEVICE_FIXTURE(olMemPrefetchTest);
+
+TEST_P(olMemPrefetchTest, SuccessHostToDevice) {
+ constexpr size_t Size = 1024;
+ void *Alloc;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+
+ std::memset(Alloc, 0x42, Size);
+
+ ASSERT_SUCCESS(olMemPrefetch(Queue, Alloc, Size,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ for (size_t I = 0; I < Size; I++)
+ ASSERT_EQ(static_cast<uint8_t *>(Alloc)[I], 0x42);
+
+ ASSERT_SUCCESS(olMemFree(Alloc));
+}
+
+TEST_P(olMemPrefetchTest, SuccessDeviceToHost) {
+ constexpr size_t Size = 1024;
+ void *Alloc;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+
+ std::memset(Alloc, 0x21, Size);
+
+ // Migrate to the device first, then bring it back.
+ ASSERT_SUCCESS(olMemPrefetch(Queue, Alloc, Size,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olMemPrefetch(Queue, Alloc, Size,
+ OL_USM_MIGRATION_FLAG_DEVICE_TO_HOST));
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ for (size_t I = 0; I < Size; I++)
+ ASSERT_EQ(static_cast<uint8_t *>(Alloc)[I], 0x21);
+
+ ASSERT_SUCCESS(olMemFree(Alloc));
+}
+
+TEST_P(olMemPrefetchTest, SuccessZeroSize) {
+ constexpr size_t Size = 1024;
+ void *Alloc;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+
+ ASSERT_SUCCESS(olMemPrefetch(Queue, Alloc, 0,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ ASSERT_SUCCESS(olMemFree(Alloc));
+}
+
+TEST_P(olMemPrefetchTest, SuccessUnsupportedAllocType) {
+ // Prefetching a non-managed allocation is not meaningful, but per the API
+ // contract the hint must be silently ignored and the call must still succeed.
+ constexpr size_t Size = 1024;
+ void *Alloc;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_DEVICE, Size, &Alloc));
+
+ ASSERT_SUCCESS(olMemPrefetch(Queue, Alloc, Size,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ ASSERT_SUCCESS(olMemFree(Alloc));
+}
+
+TEST_P(olMemPrefetchTest, InvalidFlags) {
+ constexpr size_t Size = 1024;
+ void *Alloc;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+
+ ASSERT_ERROR(OL_ERRC_INVALID_ENUMERATION,
+ olMemPrefetch(Queue, Alloc, Size, 0xdeadbeef));
+
+ ASSERT_SUCCESS(olMemFree(Alloc));
+}
More information about the llvm-commits
mailing list