Changed signatures to follow new proposed spec

Author: ianayl

sycl/include/sycl/detail/cg_types.hpp

@@ -65,6 +65,7 @@ enum class CGType : unsigned int {
   SemaphoreSignal = 25,
   ProfilingTag = 26,
   EnqueueNativeCommand = 27,
+  PrefetchUSMExpD2H = 28,
 };
 
 template <typename, typename T> struct check_fn_signature {

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

@@ -13,6 +13,7 @@
 #include <sycl/detail/common.hpp>
 #include <sycl/event.hpp>
 #include <sycl/ext/oneapi/experimental/graph.hpp>
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp>
 #include <sycl/ext/oneapi/properties/properties.hpp>
 #include <sycl/handler.hpp>
 #include <sycl/nd_range.hpp>
@@ -349,14 +350,18 @@ 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) {
+  if (type == prefetch_type::device) {
+    CGH.prefetch(Ptr, NumBytes);
+  } else {
+    CGH.ext_oneapi_prefetch_d2h(Ptr, NumBytes);
+  }
 }
 
-inline void prefetch(queue Q, void *Ptr, size_t NumBytes,
+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/USM/prefetch_exp.hpp renamed to sycl/include/sycl/ext/oneapi/experimental/enqueue_types.hpp

@@ -1,4 +1,4 @@
-//==-------- usm_prefetch_exp.hpp --- SYCL USM prefetch extensions ---------==//
+//==------ enqueue_functions.hpp ------- SYCL enqueue free functions -------==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -10,16 +10,13 @@
 
 namespace sycl {
 inline namespace _V1 {
-
 namespace ext::oneapi::experimental {
 
-/// @brief Indicates USM memory migration direction: either from host to device,
-/// or device to host.
-enum class migration_direction {
-  HOST_TO_DEVICE, /// Move data from host USM to device USM
-  DEVICE_TO_HOST  /// Move data from device USM to host USM
+enum class prefetch_type {
+  device,
+  host
 };
 
 } // namespace ext::oneapi::experimental
 } // namespace _V1
-} // namespace sycl
+} // namespace sycl

sycl/include/sycl/handler.hpp

@@ -32,8 +32,8 @@
 #include <sycl/ext/oneapi/bindless_images_mem_handle.hpp>
 #include <sycl/ext/oneapi/device_global/device_global.hpp>
 #include <sycl/ext/oneapi/device_global/properties.hpp>
-#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp>
 #include <sycl/ext/oneapi/experimental/cluster_group_prop.hpp>
+#include <sycl/ext/oneapi/experimental/enqueue_types.hpp>
 #include <sycl/ext/oneapi/experimental/graph.hpp>
 #include <sycl/ext/oneapi/experimental/raw_kernel_arg.hpp>
 #include <sycl/ext/oneapi/experimental/use_root_sync_prop.hpp>
@@ -163,9 +163,10 @@ template <class _name, class _dataT, int32_t _min_capacity, class _propertiesT,
 class pipe;
 }
 
