[SYCL][Graph] Support for Prefetch and memory advise

sycl/doc/design/CommandGraph.md

@@ -347,6 +347,8 @@ The types of commands which are unsupported, and lead to this exception are:
   This corresponds to a memory buffer write command.
 * `handler::copy(src, dest)` or `handler::memcpy(dest, src)` - Where both `src` and
    `dest` are USM pointers. This corresponds to a USM copy command.
+* `handler::prefetch()`.
+* `handler::mem_advise()`.
 
 Note that `handler::copy(src, dest)` where both `src` and `dest` are an accessor
 is supported, as a memory buffer copy command exists in the OpenCL extension.
@@ -373,6 +375,8 @@ adapter where there is matching support for each function in the list.
 | urCommandBufferAppendMemBufferWriteRectExp |  | No |
 | urCommandBufferAppendMemBufferReadRectExp |  | No |
 | urCommandBufferAppendMemBufferFillExp | clCommandFillBufferKHR | Yes |
+| urCommandBufferAppendUSMPrefetchExp |  | No |
+| urCommandBufferAppendUSMAdviseExp |  | No |
 | urCommandBufferEnqueueExp | clEnqueueCommandBufferKHR | Yes |
 |  | clCommandBarrierWithWaitListKHR | No |
 |  | clCommandCopyImageKHR | No |

sycl/include/sycl/detail/pi.def

@@ -176,6 +176,8 @@ _PI_API(piextCommandBufferMemBufferWrite)
 _PI_API(piextCommandBufferMemBufferWriteRect)
 _PI_API(piextCommandBufferMemBufferRead)
 _PI_API(piextCommandBufferMemBufferReadRect)
+_PI_API(piextCommandBufferPrefetchUSM)
+_PI_API(piextCommandBufferAdviseUSM)
 _PI_API(piextEnqueueCommandBuffer)
 
 _PI_API(piextUSMPitchedAlloc)

sycl/include/sycl/detail/pi.h

@@ -146,9 +146,10 @@
 // 14.37 Added piextUSMImportExternalPointer and piextUSMReleaseImportedPointer.
 // 14.38 Change PI_MEM_ADVICE_* values to flags for use in bitwise operations.
 // 14.39 Added PI_EXT_INTEL_DEVICE_INFO_ESIMD_SUPPORT device info query.
+// 14.40 Added piextCommandBufferPrefetchUSM and piextCommandBufferAdviseUSM
 
 #define _PI_H_VERSION_MAJOR 14
-#define _PI_H_VERSION_MINOR 39
+#define _PI_H_VERSION_MINOR 40
 
 #define _PI_STRING_HELPER(a) #a
 #define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -2463,6 +2464,38 @@ __SYCL_EXPORT pi_result piextCommandBufferMemBufferWriteRect(
     const pi_ext_sync_point *sync_point_wait_list,
     pi_ext_sync_point *sync_point);
 
+/// API to append a USM Prefetch command to the command-buffer.
+/// \param command_buffer The command-buffer to append onto.
+/// \param ptr points to the memory to migrate.
+/// \param size is the number of bytes to migrate.
+/// \param flags is a bitfield used to specify memory migration options.
+/// \param num_sync_points_in_wait_list The number of sync points in the
+/// provided wait list.
+/// \param sync_point_wait_list A list of sync points that this command must
+/// wait on.
+/// \param sync_point The sync_point associated with this memory operation.
+__SYCL_EXPORT pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer command_buffer, const void *ptr, size_t size,
+    pi_usm_migration_flags flags, pi_uint32 num_sync_points_in_wait_list,
+    const pi_ext_sync_point *sync_point_wait_list,
+    pi_ext_sync_point *sync_point);
+
+/// API to append a USM Advise command to the command-buffer.
+/// \param command_buffer The command-buffer to append onto.
+/// \param ptr is the data to be advised.
+/// \param length is the size in bytes of the memory to advise.
+/// \param advice is device specific advice.
+/// \param num_sync_points_in_wait_list The number of sync points in the
+/// provided wait list.
+/// \param sync_point_wait_list A list of sync points that this command must
+/// wait on.
+/// \param sync_point The sync_point associated with this memory operation.
+__SYCL_EXPORT pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer command_buffer, const void *ptr, size_t length,
+    pi_mem_advice advice, pi_uint32 num_sync_points_in_wait_list,
+    const pi_ext_sync_point *sync_point_wait_list,
+    pi_ext_sync_point *sync_point);
+
 /// API to submit the command-buffer to queue for execution, returns an error if
 /// the command-buffer is not finalized or another instance of the same
 /// command-buffer is currently executing.

