Change migrateMemoryToDevice to enqueue*

In line with changing from a sync op to async for memory migration
across devices in a context.

Author: Hugh Delaney

source/adapters/cuda/enqueue_native.cpp

@@ -28,22 +28,21 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueNativeCommandExp(
     ScopedStream ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
     std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
 
+    if (hQueue->getContext()->getDevices().size() > 1) {
+      for (auto i = 0u; i < NumMemsInMemList; ++i) {
+        enqueueMigrateMemoryToDeviceIfNeeded(phMemList[i], hQueue->getDevice(),
+                                             ActiveStream.getStream());
+        phMemList[i]->setLastQueueWritingToMemObj(hQueue);
+      }
+    }
+
     if (phEvent) {
       RetImplEvent =
           std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeNative(
               UR_COMMAND_ENQUEUE_NATIVE_EXP, hQueue, ActiveStream.getStream()));
       UR_CHECK_ERROR(RetImplEvent->start());
     }
 
-    if (hQueue->getContext()->getDevices().size() > 1) {
-      for (auto i = 0u; i < NumMemsInMemList; ++i) {
-        // FIXME: Update to enqueueMigrateMemory and also using
-        // setLastQueueWritingToMemObj when #1711 has merged
-        migrateMemoryToDeviceIfNeeded(phMemList[i], hQueue->getDevice());
-        phMemList[i]->setLastEventWritingToMemObj(RetImplEvent.get());
-      }
-    }
-
     pfnNativeEnqueue(hQueue, data); // This is using urQueueGetNativeHandle to
                                     // get the CUDA stream. It must be the
                                     // same stream as is used before and after

source/adapters/cuda/memory.hpp

@@ -17,8 +17,7 @@
 
 #include "common.hpp"
 #include "context.hpp"
-#include "device.hpp"
-#include "event.hpp"
+#include "queue.hpp"
 
 ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t,
                                            const ur_device_handle_t);
@@ -443,6 +442,3 @@ struct ur_mem_handle_t_ {
     }
   }
 };
-
-ur_result_t migrateMemoryToDeviceIfNeeded(ur_mem_handle_t,
-                                          const ur_device_handle_t);

source/adapters/hip/enqueue_native.cpp

@@ -12,11 +12,13 @@
 
 #include "context.hpp"
 #include "event.hpp"
+#include "memory.hpp"
 #include "queue.hpp"
 
 UR_APIEXPORT ur_result_t UR_APICALL urEnqueueNativeCommandExp(
     ur_queue_handle_t hQueue,
     ur_exp_enqueue_native_command_function_t pfnNativeEnqueue, void *data,
+    uint32_t NumMemsInMemList, const ur_mem_handle_t *phMemList,
     const ur_exp_enqueue_native_command_properties_t *,
     uint32_t NumEventsInWaitList, const ur_event_handle_t *phEventWaitList,
     ur_event_handle_t *phEvent) {
@@ -29,6 +31,14 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueNativeCommandExp(
     ScopedStream ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
     std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
 
+    if (hQueue->getContext()->getDevices().size() > 1) {
+      for (auto i = 0u; i < NumMemsInMemList; ++i) {
+        enqueueMigrateMemoryToDeviceIfNeeded(phMemList[i], hQueue->getDevice(),
+                                             ActiveStream.getStream());
+        phMemList[i]->setLastQueueWritingToMemObj(hQueue);
+      }
+    }
+
     if (phEvent) {
       RetImplEvent =
           std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeNative(

source/adapters/hip/memory.hpp

@@ -437,6 +437,3 @@ struct ur_mem_handle_t_ {
     }
   }
 };
-
-ur_result_t migrateMemoryToDeviceIfNeeded(ur_mem_handle_t,
-                                          const ur_device_handle_t);