[UR][SYCL] Implement USM prefetch from device to host in SYCL runtime and UR (#19437)

Add the ability to control USM prefetch direction (host-to-device,
device-to-host) in the enqueue_function extension:
```cpp
sycl::ext::oneapi::experimental {
    enum class prefetch_type { device, host };

    void prefetch(sycl::queue q, void* ptr, size_t numBytes,
                          prefetch_type type = prefetch_type::device);
    void prefetch(sycl::handler &h, void* ptr, size_t numBytes,
                          prefetch_type type = prefetch_type::device);
}
```

**Note:**
- ~~There is a test failure regarding a new ABI symbol: In order to not
break the ABI, I added a new handler function to represent prefetch from
device-to-host. Despite the precommit failures, this should not be an
ABI-breaking change and thus be okay to merge.~~ I modified the ABI
symbols tests as a part of this PR.

---------

Co-authored-by: Pablo Reble <[email protected]>

Author: ianayl

sycl/include/sycl/ext/oneapi/experimental/enqueue_functions.hpp

@@ -12,6 +12,7 @@
 
 #include <sycl/detail/common.hpp>
 #include <sycl/event.hpp>
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp>
 #include <sycl/ext/oneapi/experimental/graph.hpp>
 #include <sycl/ext/oneapi/properties/properties.hpp>
 #include <sycl/handler.hpp>
@@ -369,15 +370,17 @@ void fill(sycl::queue Q, T *Ptr, const T &Pattern, size_t Count,
       CodeLoc);
 }
 
