experimental/cuda2/pending_queue_actions.c - 3p/openxla/iree - Git at Google

 // Copyright 2023 The IREE Authors
 //
 // Licensed under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

 #include "experimental/cuda2/pending_queue_actions.h"

 #include <stdbool.h>

 #include "experimental/cuda2/cuda_device.h"
 #include "experimental/cuda2/cuda_dynamic_symbols.h"
 #include "experimental/cuda2/cuda_status_util.h"
 #include "experimental/cuda2/event_semaphore.h"
 #include "experimental/cuda2/graph_command_buffer.h"
 #include "iree/base/api.h"
 #include "iree/base/internal/arena.h"
 #include "iree/base/internal/synchronization.h"
 #include "iree/hal/api.h"
 #include "iree/hal/utils/deferred_command_buffer.h"
 #include "iree/hal/utils/resource_set.h"

 //===----------------------------------------------------------------------===//
 // Queue action
 //===----------------------------------------------------------------------===//

 typedef enum iree_hal_cuda2_queue_action_kind_e {
   IREE_HAL_CUDA2_QUEUE_ACTION_TYPE_EXECUTION,
   // TODO: Add support for queue alloca and dealloca.
 } iree_hal_cuda2_queue_action_kind_t;

 // A pending queue action.
 //
 // Note that this struct does not have internal synchronization; it's expected
 // to work together with the pending action queue, which synchronizes accesses.
 typedef struct iree_hal_cuda2_queue_action_t {
   // Intrusive doubly-linked list next entry pointer.
   struct iree_hal_cuda2_queue_action_t* next;
   // Intrusive doubly-linked list previous entry pointer.
   struct iree_hal_cuda2_queue_action_t* prev;

   // The owning pending actions queue. We use its allocators and pools.
   // Retained to make sure it outlives the current action.
   iree_hal_cuda2_pending_queue_actions_t* owning_actions;

   iree_hal_cuda2_queue_action_kind_t kind;
   union {
     struct {
       iree_host_size_t count;
       iree_hal_command_buffer_t* const* ptr;
     } command_buffers;
   } payload;

   // The device from which to allocate CUDA stream-based command buffers for
   // applying deferred command buffers.
   iree_hal_device_t* device;

   // The stream to launch main GPU workload.
   CUstream dispatch_cu_stream;
   // The stream to launch CUDA host function callbacks.
   CUstream callback_cu_stream;

   // Resource set to retain all associated resources by the payload.
   iree_hal_resource_set_t* resource_set;

   // Semaphore list to wait on for the payload to start on the GPU.
   iree_hal_semaphore_list_t wait_semaphore_list;
   // Semaphore list to signal after the payload completes on the GPU.
   iree_hal_semaphore_list_t signal_semaphore_list;

   // Scratch fields for analyzing whether actions are ready to issue.
   CUevent* events;
   iree_host_size_t event_count;
   bool is_pending;
 } iree_hal_cuda2_queue_action_t;

 //===----------------------------------------------------------------------===//
 // Queue action list
 //===----------------------------------------------------------------------===//

 typedef struct iree_hal_cuda2_queue_action_list_t {
   iree_hal_cuda2_queue_action_t* head;
   iree_hal_cuda2_queue_action_t* tail;
 } iree_hal_cuda2_queue_action_list_t;

 // Returns true if the action list is empty.
 static inline bool iree_hal_cuda2_queue_action_list_is_empty(
     const iree_hal_cuda2_queue_action_list_t* list) {
   return list->head == NULL;
 }

 // Pushes |action| on to the end of the given action |list|.
 static void iree_hal_cuda2_queue_action_list_push_back(
     iree_hal_cuda2_queue_action_list_t* list,
     iree_hal_cuda2_queue_action_t* action) {
   if (list->tail) {
     list->tail->next = action;
   } else {
     list->head = action;
   }
   action->next = NULL;
   action->prev = list->tail;
   list->tail = action;
 }

 // Erases |action| from |list|.
 static void iree_hal_cuda2_queue_action_list_erase(
     iree_hal_cuda2_queue_action_list_t* list,
     iree_hal_cuda2_queue_action_t* action) {
   iree_hal_cuda2_queue_action_t* next = action->next;
   iree_hal_cuda2_queue_action_t* prev = action->prev;
   if (prev) {
     prev->next = next;
     action->prev = NULL;
   } else {
     list->head = next;
   }
   if (next) {
     next->prev = prev;
     action->next = NULL;
   } else {
     list->tail = prev;
   }
 }

 // Takes all actions from |available_list| and moves them into |ready_list|.
 static void iree_hal_cuda2_queue_action_list_take_all(
     iree_hal_cuda2_queue_action_list_t* available_list,
     iree_hal_cuda2_queue_action_list_t* ready_list) {
   IREE_ASSERT(available_list != ready_list);
   ready_list->head = available_list->head;
   ready_list->tail = available_list->tail;
   available_list->head = NULL;
   available_list->tail = NULL;
 }

 //===----------------------------------------------------------------------===//
 // Pending queue actions
 //===----------------------------------------------------------------------===//

 struct iree_hal_cuda2_pending_queue_actions_t {
   // Abstract resource used for injecting reference counting and vtable;
   // must be at offset 0.
   iree_hal_resource_t resource;

   // The allocator used to create the timepoint pool.
   iree_allocator_t host_allocator;
   // The block pool to allocate resource sets from.
   iree_arena_block_pool_t* block_pool;

   // The symbols used to create and destroy CUevent objects.
   const iree_hal_cuda2_dynamic_symbols_t* symbols;

   // Non-recursive mutex guarding access to the action list.
   iree_slim_mutex_t action_mutex;

   // The double-linked list of pending actions.
   iree_hal_cuda2_queue_action_list_t action_list IREE_GUARDED_BY(action_mutex);
 };

 static const iree_hal_resource_vtable_t
     iree_hal_cuda2_pending_queue_actions_vtable;

 iree_status_t iree_hal_cuda2_pending_queue_actions_create(
     const iree_hal_cuda2_dynamic_symbols_t* symbols,
     iree_arena_block_pool_t* block_pool, iree_allocator_t host_allocator,
     iree_hal_cuda2_pending_queue_actions_t** out_actions) {
   IREE_ASSERT_ARGUMENT(symbols);
   IREE_ASSERT_ARGUMENT(block_pool);
   IREE_ASSERT_ARGUMENT(out_actions);
   IREE_TRACE_ZONE_BEGIN(z0);

   iree_hal_cuda2_pending_queue_actions_t* actions = NULL;
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0, iree_allocator_malloc(host_allocator, sizeof(*actions),
                                 (void**)&actions));
   iree_hal_resource_initialize(&iree_hal_cuda2_pending_queue_actions_vtable,
                                &actions->resource);
   actions->host_allocator = host_allocator;
   actions->block_pool = block_pool;
   actions->symbols = symbols;
   iree_slim_mutex_initialize(&actions->action_mutex);
   memset(&actions->action_list, 0, sizeof(actions->action_list));
   *out_actions = actions;

   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
 }

 static iree_hal_cuda2_pending_queue_actions_t*
 iree_hal_cuda2_pending_queue_actions_cast(iree_hal_resource_t* base_value) {
   return (iree_hal_cuda2_pending_queue_actions_t*)base_value;
 }

 void iree_hal_cuda2_pending_queue_actions_destroy(
     iree_hal_resource_t* base_actions) {
   iree_hal_cuda2_pending_queue_actions_t* actions =
       iree_hal_cuda2_pending_queue_actions_cast(base_actions);
   iree_allocator_t host_allocator = actions->host_allocator;
   IREE_TRACE_ZONE_BEGIN(z0);

   IREE_ASSERT(iree_hal_cuda2_queue_action_list_is_empty(&actions->action_list));

   iree_slim_mutex_deinitialize(&actions->action_mutex);
   iree_allocator_free(host_allocator, actions);

   IREE_TRACE_ZONE_END(z0);
 }

 static const iree_hal_resource_vtable_t
     iree_hal_cuda2_pending_queue_actions_vtable = {
         .destroy = iree_hal_cuda2_pending_queue_actions_destroy,
 };

 // Performs copy of the given |in_list| to |out_list| to retain the semaphore
 // and value list.
 static iree_status_t iree_hal_cuda2_copy_semaphore_list(
     iree_hal_semaphore_list_t in_list, iree_allocator_t host_allocator,
     iree_hal_semaphore_list_t* out_list) {
   if (in_list.count == 0) {
     memset(out_list, 0, sizeof(*out_list));
   } else {
     out_list->count = in_list.count;

     iree_host_size_t semaphore_size =
         in_list.count * sizeof(*in_list.semaphores);
     IREE_RETURN_IF_ERROR(iree_allocator_malloc(host_allocator, semaphore_size,
                                                (void**)&out_list->semaphores));
     memcpy(out_list->semaphores, in_list.semaphores, semaphore_size);

     iree_host_size_t value_size =
         in_list.count * sizeof(*in_list.payload_values);
     IREE_RETURN_IF_ERROR(iree_allocator_malloc(
         host_allocator, value_size, (void**)&out_list->payload_values));
     memcpy(out_list->payload_values, in_list.payload_values, value_size);
   }
   return iree_ok_status();
 }

 // Frees the semaphore and value list inside |semaphore_list|.
 static void iree_hal_cuda2_free_semaphore_list(
     iree_allocator_t host_allocator,
     iree_hal_semaphore_list_t* semaphore_list) {
   iree_allocator_free(host_allocator, semaphore_list->semaphores);
   iree_allocator_free(host_allocator, semaphore_list->payload_values);
 }

 iree_status_t iree_hal_cuda2_pending_queue_actions_enqueue_execution(
     iree_hal_device_t* device, CUstream dispatch_stream,
     CUstream callback_stream, iree_hal_cuda2_pending_queue_actions_t* actions,
     const iree_hal_semaphore_list_t wait_semaphore_list,
     const iree_hal_semaphore_list_t signal_semaphore_list,
     iree_host_size_t command_buffer_count,
     iree_hal_command_buffer_t* const* command_buffers) {
   IREE_ASSERT_ARGUMENT(actions);
   IREE_ASSERT_ARGUMENT(command_buffer_count == 0 || command_buffers);
   IREE_TRACE_ZONE_BEGIN(z0);

   iree_hal_cuda2_queue_action_t* action = NULL;
   IREE_RETURN_AND_END_ZONE_IF_ERROR(
       z0, iree_allocator_malloc(actions->host_allocator, sizeof(*action),
                                 (void**)&action));

   action->kind = IREE_HAL_CUDA2_QUEUE_ACTION_TYPE_EXECUTION;
   action->payload.command_buffers.count = command_buffer_count;
   action->payload.command_buffers.ptr = command_buffers;
   action->device = device;
   action->dispatch_cu_stream = dispatch_stream;
   action->callback_cu_stream = callback_stream;
   action->events = NULL;
   action->event_count = 0;
   action->is_pending = true;

   // Retain all command buffers and semaphores.
   iree_hal_resource_set_t* resource_set = NULL;
   iree_status_t status =
       iree_hal_resource_set_allocate(actions->block_pool, &resource_set);
   if (IREE_LIKELY(iree_status_is_ok(status))) {
     status = iree_hal_resource_set_insert(resource_set, command_buffer_count,
                                           command_buffers);
   }
   if (IREE_LIKELY(iree_status_is_ok(status))) {
     status =
         iree_hal_resource_set_insert(resource_set, wait_semaphore_list.count,
                                      wait_semaphore_list.semaphores);
   }
   if (IREE_LIKELY(iree_status_is_ok(status))) {
     status =
         iree_hal_resource_set_insert(resource_set, signal_semaphore_list.count,
                                      signal_semaphore_list.semaphores);
   }

   // Copy the semaphore and value list for later access.
   // TODO: avoid host allocator malloc; use some pool for the allocation.
   if (IREE_LIKELY(iree_status_is_ok(status))) {
     status = iree_hal_cuda2_copy_semaphore_list(wait_semaphore_list,
                                                 actions->host_allocator,
                                                 &action->wait_semaphore_list);
   }
   if (IREE_LIKELY(iree_status_is_ok(status))) {
     status = iree_hal_cuda2_copy_semaphore_list(signal_semaphore_list,
                                                 actions->host_allocator,
                                                 &action->signal_semaphore_list);
   }

   if (IREE_LIKELY(iree_status_is_ok(status))) {
     action->owning_actions = actions;
     iree_hal_resource_retain(actions);

     action->resource_set = resource_set;

     iree_slim_mutex_lock(&actions->action_mutex);
     iree_hal_cuda2_queue_action_list_push_back(&actions->action_list, action);
     iree_slim_mutex_unlock(&actions->action_mutex);
   } else {
     iree_hal_cuda2_free_semaphore_list(actions->host_allocator,
                                        &action->wait_semaphore_list);
     iree_hal_cuda2_free_semaphore_list(actions->host_allocator,
                                        &action->signal_semaphore_list);
     iree_hal_resource_set_free(resource_set);
     iree_allocator_free(actions->host_allocator, action);
   }

   IREE_TRACE_ZONE_END(z0);
   return status;
 }

 static void iree_hal_cuda2_pending_queue_actions_cleanup_execution(
     iree_hal_cuda2_queue_action_t* action);

 // Releases resources after action completion on the GPU and advances timeline
 // and pending actions queue.
 //
 // This is the CUDA host function callback to cudaLaunchHostFunc, invoked by a
 // CUDA driver thread.
 static void iree_hal_cuda2_execution_device_signal_host_callback(
     void* user_data) {
   iree_hal_cuda2_queue_action_t* action =
       (iree_hal_cuda2_queue_action_t*)user_data;
   iree_hal_cuda2_pending_queue_actions_t* actions = action->owning_actions;
   // Advance semaphore timelines by calling into the host signaling function.
   IREE_IGNORE_ERROR(
       iree_hal_semaphore_list_signal(action->signal_semaphore_list));
   // Destroy the current action given its done now--this also frees all retained
   // resources.
   iree_hal_cuda2_pending_queue_actions_cleanup_execution(action);
   // Try to release more pending actions to the GPU now.
   IREE_IGNORE_ERROR(iree_hal_cuda2_pending_queue_actions_issue(actions));
 }

 // Issues the given kernel dispatch |action| to the GPU.
 static iree_status_t iree_hal_cuda2_pending_queue_actions_issue_execution(
     iree_hal_cuda2_queue_action_t* action) {
   IREE_ASSERT(action->events != NULL);
   IREE_ASSERT(action->is_pending == false);
   const iree_hal_cuda2_dynamic_symbols_t* symbols =
       action->owning_actions->symbols;
   IREE_TRACE_ZONE_BEGIN(z0);

   // No need to lock given that this action is already detched from the pending
   // actions list; so only this thread is seeing it now.

   // First wait all the device CUevent in the dispatch stream.
   for (iree_host_size_t i = 0; i < action->event_count; ++i) {
     IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
         z0, symbols,
         cuStreamWaitEvent(action->dispatch_cu_stream, action->events[i],
                           CU_EVENT_WAIT_DEFAULT),
         "cuStreamWaitEvent");
   }

   // Then launch all command buffers to the dispatch stream.
   for (iree_host_size_t i = 0; i < action->payload.command_buffers.count; ++i) {
     iree_hal_command_buffer_t* command_buffer =
         action->payload.command_buffers.ptr[i];
     if (iree_hal_cuda2_graph_command_buffer_isa(command_buffer)) {
       CUgraphExec exec = iree_hal_cuda2_graph_command_buffer_handle(
           action->payload.command_buffers.ptr[i]);
       IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
           z0, symbols, cuGraphLaunch(exec, action->dispatch_cu_stream),
           "cuGraphLaunch");
     } else {
       iree_hal_command_buffer_t* stream_command_buffer = NULL;
       iree_hal_command_buffer_mode_t mode =
           IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT |
           IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION |
           IREE_HAL_COMMAND_BUFFER_MODE_UNVALIDATED;
       IREE_RETURN_AND_END_ZONE_IF_ERROR(
           z0, iree_hal_cuda2_device_create_stream_command_buffer(
                   action->device, mode, IREE_HAL_COMMAND_CATEGORY_ANY,
                   /*binding_capacity=*/0, &stream_command_buffer));
       IREE_RETURN_AND_END_ZONE_IF_ERROR(
           z0, iree_hal_resource_set_insert(action->resource_set, 1,
                                            &stream_command_buffer));
       IREE_RETURN_AND_END_ZONE_IF_ERROR(
           z0, iree_hal_deferred_command_buffer_apply(
                   command_buffer, stream_command_buffer,
                   iree_hal_buffer_binding_table_empty()));
     }
   }

   // Last record CUevent signals in the dispatch stream.
   for (iree_host_size_t i = 0; i < action->signal_semaphore_list.count; ++i) {
     // Grab a CUevent for this semaphore value signaling.
     CUevent event = NULL;
     IREE_RETURN_AND_END_ZONE_IF_ERROR(
         z0, iree_hal_cuda2_event_semaphore_acquire_timepoint_device_signal(
                 action->signal_semaphore_list.semaphores[i],
                 action->signal_semaphore_list.payload_values[i], &event));

     // Record the event signaling in the dispatch stream.
     IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
         z0, symbols, cuEventRecord(event, action->dispatch_cu_stream),
         "cuEventRecord");
     // Let the callback stream to wait on the CUevent.
     IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
         z0, symbols,
         cuStreamWaitEvent(action->callback_cu_stream, event,
                           CU_EVENT_WAIT_DEFAULT),
         "cuStreamWaitEvent");
   }

   // Now launch a host function on the callback stream to advance the semaphore
   // timeline.
   IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
       z0, symbols,
       cuLaunchHostFunc(action->callback_cu_stream,
                        iree_hal_cuda2_execution_device_signal_host_callback,
                        action),
       "cuStreamWaitEvent");

   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
 }

 // Releases resources after completing the given kernel dispatch |action|.
 static void iree_hal_cuda2_pending_queue_actions_cleanup_execution(
     iree_hal_cuda2_queue_action_t* action) {
   iree_hal_cuda2_pending_queue_actions_t* actions = action->owning_actions;
   iree_allocator_t host_allocator = actions->host_allocator;
   IREE_TRACE_ZONE_BEGIN(z0);

   iree_hal_resource_set_free(action->resource_set);
   iree_hal_cuda2_free_semaphore_list(host_allocator,
                                      &action->wait_semaphore_list);
   iree_hal_cuda2_free_semaphore_list(host_allocator,
                                      &action->signal_semaphore_list);
   iree_hal_resource_release(actions);

   iree_allocator_free(host_allocator, action);

   IREE_TRACE_ZONE_END(z0);
 }

 iree_status_t iree_hal_cuda2_pending_queue_actions_issue(
     iree_hal_cuda2_pending_queue_actions_t* actions) {
   IREE_TRACE_ZONE_BEGIN(z0);

   iree_hal_cuda2_queue_action_list_t pending_list = {NULL, NULL};
   iree_hal_cuda2_queue_action_list_t ready_list = {NULL, NULL};

   iree_slim_mutex_lock(&actions->action_mutex);

   if (iree_hal_cuda2_queue_action_list_is_empty(&actions->action_list)) {
     iree_slim_mutex_unlock(&actions->action_mutex);
     IREE_TRACE_ZONE_END(z0);
     return iree_ok_status();
   }

   // Scan through the list and categorize actions into pending and ready lists.
   for (iree_hal_cuda2_queue_action_t* action = actions->action_list.head;
        action != NULL;) {
     iree_hal_cuda2_queue_action_t* next_action = action->next;
     action->next = NULL;

     iree_host_size_t semaphore_count = action->wait_semaphore_list.count;
     iree_hal_semaphore_t** semaphores = action->wait_semaphore_list.semaphores;
     uint64_t* values = action->wait_semaphore_list.payload_values;

     // We are allocating stack space here, assuming that there won't be a lot of
     // waits and additional references to this field happens in a function call
     // from this function.
     action->events = iree_alloca(semaphore_count * sizeof(CUevent));
     action->event_count = 0;
     action->is_pending = false;

     // Look at all wait semaphores.
     for (iree_host_size_t i = 0; i < semaphore_count; ++i) {
       // If this semaphore has already signaled past the desired value, we can
       // just ignore it.
       uint64_t value = 0;
       IREE_RETURN_AND_END_ZONE_IF_ERROR(
           z0, iree_hal_semaphore_query(semaphores[i], &value));
       if (value >= values[i]) continue;

       // Try to acquire a CUevent from a device wait timepoint. If so, we can
       // use that CUevent to wait on the device. Otherwise, this action is still
       // not ready.
       CUevent event = NULL;
       IREE_RETURN_AND_END_ZONE_IF_ERROR(
           z0, iree_hal_cuda2_event_semaphore_acquire_timepoint_device_wait(
                   semaphores[i], values[i], &event));
       if (event) {
         action->events[action->event_count++] = event;
       } else {
         // Clear the scratch fields.
         action->events = NULL;
         action->event_count = 0;
         action->is_pending = true;
         break;
       }
     }

     if (action->is_pending) {
       iree_hal_cuda2_queue_action_list_push_back(&pending_list, action);
     } else {
       iree_hal_cuda2_queue_action_list_push_back(&ready_list, action);
     }

     action = next_action;
   }

   // Preserve pending timepoints.
   actions->action_list = pending_list;

   iree_slim_mutex_unlock(&actions->action_mutex);

   // Now go through the ready list and issue the actions to the GPU.
   for (iree_hal_cuda2_queue_action_t* action = ready_list.head;
        action != NULL;) {
     iree_hal_cuda2_queue_action_t* next_action = action->next;
     action->next = NULL;

     IREE_RETURN_AND_END_ZONE_IF_ERROR(
         z0, iree_hal_cuda2_pending_queue_actions_issue_execution(action));
     action->events = NULL;
     action->event_count = 0;

     action = next_action;
   }

   IREE_TRACE_ZONE_END(z0);
   return iree_ok_status();
 }
	// Copyright 2023 The IREE Authors
	//
	// Licensed under the Apache License v2.0 with LLVM Exceptions.
	// See https://llvm.org/LICENSE.txt for license information.
	// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception

	#include "experimental/cuda2/pending_queue_actions.h"

	#include <stdbool.h>

	#include "experimental/cuda2/cuda_device.h"
	#include "experimental/cuda2/cuda_dynamic_symbols.h"
	#include "experimental/cuda2/cuda_status_util.h"
	#include "experimental/cuda2/event_semaphore.h"
	#include "experimental/cuda2/graph_command_buffer.h"
	#include "iree/base/api.h"
	#include "iree/base/internal/arena.h"
	#include "iree/base/internal/synchronization.h"
	#include "iree/hal/api.h"
	#include "iree/hal/utils/deferred_command_buffer.h"
	#include "iree/hal/utils/resource_set.h"

	//===----------------------------------------------------------------------===//
	// Queue action
	//===----------------------------------------------------------------------===//

	typedef enum iree_hal_cuda2_queue_action_kind_e {
	IREE_HAL_CUDA2_QUEUE_ACTION_TYPE_EXECUTION,
	// TODO: Add support for queue alloca and dealloca.
	} iree_hal_cuda2_queue_action_kind_t;

	// A pending queue action.
	//
	// Note that this struct does not have internal synchronization; it's expected
	// to work together with the pending action queue, which synchronizes accesses.
	typedef struct iree_hal_cuda2_queue_action_t {
	// Intrusive doubly-linked list next entry pointer.
	struct iree_hal_cuda2_queue_action_t* next;
	// Intrusive doubly-linked list previous entry pointer.
	struct iree_hal_cuda2_queue_action_t* prev;

	// The owning pending actions queue. We use its allocators and pools.
	// Retained to make sure it outlives the current action.
	iree_hal_cuda2_pending_queue_actions_t* owning_actions;

	iree_hal_cuda2_queue_action_kind_t kind;
	union {
	struct {
	iree_host_size_t count;
	iree_hal_command_buffer_t* const* ptr;
	} command_buffers;
	} payload;

	// The device from which to allocate CUDA stream-based command buffers for
	// applying deferred command buffers.
	iree_hal_device_t* device;

	// The stream to launch main GPU workload.
	CUstream dispatch_cu_stream;
	// The stream to launch CUDA host function callbacks.
	CUstream callback_cu_stream;

	// Resource set to retain all associated resources by the payload.
	iree_hal_resource_set_t* resource_set;

	// Semaphore list to wait on for the payload to start on the GPU.
	iree_hal_semaphore_list_t wait_semaphore_list;
	// Semaphore list to signal after the payload completes on the GPU.
	iree_hal_semaphore_list_t signal_semaphore_list;

	// Scratch fields for analyzing whether actions are ready to issue.
	CUevent* events;
	iree_host_size_t event_count;
	bool is_pending;
	} iree_hal_cuda2_queue_action_t;

	//===----------------------------------------------------------------------===//
	// Queue action list
	//===----------------------------------------------------------------------===//

	typedef struct iree_hal_cuda2_queue_action_list_t {
	iree_hal_cuda2_queue_action_t* head;
	iree_hal_cuda2_queue_action_t* tail;
	} iree_hal_cuda2_queue_action_list_t;

	// Returns true if the action list is empty.
	static inline bool iree_hal_cuda2_queue_action_list_is_empty(
	const iree_hal_cuda2_queue_action_list_t* list) {
	return list->head == NULL;
	}

	// Pushes \|action\| on to the end of the given action \|list\|.
	static void iree_hal_cuda2_queue_action_list_push_back(
	iree_hal_cuda2_queue_action_list_t* list,
	iree_hal_cuda2_queue_action_t* action) {
	if (list->tail) {
	list->tail->next = action;
	} else {
	list->head = action;
	}
	action->next = NULL;
	action->prev = list->tail;
	list->tail = action;
	}

	// Erases \|action\| from \|list\|.
	static void iree_hal_cuda2_queue_action_list_erase(
	iree_hal_cuda2_queue_action_list_t* list,
	iree_hal_cuda2_queue_action_t* action) {
	iree_hal_cuda2_queue_action_t* next = action->next;
	iree_hal_cuda2_queue_action_t* prev = action->prev;
	if (prev) {
	prev->next = next;
	action->prev = NULL;
	} else {
	list->head = next;
	}
	if (next) {
	next->prev = prev;
	action->next = NULL;
	} else {
	list->tail = prev;
	}
	}

	// Takes all actions from \|available_list\| and moves them into \|ready_list\|.
	static void iree_hal_cuda2_queue_action_list_take_all(
	iree_hal_cuda2_queue_action_list_t* available_list,
	iree_hal_cuda2_queue_action_list_t* ready_list) {
	IREE_ASSERT(available_list != ready_list);
	ready_list->head = available_list->head;
	ready_list->tail = available_list->tail;
	available_list->head = NULL;
	available_list->tail = NULL;
	}

	//===----------------------------------------------------------------------===//
	// Pending queue actions
	//===----------------------------------------------------------------------===//

	struct iree_hal_cuda2_pending_queue_actions_t {
	// Abstract resource used for injecting reference counting and vtable;
	// must be at offset 0.
	iree_hal_resource_t resource;

	// The allocator used to create the timepoint pool.
	iree_allocator_t host_allocator;
	// The block pool to allocate resource sets from.
	iree_arena_block_pool_t* block_pool;

	// The symbols used to create and destroy CUevent objects.
	const iree_hal_cuda2_dynamic_symbols_t* symbols;

	// Non-recursive mutex guarding access to the action list.
	iree_slim_mutex_t action_mutex;

	// The double-linked list of pending actions.
	iree_hal_cuda2_queue_action_list_t action_list IREE_GUARDED_BY(action_mutex);
	};

	static const iree_hal_resource_vtable_t
	iree_hal_cuda2_pending_queue_actions_vtable;

	iree_status_t iree_hal_cuda2_pending_queue_actions_create(
	const iree_hal_cuda2_dynamic_symbols_t* symbols,
	iree_arena_block_pool_t* block_pool, iree_allocator_t host_allocator,
	iree_hal_cuda2_pending_queue_actions_t** out_actions) {
	IREE_ASSERT_ARGUMENT(symbols);
	IREE_ASSERT_ARGUMENT(block_pool);
	IREE_ASSERT_ARGUMENT(out_actions);
	IREE_TRACE_ZONE_BEGIN(z0);

	iree_hal_cuda2_pending_queue_actions_t* actions = NULL;
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_allocator_malloc(host_allocator, sizeof(*actions),
	(void**)&actions));
	iree_hal_resource_initialize(&iree_hal_cuda2_pending_queue_actions_vtable,
	&actions->resource);
	actions->host_allocator = host_allocator;
	actions->block_pool = block_pool;
	actions->symbols = symbols;
	iree_slim_mutex_initialize(&actions->action_mutex);
	memset(&actions->action_list, 0, sizeof(actions->action_list));
	*out_actions = actions;

	IREE_TRACE_ZONE_END(z0);
	return iree_ok_status();
	}

	static iree_hal_cuda2_pending_queue_actions_t*
	iree_hal_cuda2_pending_queue_actions_cast(iree_hal_resource_t* base_value) {
	return (iree_hal_cuda2_pending_queue_actions_t*)base_value;
	}

	void iree_hal_cuda2_pending_queue_actions_destroy(
	iree_hal_resource_t* base_actions) {
	iree_hal_cuda2_pending_queue_actions_t* actions =
	iree_hal_cuda2_pending_queue_actions_cast(base_actions);
	iree_allocator_t host_allocator = actions->host_allocator;
	IREE_TRACE_ZONE_BEGIN(z0);

	IREE_ASSERT(iree_hal_cuda2_queue_action_list_is_empty(&actions->action_list));

	iree_slim_mutex_deinitialize(&actions->action_mutex);
	iree_allocator_free(host_allocator, actions);

	IREE_TRACE_ZONE_END(z0);
	}

	static const iree_hal_resource_vtable_t
	iree_hal_cuda2_pending_queue_actions_vtable = {
	.destroy = iree_hal_cuda2_pending_queue_actions_destroy,
	};

	// Performs copy of the given \|in_list\| to \|out_list\| to retain the semaphore
	// and value list.
	static iree_status_t iree_hal_cuda2_copy_semaphore_list(
	iree_hal_semaphore_list_t in_list, iree_allocator_t host_allocator,
	iree_hal_semaphore_list_t* out_list) {
	if (in_list.count == 0) {
	memset(out_list, 0, sizeof(*out_list));
	} else {
	out_list->count = in_list.count;

	iree_host_size_t semaphore_size =
	in_list.count * sizeof(*in_list.semaphores);
	IREE_RETURN_IF_ERROR(iree_allocator_malloc(host_allocator, semaphore_size,
	(void**)&out_list->semaphores));
	memcpy(out_list->semaphores, in_list.semaphores, semaphore_size);

	iree_host_size_t value_size =
	in_list.count * sizeof(*in_list.payload_values);
	IREE_RETURN_IF_ERROR(iree_allocator_malloc(
	host_allocator, value_size, (void**)&out_list->payload_values));
	memcpy(out_list->payload_values, in_list.payload_values, value_size);
	}
	return iree_ok_status();
	}

	// Frees the semaphore and value list inside \|semaphore_list\|.
	static void iree_hal_cuda2_free_semaphore_list(
	iree_allocator_t host_allocator,
	iree_hal_semaphore_list_t* semaphore_list) {
	iree_allocator_free(host_allocator, semaphore_list->semaphores);
	iree_allocator_free(host_allocator, semaphore_list->payload_values);
	}

	iree_status_t iree_hal_cuda2_pending_queue_actions_enqueue_execution(
	iree_hal_device_t* device, CUstream dispatch_stream,
	CUstream callback_stream, iree_hal_cuda2_pending_queue_actions_t* actions,
	const iree_hal_semaphore_list_t wait_semaphore_list,
	const iree_hal_semaphore_list_t signal_semaphore_list,
	iree_host_size_t command_buffer_count,
	iree_hal_command_buffer_t* const* command_buffers) {
	IREE_ASSERT_ARGUMENT(actions);
	IREE_ASSERT_ARGUMENT(command_buffer_count == 0 \|\| command_buffers);
	IREE_TRACE_ZONE_BEGIN(z0);

	iree_hal_cuda2_queue_action_t* action = NULL;
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_allocator_malloc(actions->host_allocator, sizeof(*action),
	(void**)&action));

	action->kind = IREE_HAL_CUDA2_QUEUE_ACTION_TYPE_EXECUTION;
	action->payload.command_buffers.count = command_buffer_count;
	action->payload.command_buffers.ptr = command_buffers;
	action->device = device;
	action->dispatch_cu_stream = dispatch_stream;
	action->callback_cu_stream = callback_stream;
	action->events = NULL;
	action->event_count = 0;
	action->is_pending = true;

	// Retain all command buffers and semaphores.
	iree_hal_resource_set_t* resource_set = NULL;
	iree_status_t status =
	iree_hal_resource_set_allocate(actions->block_pool, &resource_set);
	if (IREE_LIKELY(iree_status_is_ok(status))) {
	status = iree_hal_resource_set_insert(resource_set, command_buffer_count,
	command_buffers);
	}
	if (IREE_LIKELY(iree_status_is_ok(status))) {
	status =
	iree_hal_resource_set_insert(resource_set, wait_semaphore_list.count,
	wait_semaphore_list.semaphores);
	}
	if (IREE_LIKELY(iree_status_is_ok(status))) {
	status =
	iree_hal_resource_set_insert(resource_set, signal_semaphore_list.count,
	signal_semaphore_list.semaphores);
	}

	// Copy the semaphore and value list for later access.
	// TODO: avoid host allocator malloc; use some pool for the allocation.
	if (IREE_LIKELY(iree_status_is_ok(status))) {
	status = iree_hal_cuda2_copy_semaphore_list(wait_semaphore_list,
	actions->host_allocator,
	&action->wait_semaphore_list);
	}
	if (IREE_LIKELY(iree_status_is_ok(status))) {
	status = iree_hal_cuda2_copy_semaphore_list(signal_semaphore_list,
	actions->host_allocator,
	&action->signal_semaphore_list);
	}

	if (IREE_LIKELY(iree_status_is_ok(status))) {
	action->owning_actions = actions;
	iree_hal_resource_retain(actions);

	action->resource_set = resource_set;

	iree_slim_mutex_lock(&actions->action_mutex);
	iree_hal_cuda2_queue_action_list_push_back(&actions->action_list, action);
	iree_slim_mutex_unlock(&actions->action_mutex);
	} else {
	iree_hal_cuda2_free_semaphore_list(actions->host_allocator,
	&action->wait_semaphore_list);
	iree_hal_cuda2_free_semaphore_list(actions->host_allocator,
	&action->signal_semaphore_list);
	iree_hal_resource_set_free(resource_set);
	iree_allocator_free(actions->host_allocator, action);
	}

	IREE_TRACE_ZONE_END(z0);
	return status;
	}

	static void iree_hal_cuda2_pending_queue_actions_cleanup_execution(
	iree_hal_cuda2_queue_action_t* action);

	// Releases resources after action completion on the GPU and advances timeline
	// and pending actions queue.
	//
	// This is the CUDA host function callback to cudaLaunchHostFunc, invoked by a
	// CUDA driver thread.
	static void iree_hal_cuda2_execution_device_signal_host_callback(
	void* user_data) {
	iree_hal_cuda2_queue_action_t* action =
	(iree_hal_cuda2_queue_action_t*)user_data;
	iree_hal_cuda2_pending_queue_actions_t* actions = action->owning_actions;
	// Advance semaphore timelines by calling into the host signaling function.
	IREE_IGNORE_ERROR(
	iree_hal_semaphore_list_signal(action->signal_semaphore_list));
	// Destroy the current action given its done now--this also frees all retained
	// resources.
	iree_hal_cuda2_pending_queue_actions_cleanup_execution(action);
	// Try to release more pending actions to the GPU now.
	IREE_IGNORE_ERROR(iree_hal_cuda2_pending_queue_actions_issue(actions));
	}

	// Issues the given kernel dispatch \|action\| to the GPU.
	static iree_status_t iree_hal_cuda2_pending_queue_actions_issue_execution(
	iree_hal_cuda2_queue_action_t* action) {
	IREE_ASSERT(action->events != NULL);
	IREE_ASSERT(action->is_pending == false);
	const iree_hal_cuda2_dynamic_symbols_t* symbols =
	action->owning_actions->symbols;
	IREE_TRACE_ZONE_BEGIN(z0);

	// No need to lock given that this action is already detched from the pending
	// actions list; so only this thread is seeing it now.

	// First wait all the device CUevent in the dispatch stream.
	for (iree_host_size_t i = 0; i < action->event_count; ++i) {
	IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
	z0, symbols,
	cuStreamWaitEvent(action->dispatch_cu_stream, action->events[i],
	CU_EVENT_WAIT_DEFAULT),
	"cuStreamWaitEvent");
	}

	// Then launch all command buffers to the dispatch stream.
	for (iree_host_size_t i = 0; i < action->payload.command_buffers.count; ++i) {
	iree_hal_command_buffer_t* command_buffer =
	action->payload.command_buffers.ptr[i];
	if (iree_hal_cuda2_graph_command_buffer_isa(command_buffer)) {
	CUgraphExec exec = iree_hal_cuda2_graph_command_buffer_handle(
	action->payload.command_buffers.ptr[i]);
	IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
	z0, symbols, cuGraphLaunch(exec, action->dispatch_cu_stream),
	"cuGraphLaunch");
	} else {
	iree_hal_command_buffer_t* stream_command_buffer = NULL;
	iree_hal_command_buffer_mode_t mode =
	IREE_HAL_COMMAND_BUFFER_MODE_ONE_SHOT \|
	IREE_HAL_COMMAND_BUFFER_MODE_ALLOW_INLINE_EXECUTION \|
	IREE_HAL_COMMAND_BUFFER_MODE_UNVALIDATED;
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_cuda2_device_create_stream_command_buffer(
	action->device, mode, IREE_HAL_COMMAND_CATEGORY_ANY,
	/binding_capacity=/0, &stream_command_buffer));
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_resource_set_insert(action->resource_set, 1,
	&stream_command_buffer));
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_deferred_command_buffer_apply(
	command_buffer, stream_command_buffer,
	iree_hal_buffer_binding_table_empty()));
	}
	}

	// Last record CUevent signals in the dispatch stream.
	for (iree_host_size_t i = 0; i < action->signal_semaphore_list.count; ++i) {
	// Grab a CUevent for this semaphore value signaling.
	CUevent event = NULL;
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_cuda2_event_semaphore_acquire_timepoint_device_signal(
	action->signal_semaphore_list.semaphores[i],
	action->signal_semaphore_list.payload_values[i], &event));

	// Record the event signaling in the dispatch stream.
	IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
	z0, symbols, cuEventRecord(event, action->dispatch_cu_stream),
	"cuEventRecord");
	// Let the callback stream to wait on the CUevent.
	IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
	z0, symbols,
	cuStreamWaitEvent(action->callback_cu_stream, event,
	CU_EVENT_WAIT_DEFAULT),
	"cuStreamWaitEvent");
	}

	// Now launch a host function on the callback stream to advance the semaphore
	// timeline.
	IREE_CUDA_RETURN_AND_END_ZONE_IF_ERROR(
	z0, symbols,
	cuLaunchHostFunc(action->callback_cu_stream,
	iree_hal_cuda2_execution_device_signal_host_callback,
	action),
	"cuStreamWaitEvent");

	IREE_TRACE_ZONE_END(z0);
	return iree_ok_status();
	}

	// Releases resources after completing the given kernel dispatch \|action\|.
	static void iree_hal_cuda2_pending_queue_actions_cleanup_execution(
	iree_hal_cuda2_queue_action_t* action) {
	iree_hal_cuda2_pending_queue_actions_t* actions = action->owning_actions;
	iree_allocator_t host_allocator = actions->host_allocator;
	IREE_TRACE_ZONE_BEGIN(z0);

	iree_hal_resource_set_free(action->resource_set);
	iree_hal_cuda2_free_semaphore_list(host_allocator,
	&action->wait_semaphore_list);
	iree_hal_cuda2_free_semaphore_list(host_allocator,
	&action->signal_semaphore_list);
	iree_hal_resource_release(actions);

	iree_allocator_free(host_allocator, action);

	IREE_TRACE_ZONE_END(z0);
	}

	iree_status_t iree_hal_cuda2_pending_queue_actions_issue(
	iree_hal_cuda2_pending_queue_actions_t* actions) {
	IREE_TRACE_ZONE_BEGIN(z0);

	iree_hal_cuda2_queue_action_list_t pending_list = {NULL, NULL};
	iree_hal_cuda2_queue_action_list_t ready_list = {NULL, NULL};

	iree_slim_mutex_lock(&actions->action_mutex);

	if (iree_hal_cuda2_queue_action_list_is_empty(&actions->action_list)) {
	iree_slim_mutex_unlock(&actions->action_mutex);
	IREE_TRACE_ZONE_END(z0);
	return iree_ok_status();
	}

	// Scan through the list and categorize actions into pending and ready lists.
	for (iree_hal_cuda2_queue_action_t* action = actions->action_list.head;
	action != NULL;) {
	iree_hal_cuda2_queue_action_t* next_action = action->next;
	action->next = NULL;

	iree_host_size_t semaphore_count = action->wait_semaphore_list.count;
	iree_hal_semaphore_t** semaphores = action->wait_semaphore_list.semaphores;
	uint64_t* values = action->wait_semaphore_list.payload_values;

	// We are allocating stack space here, assuming that there won't be a lot of
	// waits and additional references to this field happens in a function call
	// from this function.
	action->events = iree_alloca(semaphore_count * sizeof(CUevent));
	action->event_count = 0;
	action->is_pending = false;

	// Look at all wait semaphores.
	for (iree_host_size_t i = 0; i < semaphore_count; ++i) {
	// If this semaphore has already signaled past the desired value, we can
	// just ignore it.
	uint64_t value = 0;
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_semaphore_query(semaphores[i], &value));
	if (value >= values[i]) continue;

	// Try to acquire a CUevent from a device wait timepoint. If so, we can
	// use that CUevent to wait on the device. Otherwise, this action is still
	// not ready.
	CUevent event = NULL;
	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_cuda2_event_semaphore_acquire_timepoint_device_wait(
	semaphores[i], values[i], &event));
	if (event) {
	action->events[action->event_count++] = event;
	} else {
	// Clear the scratch fields.
	action->events = NULL;
	action->event_count = 0;
	action->is_pending = true;
	break;
	}
	}

	if (action->is_pending) {
	iree_hal_cuda2_queue_action_list_push_back(&pending_list, action);
	} else {
	iree_hal_cuda2_queue_action_list_push_back(&ready_list, action);
	}

	action = next_action;
	}

	// Preserve pending timepoints.
	actions->action_list = pending_list;

	iree_slim_mutex_unlock(&actions->action_mutex);

	// Now go through the ready list and issue the actions to the GPU.
	for (iree_hal_cuda2_queue_action_t* action = ready_list.head;
	action != NULL;) {
	iree_hal_cuda2_queue_action_t* next_action = action->next;
	action->next = NULL;

	IREE_RETURN_AND_END_ZONE_IF_ERROR(
	z0, iree_hal_cuda2_pending_queue_actions_issue_execution(action));
	action->events = NULL;
	action->event_count = 0;

	action = next_action;
	}

	IREE_TRACE_ZONE_END(z0);
	return iree_ok_status();
	}