// 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 "iree/hal/utils/file_transfer.h"

#include "iree/base/internal/math.h"

//===----------------------------------------------------------------------===//
// Configuration
//===----------------------------------------------------------------------===//

// TODO(benvanik): make these either compile-time configuration options so we
// can prune code paths or flags (somehow).

#if !defined(IREE_HAL_TRANSFER_WORKER_LIMIT)
// Maximum number of workers that will be used. This is something we can derive
// from the transfer size; small transfers should have 1 and we can measure to
// see how many others we need.
#define IREE_HAL_TRANSFER_WORKER_LIMIT 1
#endif  // !IREE_HAL_TRANSFER_WORKER_LIMIT

#if !defined(IREE_HAL_TRANSFER_CHUNK_SIZE)
// Bytes per worker to stage chunks of data. Larger chunks will result in less
// overhead as fewer copy operations are required.
#define IREE_HAL_TRANSFER_CHUNK_SIZE (64 * 1024 * 1024)
#endif  // !IREE_HAL_TRANSFER_CHUNK_SIZE

#if !defined(IREE_HAL_TRANSFER_CHUNKS_PER_WORKER)
// Estimated number of chunks each worker should process used to determine how
// many workers are needed as part of a transfer operation. Larger numbers will
// reduce memory overhead at the cost of latency reductions.
#define IREE_HAL_TRANSFER_CHUNKS_PER_WORKER 8
#endif  // IREE_HAL_TRANSFER_CHUNKS_PER_WORKER

//===----------------------------------------------------------------------===//
// iree_hal_transfer_operation_t
//===----------------------------------------------------------------------===//

// Maximum number of transfer workers that can be used; common usage should be
// 1-4 but on very large systems with lots of bandwidth we may be able to
// use more.
#define IREE_HAL_TRANSFER_WORKER_MAX_COUNT 64

// Each bit represents one worker within a transfer matching its ordinal in
// the operation workers array.
typedef uint64_t iree_hal_transfer_worker_bitmask_t;

// Counts the total number of workers indicated by the given worker bitmask.
#define iree_hal_transfer_worker_live_count(bitmask) \
  iree_math_count_ones_u64(bitmask)

// Describes the direction of a transfer operation.
typedef enum {
  // Transferring from the file to the buffer (read).
  IREE_HAL_TRANSFER_READ_FILE_TO_BUFFER = 0,
  // Transferring from the buffer to the file (write).
  IREE_HAL_TRANSFER_WRITE_BUFFER_TO_FILE,
} iree_hal_transfer_direction_t;

typedef struct iree_hal_transfer_operation_t iree_hal_transfer_operation_t;

// A worker greedily processing subranges of a larger transfer operation.
// Since transfers are 99% IO bound we avoid real threads: workers submit
// device operations and block on semaphore completion between chunks. When
// there are no remaining subranges the workers will exit and when the last does
// the transfer is marked complete.
typedef struct iree_hal_transfer_worker_t {
  // Parent operation this worker is a part of.
  iree_hal_transfer_operation_t* operation;
  // Used to associate tracing events with this worker.
  IREE_TRACE(int32_t trace_id;)
  // Aligned offset into the staging buffer of the worker storage.
  iree_device_size_t staging_buffer_offset;
  // Aligned length of the staging buffer storage reserved for the worker.
  iree_device_size_t staging_buffer_length;
  // Semaphore representing the timeline of the worker. The payload is a
  // monotonically increasing operation count.
  iree_hal_semaphore_t* semaphore;
  // Pending timepoint representing an in-flight operation. Upon completion of
  // the operation the semaphore will reach this value.
  uint64_t pending_timepoint;
  // Offset into the transfer operation this worker is currently processing.
  iree_device_size_t pending_transfer_offset;
  // Length of the current worker transfer; usually staging_buffer_length but
  // may be less if this worker is processing the end of the file.
  iree_device_size_t pending_transfer_length;
} iree_hal_transfer_worker_t;

