[UR][Offload] Queue flag and out-of-order queue support (#19531)

Out of order support is similar to CUDA and HIP; there's a pool of
queues, and incoming tasks are allocated to the queues round-robin
style. In addition, the memcpy helper function now creates and destroys
an ol queue, since they are cheap to create and we don't need it to
hang around.

In addition, flags are now parsed by `urQueueCreate` and can be queried
with `urQueueGetInfo`. This determines whether the queue is in-order
or out-of-order.

Author: RossBrunton

unified-runtime/source/adapters/offload/device.cpp

@@ -121,6 +121,10 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
   case UR_DEVICE_INFO_USM_SYSTEM_SHARED_SUPPORT:
     return ReturnValue(uint32_t{0});
   case UR_DEVICE_INFO_QUEUE_PROPERTIES:
+  case UR_DEVICE_INFO_QUEUE_ON_HOST_PROPERTIES:
+  case UR_DEVICE_INFO_QUEUE_ON_DEVICE_PROPERTIES:
+    return ReturnValue(
+        ur_queue_flags_t{UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE});
   case UR_DEVICE_INFO_KERNEL_LAUNCH_CAPABILITIES:
     return ReturnValue(0);
   case UR_DEVICE_INFO_SUPPORTED_PARTITIONS: {

unified-runtime/source/adapters/offload/enqueue.cpp

@@ -68,10 +68,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
   LaunchArgs.DynSharedMemory = 0;
 
   ol_event_handle_t EventOut;
+  ol_queue_handle_t Queue;
+  OL_RETURN_ON_ERR(hQueue->nextQueue(Queue));
   OL_RETURN_ON_ERR(
-      olLaunchKernel(hQueue->OffloadQueue, hQueue->OffloadDevice,
-                     hKernel->OffloadKernel, hKernel->Args.getStorage(),
-                     hKernel->Args.getStorageSize(), &LaunchArgs, &EventOut));
+      olLaunchKernel(Queue, hQueue->OffloadDevice, hKernel->OffloadKernel,
+                     hKernel->Args.getStorage(), hKernel->Args.getStorageSize(),
+                     &LaunchArgs, &EventOut));
 
   if (phEvent) {
     auto *Event = new ur_event_handle_t_(UR_COMMAND_KERNEL_LAUNCH, hQueue);
@@ -105,15 +107,20 @@ ur_result_t doMemcpy(ur_command_t Command, ur_queue_handle_t hQueue,
   (void)phEventWaitList;
   //
 
-  ol_event_handle_t EventOut = nullptr;
-
-  OL_RETURN_ON_ERR(olMemcpy(hQueue->OffloadQueue, DestPtr, DestDevice, SrcPtr,
-                            SrcDevice, size, phEvent ? &EventOut : nullptr));
-
   if (blocking) {
-    OL_RETURN_ON_ERR(olSyncQueue(hQueue->OffloadQueue));
+    OL_RETURN_ON_ERR(olMemcpy(nullptr, DestPtr, DestDevice, SrcPtr, SrcDevice,
+                              size, nullptr));
+    if (phEvent) {
+      *phEvent = ur_event_handle_t_::createEmptyEvent(Command, hQueue);
+    }
+    return UR_RESULT_SUCCESS;
   }
 
+  ol_event_handle_t EventOut = nullptr;
+  ol_queue_handle_t Queue;
+  OL_RETURN_ON_ERR(hQueue->nextQueue(Queue));
+  OL_RETURN_ON_ERR(olMemcpy(Queue, DestPtr, DestDevice, SrcPtr, SrcDevice, size,
+                            phEvent ? &EventOut : nullptr));
   if (phEvent) {
     auto *Event = new ur_event_handle_t_(Command, hQueue);
     Event->OffloadEvent = EventOut;

unified-runtime/source/adapters/offload/queue.cpp

@@ -17,22 +17,30 @@
 #include "queue.hpp"
 #include "ur2offload.hpp"
 
-UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(ur_context_handle_t hContext,
-                                                  ur_device_handle_t hDevice,
-                                                  const ur_queue_properties_t *,
-                                                  ur_queue_handle_t *phQueue) {
+UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(
+    [[maybe_unused]] ur_context_handle_t hContext, ur_device_handle_t hDevice,
+    const ur_queue_properties_t *pProps, ur_queue_handle_t *phQueue) {
 
   assert(hContext->Device == hDevice);
 
-  ur_queue_handle_t Queue = new ur_queue_handle_t_();
-  auto Res = olCreateQueue(hDevice->OffloadDevice, &Queue->OffloadQueue);
-  if (Res != OL_SUCCESS) {
-    delete Queue;
-    return offloadResultToUR(Res);
+  ur_queue_flags_t URFlags = 0;
+  if (pProps && pProps->stype == UR_STRUCTURE_TYPE_QUEUE_PROPERTIES) {
+    URFlags = pProps->flags;
   }
 
-  Queue->OffloadDevice = hDevice->OffloadDevice;
-  Queue->UrContext = hContext;
+  ur_queue_handle_t Queue =
+      new ur_queue_handle_t_(hDevice->OffloadDevice, hContext, URFlags);
+
+  // For in-order queues, create the ol queue on construction so we can report
+  // any errors earlier
+  if (!(URFlags & UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE)) {
+    [[maybe_unused]] ol_queue_handle_t InitQueue;
+    auto Res = Queue->nextQueue(InitQueue);
+    if (Res != OL_SUCCESS) {
+      delete Queue;
+      return offloadResultToUR(Res);
+    }
+  }
 
   *phQueue = Queue;
 
@@ -47,6 +55,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueGetInfo(ur_queue_handle_t hQueue,
   UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
 
   switch (propName) {
+  case UR_QUEUE_INFO_FLAGS:
+    return ReturnValue(hQueue->Flags);
   case UR_QUEUE_INFO_REFERENCE_COUNT:
     return ReturnValue(hQueue->RefCount.load());
   default:
@@ -63,15 +73,26 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueRetain(ur_queue_handle_t hQueue) {
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueRelease(ur_queue_handle_t hQueue) {
   if (--hQueue->RefCount == 0) {
-    OL_RETURN_ON_ERR(olDestroyQueue(hQueue->OffloadQueue));
+    for (auto *Q : hQueue->OffloadQueues) {
+      if (!Q) {
+        break;
+      }
+      OL_RETURN_ON_ERR(olDestroyQueue(Q));
+    }
     delete hQueue;
   }
 
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueFinish(ur_queue_handle_t hQueue) {
-  return offloadResultToUR(olSyncQueue(hQueue->OffloadQueue));
+  for (auto *Q : hQueue->OffloadQueues) {
+    if (!Q) {
+      break;
+    }
+    OL_RETURN_ON_ERR(olSyncQueue(Q));
+  }
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueGetNativeHandle(

unified-runtime/source/adapters/offload/queue.hpp

@@ -15,8 +15,42 @@
 
 #include "common.hpp"
 
+constexpr size_t OOO_QUEUE_POOL_SIZE = 32;
+
 struct ur_queue_handle_t_ : RefCounted {
-  ol_queue_handle_t OffloadQueue;
+  ur_queue_handle_t_(ol_device_handle_t Device, ur_context_handle_t UrContext,
+                     ur_queue_flags_t Flags)
+      : OffloadQueues((Flags & UR_QUEUE_FLAG_OUT_OF_ORDER_EXEC_MODE_ENABLE)
+                          ? 1
+                          : OOO_QUEUE_POOL_SIZE),
+        QueueOffset(0), OffloadDevice(Device), UrContext(UrContext),
+        Flags(Flags) {}
+
+  // In-order queues only have one element here, while out of order queues have
+  // a bank of queues to use. We rotate through them round robin instead of
+  // constantly creating new ones in case there is a long-running program that
+  // never destroys the ur queue. Out-of-order queues create ol queues when
+  // needed; any queues that are not yet created are nullptr.
+  // This is a simpler implementation of the HIP/Cuda queue pooling logic in
+  // `stream_queue_t`. In the future, if we want more performance or it
+  // simplifies the implementation of a feature, we can consider using it.
+  std::vector<ol_queue_handle_t> OffloadQueues;
+  size_t QueueOffset;
   ol_device_handle_t OffloadDevice;
   ur_context_handle_t UrContext;
+  ur_queue_flags_t Flags;
+
+  ol_result_t nextQueue(ol_queue_handle_t &Handle) {
+    auto &Slot = OffloadQueues[QueueOffset++];
+    QueueOffset %= OffloadQueues.size();
+
+    if (!Slot) {
+      if (auto Res = olCreateQueue(OffloadDevice, &Slot)) {
+        return Res;
+      }
+    }
+
+    Handle = Slot;
+    return nullptr;
+  }
 };