Adding command buffer queue affinity. (#5265)

This is needed for the CPU backend and vulkan tracing. It's a somewhat
unfortunate dependency issue and ideally we wouldn't need to track the
queue during recording however for one shot recordings we know at
compile-time what queue a stream will be submitted to. We can always
change the behavior to ignore/treat as suggestion the affinity when
recording reusable command buffers.

iree/hal/command_buffer.c

@@ -26,14 +26,15 @@ IREE_HAL_API_RETAIN_RELEASE(command_buffer);
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_command_buffer_create(
     iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(device);
   IREE_ASSERT_ARGUMENT(out_command_buffer);
   *out_command_buffer = NULL;
   IREE_TRACE_ZONE_BEGIN(z0);
   iree_status_t status =
       IREE_HAL_VTABLE_DISPATCH(device, iree_hal_device, create_command_buffer)(
-          device, mode, command_categories, out_command_buffer);
+          device, mode, command_categories, queue_affinity, out_command_buffer);
   IREE_TRACE_ZONE_END(z0);
   return status;
 }

iree/hal/command_buffer.h

@@ -64,6 +64,22 @@ enum iree_hal_command_category_e {
 };
 typedef uint32_t iree_hal_command_category_t;
 
+// A bitmask indicating affinity for a submission to use a particular set of
+// queues.
+//
+// Upon submission the queue is selected based on the flags set in
+// |command_categories| and the |queue_affinity|. As the number of available
+// queues can vary the |queue_affinity| is used to hash into the available
+// queues for the required categories. For example if 2 queues support transfer
+// commands and the affinity is 5 the resulting queue could be index hash(5)=1.
+// The affinity can thus be treated as just a way to indicate whether two
+// submissions must be placed on to the same queue. Note that the exact hashing
+// function is implementation dependent.
+typedef uint64_t iree_hal_queue_affinity_t;
+
+// Specifies that any queue may be selected.
+#define IREE_HAL_QUEUE_AFFINITY_ANY ((iree_hal_queue_affinity_t)(-1))
+
 // Bitfield specifying which execution stage a barrier should start/end at.
 //
 // Maps to VkPipelineStageFlagBits.
@@ -193,9 +209,14 @@ typedef struct iree_hal_command_buffer_s iree_hal_command_buffer_t;
 
 // Creates a command buffer ready to begin recording, possibly reusing an
 // existing one from the |device| pool.
+//
+// |queue_affinity| specifies the device queues the command buffer may be
+// submitted to. The queue affinity provided to iree_hal_device_queue_submit
+// must match or be a subset of the |queue_affinity|.
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_command_buffer_create(
     iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer);
 
 // Retains the given |command_buffer| for the caller.

iree/hal/cts/command_buffer_test.cc

@@ -41,7 +41,8 @@ TEST_P(CommandBufferTest, Create) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   EXPECT_TRUE((iree_hal_command_buffer_allowed_categories(command_buffer) &
                IREE_HAL_COMMAND_CATEGORY_DISPATCH) ==
@@ -54,7 +55,8 @@ TEST_P(CommandBufferTest, BeginEnd) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
@@ -66,7 +68,8 @@ TEST_P(CommandBufferTest, SubmitEmpty) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
@@ -81,7 +84,8 @@ TEST_P(CommandBufferTest, FillBufferWithRepeatedBytes) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_TRANSFER, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   iree_hal_buffer_t* device_buffer;
   IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer(
@@ -141,7 +145,8 @@ TEST_P(CommandBufferTest, CopyWholeBuffer) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_TRANSFER, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   // Create and fill a host buffer.
   iree_hal_buffer_t* host_buffer;
@@ -192,7 +197,8 @@ TEST_P(CommandBufferTest, CopySubBuffer) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_TRANSFER, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_TRANSFER, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   iree_hal_buffer_t* device_buffer;
   IREE_ASSERT_OK(iree_hal_allocator_allocate_buffer(

iree/hal/cts/event_test.cc

@@ -36,7 +36,8 @@ TEST_P(EventTest, SignalAndReset) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_signal_event(
@@ -60,10 +61,12 @@ TEST_P(EventTest, SubmitWithChainedCommandBuffers) {
   iree_hal_command_buffer_t* command_buffer_2;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer_1));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer_1));
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer_2));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer_2));
 
   // First command buffer signals the event when it completes.
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer_1));