// Manages an asynchronous transfer operation.
typedef struct iree_hal_transfer_operation_t {
  // Workers run synchronously within the caller's stack frame but may complete
  // the operation during setup if the transfer is small enough. The ref count
  // lets the top-level call clean up even if the operation finishes early.
  iree_atomic_ref_count_t ref_count;
  // Device this transfer operation is acting on.
  iree_hal_device_t* device;
  // Queue affinity all operations should be assigned.
  iree_hal_queue_affinity_t queue_affinity;
  // Used to associate tracing events with this worker.
  IREE_TRACE(int32_t trace_id;)

  // Direction of the operation (read file->buffer or write buffer->file).
  iree_hal_transfer_direction_t direction;
  // Retained file resource.
  iree_hal_file_t* file;
  // Offset into the file where the operation begins.
  uint64_t file_offset;
  // Retained buffer resource.
  iree_hal_buffer_t* buffer;
  // Offset into the buffer where the operation begins.
  iree_device_size_t buffer_offset;
  // Total length of the operation.
  iree_device_size_t length;

  // Sticky error status; when any worker fails this will be set to non-OK and
  // when all workers end the staging buffer will be deallocated and the signal
  // semaphores will be marked as failing.
  iree_status_t error_status;
  // Original user semaphores to signal at the end of the transfer operation.
  // Contents are stored at the end of the struct.
  iree_hal_semaphore_list_t signal_semaphore_list;

  // Shared staging buffer; contains storage for all workers.
  // We avoid a subspan buffer here to reduce overheads.
  iree_hal_buffer_t* staging_buffer;
  iree_device_size_t staging_buffer_size;

  // Offset to where the transfer head is in the operation.
  // Ranges from 0 at the start and length at the end.
  // Workers use this to consume chunks of the operation.
  iree_device_size_t transfer_head;
  // Total number of chunks remaining in the transfer.
  iree_host_size_t remaining_chunks;

  // Total number of workers participating in the operation.
  iree_host_size_t worker_count;
  // State for each worker in the operation.
  // Stored at the end of the struct.
  iree_hal_transfer_worker_t* workers;
  // One bit per worker indicating whether they are live.
  // A worker exits by not rescheduling itself and clearing this bit. When all
  // workers have exited the operation has completed.
  // When reading workers exit after enqueuing their final transfer such that
  // the final staging buffer dealloca can be asynchronously chained.
  // When writing workers exit after flushing their final chunk to the file.
  iree_hal_transfer_worker_bitmask_t live_workers;
} iree_hal_transfer_operation_t;

static void iree_hal_transfer_operation_release(
    iree_hal_transfer_operation_t* operation);
static void iree_hal_transfer_operation_destroy(
    iree_hal_transfer_operation_t* operation);

