Implement experimental 2way prefetch function

Author: ianayl

sycl/cmake/modules/FetchUnifiedRuntime.cmake

@@ -116,12 +116,12 @@ 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")
   # commit af7e275b509b41f54a66743ebf748dfb51668abf
   # Author: Maosu Zhao <maosu.zhao@intel.com>
   # Date:   Thu Oct 17 16:31:21 2024 +0800
   #     [DeviceSanitizer] Refactor the code to manage shadow memory (#2127)
-  set(UNIFIED_RUNTIME_TAG af7e275b509b41f54a66743ebf748dfb51668abf)
+  set(UNIFIED_RUNTIME_TAG 644eef9f8dfcc3931fd2883f24eace60ca98b3ed)
 
   set(UMF_BUILD_EXAMPLES OFF CACHE INTERNAL "EXAMPLES")
   # Due to the use of dependentloadflag and no installer for UMF and hwloc we need

sycl/include/sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp

@@ -0,0 +1,24 @@
+//==-------- usm_prefetch_exp.hpp --- SYCL USM prefetch extensions ---------==//
+//
+// 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 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
+  };
+
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/include/sycl/handler.hpp

@@ -35,6 +35,7 @@
 #include <sycl/ext/oneapi/experimental/cluster_group_prop.hpp>
 #include <sycl/ext/oneapi/experimental/graph.hpp>
 #include <sycl/ext/oneapi/experimental/raw_kernel_arg.hpp>
+#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp>
 #include <sycl/ext/oneapi/experimental/use_root_sync_prop.hpp>
 #include <sycl/ext/oneapi/experimental/virtual_functions.hpp>
 #include <sycl/ext/oneapi/kernel_properties/properties.hpp>
@@ -2824,6 +2825,17 @@ 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.
   ///
@@ -3253,6 +3265,9 @@ 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;

sycl/include/sycl/queue.hpp

@@ -31,6 +31,7 @@
 #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/graph.hpp>       // for command_graph...
+#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp> // for migration...
 #include <sycl/ext/oneapi/properties/properties.hpp>    // for empty_properti...
 #include <sycl/handler.hpp>                             // for handler, isDev...
 #include <sycl/id.hpp>                                  // for id
@@ -745,6 +746,24 @@ 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());
+  }
+
   /// 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.
@@ -765,6 +784,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());
+  }
+
   /// 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.
@@ -786,6 +828,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());
+  }
+
   /// Copies data from one 2D memory region to another, both pointed by
   /// USM pointers.
   /// No operations is done if \param Width or \param Height is zero. An

sycl/source/detail/cg.hpp

@@ -17,6 +17,7 @@
 #include <sycl/exception_list.hpp>  // for queue_impl
 #include <sycl/kernel.hpp>          // for kernel_impl
 #include <sycl/kernel_bundle.hpp>   // for kernel_bundle_impl
+#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp> // for migration_direction
 
 #include <assert.h> // for assert
 #include <memory>   // for shared_ptr, unique_ptr
@@ -390,14 +391,16 @@ 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,
+  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) {}
+        MDst(DstPtr), MLength(Length), MDirection(Direction) {}
   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

@@ -964,16 +964,19 @@ 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,
+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) {
   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;
   if (OutEventImpl != nullptr)
     OutEventImpl->setHostEnqueueTime();
   Adapter->call<UrApiKind::urEnqueueUSMPrefetch>(Queue->getHandleRef(), Mem,
-                                                 Length, 0, DepEvents.size(),
+                                                 Length, migration_flag, DepEvents.size(),
                                                  DepEvents.data(), OutEvent);
 }
 
@@ -1610,12 +1613,15 @@ 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,
+    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) {
   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;
   Adapter->call<UrApiKind::urCommandBufferAppendUSMPrefetchExp>(
-      CommandBuffer, Mem, Length, ur_usm_migration_flags_t(0), Deps.size(),
+      CommandBuffer, Mem, Length, migration_flag, Deps.size(),
       Deps.data(), 0, nullptr, OutSyncPoint, nullptr, nullptr);
 }
 

sycl/source/detail/memory_manager.hpp

@@ -14,6 +14,7 @@
 #include <sycl/id.hpp>
 #include <sycl/property_list.hpp>
 #include <sycl/range.hpp>
+#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp>
 
 #include <ur_api.h>
 
@@ -149,7 +150,7 @@ class MemoryManager {
                        ur_event_handle_t *OutEvent,
                        const detail::EventImplPtr &OutEventImpl);
 
