Merge pull request #1553 from hdelan/get-device-from-queue

[HIP] Get device from queue, not event
oneapi-src · May 16, 2024 · 09467b8 · 09467b8
2 parents 1c1fa63 + 4cbf210
commit 09467b8
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 23 deletions.
diff --git a/source/adapters/hip/enqueue.cpp b/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,11 @@ 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->LastEventWritingToMemObj->getQueue()->getDevice() !=
+            hQueue->getDevice()) {
+      // This event is never created with interop so getQueue is never null
+      hQueue = hBuffer->LastEventWritingToMemObj->getQueue();
+      Device = hQueue->getDevice();
       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
@@ -584,8 +587,11 @@ 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->LastEventWritingToMemObj->getQueue()->getDevice() !=
+            hQueue->getDevice()) {
+      // This event is never created with interop so getQueue is never null
+      hQueue = hBuffer->LastEventWritingToMemObj->getQueue();
+      Device = hQueue->getDevice();
       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 +1023,10 @@ 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->LastEventWritingToMemObj->getQueue()->getDevice() !=
+            hQueue->getDevice()) {
+      hQueue = hImage->LastEventWritingToMemObj->getQueue();
+      Device = hQueue->getDevice();
       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

diff --git a/source/adapters/hip/event.hpp b/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; }

diff --git a/source/adapters/hip/memory.cpp b/source/adapters/hip/memory.cpp
@@ -525,11 +525,12 @@ 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->LastEventWritingToMemObj->getQueue()->getDevice() !=
+             hDevice) {
+    UR_CHECK_ERROR(hipMemcpyDtoD(
+        Buffer.getPtr(hDevice),
+        Buffer.getPtr(Mem->LastEventWritingToMemObj->getQueue()->getDevice()),
+        Buffer.Size));
   }
   return UR_RESULT_SUCCESS;
 }
@@ -577,22 +578,24 @@ 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->LastEventWritingToMemObj->getQueue()->getDevice() !=
+             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.getArray(
+              Mem->LastEventWritingToMemObj->getQueue()->getDevice()),
           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->LastEventWritingToMemObj->getQueue()->getDevice());
       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->LastEventWritingToMemObj->getQueue()->getDevice());
       UR_CHECK_ERROR(hipDrvMemcpy3D(&CpyDesc3D));
     }
   }

diff --git a/source/adapters/hip/memory.hpp b/source/adapters/hip/memory.hpp
@@ -498,7 +498,7 @@ struct ur_mem_handle_t_ {
     LastEventWritingToMemObj = NewEvent;
     for (const auto &Device : Context->getDevices()) {
       HaveMigratedToDeviceSinceLastWrite[Device->getIndex()] =
-          Device == NewEvent->getDevice();
+          Device == NewEvent->getQueue()->getDevice();
     }
   }
 };