[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 ea0f3a1f5f15f9af7bf40bd13669afeb9ada569c
-# Merge: bb64b3e9f6d3 4a89e1c69a65
-# Author: Martin Grant <martin.morrisongrant@codeplay.com>
-# Date:   Thu Dec 19 11:26:01 2024 +0000
-#     Merge pull request #2277 from igchor/cooperative_fix
-#     [Spec] fix urKernelSuggestMaxCooperativeGroupCountExp
-set(UNIFIED_RUNTIME_TAG ea0f3a1f5f15f9af7bf40bd13669afeb9ada569c)
+# commit 0f9407c51d689c87c1835c96ffad01601c92d23c
+# Merge: ff19760f 76a96238
+# Author: Li, Ian <ian.li@intel.com>
+# Date:   Fri Dec 20 08:32:03 2024 -0800
+#     Merge branch 'main' of https://github.com/oneapi-src/unified-runtime into ianayl/2way-prefetch
+set(UNIFIED_RUNTIME_TAG 0f9407c51d689c87c1835c96ffad01601c92d23c)

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

@@ -8,26 +8,26 @@
 
 #pragma once
 
-#include <sycl/detail/array.hpp>               // for array
-#include <sycl/detail/common.hpp>              // for InitializedVal, NDLoop
-#include <sycl/detail/helpers.hpp>             // for Builder
-#include <sycl/detail/host_profiling_info.hpp> // for HostProfilingInfo
-#include <sycl/detail/item_base.hpp>           // for id
-#include <sycl/detail/kernel_desc.hpp>         // for kernel_param_kind_t
+#include <sycl/detail/array.hpp>                  // for array
+#include <sycl/detail/common.hpp>                 // for InitializedVal, NDLoop
+#include <sycl/detail/helpers.hpp>                // for Builder
+#include <sycl/detail/host_profiling_info.hpp>    // for HostProfilingInfo
+#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 <functional>  // for function
-#include <stddef.h>    // for size_t
-#include <type_traits> // for enable_if_t, false_type
-#include <utility>     // for declval
+#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
+#include <type_traits>                            // for enable_if_t, false_type
+#include <utility>                                // for declval
 
 namespace sycl {
 inline namespace _V1 {
@@ -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,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

@@ -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

@@ -8,24 +8,27 @@
 
 #pragma once
 
-#include <sycl/accessor.hpp>        // for AccessorImplHost, AccessorImplPtr
-#include <sycl/detail/cg_types.hpp> // for ArgDesc, HostTask, HostKernelBase
-#include <sycl/detail/common.hpp>   // for code_location
-#include <sycl/detail/helpers.hpp>  // for context_impl
-#include <sycl/detail/ur.hpp>       // for ur_rect_region_t, ur_rect_offset_t
-#include <sycl/event.hpp>           // for event_impl
-#include <sycl/exception_list.hpp>  // for queue_impl
+#include <sycl/accessor.hpp>           // for AccessorImplHost, AccessorImplPtr
+#include <sycl/detail/cg_types.hpp>    // for ArgDesc, HostTask, HostKernelBase
+#include <sycl/detail/common.hpp>      // for code_location
+#include <sycl/detail/helpers.hpp>     // for context_impl
+#include <sycl/detail/ur.hpp>          // for ur_rect_region_t, ur_rect_offset_t
+#include <sycl/event.hpp>              // for event_impl
+#include <sycl/exception_list.hpp>     // for queue_impl
 #include <sycl/ext/oneapi/experimental/event_mode_property.hpp>
-#include <sycl/kernel.hpp>        // for kernel_impl
-#include <sycl/kernel_bundle.hpp> // for kernel_bundle_impl
-
-#include <assert.h> // for assert
-#include <memory>   // for shared_ptr, unique_ptr
-#include <stddef.h> // for size_t
-#include <stdint.h> // for int32_t
-#include <string>   // for string
-#include <utility>  // for move
-#include <vector>   // for vector
+#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
+
+#include <assert.h>                    // for assert
+#include <memory>                      // for shared_ptr, unique_ptr
+#include <stddef.h>                    // for size_t
+#include <stdint.h>                    // for int32_t
+#include <string>                      // for string
+#include <utility>                     // for move
+#include <vector>                      // for vector
 
 namespace sycl {
 inline namespace _V1 {
@@ -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

@@ -243,6 +243,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:
@@ -636,6 +643,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 *Prefetch =
+            static_cast<sycl::detail::CGPrefetchUSMExp *>(MCommandGroup.get());
+        Stream << "Dst: " << Prefetch->getDst()
+               << " Length: " << Prefetch->getLength() << " Type: "
+               << sycl::ext::oneapi::experimental::prefetchTypeToString(
+                      Prefetch->getPrefetchType())
+               << "\\n";
+      }
+      break;
+#else
+    case sycl::detail::CGType::PrefetchUSMExpD2H:
+      Stream << "CGPrefetchUSM (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) {