-namespace ext ::oneapi ::experimental {
+namespace ext::oneapi::experimental {
 template <typename, typename> class work_group_memory;
 struct image_descriptor;
+void prefetch(handler &CGH, void *Ptr, size_t NumBytes, prefetch_type type);
 } // namespace ext::oneapi::experimental
 
 namespace ext::oneapi::experimental::detail {
@@ -2832,20 +2833,6 @@ class __SYCL_EXPORT handler {
   /// \param Count is a number of bytes to be prefetched.
   void prefetch(const void *Ptr, size_t Count);
 
-  /// Experimental implementation of prefetch supporting bidirectional USM data
-  /// migration: 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 CGH is the handler to be used for prefetching.
-  /// \param Ptr is a USM pointer to the memory to be prefetched to the
-  /// destination. \param Count is a number of bytes to be prefetched. \param
-  /// Direction indicates the direction to prefetch data to/from.
-  void ext_oneapi_prefetch_exp(
-      const void *Ptr, size_t Count,
-      ext::oneapi::experimental::migration_direction Direction =
-          ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE);
-
   /// Provides additional information to the underlying runtime about how
   /// different allocations are used.
   ///
@@ -3274,9 +3261,6 @@ class __SYCL_EXPORT handler {
   detail::code_location MCodeLoc = {};
   bool MIsFinalized = false;
   event MLastEvent;
-  /// Enum to indicate USM data migration direction
-  ext::oneapi::experimental::migration_direction MDirection =
-      ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE;
 
   // Make queue_impl class friend to be able to call finalize method.
   friend class detail::queue_impl;
@@ -3554,6 +3538,11 @@ class __SYCL_EXPORT handler {
                                       bool IsDeviceImageScoped, size_t NumBytes,
                                       size_t Offset);
 
+  // Implementation of USM prefetch, fetching from device back to host. 
+  void ext_oneapi_prefetch_d2h(const void *Ptr, size_t Count);
+  // 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

@@ -30,7 +30,6 @@
 #include <sycl/exception_list.hpp>            // for defaultAsyncHa...
 #include <sycl/ext/oneapi/device_global/device_global.hpp> // for device_global
 #include <sycl/ext/oneapi/device_global/properties.hpp> // for device_image_s...
-#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp> // for migration...
 #include <sycl/ext/oneapi/experimental/graph.hpp>    // for command_graph...
 #include <sycl/ext/oneapi/properties/properties.hpp> // for empty_properti...
 #include <sycl/handler.hpp>                          // for handler, isDev...
@@ -746,27 +745,27 @@ class __SYCL_EXPORT queue : public detail::OwnerLessBase<queue> {
                   TlsCodeLocCapture.query());
   }
 
-  /// Experimental implementation of prefetch supporting bidirectional USM data
-  /// migration: 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 Direction indicates the direction to prefetch data to/from.
-  /// \return an event representing prefetch operation.
-  event ext_oneapi_prefetch_exp(
-      const void *Ptr, size_t Count,
-      ext::oneapi::experimental::migration_direction Direction =
-          ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE,
-      const detail::code_location &CodeLoc = detail::code_location::current()) {
-    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
-    return submit(
-        [=](handler &CGH) {
-          CGH.ext_oneapi_prefetch_exp(Ptr, Count, Direction);
-        },
-        TlsCodeLocCapture.query());
-  }
+//  /// Experimental implementation of prefetch supporting bidirectional USM data
+//  /// migration: 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 Direction indicates the direction to prefetch data to/from.
+//  /// \return an event representing prefetch operation.
+//  event ext_oneapi_prefetch_exp(
+//      const void *Ptr, size_t Count,
+//      ext::oneapi::experimental::migration_direction Direction =
+//          ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE,
+//      const detail::code_location &CodeLoc = detail::code_location::current()) {
+//    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
+//    return submit(
+//        [=](handler &CGH) {
+//          CGH.ext_oneapi_prefetch_exp(Ptr, Count, Direction);
+//        },
+//        TlsCodeLocCapture.query());
+//  }
 
   /// Provides hints to the runtime library that data should be made available
   /// on a device earlier than Unified Shared Memory would normally require it
@@ -788,29 +787,29 @@ class __SYCL_EXPORT queue : public detail::OwnerLessBase<queue> {
         TlsCodeLocCapture.query());
   }
 
-  /// Experimental implementation of prefetch supporting bidirectional USM data
-  /// migration: 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 DepEvent is an event that specifies the kernel dependencies.
-  /// \param Direction indicates the direction to prefetch data to/from.
-  /// \return an event representing prefetch operation.
-  event ext_oneapi_prefetch_exp(
-      const void *Ptr, size_t Count, event DepEvent,
-      ext::oneapi::experimental::migration_direction Direction =
-          ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE,
-      const detail::code_location &CodeLoc = detail::code_location::current()) {
-    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
-    return submit(
-        [=](handler &CGH) {
-          CGH.depends_on(DepEvent);
-          CGH.ext_oneapi_prefetch_exp(Ptr, Count, Direction);
-        },
-        TlsCodeLocCapture.query());
-  }
+//   /// Experimental implementation of prefetch supporting bidirectional USM data
+//   /// migration: 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 DepEvent is an event that specifies the kernel dependencies.
+//   /// \param Direction indicates the direction to prefetch data to/from.
+//   /// \return an event representing prefetch operation.
+//   event ext_oneapi_prefetch_exp(
+//       const void *Ptr, size_t Count, event DepEvent,
+//       ext::oneapi::experimental::migration_direction Direction =
+//           ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE,
+//       const detail::code_location &CodeLoc = detail::code_location::current()) {
+//     detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
+//     return submit(
+//         [=](handler &CGH) {
+//           CGH.depends_on(DepEvent);
+//           CGH.ext_oneapi_prefetch_exp(Ptr, Count, Direction);
+//         },
+//         TlsCodeLocCapture.query());
+//   }
 
   /// Provides hints to the runtime library that data should be made available
   /// on a device earlier than Unified Shared Memory would normally require it
@@ -833,30 +832,30 @@ class __SYCL_EXPORT queue : public detail::OwnerLessBase<queue> {
         TlsCodeLocCapture.query());
   }
 
-  /// Experimental implementation of prefetch supporting bidirectional USM data
-  /// migration: 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 DepEvents is a vector of events that specifies the kernel
-  /// dependencies.
-  /// \param Direction indicates the direction to prefetch data to/from.
-  /// \return an event representing prefetch operation.
-  event ext_oneapi_prefetch_exp(
-      const void *Ptr, size_t Count, const std::vector<event> &DepEvents,
-      ext::oneapi::experimental::migration_direction Direction =
-          ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE,
-      const detail::code_location &CodeLoc = detail::code_location::current()) {
-    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
-    return submit(
-        [=](handler &CGH) {
-          CGH.depends_on(DepEvents);
-          CGH.ext_oneapi_prefetch_exp(Ptr, Count, Direction);
-        },
-        TlsCodeLocCapture.query());
-  }
+//   /// Experimental implementation of prefetch supporting bidirectional USM data
+//   /// migration: 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 DepEvents is a vector of events that specifies the kernel
+//   /// dependencies.
+//   /// \param Direction indicates the direction to prefetch data to/from.
+//   /// \return an event representing prefetch operation.
+//   event ext_oneapi_prefetch_exp(
+//       const void *Ptr, size_t Count, const std::vector<event> &DepEvents,
+//       ext::oneapi::experimental::migration_direction Direction =
+//           ext::oneapi::experimental::migration_direction::HOST_TO_DEVICE,
+//       const detail::code_location &CodeLoc = detail::code_location::current()) {
+//     detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
+//     return submit(
+//         [=](handler &CGH) {
+//           CGH.depends_on(DepEvents);
+//           CGH.ext_oneapi_prefetch_exp(Ptr, Count, Direction);
+//         },
+//         TlsCodeLocCapture.query());
+//   }
 
   /// Copies data from one 2D memory region to another, both pointed by
   /// USM pointers.

sycl/source/detail/cg.hpp

@@ -15,7 +15,6 @@
 #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/USM/prefetch_exp.hpp> // for migration_direction
 #include <sycl/kernel.hpp>                                   // for kernel_impl
 #include <sycl/kernel_bundle.hpp> // for kernel_bundle_impl
 
@@ -394,19 +393,28 @@ class CGFillUSM : public CG {
 class CGPrefetchUSM : public CG {
   void *MDst;
   size_t MLength;
-  ext::oneapi::experimental::migration_direction MDirection;
 
 public:
   CGPrefetchUSM(void *DstPtr, size_t Length, CG::StorageInitHelper CGData,
-                ext::oneapi::experimental::migration_direction Direction,
                 detail::code_location loc = {})
       : CG(CGType::PrefetchUSM, std::move(CGData), std::move(loc)),
-        MDst(DstPtr), MLength(Length), MDirection(Direction) {}
+        MDst(DstPtr), MLength(Length) {}
+  void *getDst() { return MDst; }
+  size_t getLength() { return MLength; }
+};
+
+/// "Prefetch USM" 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; }
-  ext::oneapi::experimental::migration_direction getDirection() {
-    return MDirection;
-  }
 };
 
 /// "Advise USM" command group class.

sycl/source/detail/memory_manager.cpp

@@ -966,16 +966,14 @@ void MemoryManager::fill_usm(void *Mem, QueueImplPtr Queue, size_t Length,
 
 void MemoryManager::prefetch_usm(
     void *Mem, QueueImplPtr Queue, size_t Length,
-    sycl::ext::oneapi::experimental::migration_direction Direction,
     std::vector<ur_event_handle_t> DepEvents, ur_event_handle_t *OutEvent,
-    const detail::EventImplPtr &OutEventImpl) {
+    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 =
-      UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE;
-  if (Direction ==
-      sycl::ext::oneapi::experimental::migration_direction::DEVICE_TO_HOST)
-    migration_flag = UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
+      (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>(
@@ -1617,15 +1615,13 @@ void MemoryManager::ext_oneapi_fill_cmd_buffer(
 void MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer(
     sycl::detail::ContextImplPtr Context,
     ur_exp_command_buffer_handle_t CommandBuffer, void *Mem, size_t Length,
-    sycl::ext::oneapi::experimental::migration_direction Direction,
     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 =
-      UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE;
-  if (Direction ==
-      sycl::ext::oneapi::experimental::migration_direction::DEVICE_TO_HOST)
-    migration_flag = UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
+    (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, migration_flag, Deps.size(), Deps.data(), 0,
       nullptr, OutSyncPoint, nullptr, nullptr);