static iree_status_t iree_hal_transfer_operation_create(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t signal_semaphore_list,
    iree_hal_transfer_direction_t direction, iree_hal_file_t* file,
    uint64_t file_offset, iree_hal_buffer_t* buffer,
    iree_device_size_t buffer_offset, iree_device_size_t length,
    iree_hal_file_transfer_options_t options,
    iree_hal_transfer_operation_t** out_operation) {
  IREE_ASSERT_ARGUMENT(out_operation);
  *out_operation = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_allocator_t host_allocator = iree_hal_device_host_allocator(device);

  // Determine how many workers are required and their staging reservation.
  iree_device_size_t worker_chunk_size = options.chunk_size;
  if (worker_chunk_size == IREE_HAL_FILE_TRANSFER_CHUNK_SIZE_DEFAULT) {
    worker_chunk_size = iree_min(IREE_HAL_TRANSFER_CHUNK_SIZE, length);
  }
  iree_device_size_t total_chunk_count =
      iree_device_size_ceil_div(length, worker_chunk_size);
  iree_host_size_t worker_count = options.chunk_count;
  if (worker_count == IREE_HAL_FILE_TRANSFER_CHUNK_COUNT_DEFAULT) {
    // Try to give each worker a couple chunks.
    worker_count = (iree_host_size_t)iree_device_size_ceil_div(
        total_chunk_count, IREE_HAL_TRANSFER_CHUNKS_PER_WORKER);
  }
  worker_count =
      iree_min(worker_count, iree_min(IREE_HAL_TRANSFER_WORKER_LIMIT,
                                      IREE_HAL_TRANSFER_WORKER_MAX_COUNT));

  // Calculate total size of the structure with all its associated data.
  iree_hal_transfer_operation_t* operation = NULL;
  iree_host_size_t total_size = sizeof(*operation);
  iree_host_size_t semaphores_offset =
      iree_host_align(total_size, iree_max_align_t);
  total_size = semaphores_offset + sizeof(signal_semaphore_list.semaphores[0]) *
                                       signal_semaphore_list.count;
  iree_host_size_t payload_values_offset =
      iree_host_align(total_size, iree_max_align_t);
  total_size =
      payload_values_offset + sizeof(signal_semaphore_list.payload_values[0]) *
                                  signal_semaphore_list.count;
  iree_host_size_t worker_offset =
      iree_host_align(total_size, iree_max_align_t);
  total_size = worker_offset + sizeof(operation->workers[0]) * worker_count;

  // Allocate and initialize the struct.
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0,
      iree_allocator_malloc(host_allocator, total_size, (void**)&operation));
  iree_atomic_ref_count_init(&operation->ref_count);
  operation->device = device;
  iree_hal_device_retain(device);
  operation->queue_affinity = queue_affinity;
  operation->direction = direction;
  operation->file = file;
  iree_hal_file_retain(file);
  operation->file_offset = file_offset;
  operation->buffer = buffer;
  iree_hal_buffer_retain(buffer);
  operation->buffer_offset = buffer_offset;
  operation->length = length;
  operation->staging_buffer_size = worker_count * worker_chunk_size;
  operation->transfer_head = 0;
  operation->remaining_chunks = (iree_host_size_t)total_chunk_count;
  operation->worker_count = worker_count;

  // Assign all pointers to the struct suffix storage.
  // We do this first so that if we have to free the struct we have valid
  // pointers.
  operation->signal_semaphore_list.count = signal_semaphore_list.count;
  operation->signal_semaphore_list.semaphores =
      (iree_hal_semaphore_t**)((uintptr_t)operation + semaphores_offset);
  operation->signal_semaphore_list.payload_values =
      (uint64_t*)((uintptr_t)operation + payload_values_offset);
  operation->workers =
      (iree_hal_transfer_worker_t*)((uintptr_t)operation + worker_offset);

  // Assign a unique ID we'll use to make it easier to track what individual
  // steps are part of this transfer.
  IREE_TRACE({
    static iree_atomic_int32_t next_trace_id = IREE_ATOMIC_VAR_INIT(0);
    operation->trace_id =
        iree_atomic_fetch_add(&next_trace_id, 1, iree_memory_order_seq_cst);
    IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, operation->trace_id);
  });

  // Retain each signal semaphore for ourselves as we don't know if the caller
  // will hold them for the lifetime of the operation.
  memcpy(operation->signal_semaphore_list.semaphores,
         signal_semaphore_list.semaphores,
         sizeof(signal_semaphore_list.semaphores[0]) *
             signal_semaphore_list.count);
  memcpy(operation->signal_semaphore_list.payload_values,
         signal_semaphore_list.payload_values,
         sizeof(signal_semaphore_list.payload_values[0]) *
             signal_semaphore_list.count);
  for (iree_host_size_t i = 0; i < signal_semaphore_list.count; ++i) {
    iree_hal_semaphore_retain(signal_semaphore_list.semaphores[i]);
  }

  // Initialize all workers.
  iree_status_t status = iree_ok_status();
  for (iree_host_size_t i = 0; i < worker_count; ++i) {
    iree_hal_transfer_worker_t* worker = &operation->workers[i];
    worker->operation = operation;

    // Assign a unique ID we'll use to make it easier to track what individual
    // steps are part of this worker. It only needs to be unique within the
    // operation.
    IREE_TRACE(worker->trace_id = (int64_t)i);

    // View into the staging buffer where the worker keeps its memory.
    worker->staging_buffer_offset = worker_chunk_size * i;
    worker->staging_buffer_length = worker_chunk_size;

    // Create semaphore for tracking worker progress.
    worker->pending_timepoint = 0ull;
    status = iree_hal_semaphore_create(
        device, IREE_HAL_QUEUE_AFFINITY_ANY, worker->pending_timepoint,
        IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &worker->semaphore);
    if (!iree_status_is_ok(status)) break;
  }

  if (iree_status_is_ok(status)) {
    IREE_TRACE_ZONE_APPEND_TEXT(z0, "worker count: ");
    IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)worker_count);
    IREE_TRACE_ZONE_APPEND_TEXT(z0, "worker chunk size: ");
    IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)worker_chunk_size);
    *out_operation = operation;
  } else {
    iree_hal_transfer_operation_release(operation);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static void iree_hal_transfer_operation_retain(
    iree_hal_transfer_operation_t* operation) {
  if (IREE_LIKELY(operation)) {
    iree_atomic_ref_count_inc(&operation->ref_count);
  }
}

static void iree_hal_transfer_operation_release(
    iree_hal_transfer_operation_t* operation) {
  if (IREE_LIKELY(operation) &&
      iree_atomic_ref_count_dec(&operation->ref_count) == 1) {
    iree_hal_transfer_operation_destroy(operation);
  }
}

static void iree_hal_transfer_operation_destroy(
    iree_hal_transfer_operation_t* operation) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, operation->trace_id);
  iree_allocator_t host_allocator =
      iree_hal_device_host_allocator(operation->device);

  // All workers must have exited before we free the operation memory.
  IREE_ASSERT(operation->live_workers == 0, "all workers must have exited");

  for (iree_host_size_t i = 0; i < operation->worker_count; ++i) {
    iree_hal_semaphore_release(operation->workers[i].semaphore);
  }
  iree_hal_buffer_release(operation->staging_buffer);
  for (iree_host_size_t i = 0; i < operation->signal_semaphore_list.count;
       ++i) {
    iree_hal_semaphore_release(operation->signal_semaphore_list.semaphores[i]);
  }
  iree_hal_buffer_release(operation->buffer);
  iree_hal_file_release(operation->file);
  iree_hal_device_release(operation->device);
  iree_status_ignore(operation->error_status);

  iree_allocator_free(host_allocator, operation);

  IREE_TRACE_ZONE_END(z0);
}