iree/hal/cts/semaphore_submission_test.cc

@@ -57,7 +57,8 @@ TEST_P(SemaphoreSubmissionTest, SubmitAndSignal) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
@@ -94,7 +95,8 @@ TEST_P(SemaphoreSubmissionTest, SubmitWithWait) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));
 
@@ -142,7 +144,8 @@ TEST_P(SemaphoreSubmissionTest, SubmitWithMultipleSemaphores) {
   iree_hal_command_buffer_t* command_buffer;
   IREE_ASSERT_OK(iree_hal_command_buffer_create(
       device_, IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT,
-      IREE_HAL_COMMAND_CATEGORY_DISPATCH, &command_buffer));
+      IREE_HAL_COMMAND_CATEGORY_DISPATCH, IREE_HAL_QUEUE_AFFINITY_ANY,
+      &command_buffer));
 
   IREE_ASSERT_OK(iree_hal_command_buffer_begin(command_buffer));
   IREE_ASSERT_OK(iree_hal_command_buffer_end(command_buffer));

iree/hal/cuda/cuda_device.c

@@ -155,10 +155,12 @@ static iree_hal_allocator_t* iree_hal_cuda_device_allocator(
 static iree_status_t iree_hal_cuda_device_create_command_buffer(
     iree_hal_device_t* base_device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device);
   return iree_hal_cuda_graph_command_buffer_allocate(
-      &device->context_wrapper, mode, command_categories, out_command_buffer);
+      &device->context_wrapper, mode, command_categories, queue_affinity,
+      out_command_buffer);
 }
 
 static iree_status_t iree_hal_cuda_device_create_descriptor_set(
@@ -218,8 +220,9 @@ static iree_status_t iree_hal_cuda_device_create_semaphore(
 
 static iree_status_t iree_hal_cuda_device_queue_submit(
     iree_hal_device_t* base_device,
-    iree_hal_command_category_t command_categories, uint64_t queue_affinity,
-    iree_host_size_t batch_count, const iree_hal_submission_batch_t* batches) {
+    iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
+    const iree_hal_submission_batch_t* batches) {
   iree_hal_cuda_device_t* device = iree_hal_cuda_device_cast(base_device);
   for (int i = 0; i < batch_count; i++) {
     for (int j = 0; j < batches[i].command_buffer_count; j++) {

iree/hal/cuda/graph_command_buffer.c

@@ -28,6 +28,7 @@ typedef struct {
   iree_hal_cuda_context_wrapper_t* context;
   iree_hal_command_buffer_mode_t mode;
   iree_hal_command_category_t allowed_categories;
+  iree_hal_queue_affinity_t queue_affinity;
   CUgraph graph;
   CUgraphExec exec;
   // Keep track of the last node added to the command buffer as we are currently
@@ -52,6 +53,7 @@ iree_status_t iree_hal_cuda_graph_command_buffer_allocate(
     iree_hal_cuda_context_wrapper_t* context,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(context);
   IREE_ASSERT_ARGUMENT(out_command_buffer);
@@ -72,6 +74,7 @@ iree_status_t iree_hal_cuda_graph_command_buffer_allocate(
     command_buffer->context = context;
     command_buffer->mode = mode;
     command_buffer->allowed_categories = command_categories;
+    command_buffer->queue_affinity = queue_affinity;
     command_buffer->graph = graph;
     command_buffer->exec = NULL;
     command_buffer->last_node = NULL;

iree/hal/cuda/graph_command_buffer.h

@@ -29,6 +29,7 @@ iree_status_t iree_hal_cuda_graph_command_buffer_allocate(
     iree_hal_cuda_context_wrapper_t* context,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer);
 
 // Returns the native cuda graph associated to the command buffer.

iree/hal/device.c

@@ -42,7 +42,7 @@ iree_hal_device_allocator(iree_hal_device_t* device) {
 
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_device_queue_submit(
     iree_hal_device_t* device, iree_hal_command_category_t command_categories,
-    uint64_t queue_affinity, iree_host_size_t batch_count,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
     const iree_hal_submission_batch_t* batches) {
   IREE_ASSERT_ARGUMENT(device);
   IREE_ASSERT_ARGUMENT(!batch_count || batches);

iree/hal/device.h

@@ -171,7 +171,7 @@ iree_hal_device_allocator(iree_hal_device_t* device);
 // https://www.khronos.org/registry/vulkan/specs/1.2-extensions/man/html/vkQueueSubmit.html
 IREE_API_EXPORT iree_status_t IREE_API_CALL iree_hal_device_queue_submit(
     iree_hal_device_t* device, iree_hal_command_category_t command_categories,
-    uint64_t queue_affinity, iree_host_size_t batch_count,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
     const iree_hal_submission_batch_t* batches);
 
 // Blocks the caller until the semaphores reach or exceed the specified payload
@@ -242,6 +242,7 @@ typedef struct {
   iree_status_t(IREE_API_PTR* create_command_buffer)(
       iree_hal_device_t* device, iree_hal_command_buffer_mode_t mode,
       iree_hal_command_category_t command_categories,
+      iree_hal_queue_affinity_t queue_affinity,
       iree_hal_command_buffer_t** out_command_buffer);
 
   iree_status_t(IREE_API_PTR* create_descriptor_set)(
@@ -276,7 +277,7 @@ typedef struct {
 
   iree_status_t(IREE_API_PTR* queue_submit)(
       iree_hal_device_t* device, iree_hal_command_category_t command_categories,
-      uint64_t queue_affinity, iree_host_size_t batch_count,
+      iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
       const iree_hal_submission_batch_t* batches);
 
   iree_status_t(IREE_API_PTR* wait_semaphores_with_deadline)(

iree/hal/local/task_command_buffer.c

@@ -42,6 +42,7 @@ typedef struct {
   iree_task_scope_t* scope;
   iree_hal_command_buffer_mode_t mode;
   iree_hal_command_category_t allowed_categories;
+  iree_hal_queue_affinity_t queue_affinity;
 
   // Arena used for all allocations; references the shared device block pool.
   iree_arena_allocator_t arena;
@@ -108,6 +109,7 @@ iree_status_t iree_hal_task_command_buffer_create(
     iree_hal_device_t* device, iree_task_scope_t* scope,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_arena_block_pool_t* block_pool,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(device);
@@ -140,6 +142,7 @@ iree_status_t iree_hal_task_command_buffer_create(
     command_buffer->scope = scope;
     command_buffer->mode = mode;
     command_buffer->allowed_categories = command_categories;
+    command_buffer->queue_affinity = queue_affinity;
     iree_arena_initialize(block_pool, &command_buffer->arena);
     iree_task_list_initialize(&command_buffer->root_tasks);
     iree_task_list_initialize(&command_buffer->leaf_tasks);

iree/hal/local/task_command_buffer.h

@@ -29,6 +29,7 @@ iree_status_t iree_hal_task_command_buffer_create(
     iree_hal_device_t* device, iree_task_scope_t* scope,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_arena_block_pool_t* block_pool,
     iree_hal_command_buffer_t** out_command_buffer);
 

iree/hal/local/task_device.c

@@ -197,17 +197,31 @@ static iree_hal_allocator_t* iree_hal_task_device_allocator(
   return device->device_allocator;
 }
 
+// Returns the queue index to submit work to based on the |queue_affinity|.
+//
+// If we wanted to have dedicated transfer queues we'd fork off based on
+// command_categories. For now all queues are general purpose.
+static iree_host_size_t iree_hal_task_device_select_queue(
+    iree_hal_task_device_t* device,
+    iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity) {
+  // TODO(benvanik): evaluate if we want to obscure this mapping a bit so that
+  // affinity really means "equivalent affinities map to equivalent queues" and
+  // not a specific queue index.
+  return queue_affinity % device->queue_count;
+}
+
 static iree_status_t iree_hal_task_device_create_command_buffer(
     iree_hal_device_t* base_device, iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree_hal_command_buffer_t** out_command_buffer) {
   iree_hal_task_device_t* device = iree_hal_task_device_cast(base_device);
-  // TODO(benvanik): prevent the need for taking a scope here. We need it to
-  // construct the tasks as we record but unfortunately then that means we would
-  // need to know which queue we'd be submitting against ahead of time.
+  iree_host_size_t queue_index = iree_hal_task_device_select_queue(
+      device, command_categories, queue_affinity);
   return iree_hal_task_command_buffer_create(
-      base_device, &device->queues[0].scope, mode, command_categories,
-      &device->large_block_pool, out_command_buffer);
+      base_device, &device->queues[queue_index].scope, mode, command_categories,
+      queue_affinity, &device->large_block_pool, out_command_buffer);
 }
 
 static iree_status_t iree_hal_task_device_create_descriptor_set(
@@ -264,26 +278,14 @@ static iree_status_t iree_hal_task_device_create_semaphore(
                                         device->host_allocator, out_semaphore);
 }
 
-// Returns the queue index to submit work to based on the |queue_affinity|.
-//
-// If we wanted to have dedicated transfer queues we'd fork off based on
-// command_categories. For now all queues are general purpose.
-static iree_host_size_t iree_hal_device_select_queue(
-    iree_hal_task_device_t* device,
-    iree_hal_command_category_t command_categories, uint64_t queue_affinity) {
-  // TODO(benvanik): evaluate if we want to obscure this mapping a bit so that
-  // affinity really means "equivalent affinities map to equivalent queues" and
-  // not a specific queue index.
-  return queue_affinity % device->queue_count;
-}
-
 static iree_status_t iree_hal_task_device_queue_submit(
     iree_hal_device_t* base_device,
-    iree_hal_command_category_t command_categories, uint64_t queue_affinity,
-    iree_host_size_t batch_count, const iree_hal_submission_batch_t* batches) {
+    iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity, iree_host_size_t batch_count,
+    const iree_hal_submission_batch_t* batches) {
   iree_hal_task_device_t* device = iree_hal_task_device_cast(base_device);
-  iree_host_size_t queue_index =
-      iree_hal_device_select_queue(device, command_categories, queue_affinity);
+  iree_host_size_t queue_index = iree_hal_task_device_select_queue(
+      device, command_categories, queue_affinity);
   return iree_hal_task_queue_submit(&device->queues[queue_index], batch_count,
                                     batches);
 }

iree/hal/vulkan/direct_command_buffer.cc

@@ -35,6 +35,7 @@ typedef struct {
   VkDeviceHandle* logical_device;
   iree_hal_command_buffer_mode_t mode;
   iree_hal_command_category_t allowed_categories;
+  iree_hal_queue_affinity_t queue_affinity;
 
   VkCommandPoolHandle* command_pool;
   VkCommandBuffer handle;
@@ -66,6 +67,7 @@ iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     iree::hal::vulkan::VkCommandPoolHandle* command_pool,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
     iree_hal_command_buffer_t** out_command_buffer) {
   IREE_ASSERT_ARGUMENT(logical_device);
@@ -95,6 +97,7 @@ iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     command_buffer->logical_device = logical_device;
     command_buffer->mode = mode;
     command_buffer->allowed_categories = command_categories;
+    command_buffer->queue_affinity = queue_affinity;
     command_buffer->command_pool = command_pool;
     command_buffer->handle = handle;
     command_buffer->syms = logical_device->syms().get();

iree/hal/vulkan/direct_command_buffer.h

@@ -29,6 +29,7 @@ iree_status_t iree_hal_vulkan_direct_command_buffer_allocate(
     iree::hal::vulkan::VkCommandPoolHandle* command_pool,
     iree_hal_command_buffer_mode_t mode,
     iree_hal_command_category_t command_categories,
+    iree_hal_queue_affinity_t queue_affinity,
     iree::hal::vulkan::DescriptorPoolCache* descriptor_pool_cache,
     iree_hal_command_buffer_t** out_command_buffer);