[UR][SYCL] Implement USM prefetch from device to host in SYCL runtime and UR

sycl/include/sycl/detail/cg_types.hpp

@@ -67,6 +67,11 @@ enum class CGType : unsigned int {
   EnqueueNativeCommand = 27,
   AsyncAlloc = 28,
   AsyncFree = 29,
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  PrefetchUSMExp = 30,
+#else
+  PrefetchUSMExpD2H = 30,
+#endif
 };
 
 template <typename, typename T> struct check_fn_signature {

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,25 @@ 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) {
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  CGH.ext_oneapi_prefetch_exp(Ptr, NumBytes, Type);
+#else
+  if (Type == prefetch_type::device) {
+    CGH.prefetch(Ptr, NumBytes);
+  } else {
+    CGH.ext_oneapi_prefetch_d2h(Ptr, NumBytes);
+  }
+#endif
 }
 
 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,9 @@ 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;
+void prefetch(handler &CGH, void *Ptr, size_t NumBytes, prefetch_type 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);
@@ -3188,6 +3191,10 @@ class __SYCL_EXPORT handler {
   void *MDstPtr = nullptr;
   /// Length to copy or fill (for USM operations).
   size_t MLength = 0;
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  // Prefetch direction for ext_oneapi_prefetch_exp
+  ext::oneapi::experimental::prefetch_type MPrefetchType;
+#endif
   /// Pattern that is used to fill memory object in case command type is fill.
   std::vector<unsigned char> MPattern;
   /// Storage for a lambda or function object.
@@ -3454,6 +3461,27 @@ class __SYCL_EXPORT handler {
   void ext_oneapi_memset2d_impl(void *Dest, size_t DestPitch, int Value,
                                 size_t Width, size_t Height);
 
+// Implementation of enqueue_functions extension's USM prefetch, allowing for
+// prefetching memory from both host to device and vice versa.
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  // Prefetch implementation that accounts for prefetching both directions, but
+  // introduces a "prefetch type" field to handler/CG nodes: this results in an
+  // ABI break.
+  void ext_oneapi_prefetch_exp(const void *Ptr, size_t Count,
+                               ext::oneapi::experimental::prefetch_type Type);
+#else
+  // Non-ABI breaking implementation that implements prefetching from device to
+  // host as a separate function.
+  void ext_oneapi_prefetch_d2h(const void *Ptr, size_t Count);
+  // TODO upon next ABI-breaking cycle, decide which approach to go with.
+#endif
+
+  // Enqueue_functions extension's prefetch function is friended in order to
+  // call private handler function ext_oneapi_prefetch_d2h.
+  friend void ext::oneapi::experimental::prefetch(
+      handler &CGH, void *Ptr, size_t NumBytes,
+      ext::oneapi::experimental::prefetch_type Type);
+
   // Implementation of memcpy to device_global.
   void memcpyToDeviceGlobal(const void *DeviceGlobalPtr, const void *Src,
                             bool IsDeviceImageScoped, size_t NumBytes,

sycl/source/detail/cg.hpp

@@ -408,6 +408,41 @@ class CGPrefetchUSM : public CG {
   size_t getLength() { return MLength; }
 };
 
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+/// Enqueue_functions extension USM Prefetch command group class
+class CGPrefetchUSMExp : public CG {
+  void *MDst;
+  size_t MLength;
+  ext::oneapi::experimental::prefetch_type MPrefetchType;
+
+public:
+  CGPrefetchUSMExp(void *DstPtr, size_t Length, CG::StorageInitHelper CGData,
+                   ext::oneapi::experimental::prefetch_type PrefetchType,
+                   detail::code_location loc = {})
+      : CG(CGType::PrefetchUSMExp, std::move(CGData), std::move(loc)),
+        MDst(DstPtr), MLength(Length), MPrefetchType(PrefetchType) {}
+  void *getDst() { return MDst; }
+  size_t getLength() { return MLength; }
+  ext::oneapi::experimental::prefetch_type getPrefetchType() {
+    return MPrefetchType;
+  }
+};
+#else
+/// Enqueue_functions USM device-to-host prefetch command group class
+class CGPrefetchUSMExpD2H : public CG {
+  void *MDst;
+  size_t MLength;
+
+public:
+  CGPrefetchUSMExpD2H(void *DstPtr, size_t Length, CG::StorageInitHelper CGData,
+                      detail::code_location loc = {})
+      : CG(CGType::PrefetchUSMExpD2H, std::move(CGData), std::move(loc)),
+        MDst(DstPtr), MLength(Length) {}
+  void *getDst() { return MDst; }
+  size_t getLength() { return MLength; }
+};
+#endif
+
 /// "Advise USM" command group class.
 class CGAdviseUSM : public CG {
   void *MDst;

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>
@@ -55,6 +56,11 @@ inline node_type getNodeTypeFromCG(sycl::detail::CGType CGType) {
   case sycl::detail::CGType::FillUSM:
     return node_type::memfill;
   case sycl::detail::CGType::PrefetchUSM:
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  case sycl::detail::CGType::PrefetchUSMExp:
+#else
+  case sycl::detail::CGType::PrefetchUSMExpD2H:
+#endif
     return node_type::prefetch;
   case sycl::detail::CGType::AdviseUSM:
     return node_type::memadvise;
@@ -247,6 +253,13 @@ class node_impl : public std::enable_shared_from_this<node_impl> {
       return createCGCopy<sycl::detail::CGFillUSM>();
     case sycl::detail::CGType::PrefetchUSM:
       return createCGCopy<sycl::detail::CGPrefetchUSM>();
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+    case sycl::detail::CGType::PrefetchUSMExp:
+      return createCGCopy<sycl::detail::CGPrefetchUSMExp>();
+#else
+    case sycl::detail::CGType::PrefetchUSMExpD2H:
+      return createCGCopy<sycl::detail::CGPrefetchUSMExpD2H>();
+#endif
     case sycl::detail::CGType::AdviseUSM:
       return createCGCopy<sycl::detail::CGAdviseUSM>();
     case sycl::detail::CGType::Copy2DUSM:
@@ -658,6 +671,31 @@ class node_impl : public std::enable_shared_from_this<node_impl> {
                << " Length: " << Prefetch->getLength() << "\\n";
       }
       break;
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+    case sycl::detail::CGType::PrefetchUSMExp:
+      Stream << "CGPrefetchUSMExp \\n";
+      if (Verbose) {
+        sycl::detail::CGPrefetchUSMExp *PrefetchExp =
+            static_cast<sycl::detail::CGPrefetchUSMExp *>(MCommandGroup.get());
+        Stream << "Dst: " << PrefetchExp->getDst()
+               << " Length: " << PrefetchExp->getLength() << " PrefetchType: "
+               << sycl::ext::oneapi::experimental::prefetchTypeToString(
+                      PrefetchExp->getPrefetchType())
+               << "\\n";
+      }
+      break;
+#else
+    case sycl::detail::CGType::PrefetchUSMExpD2H:
+      Stream << "CGPrefetchUSMExpD2H (Experimental, Device to host) \\n";
+      if (Verbose) {
+        sycl::detail::CGPrefetchUSMExpD2H *Prefetch =
+            static_cast<sycl::detail::CGPrefetchUSMExpD2H *>(
+                MCommandGroup.get());
+        Stream << "Dst: " << Prefetch->getDst()
+               << " Length: " << Prefetch->getLength() << "\\n";
+      }
+      break;
+#endif
     case sycl::detail::CGType::AdviseUSM:
       Stream << "CGAdviseUSM \\n";
       if (Verbose) {

sycl/source/detail/memory_manager.cpp

@@ -922,13 +922,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,
@@ -1539,11 +1544,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,