// Notifies listeners that the operation has completed and releases its memory.
// If this was a read then the staging buffer dealloca will be chained to the
// last asynchronous copies. In writes the last flush to the file happened
// synchronously so the dealloca happens immediately.
//
// Pre-condition: all workers have exited and there are no operations in flight.
// Post-condition: the operation is freed.
static void iree_hal_transfer_operation_notify_completion(
    iree_hal_transfer_operation_t* operation) {
  IREE_ASSERT_ARGUMENT(operation);
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, operation->trace_id);

  // We can only free the operation if no workers have pending work.
  IREE_ASSERT(operation->live_workers == 0, "no workers can be live");

  // Deallocating the staging buffer can only happen after all workers have
  // completed copies into/out-of it. In reads it's expected there are copies
  // in-flight and we can wait on all worker semaphores. In writes the last
  // flush to the file happened synchronously so we don't need to wait at all.
  iree_hal_semaphore_list_t wait_semaphore_list =
      iree_hal_semaphore_list_empty();
  if (operation->direction == IREE_HAL_TRANSFER_READ_FILE_TO_BUFFER) {
    wait_semaphore_list.count = operation->worker_count;
    wait_semaphore_list.semaphores = (iree_hal_semaphore_t**)iree_alloca(
        wait_semaphore_list.count * sizeof(wait_semaphore_list.semaphores[0]));
    wait_semaphore_list.payload_values =
        (uint64_t*)iree_alloca(wait_semaphore_list.count *
                               sizeof(wait_semaphore_list.payload_values[0]));
    for (iree_host_size_t i = 0; i < operation->worker_count; ++i) {
      iree_hal_transfer_worker_t* worker = &operation->workers[i];
      wait_semaphore_list.semaphores[i] = worker->semaphore;
      wait_semaphore_list.payload_values[i] = worker->pending_timepoint;
    }
  }

  // When the dealloca completes signal the original semaphores passed in to the
  // operation. If the transfer failed then we need to signal them all to
  // failure after the dealloca so we use the same semaphores but set the
  // failure payload.
  iree_hal_semaphore_list_t signal_semaphore_list =
      operation->signal_semaphore_list;
  if (!iree_status_is_ok(operation->error_status)) {
    for (iree_host_size_t i = 0; i < signal_semaphore_list.count; ++i) {
      signal_semaphore_list.payload_values[i] =
          IREE_HAL_SEMAPHORE_FAILURE_VALUE;
    }
  }

  // Schedule staging buffer deallocation.
  // Note that we need to do this even if the operation failed and we want it to
  // be scheduled after any copies that may be in-flight (say worker 4 failed
  // on a chunk but workers 0-3 succeeded).
  iree_status_t status = iree_hal_device_queue_dealloca(
      operation->device, operation->queue_affinity, wait_semaphore_list,
      signal_semaphore_list, operation->staging_buffer,
      IREE_HAL_DEALLOCA_FLAG_NONE);

  // If the dealloca failed we don't have a great way of letting anyone know.
  // We'll just drop it on the floor for now and let the buffer be freed by
  // reference counting.
  iree_status_ignore(status);

  IREE_TRACE_ZONE_END(z0);
}

