[llvm] [offload] Add olMemPrefetch to liboffload (PR #206752)
Łukasz Plewa via llvm-commits
llvm-commits at lists.llvm.org
Thu Jul 2 06:00:38 PDT 2026
https://github.com/lplewa updated https://github.com/llvm/llvm-project/pull/206752
>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 1/3] [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));
+}
>From b000911896a9ee2431a4fe5a793566908ac4eff4 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Plewa?= <lukasz.plewa at intel.com>
Date: Tue, 30 Jun 2026 17:56:23 +0200
Subject: [PATCH 2/3] clang format unittest
---
.../OffloadAPI/memory/olMemPrefetch.cpp | 20 +++++++++----------
1 file changed, 10 insertions(+), 10 deletions(-)
diff --git a/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp b/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
index ef077de46dace..d801836bd85b0 100644
--- a/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
+++ b/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
@@ -20,8 +20,8 @@ TEST_P(olMemPrefetchTest, SuccessHostToDevice) {
std::memset(Alloc, 0x42, Size);
- ASSERT_SUCCESS(olMemPrefetch(Queue, Alloc, Size,
- OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ 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++)
@@ -38,10 +38,10 @@ TEST_P(olMemPrefetchTest, SuccessDeviceToHost) {
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(
+ 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++)
@@ -55,8 +55,8 @@ TEST_P(olMemPrefetchTest, SuccessZeroSize) {
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(
+ olMemPrefetch(Queue, Alloc, 0, OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
ASSERT_SUCCESS(olSyncQueue(Queue));
ASSERT_SUCCESS(olMemFree(Alloc));
@@ -69,8 +69,8 @@ TEST_P(olMemPrefetchTest, SuccessUnsupportedAllocType) {
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(
+ olMemPrefetch(Queue, Alloc, Size, OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
ASSERT_SUCCESS(olSyncQueue(Queue));
ASSERT_SUCCESS(olMemFree(Alloc));
>From 59c8f5bdd57a27e35e8c8de8ec1fb5729f1017fd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Plewa?= <lukasz.plewa at intel.com>
Date: Thu, 2 Jul 2026 14:59:31 +0200
Subject: [PATCH 3/3] [offload] change prefetch to accept multiple ptrs
---
offload/liboffload/API/Memory.td | 13 ++-
offload/liboffload/src/OffloadImpl.cpp | 11 ++-
.../common/include/PluginInterface.h | 14 ++--
.../common/src/PluginInterface.cpp | 5 +-
offload/plugins-nextgen/cuda/src/rtl.cpp | 42 ++++++----
.../level_zero/include/L0Device.h | 3 +-
.../level_zero/src/L0Device.cpp | 10 ++-
.../OffloadAPI/memory/olMemPrefetch.cpp | 83 ++++++++++++++++---
8 files changed, 133 insertions(+), 48 deletions(-)
diff --git a/offload/liboffload/API/Memory.td b/offload/liboffload/API/Memory.td
index 7eac24caf7ea1..a7470e859c46e 100644
--- a/offload/liboffload/API/Memory.td
+++ b/offload/liboffload/API/Memory.td
@@ -237,16 +237,23 @@ def olMemPrefetch : Function {
"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.",
+ "`Mems` and `Sizes` must both contain `Count` elements; element `i` of `Sizes` is the "
+ "size in bytes of the region pointed to by element `i` of `Mems`.",
];
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<"size_t", "Count", "number of entries in `Mems` and `Sizes`. Zero is allowed and is a no-op.", PARAM_IN>,
+ Param<"const void**", "Mems", "array of `Count` pointers, each pointing into a USM allocation to be prefetched. May be NULL if `Count` is 0.", PARAM_IN_OPTIONAL>,
+ Param<"const size_t*", "Sizes", "array of `Count` sizes in bytes. A zero entry is allowed and is a no-op for that element. May be NULL if `Count` is 0.", PARAM_IN_OPTIONAL>,
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_NULL_POINTER", [
+ "`Count > 0 && Mems == NULL`",
+ "`Count > 0 && Sizes == NULL`"
+ ]>,
Return<"OL_ERRC_INVALID_ENUMERATION", [
- "Flags contains bits other than the defined OL_USM_MIGRATION_FLAG_* values"
+ "`(Flags & ~(OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE | OL_USM_MIGRATION_FLAG_DEVICE_TO_HOST)) != 0`"
]>,
];
}
diff --git a/offload/liboffload/src/OffloadImpl.cpp b/offload/liboffload/src/OffloadImpl.cpp
index 3754ae7d2b5b2..bd5236bf10d21 100644
--- a/offload/liboffload/src/OffloadImpl.cpp
+++ b/offload/liboffload/src/OffloadImpl.cpp
@@ -1030,15 +1030,14 @@ 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,
+Error olMemPrefetch_impl(ol_queue_handle_t Queue, size_t Count,
+ const void **Mems, const size_t *Sizes,
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);
+ if (Count == 0)
+ return Error::success();
bool ToHost = (Flags & OL_USM_MIGRATION_FLAG_DEVICE_TO_HOST) != 0;
- return Queue->Device->Device->dataPrefetch(Mem, Size, ToHost,
+ return Queue->Device->Device->dataPrefetch(Count, Mems, Sizes, ToHost,
Queue->AsyncInfo);
}
diff --git a/offload/plugins-nextgen/common/include/PluginInterface.h b/offload/plugins-nextgen/common/include/PluginInterface.h
index e299b710aaf89..ffd1457c15add 100644
--- a/offload/plugins-nextgen/common/include/PluginInterface.h
+++ b/offload/plugins-nextgen/common/include/PluginInterface.h
@@ -1043,13 +1043,13 @@ 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,
+ /// Prefetch a batch of memory ranges. Mems[i] has size Sizes[i].
+ /// ToHost = true migrates towards the host, false towards the device.
+ /// Backends without prefetch support treat the call as a no-op.
+ Error dataPrefetch(size_t Count, const void **Mems, const size_t *Sizes,
+ bool ToHost, __tgt_async_info *AsyncInfo);
+ virtual Error dataPrefetchImpl(size_t Count, const void **Mems,
+ const size_t *Sizes, bool ToHost,
AsyncInfoWrapperTy &AsyncInfoWrapper) {
return Plugin::success();
}
diff --git a/offload/plugins-nextgen/common/src/PluginInterface.cpp b/offload/plugins-nextgen/common/src/PluginInterface.cpp
index 5f1f68653dce9..2f8a8d312aec8 100644
--- a/offload/plugins-nextgen/common/src/PluginInterface.cpp
+++ b/offload/plugins-nextgen/common/src/PluginInterface.cpp
@@ -1155,10 +1155,11 @@ Error GenericDeviceTy::dataFill(void *TgtPtr, const void *PatternPtr,
return Err;
}
-Error GenericDeviceTy::dataPrefetch(const void *Mem, int64_t Size, bool ToHost,
+Error GenericDeviceTy::dataPrefetch(size_t Count, const void **Mems,
+ const size_t *Sizes, bool ToHost,
__tgt_async_info *AsyncInfo) {
AsyncInfoWrapperTy AsyncInfoWrapper(*this, AsyncInfo);
- auto Err = dataPrefetchImpl(Mem, Size, ToHost, AsyncInfoWrapper);
+ auto Err = dataPrefetchImpl(Count, Mems, Sizes, ToHost, AsyncInfoWrapper);
AsyncInfoWrapper.finalize(Err);
return Err;
}
diff --git a/offload/plugins-nextgen/cuda/src/rtl.cpp b/offload/plugins-nextgen/cuda/src/rtl.cpp
index 9d04e3d433296..34b8b6994b4e8 100644
--- a/offload/plugins-nextgen/cuda/src/rtl.cpp
+++ b/offload/plugins-nextgen/cuda/src/rtl.cpp
@@ -928,17 +928,17 @@ struct CUDADeviceTy : public GenericDeviceTy {
}
/// Prefetch managed memory to the device or back to the host.
- Error dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
+ // TODO: switch to cuMemPrefetchBatchAsync once the minimum supported CUDA
+ // driver is 13 or newer. That entry point takes the (Mems, Sizes, Count)
+ // arrays directly and lets the driver batch the migration.
+ Error dataPrefetchImpl(size_t Count, const void **Mems, const size_t *Sizes,
+ bool ToHost,
AsyncInfoWrapperTy &AsyncInfoWrapper) override {
- if (Size == 0)
+ if (Count == 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
@@ -949,17 +949,29 @@ struct CUDADeviceTy : public GenericDeviceTy {
!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();
+ CUstream Stream;
+ if (auto Err = getStream(AsyncInfoWrapper, Stream))
+ return Err;
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");
+ for (size_t I = 0; I < Count; I++) {
+ if (Sizes[I] == 0)
+ continue;
+
+ // Prefetch only works with USM (managed) memory; ignore the hint
+ // otherwise.
+ unsigned int IsManaged = 0;
+ if (cuPointerGetAttribute(&IsManaged, CU_POINTER_ATTRIBUTE_IS_MANAGED,
+ (CUdeviceptr)Mems[I]) != CUDA_SUCCESS ||
+ !IsManaged)
+ continue;
+
+ CUresult Res =
+ cuMemPrefetchAsync((CUdeviceptr)Mems[I], Sizes[I], Dst, Stream);
+ if (auto Err = Plugin::check(Res, "error in cuMemPrefetchAsync: %s"))
+ return Err;
+ }
+ return Plugin::success();
}
/// Initialize the async info for interoperability purposes.
diff --git a/offload/plugins-nextgen/level_zero/include/L0Device.h b/offload/plugins-nextgen/level_zero/include/L0Device.h
index 9f55b1e55b342..408fb452a0dc8 100644
--- a/offload/plugins-nextgen/level_zero/include/L0Device.h
+++ b/offload/plugins-nextgen/level_zero/include/L0Device.h
@@ -507,7 +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,
+ Error dataPrefetchImpl(size_t Count, const void **Mems, const size_t *Sizes,
+ bool ToHost,
AsyncInfoWrapperTy &AsyncInfoWrapper) override;
Error synchronizeImpl(__tgt_async_info &AsyncInfo,
bool ReleaseQueue) override;
diff --git a/offload/plugins-nextgen/level_zero/src/L0Device.cpp b/offload/plugins-nextgen/level_zero/src/L0Device.cpp
index d18741a19e75c..e0dfe0596b5ce 100644
--- a/offload/plugins-nextgen/level_zero/src/L0Device.cpp
+++ b/offload/plugins-nextgen/level_zero/src/L0Device.cpp
@@ -685,9 +685,10 @@ Error L0DeviceTy::dataFillImpl(void *TgtPtr, const void *PatternPtr,
AsyncInfoWrapper);
}
-Error L0DeviceTy::dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
+Error L0DeviceTy::dataPrefetchImpl(size_t Count, const void **Mems,
+ const size_t *Sizes, bool ToHost,
AsyncInfoWrapperTy &AsyncInfoWrapper) {
- if (Size == 0)
+ if (Count == 0)
return Plugin::success();
// Level Zero only supports prefetching memory to the device. A prefetch
@@ -699,7 +700,10 @@ Error L0DeviceTy::dataPrefetchImpl(const void *Mem, int64_t Size, bool ToHost,
auto QueueOrErr = getOrCreateQueue(AsyncInfo);
if (!QueueOrErr)
return QueueOrErr.takeError();
- return (*QueueOrErr)->memoryPrefetch(Mem, Size);
+ for (size_t I = 0; I < Count; I++)
+ if (auto Err = (*QueueOrErr)->memoryPrefetch(Mems[I], Sizes[I]))
+ return Err;
+ return Plugin::success();
}
Expected<void *> L0DeviceTy::dataAlloc(size_t Size, size_t Align, int32_t Kind,
diff --git a/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp b/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
index d801836bd85b0..92762ca554c35 100644
--- a/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
+++ b/offload/unittests/OffloadAPI/memory/olMemPrefetch.cpp
@@ -20,8 +20,10 @@ TEST_P(olMemPrefetchTest, SuccessHostToDevice) {
std::memset(Alloc, 0x42, Size);
- ASSERT_SUCCESS(
- olMemPrefetch(Queue, Alloc, Size, OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ const void *Mems[] = {Alloc};
+ const size_t Sizes[] = {Size};
+ ASSERT_SUCCESS(olMemPrefetch(Queue, 1, Mems, Sizes,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
ASSERT_SUCCESS(olSyncQueue(Queue));
for (size_t I = 0; I < Size; I++)
@@ -37,11 +39,13 @@ TEST_P(olMemPrefetchTest, SuccessDeviceToHost) {
std::memset(Alloc, 0x21, Size);
+ const void *Mems[] = {Alloc};
+ const size_t Sizes[] = {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(olMemPrefetch(Queue, 1, Mems, Sizes,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olMemPrefetch(Queue, 1, Mems, Sizes,
+ OL_USM_MIGRATION_FLAG_DEVICE_TO_HOST));
ASSERT_SUCCESS(olSyncQueue(Queue));
for (size_t I = 0; I < Size; I++)
@@ -50,13 +54,46 @@ TEST_P(olMemPrefetchTest, SuccessDeviceToHost) {
ASSERT_SUCCESS(olMemFree(Alloc));
}
+TEST_P(olMemPrefetchTest, SuccessMultiple) {
+ constexpr size_t NumAllocs = 3;
+ constexpr size_t Size = 1024;
+ void *Allocs[NumAllocs];
+ const void *Mems[NumAllocs];
+ size_t Sizes[NumAllocs];
+ for (size_t I = 0; I < NumAllocs; I++) {
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Allocs[I]));
+ std::memset(Allocs[I], static_cast<int>(0x10 + I), Size);
+ Mems[I] = Allocs[I];
+ Sizes[I] = Size;
+ }
+
+ ASSERT_SUCCESS(olMemPrefetch(Queue, NumAllocs, Mems, Sizes,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+
+ for (size_t I = 0; I < NumAllocs; I++) {
+ for (size_t J = 0; J < Size; J++)
+ ASSERT_EQ(static_cast<uint8_t *>(Allocs[I])[J],
+ static_cast<uint8_t>(0x10 + I));
+ ASSERT_SUCCESS(olMemFree(Allocs[I]));
+ }
+}
+
+TEST_P(olMemPrefetchTest, SuccessZeroCount) {
+ ASSERT_SUCCESS(olMemPrefetch(Queue, 0, nullptr, nullptr,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+ ASSERT_SUCCESS(olSyncQueue(Queue));
+}
+
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));
+ const void *Mems[] = {Alloc};
+ const size_t Sizes[] = {0};
+ ASSERT_SUCCESS(olMemPrefetch(Queue, 1, Mems, Sizes,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
ASSERT_SUCCESS(olSyncQueue(Queue));
ASSERT_SUCCESS(olMemFree(Alloc));
@@ -69,8 +106,10 @@ TEST_P(olMemPrefetchTest, SuccessUnsupportedAllocType) {
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));
+ const void *Mems[] = {Alloc};
+ const size_t Sizes[] = {Size};
+ ASSERT_SUCCESS(olMemPrefetch(Queue, 1, Mems, Sizes,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
ASSERT_SUCCESS(olSyncQueue(Queue));
ASSERT_SUCCESS(olMemFree(Alloc));
@@ -81,8 +120,30 @@ TEST_P(olMemPrefetchTest, InvalidFlags) {
void *Alloc;
ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+ const void *Mems[] = {Alloc};
+ const size_t Sizes[] = {Size};
ASSERT_ERROR(OL_ERRC_INVALID_ENUMERATION,
- olMemPrefetch(Queue, Alloc, Size, 0xdeadbeef));
+ olMemPrefetch(Queue, 1, Mems, Sizes, 0xdeadbeef));
+
+ ASSERT_SUCCESS(olMemFree(Alloc));
+}
+
+TEST_P(olMemPrefetchTest, InvalidNullMems) {
+ const size_t Sizes[] = {0};
+ ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+ olMemPrefetch(Queue, 1, nullptr, Sizes,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
+}
+
+TEST_P(olMemPrefetchTest, InvalidNullSizes) {
+ constexpr size_t Size = 1024;
+ void *Alloc;
+ ASSERT_SUCCESS(olMemAlloc(Device, OL_ALLOC_TYPE_MANAGED, Size, &Alloc));
+
+ const void *Mems[] = {Alloc};
+ ASSERT_ERROR(OL_ERRC_INVALID_NULL_POINTER,
+ olMemPrefetch(Queue, 1, Mems, nullptr,
+ OL_USM_MIGRATION_FLAG_HOST_TO_DEVICE));
ASSERT_SUCCESS(olMemFree(Alloc));
}
More information about the llvm-commits
mailing list