Get device from queue, not event

If an event was constructed using some interop method it doesn't have an
associated queue. Implicit migration of buffers was working on the
assumption that each event has an associated device, which can be
obtained from the queue. This patch removes these assumptions.

Author: Hugh Delaney

source/adapters/hip/enqueue.cpp

@@ -20,7 +20,7 @@
 
 extern size_t imageElementByteSize(hipArray_Format ArrayFormat);
 
-ur_result_t enqueueEventsWait(ur_queue_handle_t, hipStream_t Stream,
+ur_result_t enqueueEventsWait(ur_queue_handle_t Queue, hipStream_t Stream,
                               uint32_t NumEventsInWaitList,
                               const ur_event_handle_t *EventWaitList) {
   if (!EventWaitList) {
@@ -29,8 +29,8 @@ ur_result_t enqueueEventsWait(ur_queue_handle_t, hipStream_t Stream,
   try {
     auto Result = forLatestEvents(
         EventWaitList, NumEventsInWaitList,
-        [Stream](ur_event_handle_t Event) -> ur_result_t {
-          ScopedContext Active(Event->getDevice());
+        [Stream, Queue](ur_event_handle_t Event) -> ur_result_t {
+          ScopedContext Active(Queue->getDevice());
           if (Event->isCompleted() || Event->getStream() == Stream) {
             return UR_RESULT_SUCCESS;
           } else {
@@ -218,8 +218,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferRead(
     // last queue to write to the MemBuffer, meaning we must perform the copy
     // from a different device
     if (hBuffer->LastEventWritingToMemObj &&
-        hBuffer->LastEventWritingToMemObj->getDevice() != hQueue->getDevice()) {
-      Device = hBuffer->LastEventWritingToMemObj->getDevice();
+        hBuffer->LastDeviceWritingToMemObj != hQueue->getDevice()) {
+      Device = hBuffer->LastDeviceWritingToMemObj;
       ScopedContext Active(Device);
       HIPStream = hipStream_t{0}; // Default stream for different device
       // We may have to wait for an event on another queue if it is the last
@@ -367,7 +367,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
         // if it has been written to
         if (phEvent && (MemArg.AccessFlags &
                         (UR_MEM_FLAG_READ_WRITE | UR_MEM_FLAG_WRITE_ONLY))) {
-          MemArg.Mem->setLastEventWritingToMemObj(RetImplEvent.get());
+          MemArg.Mem->setLastEventWritingToMemObj(RetImplEvent.get(),
+                                                  hQueue->getDevice());
         }
       }
       // We can release the MemoryMigrationMutexes now
@@ -584,8 +585,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferReadRect(
     // last queue to write to the MemBuffer, meaning we must perform the copy
     // from a different device
     if (hBuffer->LastEventWritingToMemObj &&
-        hBuffer->LastEventWritingToMemObj->getDevice() != hQueue->getDevice()) {
-      Device = hBuffer->LastEventWritingToMemObj->getDevice();
+        hBuffer->LastDeviceWritingToMemObj != hQueue->getDevice()) {
+      Device = hBuffer->LastDeviceWritingToMemObj;
       ScopedContext Active(Device);
       HIPStream = hipStream_t{0}; // Default stream for different device
       // We may have to wait for an event on another queue if it is the last
@@ -1017,8 +1018,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemImageRead(
     // last queue to write to the MemBuffer, meaning we must perform the copy
     // from a different device
     if (hImage->LastEventWritingToMemObj &&
-        hImage->LastEventWritingToMemObj->getDevice() != hQueue->getDevice()) {
-      Device = hImage->LastEventWritingToMemObj->getDevice();
+        hImage->LastDeviceWritingToMemObj != hQueue->getDevice()) {
+      Device = hImage->LastDeviceWritingToMemObj;
       ScopedContext Active(Device);
       HIPStream = hipStream_t{0}; // Default stream for different device
       // We may have to wait for an event on another queue if it is the last

source/adapters/hip/event.hpp

@@ -28,8 +28,6 @@ struct ur_event_handle_t_ {
 
   ur_queue_handle_t getQueue() const noexcept { return Queue; }
 
-  ur_device_handle_t getDevice() const noexcept { return Queue->getDevice(); }
-
   hipStream_t getStream() const noexcept { return Stream; }
 
   uint32_t getComputeStreamToken() const noexcept { return StreamToken; }

source/adapters/hip/memory.cpp

@@ -525,11 +525,10 @@ inline ur_result_t migrateBufferToDevice(ur_mem_handle_t Mem,
       UR_CHECK_ERROR(
           hipMemcpyHtoD(Buffer.getPtr(hDevice), Buffer.HostPtr, Buffer.Size));
     }
-  } else if (Mem->LastEventWritingToMemObj->getDevice() != hDevice) {
-    UR_CHECK_ERROR(
-        hipMemcpyDtoD(Buffer.getPtr(hDevice),
-                      Buffer.getPtr(Mem->LastEventWritingToMemObj->getDevice()),
-                      Buffer.Size));
+  } else if (Mem->LastDeviceWritingToMemObj != hDevice) {
+    UR_CHECK_ERROR(hipMemcpyDtoD(Buffer.getPtr(hDevice),
+                                 Buffer.getPtr(Mem->LastDeviceWritingToMemObj),
+                                 Buffer.Size));
   }
   return UR_RESULT_SUCCESS;
 }
@@ -577,22 +576,19 @@ inline ur_result_t migrateImageToDevice(ur_mem_handle_t Mem,
       CpyDesc3D.srcHost = Image.HostPtr;
       UR_CHECK_ERROR(hipDrvMemcpy3D(&CpyDesc3D));
     }
-  } else if (Mem->LastEventWritingToMemObj->getDevice() != hDevice) {
+  } else if (Mem->LastDeviceWritingToMemObj != hDevice) {
     if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE1D) {
       // FIXME: 1D memcpy from DtoD going through the host.
       UR_CHECK_ERROR(hipMemcpyAtoH(
-          Image.HostPtr,
-          Image.getArray(Mem->LastEventWritingToMemObj->getDevice()),
+          Image.HostPtr, Image.getArray(Mem->LastDeviceWritingToMemObj),
           0 /*srcOffset*/, ImageSizeBytes));
       UR_CHECK_ERROR(
           hipMemcpyHtoA(ImageArray, 0, Image.HostPtr, ImageSizeBytes));
     } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE2D) {
-      CpyDesc2D.srcArray =
-          Image.getArray(Mem->LastEventWritingToMemObj->getDevice());
+      CpyDesc2D.srcArray = Image.getArray(Mem->LastDeviceWritingToMemObj);
       UR_CHECK_ERROR(hipMemcpyParam2D(&CpyDesc2D));
     } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE3D) {
-      CpyDesc3D.srcArray =
-          Image.getArray(Mem->LastEventWritingToMemObj->getDevice());
+      CpyDesc3D.srcArray = Image.getArray(Mem->LastDeviceWritingToMemObj);
       UR_CHECK_ERROR(hipDrvMemcpy3D(&CpyDesc3D));
     }
   }

source/adapters/hip/memory.hpp

@@ -393,6 +393,10 @@ struct ur_mem_handle_t_ {
   // We should wait on this event prior to migrating memory across allocations
   // in this ur_mem_handle_t_
   ur_event_handle_t LastEventWritingToMemObj{nullptr};
+  // Since the event may not contain device info (if using interop, which
+  // doesn't take a queue) we should use this member var to keep track of which
+  // device has most recent view of data
+  ur_device_handle_t LastDeviceWritingToMemObj{nullptr};
 
   // Enumerates all possible types of accesses.
   enum access_mode_t { unknown, read_write, read_only, write_only };
@@ -487,18 +491,23 @@ struct ur_mem_handle_t_ {
 
   uint32_t getReferenceCount() const noexcept { return RefCount; }
 
-  void setLastEventWritingToMemObj(ur_event_handle_t NewEvent) {
+  void setLastEventWritingToMemObj(ur_event_handle_t NewEvent,
+                                   ur_device_handle_t RecentDevice) {
     assert(NewEvent && "Invalid event!");
     // This entry point should only ever be called when using multi device ctx
     assert(Context->Devices.size() > 1);
+    urEventRetain(NewEvent);
+    urDeviceRetain(RecentDevice);
     if (LastEventWritingToMemObj != nullptr) {
       urEventRelease(LastEventWritingToMemObj);
     }
-    urEventRetain(NewEvent);
+    if (LastDeviceWritingToMemObj != nullptr) {
+      urDeviceRelease(LastDeviceWritingToMemObj);
+    }
     LastEventWritingToMemObj = NewEvent;
     for (const auto &Device : Context->getDevices()) {
       HaveMigratedToDeviceSinceLastWrite[Device->getIndex()] =
-          Device == NewEvent->getDevice();
+          Device == RecentDevice;
     }
   }
 };