// Exits the |worker| indicating it will process no more chunks.
// If this is the last worker to exit the transfer is considered completed.
// The first non-OK |status| provided will be set as a stick error on the
// |operation| and all workers will check it and exit themselves asynchronously.
//
// NOTE: this may end the entire operation and free the operation (and worker)
// memory. Callers must not touch either |operation| or |worker| after calling
// this method.
static iree_status_t iree_hal_transfer_worker_exit(
    iree_hal_transfer_operation_t* operation,
    iree_hal_transfer_worker_t* worker, iree_status_t status) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)operation->trace_id);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)worker->trace_id);

  // Check if this is the first failure of an operation and set the error bit.
  // Otherwise we ignore the error here as it's probably just telling us the
  // worker has aborted.
  if (iree_status_is_ok(operation->error_status) &&
      !iree_status_is_ok(status)) {
    operation->error_status = status;
  } else {
    iree_status_ignore(status);
  }

  // Clear the worker live bit and see if there are any more workers live. So
  // long as there is at least one we need to keep the operation running.
  iree_host_size_t worker_index =
      (iree_host_size_t)(worker - operation->workers);
  operation->live_workers &= ~(1ull << worker_index);
  if (operation->live_workers > 0) {
    // Other workers are still live - this is just one worker exiting by not
    // rescheduling itself.
    iree_hal_transfer_operation_release(operation);
    IREE_TRACE_ZONE_END(z0);
    return iree_ok_status();
  }

  // This was the last worker; the operation has completed!
  iree_hal_transfer_operation_notify_completion(operation);

  // Free the operation - the dealloca (and maybe even some copies) may still be
  // in-flight but all on the device side and not using any resources on the
  // operation.
  iree_hal_transfer_operation_release(operation);
  // NOTE: at this point the worker may have freed itself and its memory can
  // no longer be used!

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();  // always ok; just for convenience.
}

// Runs a single read worker to completion: reads chunks from the file into the
// staging buffer and issues device copies from staging into the target buffer.
// Each iteration waits for the previous copy to complete before reading the
// next chunk. The worker exits when there are no remaining chunks or when
// enough workers are live to cover the remaining work.
static void iree_hal_transfer_worker_run_read(
    iree_hal_transfer_operation_t* operation,
    iree_hal_transfer_worker_t* worker) {
  // Wait for the staging buffer allocation to complete.
  iree_status_t status = iree_hal_semaphore_wait(
      worker->semaphore, worker->pending_timepoint, iree_infinite_timeout(),
      IREE_ASYNC_WAIT_FLAG_NONE);

  while (iree_status_is_ok(status)) {
    // Bail if the operation has failed (another worker hit an error).
    if (!iree_status_is_ok(operation->error_status)) break;

    // Exit if we're out of chunks to process.
    if (operation->remaining_chunks == 0) break;

    // Grab a piece of the transfer to operate on.
    --operation->remaining_chunks;
    iree_device_size_t transfer_offset = operation->transfer_head;
    iree_device_size_t transfer_length = iree_min(
        operation->length - transfer_offset, worker->staging_buffer_length);
    IREE_ASSERT(transfer_length > 0,
                "should not have ticked if there was no work to do");
    operation->transfer_head += transfer_length;

    // Timeline increments by one.
    uint64_t wait_timepoint = worker->pending_timepoint;
    iree_hal_semaphore_list_t wait_semaphore_list = {
        .count = 1,
        .semaphores = &worker->semaphore,
        .payload_values = &wait_timepoint,
    };
    uint64_t signal_timepoint = ++worker->pending_timepoint;
    iree_hal_semaphore_list_t signal_semaphore_list = {
        .count = 1,
        .semaphores = &worker->semaphore,
        .payload_values = &signal_timepoint,
    };

    // Track the pending copy operation so we know where to place it in the
    // buffer.
    worker->pending_transfer_offset = transfer_offset;
    worker->pending_transfer_length = transfer_length;

    // Synchronously copy the contents from the file to the staging buffer.
    status = iree_hal_file_read(
        operation->file,
        operation->file_offset + worker->pending_transfer_offset,
        operation->staging_buffer, worker->staging_buffer_offset,
        worker->pending_transfer_length);
    if (!iree_status_is_ok(status)) break;

    // Issue asynchronous copy from the staging buffer into the target buffer.
    status = iree_hal_device_queue_copy(
        operation->device, operation->queue_affinity, wait_semaphore_list,
        signal_semaphore_list, operation->staging_buffer,
        worker->staging_buffer_offset, operation->buffer,
        operation->buffer_offset + transfer_offset, transfer_length,
        IREE_HAL_COPY_FLAG_NONE);
    if (!iree_status_is_ok(status)) break;

    // If there are enough live workers to cover all remaining chunks we can
    // exit now and avoid an additional host wake (+ latency) by the wait.
    if (iree_hal_transfer_worker_live_count(operation->live_workers) >
        operation->remaining_chunks) {
      break;
    }

    // Wait for the copy to complete before processing the next chunk.
    status = iree_hal_semaphore_wait(
        worker->semaphore, worker->pending_timepoint, iree_infinite_timeout(),
        IREE_ASYNC_WAIT_FLAG_NONE);
  }

  iree_status_ignore(iree_hal_transfer_worker_exit(operation, worker, status));
}