-  static void prefetch_usm(void *Ptr, QueueImplPtr Queue, size_t Len,
+  static void prefetch_usm(void *Ptr, QueueImplPtr Queue, size_t Len, sycl::ext::oneapi::experimental::migration_direction Direction,
                            std::vector<ur_event_handle_t> DepEvents,
                            ur_event_handle_t *OutEvent,
                            const detail::EventImplPtr &OutEventImpl);
@@ -250,6 +251,7 @@ class MemoryManager {
   static void 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);
 

sycl/source/detail/scheduler/commands.cpp

@@ -2867,7 +2867,7 @@ ur_result_t ExecCGCommand::enqueueImpCommandBuffer() {
     CGPrefetchUSM *Prefetch = (CGPrefetchUSM *)MCommandGroup.get();
     MemoryManager::ext_oneapi_prefetch_usm_cmd_buffer(
         MQueue->getContextImplPtr(), MCommandBuffer, Prefetch->getDst(),
-        Prefetch->getLength(), std::move(MSyncPointDeps), &OutSyncPoint);
+        Prefetch->getLength(), Prefetch->getDirection(), std::move(MSyncPointDeps), &OutSyncPoint);
     MEvent->setSyncPoint(OutSyncPoint);
     return UR_RESULT_SUCCESS;
   }
@@ -3044,7 +3044,7 @@ ur_result_t ExecCGCommand::enqueueImpQueue() {
   case CGType::PrefetchUSM: {
     CGPrefetchUSM *Prefetch = (CGPrefetchUSM *)MCommandGroup.get();
     MemoryManager::prefetch_usm(Prefetch->getDst(), MQueue,
-                                Prefetch->getLength(), std::move(RawEvents),
+                                Prefetch->getLength(), Prefetch->getDirection(), std::move(RawEvents),
                                 Event, MEvent);
     if (Event)
       MEvent->setHandle(*Event);

sycl/source/handler.cpp

@@ -377,7 +377,7 @@ event handler::finalize() {
     break;
   case detail::CGType::PrefetchUSM:
     CommandGroup.reset(new detail::CGPrefetchUSM(
-        MDstPtr, MLength, std::move(impl->CGData), MCodeLoc));
+        MDstPtr, MLength, std::move(impl->CGData), MDirection, MCodeLoc));
     break;
   case detail::CGType::AdviseUSM:
     CommandGroup.reset(new detail::CGAdviseUSM(MDstPtr, MLength, impl->MAdvice,
@@ -970,6 +970,14 @@ void handler::prefetch(const void *Ptr, size_t Count) {
   setType(detail::CGType::PrefetchUSM);
 }
 
+void handler::ext_oneapi_prefetch_exp(const void* ptr, size_t Count, ext::oneapi::experimental::migration_direction Direction) {
+  throwIfActionIsCreated();
+  MDstPtr = const_cast<void *>(ptr);
+  MLength = Count;
+  MDirection = Direction;
+  setType(sycl::detail::CGType::PrefetchUSM);
+}
+
 void handler::mem_advise(const void *Ptr, size_t Count, int Advice) {
   throwIfActionIsCreated();
   MDstPtr = const_cast<void *>(Ptr);

sycl/test-e2e/USM/prefetch_exp.cpp

@@ -0,0 +1,73 @@
+//==---- prefetch_exp.cpp - Experimental 2-way USM prefetch test ------------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: %{build} -o %t1.out
+// RUN: %{run} %t1.out
+
+#include <sycl/detail/core.hpp>
+#include <sycl/usm.hpp>
+#include <sycl/ext/oneapi/experimental/USM/prefetch_exp.hpp>
+
+using namespace sycl;
+
+static constexpr int count = 100;
+
+int main() {
+  queue q([](exception_list el) {
+    for (auto &e : el)
+      throw e;
+  });
+  if (q.get_device().get_info<info::device::usm_shared_allocations>()) {
+    float *src = (float *)malloc_shared(sizeof(float) * count, q.get_device(),
+                                        q.get_context());
+    float *dest = (float *)malloc_shared(sizeof(float) * count, q.get_device(),
+                                         q.get_context());
+    for (int i = 0; i < count; i++)
+      src[i] = i;
+
+    // Test host to device prefetch_exp(handler &CGH, ..)
+    {
+      event init_prefetch = q.submit(
+          [&](handler &cgh) { cgh.ext_oneapi_prefetch_exp(src, sizeof(float) * count); });
+
+      q.submit([&](handler &cgh) {
+        cgh.depends_on(init_prefetch);
+        cgh.single_task<class double_dest>([=]() {
+          for (int i = 0; i < count; i++)
+            dest[i] = 2 * src[i];
+        });
+      }
+      q.wait_and_throw();
+
+      for (int i = 0; i < count; i++) {
+        assert(dest[i] == i * 2);
+      }
+    }
+
+    // Test queue::prefetch
+    {
+      event init_prefetch = q.ext_oneapi_prefetch_exp(src, sizeof(float) * count);
+
+      q.submit([&](handler &cgh) {
+        cgh.depends_on(init_prefetch);
+        cgh.single_task<class double_dest3>([=]() {
+          for (int i = 0; i < count; i++)
+            dest[i] = 3 * src[i];
+        });
+      });
+      q.wait_and_throw();
+
+      for (int i = 0; i < count; i++) {
+        assert(dest[i] == i * 3);
+      }
+    }
+    free(src, q);
+    free(dest, q);
+  }
+  return 0;
+}