-inline void prefetch(handler &CGH, void *Ptr, size_t NumBytes) {
-  CGH.prefetch(Ptr, NumBytes);
+inline void prefetch(handler &CGH, void *Ptr, size_t NumBytes,
+                     prefetch_type Type = prefetch_type::device) {
+  CGH.prefetch(Ptr, NumBytes, Type);
 }
 
 inline void prefetch(queue Q, void *Ptr, size_t NumBytes,
+                     prefetch_type Type = prefetch_type::device,
                      const sycl::detail::code_location &CodeLoc =
                          sycl::detail::code_location::current()) {
   submit(
-      std::move(Q), [&](handler &CGH) { prefetch(CGH, Ptr, NumBytes); },
+      std::move(Q), [&](handler &CGH) { prefetch(CGH, Ptr, NumBytes, Type); },
       CodeLoc);
 }
 

sycl/include/sycl/ext/oneapi/experimental/enqueue_types.hpp

@@ -0,0 +1,33 @@
+//==--------------- enqueue_types.hpp ---- SYCL enqueue types --------------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <string>
+
+namespace sycl {
+inline namespace _V1 {
+namespace ext::oneapi::experimental {
+
+/// @brief Indicates the destination device for USM data to be prefetched to.
+enum class prefetch_type { device, host };
+
+inline std::string prefetchTypeToString(prefetch_type value) {
+  switch (value) {
+  case sycl::ext::oneapi::experimental::prefetch_type::device:
+    return "prefetch_type::device";
+  case sycl::ext::oneapi::experimental::prefetch_type::host:
+    return "prefetch_type::host";
+  default:
+    return "prefetch_type::unknown";
+  }
+}
+
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/handler.hpp

@@ -149,6 +149,8 @@ namespace ext ::oneapi ::experimental {
 template <typename, typename> class work_group_memory;
 template <typename, typename> class dynamic_work_group_memory;
 struct image_descriptor;
+enum class prefetch_type;
+
 __SYCL_EXPORT void async_free(sycl::handler &h, void *ptr);
 __SYCL_EXPORT void *async_malloc(sycl::handler &h, sycl::usm::alloc kind,
                                  size_t size);
@@ -2627,6 +2629,16 @@ class __SYCL_EXPORT handler {
   /// \param Count is a number of bytes to be prefetched.
   void prefetch(const void *Ptr, size_t Count);
 
+  /// Provides hints to the runtime library that data should be made available
+  /// on a device earlier than Unified Shared Memory would normally require it
+  /// to be available.
+  ///
+  /// \param Ptr is a USM pointer to the memory to be prefetched to the device.
+  /// \param Count is a number of bytes to be prefetched.
+  /// \param Type is type of prefetch, i.e. fetch to device or fetch to host.
+  void prefetch(const void *Ptr, size_t Count,
+                ext::oneapi::experimental::prefetch_type Type);
+
   /// Provides additional information to the underlying runtime about how
   /// different allocations are used.
   ///

sycl/source/detail/cg.hpp

@@ -397,14 +397,19 @@ class CGFillUSM : public CG {
 class CGPrefetchUSM : public CG {
   void *MDst;
   size_t MLength;
+  ext::oneapi::experimental::prefetch_type MPrefetchType;
 
 public:
   CGPrefetchUSM(void *DstPtr, size_t Length, CG::StorageInitHelper CGData,
+                ext::oneapi::experimental::prefetch_type PrefetchType,
                 detail::code_location loc = {})
       : CG(CGType::PrefetchUSM, std::move(CGData), std::move(loc)),
-        MDst(DstPtr), MLength(Length) {}
-  void *getDst() { return MDst; }
-  size_t getLength() { return MLength; }
+        MDst(DstPtr), MLength(Length), MPrefetchType(PrefetchType) {}
+  void *getDst() const { return MDst; }
+  size_t getLength() const { return MLength; }
+  ext::oneapi::experimental::prefetch_type getPrefetchType() const {
+    return MPrefetchType;
+  }
 };
 
 /// "Advise USM" command group class.

sycl/source/detail/graph/node_impl.hpp

@@ -15,7 +15,8 @@
 #include <sycl/detail/cg_types.hpp>    // for CGType
 #include <sycl/detail/kernel_desc.hpp> // for kernel_param_kind_t
 
-#include <sycl/ext/oneapi/experimental/graph/node.hpp> // for node
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp> // for prefetchType
+#include <sycl/ext/oneapi/experimental/graph/node.hpp>    // for node
 
 #include <cstring>
 #include <fstream>
@@ -655,7 +656,10 @@ class node_impl : public std::enable_shared_from_this<node_impl> {
         sycl::detail::CGPrefetchUSM *Prefetch =
             static_cast<sycl::detail::CGPrefetchUSM *>(MCommandGroup.get());
         Stream << "Dst: " << Prefetch->getDst()
-               << " Length: " << Prefetch->getLength() << "\\n";
+               << " Length: " << Prefetch->getLength() << " PrefetchType: "
+               << sycl::ext::oneapi::experimental::prefetchTypeToString(
+                      Prefetch->getPrefetchType())
+               << "\\n";
       }
       break;
     case sycl::detail::CGType::AdviseUSM:

sycl/source/detail/handler_impl.hpp

@@ -12,6 +12,7 @@
 #include <detail/cg.hpp>
 #include <detail/kernel_bundle_impl.hpp>
 #include <memory>
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp>
 
 namespace sycl {
 inline namespace _V1 {
@@ -91,6 +92,10 @@ class handler_impl {
   /// property.
   bool MIsDeviceImageScoped = false;
 
+  /// Direction of USM prefetch / destination device.
+  sycl::ext::oneapi::experimental::prefetch_type MPrefetchType =
+      sycl::ext::oneapi::experimental::prefetch_type::device;
+
   // Program scope pipe information.
 
   // Pipe name that uniquely identifies a pipe.

sycl/source/detail/memory_manager.cpp

@@ -925,13 +925,18 @@ void MemoryManager::fill_usm(void *Mem, queue_impl &Queue, size_t Length,
       DepEvents.size(), DepEvents.data(), OutEvent);
 }
 
-void MemoryManager::prefetch_usm(void *Mem, queue_impl &Queue, size_t Length,
-                                 std::vector<ur_event_handle_t> DepEvents,
-                                 ur_event_handle_t *OutEvent) {
+void MemoryManager::prefetch_usm(
+    void *Mem, queue_impl &Queue, size_t Length,
+    std::vector<ur_event_handle_t> DepEvents, ur_event_handle_t *OutEvent,
+    sycl::ext::oneapi::experimental::prefetch_type Dest) {
   adapter_impl &Adapter = Queue.getAdapter();
-  Adapter.call<UrApiKind::urEnqueueUSMPrefetch>(Queue.getHandleRef(), Mem,
-                                                Length, 0u, DepEvents.size(),
-                                                DepEvents.data(), OutEvent);
+  ur_usm_migration_flags_t MigrationFlag =
+      (Dest == sycl::ext::oneapi::experimental::prefetch_type::device)
+          ? UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE
+          : UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
+  Adapter.call<UrApiKind::urEnqueueUSMPrefetch>(
+      Queue.getHandleRef(), Mem, Length, MigrationFlag, DepEvents.size(),
+      DepEvents.data(), OutEvent);
 }
 
 void MemoryManager::advise_usm(const void *Mem, queue_impl &Queue,
@@ -1542,11 +1547,16 @@ void MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer(
     sycl::detail::context_impl *Context,
     ur_exp_command_buffer_handle_t CommandBuffer, void *Mem, size_t Length,
     std::vector<ur_exp_command_buffer_sync_point_t> Deps,
-    ur_exp_command_buffer_sync_point_t *OutSyncPoint) {
+    ur_exp_command_buffer_sync_point_t *OutSyncPoint,
+    sycl::ext::oneapi::experimental::prefetch_type Dest) {
   adapter_impl &Adapter = Context->getAdapter();
+  ur_usm_migration_flags_t MigrationFlag =
+      (Dest == sycl::ext::oneapi::experimental::prefetch_type::device)
+          ? UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE
+          : UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
   Adapter.call<UrApiKind::urCommandBufferAppendUSMPrefetchExp>(
-      CommandBuffer, Mem, Length, ur_usm_migration_flags_t(0), Deps.size(),
-      Deps.data(), 0u, nullptr, OutSyncPoint, nullptr, nullptr);
+      CommandBuffer, Mem, Length, MigrationFlag, Deps.size(), Deps.data(), 0,
+      nullptr, OutSyncPoint, nullptr, nullptr);
 }
 
 void MemoryManager::ext_oneapi_advise_usm_cmd_buffer(

sycl/source/detail/memory_manager.hpp

@@ -11,6 +11,7 @@
 #include <detail/sycl_mem_obj_i.hpp>
 #include <sycl/access/access.hpp>
 #include <sycl/detail/export.hpp>
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp> // for prefetch_type
 #include <sycl/id.hpp>
 #include <sycl/property_list.hpp>
 #include <sycl/range.hpp>
@@ -146,9 +147,12 @@ class MemoryManager {
                        std::vector<ur_event_handle_t> DepEvents,
                        ur_event_handle_t *OutEvent);
 
-  static void prefetch_usm(void *Ptr, queue_impl &Queue, size_t Len,
-                           std::vector<ur_event_handle_t> DepEvents,
-                           ur_event_handle_t *OutEvent);
+  static void
+  prefetch_usm(void *Ptr, queue_impl &Queue, size_t Len,
+               std::vector<ur_event_handle_t> DepEvents,
+               ur_event_handle_t *OutEvent,
+               sycl::ext::oneapi::experimental::prefetch_type Dest =
+                   sycl::ext::oneapi::experimental::prefetch_type::device);
 
   static void advise_usm(const void *Ptr, queue_impl &Queue, size_t Len,
                          ur_usm_advice_flags_t Advice,
@@ -245,7 +249,9 @@ class MemoryManager {
       sycl::detail::context_impl *Context,
       ur_exp_command_buffer_handle_t CommandBuffer, void *Mem, size_t Length,
       std::vector<ur_exp_command_buffer_sync_point_t> Deps,
-      ur_exp_command_buffer_sync_point_t *OutSyncPoint);
+      ur_exp_command_buffer_sync_point_t *OutSyncPoint,
+      sycl::ext::oneapi::experimental::prefetch_type Dest =
+          sycl::ext::oneapi::experimental::prefetch_type::device);
 
   static void ext_oneapi_advise_usm_cmd_buffer(
       sycl::detail::context_impl *Context,

sycl/source/detail/scheduler/commands.cpp

@@ -2980,7 +2980,8 @@ ur_result_t ExecCGCommand::enqueueImpCommandBuffer() {
     if (auto Result = callMemOpHelper(
             MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer,
             &MQueue->getContextImpl(), MCommandBuffer, Prefetch->getDst(),
-            Prefetch->getLength(), std::move(MSyncPointDeps), &OutSyncPoint);
+            Prefetch->getLength(), std::move(MSyncPointDeps), &OutSyncPoint,
+            Prefetch->getPrefetchType());
         Result != UR_RESULT_SUCCESS)
       return Result;
 
@@ -3291,7 +3292,8 @@ ur_result_t ExecCGCommand::enqueueImpQueue() {
     CGPrefetchUSM *Prefetch = (CGPrefetchUSM *)MCommandGroup.get();
     if (auto Result = callMemOpHelper(
             MemoryManager::prefetch_usm, Prefetch->getDst(), *MQueue,
-            Prefetch->getLength(), std::move(RawEvents), Event);
+            Prefetch->getLength(), std::move(RawEvents), Event,
+            Prefetch->getPrefetchType());
         Result != UR_RESULT_SUCCESS)
       return Result;
 

sycl/source/handler.cpp

@@ -36,6 +36,7 @@
 #include <sycl/stream.hpp>
 
 #include <sycl/ext/oneapi/bindless_images_memory.hpp>
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp>
 #include <sycl/ext/oneapi/experimental/graph.hpp>
 #include <sycl/ext/oneapi/experimental/work_group_memory.hpp>
 #include <sycl/ext/oneapi/memcpy2d.hpp>
@@ -745,8 +746,9 @@ event handler::finalize() {
                                              MCodeLoc));
     break;
   case detail::CGType::PrefetchUSM:
-    CommandGroup.reset(new detail::CGPrefetchUSM(
-        MDstPtr, MLength, std::move(impl->CGData), MCodeLoc));
+    CommandGroup.reset(
+        new detail::CGPrefetchUSM(MDstPtr, MLength, std::move(impl->CGData),
+                                  impl->MPrefetchType, MCodeLoc));
     break;
   case detail::CGType::AdviseUSM:
     CommandGroup.reset(new detail::CGAdviseUSM(MDstPtr, MLength, impl->MAdvice,
@@ -1499,6 +1501,16 @@ void handler::prefetch(const void *Ptr, size_t Count) {
   throwIfActionIsCreated();
   MDstPtr = const_cast<void *>(Ptr);
   MLength = Count;
+  impl->MPrefetchType = ext::oneapi::experimental::prefetch_type::device;
+  setType(detail::CGType::PrefetchUSM);
+}
+
+void handler::prefetch(const void *Ptr, size_t Count,
+                       ext::oneapi::experimental::prefetch_type Type) {
+  throwIfActionIsCreated();
+  MDstPtr = const_cast<void *>(Ptr);
+  MLength = Count;
+  impl->MPrefetchType = Type;
   setType(detail::CGType::PrefetchUSM);
 }