// Begins the read transfer operation after |wait_semaphore_list| is satisfied.
// Allocates a staging buffer then runs each worker to completion sequentially.
static iree_status_t iree_hal_transfer_operation_launch_read(
    iree_hal_transfer_operation_t* operation,
    iree_hal_semaphore_list_t wait_semaphore_list) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)operation->trace_id);
  IREE_ASSERT(operation->direction == IREE_HAL_TRANSFER_READ_FILE_TO_BUFFER);

  // Staging buffers get allocated based on the direction we are transferring.
  // This optimizes for access patterns such as sequential writes from the host
  // when staging into the buffer and sequential cached reads from the host when
  // staging out of the buffer.
  iree_hal_buffer_params_t staging_buffer_params = {
      .access = IREE_HAL_MEMORY_ACCESS_ALL,
      // TODO(benvanik): make staging alignment an option/device query?
      .min_alignment = 64,
      .queue_affinity = operation->queue_affinity,
      .type = IREE_HAL_MEMORY_TYPE_OPTIMAL_FOR_HOST |
              IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE,
      .usage = IREE_HAL_BUFFER_USAGE_TRANSFER |
               IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED |
               IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_SEQUENTIAL_WRITE,
  };

  // Queue the staging buffer allocation.
  // When it completes we'll do the first host->device copy via mapping.
  iree_hal_semaphore_list_t alloca_semaphore_list = {
      .count = operation->worker_count,
      .semaphores =
          iree_alloca(sizeof(iree_hal_semaphore_t*) * operation->worker_count),
      .payload_values = iree_alloca(sizeof(uint64_t) * operation->worker_count),
  };
  for (iree_host_size_t i = 0; i < operation->worker_count; ++i) {
    iree_hal_transfer_worker_t* worker = &operation->workers[i];
    alloca_semaphore_list.semaphores[i] = worker->semaphore;
    uint64_t signal_timepoint = ++worker->pending_timepoint;
    alloca_semaphore_list.payload_values[i] = signal_timepoint;
  }
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_device_queue_alloca(
              operation->device, operation->queue_affinity, wait_semaphore_list,
              alloca_semaphore_list, /*pool=*/NULL, staging_buffer_params,
              operation->staging_buffer_size, IREE_HAL_ALLOCA_FLAG_NONE,
              &operation->staging_buffer));

  // Run each worker to completion. Workers greedily consume chunks and exit
  // when there are no more chunks or enough workers are already covering the
  // remaining work.
  for (iree_host_size_t worker_index = 0;
       worker_index < operation->worker_count; ++worker_index) {
    iree_hal_transfer_worker_t* worker = &operation->workers[worker_index];
    operation->live_workers |= 1ull << worker_index;
    iree_hal_transfer_operation_retain(operation);

    iree_hal_transfer_worker_run_read(operation, worker);

    // The worker may have consumed all remaining chunks.
    if (operation->remaining_chunks == 0) break;
  }

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

// Runs a single write worker to completion: issues device copies from the
// source buffer into the staging buffer, waits for completion, then writes
// the staged data to the file. Each iteration processes one chunk.
// The worker exits when there are no remaining chunks or on error.
static void iree_hal_transfer_worker_run_write(
    iree_hal_transfer_operation_t* operation,
    iree_hal_transfer_worker_t* worker) {
  iree_status_t status = iree_ok_status();

  while (iree_status_is_ok(status)) {
    // Bail if the operation has failed (another worker hit an error).
    if (!iree_status_is_ok(operation->error_status)) break;

    // Exit if we're out of chunks to process.
    if (operation->remaining_chunks == 0) break;

    // Grab a piece of the transfer to operate on.
    --operation->remaining_chunks;
    iree_device_size_t transfer_offset = operation->transfer_head;
    iree_device_size_t transfer_length = iree_min(
        operation->length - transfer_offset, worker->staging_buffer_length);
    IREE_ASSERT(transfer_length > 0,
                "should not have ticked if there was no work to do");
    operation->transfer_head += transfer_length;

    // Timeline increments by one.
    uint64_t wait_timepoint = worker->pending_timepoint;
    iree_hal_semaphore_list_t wait_semaphore_list = {
        .count = 1,
        .semaphores = &worker->semaphore,
        .payload_values = &wait_timepoint,
    };
    uint64_t signal_timepoint = ++worker->pending_timepoint;
    iree_hal_semaphore_list_t signal_semaphore_list = {
        .count = 1,
        .semaphores = &worker->semaphore,
        .payload_values = &signal_timepoint,
    };

    // Track the pending copy operation so we know where to place it in the
    // file.
    worker->pending_transfer_offset = transfer_offset;
    worker->pending_transfer_length = transfer_length;

    // Issue an asynchronous copy from the source buffer to the staging buffer.
    status = iree_hal_device_queue_copy(
        operation->device, operation->queue_affinity, wait_semaphore_list,
        signal_semaphore_list, operation->buffer,
        operation->buffer_offset + transfer_offset, operation->staging_buffer,
        worker->staging_buffer_offset, transfer_length,
        IREE_HAL_COPY_FLAG_NONE);
    if (!iree_status_is_ok(status)) break;

    // Wait for the copy to complete so we can write the staging data to the
    // file.
    status = iree_hal_semaphore_wait(worker->semaphore, signal_timepoint,
                                     iree_infinite_timeout(),
                                     IREE_ASYNC_WAIT_FLAG_NONE);
    if (!iree_status_is_ok(status)) break;

    // Synchronously write the staged data to the file.
    status = iree_hal_file_write(
        operation->file,
        operation->file_offset + worker->pending_transfer_offset,
        operation->staging_buffer, worker->staging_buffer_offset,
        worker->pending_transfer_length);
  }

  iree_status_ignore(iree_hal_transfer_worker_exit(operation, worker, status));
}