sycl/plugins/cuda/pi_cuda.cpp

@@ -1137,6 +1137,24 @@ pi_result piextCommandBufferMemBufferWriteRect(
       NumSyncPointsInWaitList, SyncPointWaitList, SyncPoint);
 }
 
+pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Size,
+    pi_usm_migration_flags Flags, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferPrefetchUSM(CommandBuffer, Ptr, Size, Flags,
+                                              NumSyncPointsInWaitList,
+                                              SyncPointWaitList, SyncPoint);
+}
+
+pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Length,
+    pi_mem_advice Advice, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferAdviseUSM(CommandBuffer, Ptr, Length, Advice,
+                                            NumSyncPointsInWaitList,
+                                            SyncPointWaitList, SyncPoint);
+}
+
 pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
                                     pi_queue Queue,
                                     pi_uint32 NumEventsInWaitList,

sycl/plugins/hip/pi_hip.cpp

@@ -1145,6 +1145,24 @@ pi_result piextCommandBufferMemBufferWriteRect(
       NumSyncPointsInWaitList, SyncPointWaitList, SyncPoint);
 }
 
+pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Size,
+    pi_usm_migration_flags Flags, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferPrefetchUSM(CommandBuffer, Ptr, Size, Flags,
+                                              NumSyncPointsInWaitList,
+                                              SyncPointWaitList, SyncPoint);
+}
+
+pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Length,
+    pi_mem_advice Advice, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferAdviseUSM(CommandBuffer, Ptr, Length, Advice,
+                                            NumSyncPointsInWaitList,
+                                            SyncPointWaitList, SyncPoint);
+}
+
 pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
                                     pi_queue Queue,
                                     pi_uint32 NumEventsInWaitList,

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -1304,6 +1304,24 @@ pi_result piextCommandBufferMemBufferWriteRect(
       NumSyncPointsInWaitList, SyncPointWaitList, SyncPoint);
 }
 
+pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Size,
+    pi_usm_migration_flags Flags, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferPrefetchUSM(CommandBuffer, Ptr, Size, Flags,
+                                              NumSyncPointsInWaitList,
+                                              SyncPointWaitList, SyncPoint);
+}
+
+pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Length,
+    pi_mem_advice Advice, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferAdviseUSM(CommandBuffer, Ptr, Length, Advice,
+                                            NumSyncPointsInWaitList,
+                                            SyncPointWaitList, SyncPoint);
+}
+
 pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
                                     pi_queue Queue,
                                     pi_uint32 NumEventsInWaitList,

sycl/plugins/opencl/pi_opencl.cpp

@@ -1081,6 +1081,24 @@ pi_result piextCommandBufferMemBufferWriteRect(
       NumSyncPointsInWaitList, SyncPointWaitList, SyncPoint);
 }
 
+pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Size,
+    pi_usm_migration_flags Flags, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferPrefetchUSM(CommandBuffer, Ptr, Size, Flags,
+                                              NumSyncPointsInWaitList,
+                                              SyncPointWaitList, SyncPoint);
+}
+
+pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Length,
+    pi_mem_advice Advice, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferAdviseUSM(CommandBuffer, Ptr, Length, Advice,
+                                            NumSyncPointsInWaitList,
+                                            SyncPointWaitList, SyncPoint);
+}
+
 pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
                                     pi_queue Queue,
                                     pi_uint32 NumEventsInWaitList,

sycl/plugins/unified_runtime/pi2ur.hpp

