[UR] Add support for bidirectional USM prefetching

include/ur_api.h

@@ -8279,18 +8279,20 @@ typedef enum ur_map_flag_t {
 #define UR_MAP_FLAGS_MASK 0xfffffff8
 
 ///////////////////////////////////////////////////////////////////////////////
-/// @brief Map flags
+/// @brief USM migration flags, indicating the direction data is migrated in
 typedef uint32_t ur_usm_migration_flags_t;
 typedef enum ur_usm_migration_flag_t {
-  /// Default migration TODO: Add more enums!
-  UR_USM_MIGRATION_FLAG_DEFAULT = UR_BIT(0),
+  /// Migrate data from host to device
+  UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE = UR_BIT(0),
+  /// Migrate data from device to host
+  UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST = UR_BIT(1),
   /// @cond
   UR_USM_MIGRATION_FLAG_FORCE_UINT32 = 0x7fffffff
   /// @endcond
 
 } ur_usm_migration_flag_t;
 /// @brief Bit Mask for validating ur_usm_migration_flags_t
-#define UR_USM_MIGRATION_FLAGS_MASK 0xfffffffe
+#define UR_USM_MIGRATION_FLAGS_MASK 0xfffffffc
 
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Enqueue a command to map a region of the buffer object into the host
@@ -8549,7 +8551,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMPrefetch(
     const void *pMem,
     /// [in] size in bytes to be fetched
     size_t size,
-    /// [in] USM prefetch flags
+    /// [in] USM migration flags
     ur_usm_migration_flags_t flags,
     /// [in] size of the event wait list
     uint32_t numEventsInWaitList,
@@ -10909,7 +10911,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferAppendUSMPrefetchExp(
     const void *pMemory,
     /// [in] size in bytes to be fetched.
     size_t size,
-    /// [in] USM prefetch flags
+    /// [in] USM migration flags
     ur_usm_migration_flags_t flags,
     /// [in] The number of sync points in the provided dependency list.
     uint32_t numSyncPointsInWaitList,

include/ur_print.hpp

@@ -10300,8 +10300,11 @@ inline ur_result_t printFlag<ur_map_flag_t>(std::ostream &os, uint32_t flag) {
 inline std::ostream &operator<<(std::ostream &os,
                                 enum ur_usm_migration_flag_t value) {
   switch (value) {
-  case UR_USM_MIGRATION_FLAG_DEFAULT:
-    os << "UR_USM_MIGRATION_FLAG_DEFAULT";
+  case UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE:
+    os << "UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE";
+    break;
+  case UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST:
+    os << "UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST";
     break;
   default:
     os << "unknown enumerator";
@@ -10319,15 +10322,26 @@ inline ur_result_t printFlag<ur_usm_migration_flag_t>(std::ostream &os,
   uint32_t val = flag;
   bool first = true;
 
-  if ((val & UR_USM_MIGRATION_FLAG_DEFAULT) ==
-      (uint32_t)UR_USM_MIGRATION_FLAG_DEFAULT) {
-    val ^= (uint32_t)UR_USM_MIGRATION_FLAG_DEFAULT;
+  if ((val & UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE) ==
+      (uint32_t)UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE) {
+    val ^= (uint32_t)UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE;
+    if (!first) {
+      os << " | ";
+    } else {
+      first = false;
+    }
+    os << UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE;
+  }
+
+  if ((val & UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST) ==
+      (uint32_t)UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST) {
+    val ^= (uint32_t)UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
     if (!first) {
       os << " | ";
     } else {
       first = false;
     }
-    os << UR_USM_MIGRATION_FLAG_DEFAULT;
+    os << UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST;
   }
   if (val != 0) {
     std::bitset<32> bits(val);

scripts/core/enqueue.yml

@@ -849,13 +849,16 @@ etors:
       value: "$X_BIT(2)"
 --- #--------------------------------------------------------------------------
 type: enum
-desc: "Map flags"
-class: $xDevice
+desc: "USM migration flags, indicating the direction data is migrated in"
+class: $xEnqueue
 name: $x_usm_migration_flags_t
 etors:
-    - name: DEFAULT
-      desc: "Default migration TODO: Add more enums! "
+    - name: HOST_TO_DEVICE
+      desc: "Migrate data from host to device"
       value: "$X_BIT(0)"
+    - name: DEVICE_TO_HOST
+      desc: "Migrate data from device to host"
+      value: "$X_BIT(1)"
 --- #--------------------------------------------------------------------------
 type: function
 desc: "Enqueue a command to map a region of the buffer object into the host address space and return a pointer to the mapped region"
@@ -1144,7 +1147,7 @@ params:
       desc: "[in] size in bytes to be fetched"
     - type: $x_usm_migration_flags_t
       name: flags
-      desc: "[in] USM prefetch flags"
+      desc: "[in] USM migration flags"
     - type: uint32_t
       name: numEventsInWaitList
       desc: "[in] size of the event wait list"

scripts/core/exp-command-buffer.yml

@@ -1024,7 +1024,7 @@ params:
       desc: "[in] size in bytes to be fetched."
     - type: $x_usm_migration_flags_t
       name: flags
-      desc: "[in] USM prefetch flags"
+      desc: "[in] USM migration flags"
     - type: uint32_t
       name: numSyncPointsInWaitList
       desc: "[in] The number of sync points in the provided dependency list."

source/adapters/cuda/enqueue.cpp

@@ -1598,13 +1598,25 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMPrefetch(
     ur_queue_handle_t hQueue, const void *pMem, size_t size,
     ur_usm_migration_flags_t flags, uint32_t numEventsInWaitList,
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
-  std::ignore = flags;
+  ur_device_handle_t Device = hQueue->getDevice();
+  CUdevice TargetDevice;
+  switch (flags) {
+  case UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE:
+    TargetDevice = Device->get();
+    break;
+  case UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST:
+    TargetDevice = CU_DEVICE_CPU;
+    break;
+  default:
+    setErrorMessage("Invalid USM migration flag",
+                    UR_RESULT_ERROR_INVALID_ENUMERATION);
+    return UR_RESULT_ERROR_INVALID_ENUMERATION;
+  }
 
   size_t PointerRangeSize = 0;
   UR_CHECK_ERROR(cuPointerGetAttribute(
       &PointerRangeSize, CU_POINTER_ATTRIBUTE_RANGE_SIZE, (CUdeviceptr)pMem));
   UR_ASSERT(size <= PointerRangeSize, UR_RESULT_ERROR_INVALID_SIZE);
-  ur_device_handle_t Device = hQueue->getDevice();
 
   // Certain cuda devices and Windows do not have support for some Unified
   // Memory features. cuMemPrefetchAsync requires concurrent memory access
@@ -1641,7 +1653,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMPrefetch(
       UR_CHECK_ERROR(EventPtr->start());
     }
     UR_CHECK_ERROR(
-        cuMemPrefetchAsync((CUdeviceptr)pMem, size, Device->get(), CuStream));
+        cuMemPrefetchAsync((CUdeviceptr)pMem, size, TargetDevice, CuStream));
     if (phEvent) {
       UR_CHECK_ERROR(EventPtr->record());
       *phEvent = EventPtr.release();

source/adapters/hip/enqueue.cpp

@@ -1365,10 +1365,24 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMPrefetch(
     ur_queue_handle_t hQueue, const void *pMem, size_t size,
     ur_usm_migration_flags_t flags, uint32_t numEventsInWaitList,
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
-  std::ignore = flags;
 
   void *HIPDevicePtr = const_cast<void *>(pMem);
   ur_device_handle_t Device = hQueue->getDevice();
+  hipDevice_t TargetDevice;
+  switch (flags) {
+  case UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE:
+    TargetDevice = Device->get();
+    break;
+  case UR_USM_MIGRATION_FLAG_DEVICE_TO_HOST:
+    // HIP doesn't have a constant for host like CUDA does; -1 is used instead
+    // https://github.com/ROCm/HIP/blob/3d60bd3a6415c2/docs/how-to/unified_memory.rst#L376
+    TargetDevice = -1;
+    break;
+  default:
+    setErrorMessage("Invalid USM migration flag",
+                    UR_RESULT_ERROR_INVALID_ENUMERATION);
+    return UR_RESULT_ERROR_INVALID_ENUMERATION;
+  }
 
 // HIP_POINTER_ATTRIBUTE_RANGE_SIZE is not an attribute in ROCM < 5,
 // so we can't perform this check for such cases.
@@ -1428,8 +1442,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMPrefetch(
       return UR_RESULT_ERROR_ADAPTER_SPECIFIC;
     }
 
-    UR_CHECK_ERROR(
-        hipMemPrefetchAsync(pMem, size, hQueue->getDevice()->get(), HIPStream));
+    UR_CHECK_ERROR(hipMemPrefetchAsync(pMem, size, TargetDevice, HIPStream));
     releaseEvent();
   } catch (ur_result_t Err) {
     return Err;

source/adapters/level_zero/command_buffer.cpp

@@ -1304,13 +1304,18 @@ ur_result_t urCommandBufferAppendUSMPrefetchExp(
   std::ignore = EventWaitList;
   std::ignore = Event;
   std::ignore = Command;
-  std::ignore = Flags;
 
   if (CommandBuffer->IsInOrderCmdList) {
-    // Add the prefetch command to the command-buffer.
-    // Note that L0 does not handle migration flags.
-    ZE2UR_CALL(zeCommandListAppendMemoryPrefetch,
-               (CommandBuffer->ZeComputeCommandList, Mem, Size));
+    if (Flags == UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE) {
+      // Add the prefetch command to the command-buffer.
+      ZE2UR_CALL(zeCommandListAppendMemoryPrefetch,
+                 (CommandBuffer->ZeComputeCommandList, Mem, Size));
+    } else {
+      // L0 currently does not handle migration flags -- All other migration
+      // behavior is ignored:
+      logger::warning("USM migration from device to host is not currently "
+                      "supported by level zero.");
+    }
   } else {
     std::vector<ze_event_handle_t> ZeEventList;
     ze_event_handle_t ZeLaunchEvent = nullptr;
@@ -1324,10 +1329,16 @@ ur_result_t urCommandBufferAppendUSMPrefetchExp(
                   ZeEventList.data()));
     }
 
-    // Add the prefetch command to the command-buffer.
-    // Note that L0 does not handle migration flags.
-    ZE2UR_CALL(zeCommandListAppendMemoryPrefetch,
-               (CommandBuffer->ZeComputeCommandList, Mem, Size));
+    if (Flags == UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE) {
+      // Add the prefetch command to the command buffer.
+      ZE2UR_CALL(zeCommandListAppendMemoryPrefetch,
+                 (CommandBuffer->ZeComputeCommandList, Mem, Size));
+    } else {
+      // L0 currently does not handle migration flags -- All other migration
+      // behavior is ignored:
+      logger::warning("USM migration from device to host is not currently "
+                      "supported by level zero.");
+    }
 
     // Level Zero does not have a completion "event" with the prefetch API,
     // so manually add command to signal our event.

source/adapters/level_zero/memory.cpp

@@ -1290,7 +1290,6 @@ ur_result_t urEnqueueUSMPrefetch(
     /// [in,out][optional] return an event object that identifies this
     /// particular command instance.
     ur_event_handle_t *OutEvent) {
-  std::ignore = Flags;
   // Lock automatically releases when this goes out of scope.
   std::scoped_lock<ur_shared_mutex> lock(Queue->Mutex);
 
@@ -1330,8 +1329,14 @@ ur_result_t urEnqueueUSMPrefetch(
     ZE2UR_CALL(zeCommandListAppendWaitOnEvents,
                (ZeCommandList, WaitList.Length, WaitList.ZeEventList));
   }
-  // TODO: figure out how to translate "flags"
-  ZE2UR_CALL(zeCommandListAppendMemoryPrefetch, (ZeCommandList, Mem, Size));
+
+  if (Flags == UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE) {
+    ZE2UR_CALL(zeCommandListAppendMemoryPrefetch, (ZeCommandList, Mem, Size));
+  } else {
+    // L0 does not suppot migrating from device to host yet: skip procedure
+    logger::warning("urEnqueueUSMPrefetch: Prefetch from device to host not yet"
+                    " supported by level zero");
+  }
 
   // TODO: Level Zero does not have a completion "event" with the prefetch API,
   // so manually add command to signal our event.

source/adapters/level_zero/v2/queue_immediate_in_order.cpp

@@ -629,8 +629,6 @@ ur_result_t ur_queue_immediate_in_order_t::enqueueUSMPrefetch(
     ur_event_handle_t *phEvent) {
   TRACK_SCOPE_LATENCY("ur_queue_immediate_in_order_t::enqueueUSMPrefetch");
 
-  std::ignore = flags;
-
   std::scoped_lock<ur_shared_mutex> lock(this->Mutex);
 
   auto zeSignalEvent = getSignalEvent(phEvent, UR_COMMAND_USM_PREFETCH);
@@ -643,12 +641,19 @@ ur_result_t ur_queue_immediate_in_order_t::enqueueUSMPrefetch(
         zeCommandListAppendWaitOnEvents,
         (commandListManager.getZeCommandList(), numWaitEvents, pWaitEvents));
   }
-  // TODO: figure out how to translate "flags"
-  ZE2UR_CALL(zeCommandListAppendMemoryPrefetch,
-             (commandListManager.getZeCommandList(), pMem, size));
-  if (zeSignalEvent) {
-    ZE2UR_CALL(zeCommandListAppendSignalEvent,
-               (commandListManager.getZeCommandList(), zeSignalEvent));
+
+  if (flags == UR_USM_MIGRATION_FLAG_HOST_TO_DEVICE) {
+    ZE2UR_CALL(zeCommandListAppendMemoryPrefetch,
+               (commandListManager.getZeCommandList(), pMem, size));
+    if (zeSignalEvent) {
+      ZE2UR_CALL(zeCommandListAppendSignalEvent,
+                 (commandListManager.getZeCommandList(), zeSignalEvent));
+    }
+  } else {
+    // L0 does not suppot migrating from device to host yet: skip procedure
+    setErrorMessage("Prefetch from device to host not yet supported by level "
+                    "zero.",
+                    UR_RESULT_SUCCESS);
   }
 
   return UR_RESULT_SUCCESS;

source/adapters/mock/ur_mockddi.cpp

@@ -6720,7 +6720,7 @@ __urdlllocal ur_result_t UR_APICALL urEnqueueUSMPrefetch(
     const void *pMem,
     /// [in] size in bytes to be fetched
     size_t size,
-    /// [in] USM prefetch flags
+    /// [in] USM migration flags
     ur_usm_migration_flags_t flags,
     /// [in] size of the event wait list
     uint32_t numEventsInWaitList,
@@ -9584,7 +9584,7 @@ __urdlllocal ur_result_t UR_APICALL urCommandBufferAppendUSMPrefetchExp(
     const void *pMemory,
     /// [in] size in bytes to be fetched.
     size_t size,
-    /// [in] USM prefetch flags
+    /// [in] USM migration flags
     ur_usm_migration_flags_t flags,
     /// [in] The number of sync points in the provided dependency list.
     uint32_t numSyncPointsInWaitList,

source/adapters/opencl/common.cpp

@@ -13,15 +13,15 @@
 namespace cl_adapter {
 
 /* Global variables for urAdapterGetLastError() */
-thread_local int32_t ErrorMessageCode = 0;
+thread_local ur_result_t ErrorMessageCode = UR_RESULT_SUCCESS;
 thread_local char ErrorMessage[MaxMessageSize]{};
 
-[[maybe_unused]] void setErrorMessage(const char *Message, int32_t ErrorCode) {
+[[maybe_unused]] void setErrorMessage(const char *Message,
+                                      ur_result_t ErrorCode) {
   assert(strlen(Message) < cl_adapter::MaxMessageSize);
   // Copy at most MaxMessageSize - 1 bytes to ensure the resultant string is
   // always null terminated.
   strncpy(cl_adapter::ErrorMessage, Message, MaxMessageSize - 1);
-
   ErrorMessageCode = ErrorCode;
 }
 } // namespace cl_adapter

source/adapters/opencl/common.hpp

@@ -150,7 +150,7 @@ inline const OpenCLVersion V3_0(3, 0);
 
 namespace cl_adapter {
 constexpr size_t MaxMessageSize = 256;
-extern thread_local int32_t ErrorMessageCode;
+extern thread_local ur_result_t ErrorMessageCode;
 extern thread_local char ErrorMessage[MaxMessageSize];
 
 // Utility function for setting a message and warning
@@ -203,6 +203,7 @@ CONSTFIX char EnqueueWriteGlobalVariableName[] =
     "clEnqueueWriteGlobalVariableINTEL";
 CONSTFIX char EnqueueReadGlobalVariableName[] =
     "clEnqueueReadGlobalVariableINTEL";
+CONSTFIX char EnqueueMigrateMemName[] = "clEnqueueMigrateMemINTEL";
 // Names of host pipe functions queried from OpenCL
 CONSTFIX char EnqueueReadHostPipeName[] = "clEnqueueReadHostPipeINTEL";
 CONSTFIX char EnqueueWriteHostPipeName[] = "clEnqueueWriteHostPipeINTEL";

source/adapters/opencl/extension_functions.def

@@ -8,6 +8,7 @@ CL_EXTENSION_FUNC(clMemBlockingFreeINTEL)
 CL_EXTENSION_FUNC(clSetKernelArgMemPointerINTEL)
 CL_EXTENSION_FUNC(clEnqueueMemFillINTEL)
 CL_EXTENSION_FUNC(clEnqueueMemcpyINTEL)
+CL_EXTENSION_FUNC(clEnqueueMigrateMemINTEL)
 CL_EXTENSION_FUNC(clGetMemAllocInfoINTEL)
 CL_EXTENSION_FUNC(clEnqueueWriteGlobalVariable)
 CL_EXTENSION_FUNC(clEnqueueReadGlobalVariable)