[UR][SYCL] Add support for bidirectional USM prefetching

sycl/cmake/modules/FetchUnifiedRuntime.cmake

@@ -116,7 +116,7 @@ if(SYCL_UR_USE_FETCH_CONTENT)
       CACHE PATH "Path to external '${name}' adapter source dir" FORCE)
   endfunction()
 
-  set(UNIFIED_RUNTIME_REPO "https://github.com/oneapi-src/unified-runtime.git")
+  set(UNIFIED_RUNTIME_REPO "https://github.com/ianayl/unified-runtime.git")
   include(${CMAKE_CURRENT_SOURCE_DIR}/cmake/modules/UnifiedRuntimeTag.cmake)
 
   set(UMF_BUILD_EXAMPLES OFF CACHE INTERNAL "EXAMPLES")

sycl/cmake/modules/UnifiedRuntimeTag.cmake

@@ -1,7 +1,6 @@
-# commit 39df0317814c164f5242eda8d6f08550f6268492
-# Merge: 68d93efd be27d8f0
-# Author: Kenneth Benzie (Benie) <k.benzie@codeplay.com>
-# Date:   Mon Dec 16 13:53:13 2024 +0000
-#     Merge pull request #2467 from nrspruit/fix_external_import_function_call
-#     [L0] Fix external semaphore import function calls to match the header
-set(UNIFIED_RUNTIME_TAG 39df0317814c164f5242eda8d6f08550f6268492)
+# commit 8e819c2d4d3cd7e8a5d8b5ea64ed22dab575737f
+# Merge: 65ca7748 39df0317
+# Author: Li, Ian <ian.li@intel.com>
+# Date:   Mon Dec 16 10:05:24 2024 -0800
+#     Merge branch 'main' of https://github.com/oneapi-src/unified-runtime into ianayl/2way-prefetch
+set(UNIFIED_RUNTIME_TAG 8e819c2d4d3cd7e8a5d8b5ea64ed22dab575737f)

sycl/doc/extensions/experimental/sycl_ext_oneapi_enqueue_functions.asciidoc

@@ -633,14 +633,22 @@ a!
 ----
 namespace sycl::ext::oneapi::experimental {
 
-void prefetch(sycl::queue q, void* ptr, size_t numBytes);
+enum class prefetch_type { device, host };
 
-void prefetch(sycl::handler &h, void* ptr, size_t numBytes);
+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);
 
 }
 ----
 !====
-_Effects_: Enqueues a `prefetch` to the `sycl::queue` or `sycl::handler`.
+_Effects_: Enqueues a `prefetch` to the `sycl::queue` or `sycl::handler`.  The
+`type` parameter tells the direction of the prefetch operation: When the value
+is `prefetch_type::device`, the memory is prefetched _to_ the device associated
+with the queue. When the value is `prefetch_type::host`, the memory is
+prefetched _to_ the host, regardless of the device on which it currently resides.
 
 a|
 [frame=all,grid=none]

sycl/include/sycl/detail/cg_types.hpp

@@ -15,14 +15,14 @@
 #include <sycl/detail/item_base.hpp>           // for id
 #include <sycl/detail/kernel_desc.hpp>         // for kernel_param_kind_t
 #include <sycl/exception.hpp>
-#include <sycl/group.hpp>                      // for group
-#include <sycl/h_item.hpp>                     // for h_item
-#include <sycl/id.hpp>                         // for id
-#include <sycl/item.hpp>                       // for item
-#include <sycl/kernel_handler.hpp>             // for kernel_handler
-#include <sycl/nd_item.hpp>                    // for nd_item
-#include <sycl/nd_range.hpp>                   // for nd_range
-#include <sycl/range.hpp>                      // for range, operator*
+#include <sycl/group.hpp>          // for group
+#include <sycl/h_item.hpp>         // for h_item
+#include <sycl/id.hpp>             // for id
+#include <sycl/item.hpp>           // for item
+#include <sycl/kernel_handler.hpp> // for kernel_handler
+#include <sycl/nd_item.hpp>        // for nd_item
+#include <sycl/nd_range.hpp>       // for nd_range
+#include <sycl/range.hpp>          // for range, operator*
 
 #include <functional>  // for function
 #include <stddef.h>    // for size_t