// Begins the write transfer operation after |wait_semaphore_list| is satisfied.
// Allocates a staging buffer then runs each worker to completion sequentially.
static iree_status_t iree_hal_transfer_operation_launch_write(
    iree_hal_transfer_operation_t* operation,
    iree_hal_semaphore_list_t wait_semaphore_list) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_VALUE_I64(z0, (int64_t)operation->trace_id);
  IREE_ASSERT(operation->direction == IREE_HAL_TRANSFER_WRITE_BUFFER_TO_FILE);

  // Staging buffers get allocated based on the direction we are transferring.
  // This optimizes for access patterns such as sequential writes from the host
  // when staging into the buffer and sequential cached reads from the host when
  // staging out of the buffer.
  iree_hal_buffer_params_t staging_buffer_params = {
      .access = IREE_HAL_MEMORY_ACCESS_ALL,
      // TODO(benvanik): make staging alignment an option/device query?
      .min_alignment = 64,
      .queue_affinity = operation->queue_affinity,
      .type = IREE_HAL_MEMORY_TYPE_OPTIMAL_FOR_HOST |
              IREE_HAL_MEMORY_TYPE_HOST_CACHED |
              IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE,
      .usage = IREE_HAL_BUFFER_USAGE_TRANSFER |
               IREE_HAL_BUFFER_USAGE_MAPPING_SCOPED |
               IREE_HAL_BUFFER_USAGE_MAPPING_ACCESS_RANDOM,
  };

  // Queue the staging buffer allocation.
  // When it completes we'll signal each worker to start its first transfer.
  iree_hal_semaphore_list_t alloca_semaphore_list = {
      .count = operation->worker_count,
      .semaphores =
          iree_alloca(sizeof(iree_hal_semaphore_t*) * operation->worker_count),
      .payload_values = iree_alloca(sizeof(uint64_t) * operation->worker_count),
  };
  for (iree_host_size_t i = 0; i < operation->worker_count; ++i) {
    iree_hal_transfer_worker_t* worker = &operation->workers[i];
    alloca_semaphore_list.semaphores[i] = worker->semaphore;
    uint64_t signal_timepoint = ++worker->pending_timepoint;
    alloca_semaphore_list.payload_values[i] = signal_timepoint;
  }
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_device_queue_alloca(
              operation->device, operation->queue_affinity, wait_semaphore_list,
              alloca_semaphore_list, /*pool=*/NULL, staging_buffer_params,
              operation->staging_buffer_size, IREE_HAL_ALLOCA_FLAG_NONE,
              &operation->staging_buffer));

  // Run each worker to completion. Workers greedily consume chunks and exit
  // when there are no more chunks remaining.
  for (iree_host_size_t worker_index = 0;
       worker_index < operation->worker_count; ++worker_index) {
    iree_hal_transfer_worker_t* worker = &operation->workers[worker_index];
    operation->live_workers |= 1ull << worker_index;
    iree_hal_transfer_operation_retain(operation);

    iree_hal_transfer_worker_run_write(operation, worker);

    // The worker may have consumed all remaining chunks.
    if (operation->remaining_chunks == 0) break;
  }

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Streaming file IO API
//===----------------------------------------------------------------------===//

