// Copyright 2026 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/local/inline_dispatch.h"

#include "iree/hal/local/local_executable.h"

iree_status_t iree_hal_local_executable_dispatch_inline(
    iree_hal_executable_t* executable, iree_hal_executable_function_t function,
    const iree_hal_dispatch_config_t config, iree_const_byte_span_t constants,
    const iree_hal_buffer_ref_t* bindings, iree_host_size_t binding_count,
    iree_hal_dispatch_flags_t flags) {
  if (iree_hal_dispatch_uses_custom_arguments(flags)) {
    return iree_make_status(IREE_STATUS_UNIMPLEMENTED,
                            "direct/indirect arguments are not supported on "
                            "the inline CPU dispatch path");
  }
  const bool uses_indirect_parameters =
      iree_hal_dispatch_uses_indirect_parameters(flags);
  iree_hal_local_executable_t* local_executable =
      iree_hal_local_executable_cast(executable);
  if (!iree_hal_executable_function_is_index_in_range(
          function, local_executable->export_count)) {
    return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                            "function id %" PRIu64
                            " out of range (count: %" PRIhsz ")",
                            function.value, local_executable->export_count);
  }
  const uint32_t function_index = iree_hal_executable_function_index(function);

  IREE_TRACE_ZONE_BEGIN(z0);

  // Stack-allocate mapping and pointer arrays for the bindings.
  iree_hal_buffer_mapping_t* mappings = NULL;
  void** binding_ptrs = NULL;
  size_t* binding_lengths = NULL;
  if (binding_count > 0) {
    mappings = (iree_hal_buffer_mapping_t*)iree_alloca(binding_count *
                                                       sizeof(*mappings));
    memset(mappings, 0, binding_count * sizeof(*mappings));
    binding_ptrs = (void**)iree_alloca(binding_count * sizeof(*binding_ptrs));
    binding_lengths =
        (size_t*)iree_alloca(binding_count * sizeof(*binding_lengths));
  }

  // Map all binding buffers.
  iree_status_t status = iree_ok_status();
  iree_host_size_t mapped_count = 0;
  for (iree_host_size_t i = 0; i < binding_count && iree_status_is_ok(status);
       ++i) {
    status = iree_hal_buffer_map_range(
        bindings[i].buffer, IREE_HAL_MAPPING_MODE_SCOPED,
        IREE_HAL_MEMORY_ACCESS_ANY, bindings[i].offset, bindings[i].length,
        &mappings[i]);
    if (iree_status_is_ok(status)) {
      binding_ptrs[i] = mappings[i].contents.data;
      binding_lengths[i] = mappings[i].contents.data_length;
      ++mapped_count;
    }
  }

  iree_hal_buffer_mapping_t parameter_mapping = {{0}};
  bool has_parameter_mapping = false;
  if (iree_status_is_ok(status) && uses_indirect_parameters) {
    status = iree_hal_buffer_map_range(
        config.workgroup_count_ref.buffer, IREE_HAL_MAPPING_MODE_SCOPED,
        IREE_HAL_MEMORY_ACCESS_READ, config.workgroup_count_ref.offset,
        3 * sizeof(uint32_t), &parameter_mapping);
    has_parameter_mapping = iree_status_is_ok(status);
  }

  // Execute all workgroups inline.
  if (iree_status_is_ok(status)) {
    iree_hal_executable_dispatch_state_v0_t dispatch_state = {
        .workgroup_size_x =
            config.workgroup_size[0] ? config.workgroup_size[0] : 1,
        .workgroup_size_y =
            config.workgroup_size[1] ? config.workgroup_size[1] : 1,
        .workgroup_size_z =
            (uint16_t)(config.workgroup_size[2] ? config.workgroup_size[2] : 1),
        .constant_count = (uint16_t)(constants.data_length / sizeof(uint32_t)),
        .workgroup_count_x = config.workgroup_count[0],
        .workgroup_count_y = config.workgroup_count[1],
        .workgroup_count_z = (uint16_t)config.workgroup_count[2],
        .max_concurrency = 1,
        .binding_count = (uint8_t)binding_count,
        .constants = (const uint32_t*)constants.data,
        .binding_ptrs = binding_ptrs,
        .binding_lengths = binding_lengths,
    };
    if (uses_indirect_parameters) {
      const uint32_t* workgroup_count =
          (const uint32_t*)parameter_mapping.contents.data;
      dispatch_state.workgroup_count_x = workgroup_count[0];
      dispatch_state.workgroup_count_y = workgroup_count[1];
      dispatch_state.workgroup_count_z = (uint16_t)workgroup_count[2];
    }
    status = iree_hal_local_executable_issue_dispatch_inline(
        local_executable, function_index, &dispatch_state, /*processor_id=*/0,
        iree_byte_span_empty());
  }

  // Unmap all buffers (even on failure — mappings hold refs).
  if (has_parameter_mapping) {
    status = iree_status_join(status,
                              iree_hal_buffer_unmap_range(&parameter_mapping));
  }
  for (iree_host_size_t i = 0; i < mapped_count; ++i) {
    status =
        iree_status_join(status, iree_hal_buffer_unmap_range(&mappings[i]));
  }

  IREE_TRACE_ZONE_END(z0);
  return status;
}