@@ -4568,6 +4568,41 @@ inline pi_result piextCommandBufferMemBufferWrite(
   return PI_SUCCESS;
 }
 
+inline pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Size,
+    pi_usm_migration_flags Flags, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+
+  // flags is currently unused so fail if set
+  PI_ASSERT(Flags == 0, PI_ERROR_INVALID_VALUE);
+
+  ur_exp_command_buffer_handle_t UrCommandBuffer =
+      reinterpret_cast<ur_exp_command_buffer_handle_t>(CommandBuffer);
+
+  // TODO: to map from pi_usm_migration_flags to
+  // ur_usm_migration_flags_t
+  // once we have those defined
+  ur_usm_migration_flags_t UrFlags{};
+  HANDLE_ERRORS(urCommandBufferAppendUSMPrefetchExp(
+      UrCommandBuffer, Ptr, Size, UrFlags, NumSyncPointsInWaitList,
+      SyncPointWaitList, SyncPoint));
+  return PI_SUCCESS;
+}
+
+inline pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Length,
+    pi_mem_advice Advice, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  ur_exp_command_buffer_handle_t UrCommandBuffer =
+      reinterpret_cast<ur_exp_command_buffer_handle_t>(CommandBuffer);
+
+  ur_usm_advice_flags_t UrAdvice{};
+  HANDLE_ERRORS(urCommandBufferAppendUSMAdviseExp(
+      UrCommandBuffer, Ptr, Length, UrAdvice, NumSyncPointsInWaitList,
+      SyncPointWaitList, SyncPoint));
+  return PI_SUCCESS;
+}
+
 inline pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
                                            pi_queue Queue,
                                            pi_uint32 NumEventsInWaitList,

sycl/plugins/unified_runtime/pi_unified_runtime.cpp

@@ -1102,6 +1102,24 @@ pi_result piextCommandBufferMemBufferWriteRect(
       NumSyncPointsInWaitList, SyncPointWaitList, SyncPoint);
 }
 
+pi_result piextCommandBufferPrefetchUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Size,
+    pi_usm_migration_flags Flags, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferPrefetchUSM(CommandBuffer, Ptr, Size, Flags,
+                                              NumSyncPointsInWaitList,
+                                              SyncPointWaitList, SyncPoint);
+}
+
+pi_result piextCommandBufferAdviseUSM(
+    pi_ext_command_buffer CommandBuffer, const void *Ptr, size_t Length,
+    pi_mem_advice Advice, pi_uint32 NumSyncPointsInWaitList,
+    const pi_ext_sync_point *SyncPointWaitList, pi_ext_sync_point *SyncPoint) {
+  return pi2ur::piextCommandBufferAdviseUSM(CommandBuffer, Ptr, Length, Advice,
+                                            NumSyncPointsInWaitList,
+                                            SyncPointWaitList, SyncPoint);
+}
+
 pi_result piextEnqueueCommandBuffer(pi_ext_command_buffer CommandBuffer,
                                     pi_queue Queue,
                                     pi_uint32 NumEventsInWaitList,

sycl/source/detail/memory_manager.cpp

@@ -194,15 +194,15 @@ void memBufferMapHelper(const PluginPtr &Plugin, pi_queue Queue, pi_mem Buffer,
   // We only want to instrument piEnqueueMemBufferMap
 
 #ifdef XPTI_ENABLE_INSTRUMENTATION
-    CorrID = emitMemAllocBeginTrace(MemObjID, Size, 0 /* guard zone */);
-    xpti::utils::finally _{[&] {
-      emitMemAllocEndTrace(MemObjID, (uintptr_t)(*RetMap), Size,
-                           0 /* guard zone */, CorrID);
-    }};
+  CorrID = emitMemAllocBeginTrace(MemObjID, Size, 0 /* guard zone */);
+  xpti::utils::finally _{[&] {
+    emitMemAllocEndTrace(MemObjID, (uintptr_t)(*RetMap), Size,
+                         0 /* guard zone */, CorrID);
+  }};
 #endif
-    Plugin->call<PiApiKind::piEnqueueMemBufferMap>(
-        Queue, Buffer, Blocking, Flags, Offset, Size, NumEvents, WaitList,
-        Event, RetMap);
+  Plugin->call<PiApiKind::piEnqueueMemBufferMap>(Queue, Buffer, Blocking, Flags,
+                                                 Offset, Size, NumEvents,
+                                                 WaitList, Event, RetMap);
 }
 
 void memUnmapHelper(const PluginPtr &Plugin, pi_queue Queue, pi_mem Mem,
@@ -1666,6 +1666,33 @@ void MemoryManager::ext_oneapi_copy_usm_cmd_buffer(
   }
 }
 
+void MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer(
+    sycl::detail::ContextImplPtr Context,
+    sycl::detail::pi::PiExtCommandBuffer CommandBuffer, void *Mem,
+    size_t Length, std::vector<sycl::detail::pi::PiExtSyncPoint> Deps,
+    sycl::detail::pi::PiExtSyncPoint *OutSyncPoint) {
+  assert(!Context->is_host() && "Host queue not supported in prefetch_usm.");
+
+  const PluginPtr &Plugin = Context->getPlugin();
+  Plugin->call<PiApiKind::piextCommandBufferPrefetchUSM>(
+      CommandBuffer, Mem, Length, _pi_usm_migration_flags(0), Deps.size(),
+      Deps.data(), OutSyncPoint);
+}
+
+void MemoryManager::ext_oneapi_advise_usm_cmd_buffer(
+    sycl::detail::ContextImplPtr Context,
+    sycl::detail::pi::PiExtCommandBuffer CommandBuffer, const void *Mem,
+    size_t Length, pi_mem_advice Advice,
+    std::vector<sycl::detail::pi::PiExtSyncPoint> Deps,
+    sycl::detail::pi::PiExtSyncPoint *OutSyncPoint) {
+  assert(!Context->is_host() && "Host queue not supported in advise_usm.");
+
+  const PluginPtr &Plugin = Context->getPlugin();
+  Plugin->call<PiApiKind::piextCommandBufferAdviseUSM>(
+      CommandBuffer, Mem, Length, Advice, Deps.size(), Deps.data(),
+      OutSyncPoint);
+}
+
 void MemoryManager::copy_image_bindless(
     void *Src, QueueImplPtr Queue, void *Dst,
     const sycl::detail::pi::PiMemImageDesc &Desc,

sycl/source/detail/memory_manager.hpp

@@ -316,6 +316,19 @@ class __SYCL_EXPORT MemoryManager {
       void *DstMem, std::vector<sycl::detail::pi::PiExtSyncPoint> Deps,
       sycl::detail::pi::PiExtSyncPoint *OutSyncPoint);
 
+  static void ext_oneapi_prefetch_usm_cmd_buffer(
+      sycl::detail::ContextImplPtr Context,
+      sycl::detail::pi::PiExtCommandBuffer CommandBuffer, void *Mem,
+      size_t Length, std::vector<sycl::detail::pi::PiExtSyncPoint> Deps,
+      sycl::detail::pi::PiExtSyncPoint *OutSyncPoint);
+
+  static void ext_oneapi_advise_usm_cmd_buffer(
+      sycl::detail::ContextImplPtr Context,
+      sycl::detail::pi::PiExtCommandBuffer CommandBuffer, const void *Mem,
+      size_t Length, pi_mem_advice Advice,
+      std::vector<sycl::detail::pi::PiExtSyncPoint> Deps,
+      sycl::detail::pi::PiExtSyncPoint *OutSyncPoint);
+
   static void
   copy_image_bindless(void *Src, QueueImplPtr Queue, void *Dst,
                       const sycl::detail::pi::PiMemImageDesc &Desc,

sycl/source/detail/queue_impl.cpp

@@ -227,6 +227,17 @@ event queue_impl::mem_advise(const std::shared_ptr<detail::queue_impl> &Self,
                              const void *Ptr, size_t Length,
                              pi_mem_advice Advice,
                              const std::vector<event> &DepEvents) {
+  // If we have a command graph set we need to capture the advise through normal
+  // queue submission.
+  if (MGraph.lock()) {
+    return submit(
+        [&](handler &CGH) {
+          CGH.depends_on(DepEvents);
+          CGH.mem_advise(Ptr, Length, Advice);
+        },
+        Self, {});
+  }
+
   if (MHasDiscardEventsSupport) {
     MemoryManager::advise_usm(Ptr, Self, Length, Advice,
                               getOrWaitEvents(DepEvents, MContext), nullptr);

sycl/source/detail/scheduler/commands.cpp

@@ -2764,6 +2764,23 @@ pi_int32 ExecCGCommand::enqueueImpCommandBuffer() {
     MEvent->setSyncPoint(OutSyncPoint);
     return PI_SUCCESS;
   }
+  case CG::CGTYPE::PrefetchUSM: {
+    CGPrefetchUSM *Prefetch = (CGPrefetchUSM *)MCommandGroup.get();
+    MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer(
+        MQueue->getContextImplPtr(), MCommandBuffer, Prefetch->getDst(),
+        Prefetch->getLength(), std::move(MSyncPointDeps), &OutSyncPoint);
+    MEvent->setSyncPoint(OutSyncPoint);
+    return PI_SUCCESS;
+  }
+  case CG::CGTYPE::AdviseUSM: {
+    CGAdviseUSM *Advise = (CGAdviseUSM *)MCommandGroup.get();
+    MemoryManager::ext_oneapi_advise_usm_cmd_buffer(
+        MQueue->getContextImplPtr(), MCommandBuffer, Advise->getDst(),
+        Advise->getLength(), Advise->getAdvice(), std::move(MSyncPointDeps),
+        &OutSyncPoint);
+    MEvent->setSyncPoint(OutSyncPoint);
+    return PI_SUCCESS;
+  }
   default:
     throw runtime_error("CG type not implemented for command buffers.",
                         PI_ERROR_INVALID_OPERATION);

sycl/test-e2e/Graph/Explicit/memadvise.cpp

@@ -0,0 +1,32 @@
+// RUN: %{build} -o %t.out
+// RUN: %if linux && (level_zero || cuda) %{ env SYCL_PI_TRACE=2 %{run} %t.out 2>&1 FileCheck %s %} %else %{ %{run} %t.out %}
+
+// Since Mem advise is only a memory hint that doesn't
+// impact results but only performances, we verify
+// that a node is correctly added by checking PI function calls.
+
+// CHECK: piextCommandBufferAdviseUSM
+// CHECK-NEXT: <unknown> : 0x[[#%x,COMMAND_BUFFER:]]
+// CHECK-NEXT: <unknown> : 0x[[#%x,PTR:]]
+// CHECK-NEXT: <unknown> : 400
+// CHECK-NEXT: <unknown> : 0
+// CHECK-NEXT: <unknown> : 0
+// CHECK-NEXT: <unknown> : 0
+// CHECK-NEXT: <unknown> : 0x[[#%x,ADVISE_SYNC_POINT:]]
+// CHECK: pi_result : PI_SUCCESS
+
+// CHECK: piextCommandBufferNDRangeKernel(
+// CHECK-NEXT: <unknown> : 0x[[#COMMAND_BUFFER]]
+// CHECK-NEXT: <unknown> : 0x[[#%x,KERNEL:]]
+// CHECK-NEXT: <unknown> : 1
+// CHECK-NEXT: <unknown> : 0x[[#%x,GLOBAL_WORK_OFFSET:]]
+// CHECK-NEXT: <unknown> : 0x[[#%x,GLOBAL_WORK_SIZE:]]
+// CHECK-NEXT: <unknown> : 0
+// CHECK-NEXT: <unknown> : 1
+// CHECK-NEXT: <unknown> : 0x[[#%x,SYNC_POINT_WAIT_LIST:]]
+// CHECK-NEXT: <unknown> : 0x[[#%x,KERNEL_SYNC_POINT:]]
+// CHECK: pi_result : PI_SUCCESS
+
+#define GRAPH_E2E_EXPLICIT
+
+#include "../Inputs/memadvise.cpp"