IREE_API_EXPORT iree_status_t iree_hal_device_queue_read_streaming(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t wait_semaphore_list,
    const iree_hal_semaphore_list_t signal_semaphore_list,
    iree_hal_file_t* source_file, uint64_t source_offset,
    iree_hal_buffer_t* target_buffer, iree_device_size_t target_offset,
    iree_device_size_t length, iree_hal_read_flags_t flags,
    iree_hal_file_transfer_options_t options) {
  IREE_RETURN_IF_ERROR(
      iree_hal_file_validate_access(source_file, IREE_HAL_MEMORY_ACCESS_READ));

  // If the file has a device-visible storage buffer we can issue a direct
  // device copy without staging. HOST_LOCAL-only buffers (e.g. heap-wrapped
  // host allocations that failed device import) must fall through to the
  // streaming path which handles host↔device staging.
  iree_hal_buffer_t* storage_buffer = iree_hal_file_storage_buffer(source_file);
  if (storage_buffer &&
      iree_all_bits_set(iree_hal_buffer_memory_type(storage_buffer),
                        IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
    return iree_hal_device_queue_copy(
        device, queue_affinity, wait_semaphore_list, signal_semaphore_list,
        storage_buffer, (iree_device_size_t)source_offset, target_buffer,
        target_offset, length, IREE_HAL_COPY_FLAG_NONE);
  }

  // This host-side transfer utility requires synchronous I/O.
  // HAL implementations are expected to handle asynchronous files themselves.
  if (!iree_hal_file_supports_synchronous_io(source_file)) {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "provided source file does not support synchronous I/O and cannot be "
        "used with streaming file transfer");
  }

  // Allocate full transfer operation.
  iree_hal_transfer_operation_t* operation = NULL;
  IREE_RETURN_IF_ERROR(iree_hal_transfer_operation_create(
      device, queue_affinity, signal_semaphore_list,
      IREE_HAL_TRANSFER_READ_FILE_TO_BUFFER, source_file, source_offset,
      target_buffer, target_offset, length, options, &operation));

  // Run the streaming transfer synchronously. Workers submit device operations
  // and wait for semaphore completion between transfer chunks.
  iree_status_t status =
      iree_hal_transfer_operation_launch_read(operation, wait_semaphore_list);

  iree_hal_transfer_operation_release(operation);

  return status;
}

IREE_API_EXPORT iree_status_t iree_hal_device_queue_write_streaming(
    iree_hal_device_t* device, iree_hal_queue_affinity_t queue_affinity,
    const iree_hal_semaphore_list_t wait_semaphore_list,
    const iree_hal_semaphore_list_t signal_semaphore_list,
    iree_hal_buffer_t* source_buffer, iree_device_size_t source_offset,
    iree_hal_file_t* target_file, uint64_t target_offset,
    iree_device_size_t length, iree_hal_write_flags_t flags,
    iree_hal_file_transfer_options_t options) {
  // EXPERIMENTAL: assume memory files only today (as that's all we have).
  IREE_RETURN_IF_ERROR(
      iree_hal_file_validate_access(target_file, IREE_HAL_MEMORY_ACCESS_WRITE));

  // If the file has a device-visible storage buffer we can issue a direct
  // device copy without staging. HOST_LOCAL-only buffers (e.g. heap-wrapped
  // host allocations that failed device import) must fall through to the
  // streaming path which handles host↔device staging.
  iree_hal_buffer_t* storage_buffer = iree_hal_file_storage_buffer(target_file);
  if (storage_buffer &&
      iree_all_bits_set(iree_hal_buffer_memory_type(storage_buffer),
                        IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE)) {
    return iree_hal_device_queue_copy(
        device, queue_affinity, wait_semaphore_list, signal_semaphore_list,
        source_buffer, source_offset, storage_buffer,
        (iree_device_size_t)target_offset, length, IREE_HAL_COPY_FLAG_NONE);
  }

  // This host-side transfer utility requires synchronous I/O.
  // HAL implementations are expected to handle asynchronous files themselves.
  if (!iree_hal_file_supports_synchronous_io(target_file)) {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "provided target file does not support synchronous I/O and cannot be "
        "used with streaming file transfer");
  }

  // Allocate full transfer operation.
  iree_hal_transfer_operation_t* operation = NULL;
  IREE_RETURN_IF_ERROR(iree_hal_transfer_operation_create(
      device, queue_affinity, signal_semaphore_list,
      IREE_HAL_TRANSFER_WRITE_BUFFER_TO_FILE, target_file, target_offset,
      source_buffer, source_offset, length, options, &operation));

  // Run the streaming transfer synchronously. Workers submit device operations
  // and wait for semaphore completion between transfer chunks.
  iree_status_t status =
      iree_hal_transfer_operation_launch_write(operation, wait_semaphore_list);

  iree_hal_transfer_operation_release(operation);

  return status;
}