@@ -65,6 +65,11 @@ enum class CGType : unsigned int {
   SemaphoreSignal = 25,
   ProfilingTag = 26,
   EnqueueNativeCommand = 27,
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  PrefetchUSMExp = 28,
+#else
+  PrefetchUSMExpD2H = 28,
+#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>
@@ -356,14 +357,24 @@ void fill(sycl::queue Q, T *Ptr, const T &Pattern, size_t Count,
   submit(Q, [&](handler &CGH) { fill<T>(CGH, Ptr, Pattern, 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(Q, [&](handler &CGH) { prefetch(CGH, Ptr, NumBytes); }, CodeLoc);
+  submit(Q, [&](handler &CGH) { prefetch(CGH, Ptr, NumBytes, type); }, CodeLoc);
 }
 
 inline void mem_advise(handler &CGH, void *Ptr, size_t NumBytes, int Advice) {

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

@@ -0,0 +1,20 @@
+//==--------------- 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
+
+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 };
+
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/handler.hpp

@@ -151,10 +151,11 @@ template <class _name, class _dataT, int32_t _min_capacity, class _propertiesT,
 class pipe;
 }
 
-namespace ext ::oneapi ::experimental {
-template <typename, typename>
-class work_group_memory;
+namespace ext::oneapi::experimental {
+template <typename, typename> class work_group_memory;
 struct image_descriptor;
+enum class prefetch_type;
+void prefetch(handler &CGH, void *Ptr, size_t NumBytes, prefetch_type type);
 } // namespace ext::oneapi::experimental
 
 namespace ext::oneapi::experimental::detail {
@@ -482,7 +483,8 @@ class __SYCL_EXPORT handler {
 
   /// Saves the location of user's code passed in \p CodeLoc for future usage in
   /// finalize() method.
-  /// TODO: remove the first version of this func (the one without the IsTopCodeLoc arg)
+  /// TODO: remove the first version of this func (the one without the
+  /// IsTopCodeLoc arg)
   ///   at the next ABI breaking window since removing it breaks ABI on windows.
   void saveCodeLoc(detail::code_location CodeLoc);
   void saveCodeLoc(detail::code_location CodeLoc, bool IsTopCodeLoc);
@@ -692,8 +694,9 @@ class __SYCL_EXPORT handler {
         detail::KernelLambdaHasKernelHandlerArgT<KernelType,
                                                  LambdaArgType>::value;
 
-    MHostKernel = std::make_unique<
-        detail::HostKernel<KernelType, LambdaArgType, Dims>>(KernelFunc);
+    MHostKernel =
+        std::make_unique<detail::HostKernel<KernelType, LambdaArgType, Dims>>(
+            KernelFunc);
 
     constexpr bool KernelHasName =
         detail::getKernelName<KernelName>() != nullptr &&
@@ -3218,6 +3221,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
+  /// USM prefetch direction.
+  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.
@@ -3502,6 +3509,28 @@ class __SYCL_EXPORT handler {
                                       bool IsDeviceImageScoped, size_t NumBytes,
                                       size_t Offset);
 
+// Enqueue function extension's implementation USM prefetch, enabling USM
+// prefetch from both host to device, and device to host.
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+  // With breaking changes enabled, the handler/CG Nodes have their fields
+  // modified. This function updates that field in the CG Node.
+  void ext_oneapi_prefetch_exp(const void *Ptr, size_t Count,
+                               ext::oneapi::experimental::prefetch_type Type);
+  // This is a separate function to keep the current handler.prefetch function
+  // the same.
+#else
+  // Without breaking changes, the handler/CG Nodes fields cannot be modified,
+  // meaning 1 CG node type cannot indicate both prefetch directions. Thus, the
+  // default handler.prefetch indicates host to device, and this function serves
+  // as device to host.
+  void ext_oneapi_prefetch_d2h(const void *Ptr, size_t Count);
+#endif
+  // Friend prefetch from the enqueue functions extension to allow call to
+  // private function ext_oneapi_prefetch_d2h
+  friend void sycl::ext::oneapi::experimental::prefetch(
+      handler &CGH, void *Ptr, size_t NumBytes,
+      sycl::ext::oneapi::experimental::prefetch_type type);
+
   // Changing values in this will break ABI/API.
   enum class StableKernelCacheConfig : int32_t {
     Default = 0,

sycl/include/sycl/queue.hpp

@@ -2747,7 +2747,8 @@ class __SYCL_EXPORT queue : public detail::OwnerLessBase<queue> {
   queue(std::shared_ptr<detail::queue_impl> impl) : impl(impl) {}
 
   template <class Obj>
-  friend const decltype(Obj::impl)& detail::getSyclObjImpl(const Obj &SyclObject);
+  friend const decltype(Obj::impl) &
+  detail::getSyclObjImpl(const Obj &SyclObject);
   template <class T>
   friend T detail::createSyclObjFromImpl(decltype(T::impl) ImplObj);
 

sycl/source/detail/cg.hpp

@@ -16,6 +16,9 @@
 #include <sycl/event.hpp>           // for event_impl
 #include <sycl/exception_list.hpp>  // for queue_impl
 #include <sycl/ext/oneapi/experimental/event_mode_property.hpp>
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp> // for prefetch_type
+#endif
 #include <sycl/kernel.hpp>        // for kernel_impl
 #include <sycl/kernel_bundle.hpp> // for kernel_bundle_impl
 
@@ -406,6 +409,41 @@ class CGPrefetchUSM : public CG {
   size_t getLength() { return MLength; }
 };
 
+/// Command group class for experimental USM prefetch provided in the enqueue
+/// functions extension.
+#ifdef __INTEL_PREVIEW_BREAKING_CHANGES
+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 Type,
+                   detail::code_location loc = {})
+      : CG(CGType::PrefetchUSMExp, std::move(CGData), std::move(loc)),
+        MDst(DstPtr), MLength(Length), MPrefetchType(Type) {}
+  void *getDst() { return MDst; }
+  size_t getLength() { return MLength; }
+  ext::oneapi::experimental::prefetch_type getPrefetchType() {
+    return MPrefetchType;
+  }
+};
+#else
+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_impl.hpp

@@ -242,6 +242,19 @@ 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:
+      throw sycl::exception(sycl::make_error_code(errc::feature_not_supported),
+                            "Prefetch as a part of the experimental enqueue "
+                            "function extension is currently not supported by "
+                            "SYCL Graph extension.");
+      // return createCGCopy<sycl::detail::CGPrefetchUSMExp>();
+#else
+    case sycl::detail::CGType::PrefetchUSMExpD2H:
+      throw sycl::exception(sycl::make_error_code(errc::feature_not_supported),
+                            "Prefetch from device to host is currently not "
+                            "supported by SYCL Graph extension.");
+#endif
     case sycl::detail::CGType::AdviseUSM:
       return createCGCopy<sycl::detail::CGAdviseUSM>();
     case sycl::detail::CGType::Copy2DUSM:

sycl/source/detail/memory_manager.cpp

@@ -964,17 +964,22 @@ void MemoryManager::fill_usm(void *Mem, QueueImplPtr Queue, size_t Length,
       DepEvents.size(), DepEvents.data(), OutEvent);
 }
 
-void MemoryManager::prefetch_usm(void *Mem, QueueImplPtr Queue, size_t Length,
-                                 std::vector<ur_event_handle_t> DepEvents,
-                                 ur_event_handle_t *OutEvent,
-                                 const detail::EventImplPtr &OutEventImpl) {
+void MemoryManager::prefetch_usm(
+    void *Mem, QueueImplPtr Queue, size_t Length,
+    std::vector<ur_event_handle_t> DepEvents, ur_event_handle_t *OutEvent,
+    const detail::EventImplPtr &OutEventImpl,
+    sycl::ext::oneapi::experimental::prefetch_type Dest) {
   assert(Queue && "USM prefetch must be called with a valid device queue");
   const AdapterPtr &Adapter = Queue->getAdapter();
+  ur_usm_migration_flags_t migration_flag =
+      (Dest == sycl::ext::oneapi::experimental::prefetch_type::device)
+          ? UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE
+          : UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
   if (OutEventImpl != nullptr)
     OutEventImpl->setHostEnqueueTime();
-  Adapter->call<UrApiKind::urEnqueueUSMPrefetch>(Queue->getHandleRef(), Mem,
-                                                 Length, 0, DepEvents.size(),
-                                                 DepEvents.data(), OutEvent);
+  Adapter->call<UrApiKind::urEnqueueUSMPrefetch>(
+      Queue->getHandleRef(), Mem, Length, migration_flag, DepEvents.size(),
+      DepEvents.data(), OutEvent);
 }
 
 void MemoryManager::advise_usm(const void *Mem, QueueImplPtr Queue,
@@ -1612,11 +1617,16 @@ void MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer(
     sycl::detail::ContextImplPtr 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) {
   const AdapterPtr &Adapter = Context->getAdapter();
+  ur_usm_migration_flags_t migration_flag =
+      (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(), 0, nullptr, OutSyncPoint, nullptr, nullptr);
+      CommandBuffer, Mem, Length, migration_flag, 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>
@@ -149,10 +150,13 @@ class MemoryManager {
                        ur_event_handle_t *OutEvent,
                        const detail::EventImplPtr &OutEventImpl);
 
-  static void prefetch_usm(void *Ptr, QueueImplPtr Queue, size_t Len,
-                           std::vector<ur_event_handle_t> DepEvents,
-                           ur_event_handle_t *OutEvent,
-                           const detail::EventImplPtr &OutEventImpl);
+  static void
+  prefetch_usm(void *Ptr, QueueImplPtr Queue, size_t Len,
+               std::vector<ur_event_handle_t> DepEvents,
+               ur_event_handle_t *OutEvent,
+               const detail::EventImplPtr &OutEventImpl,
+               sycl::ext::oneapi::experimental::prefetch_type Dest =
+                   sycl::ext::oneapi::experimental::prefetch_type::device);
 
   static void advise_usm(const void *Ptr, QueueImplPtr Queue, size_t Len,
                          ur_usm_advice_flags_t Advice,
@@ -251,7 +255,9 @@ class MemoryManager {
       sycl::detail::ContextImplPtr 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::ContextImplPtr Context,