// Copyright 2025 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/drivers/amdgpu/registration/driver_module.h"

#include "iree/base/api.h"
#include "iree/base/tooling/flags.h"
#include "iree/hal/drivers/amdgpu/api.h"

IREE_FLAG_LIST(string, amdgpu_libhsa_search_path,
               "Search path (directory or file) for the ROCR-Runtime library "
               "(`libhsa-runtime64.so`, etc).");

IREE_FLAG(int64_t, amdgpu_host_block_pool_small_size, 0,
          "Size in bytes of a small host block in the pool. Must be a power of "
          "two or 0 for the default.");
IREE_FLAG(int64_t, amdgpu_host_block_pool_large_size, 0,
          "Size in bytes of a large host block in the pool. Must be a power of "
          "two or 0 for the default.");
IREE_FLAG(int64_t, amdgpu_host_block_pool_command_buffer_size, 0,
          "Usable size in bytes of a command-buffer recording block in the "
          "host block pool. Must be a power of two or 0 for the default.");

IREE_FLAG(int64_t, amdgpu_device_block_pool_small_size, 0,
          "Size in bytes of a small device block in the pool. Must be a power "
          "of two or 0 for the default.");
IREE_FLAG(int64_t, amdgpu_device_block_pool_small_capacity, 0,
          "Initial small block pool block allocation count in blocks or 0 for "
          "the default.");
IREE_FLAG(int64_t, amdgpu_device_block_pool_large_size, 0,
          "Size in bytes of a large device block in the pool. Must be a power "
          "of two or 0 for the default.");
IREE_FLAG(int64_t, amdgpu_device_block_pool_large_capacity, 0,
          "Initial large block pool block allocation count in blocks or 0 for "
          "the default.");

IREE_FLAG(int64_t, amdgpu_default_pool_range_length, 0,
          "Logical byte length of the default TLSF queue-allocation pool per "
          "physical device or 0 for the default.");
IREE_FLAG(int64_t, amdgpu_default_pool_alignment, 0,
          "Minimum byte alignment for default-pool reservations. Must be a "
          "power of two or 0 for the default.");
IREE_FLAG(int32_t, amdgpu_default_pool_frontier_capacity, 0,
          "Maximum death-frontier entry count stored per free default-pool "
          "block or 0 for the default.");

IREE_FLAG(string, amdgpu_queue_placement, "any",
          "Device queue placement: 'any' (currently host), 'host', or "
          "'device' (reserved and currently unsupported).");

IREE_FLAG(string, amdgpu_command_buffer_mode, "aql",
          "Command-buffer execution mode: 'aql', 'pm4', or 'auto'. PM4 is an "
          "experimental gfx1100 dispatch-only reusable-command-buffer path.");
IREE_FLAG(
    string, amdgpu_pm4_command_buffer_publication_mode, "host-copy",
    "PM4 command-buffer resident publication mode: 'host-copy' writes host "
    "staging builders and publishes populated segments with hsa_memory_copy; "
    "'host-async-copy' publishes one contiguous staging image with "
    "hsa_amd_memory_async_copy and waits in end(); "
    "'host-async-copy-nonblocking' publishes the same staging image and makes "
    "queue execution wait on publication completion.");

IREE_FLAG(bool, amdgpu_preallocate_pools, true,
          "Preallocates a reasonable number of resources in pools to reduce "
          "initial execution latency.");

IREE_FLAG(bool, amdgpu_exclusive_execution, false,
          "Reserved for exclusive queue scheduling; currently unsupported.");

IREE_FLAG(
    bool, amdgpu_force_wait_barrier_defer, false,
    "Forces cross-queue wait barriers through the software deferral path "
    "instead of using the device-side strategy selected from the GPU ISA.");

IREE_FLAG(bool, amdgpu_experimental_pm4_command_buffers, false,
          "Enables PM4 dispatch command-buffer capabilities on unvalidated "
          "gfx9-gfx12 "
          "targets. This is for hardware bring-up only; default automatic PM4 "
          "selection remains limited to validated GPU ISAs.");

IREE_FLAG(int64_t, amdgpu_wait_active_for_ns, 0,
          "Reserved for future HSA active-wait tuning. Must be 0 today.");

static iree_status_t iree_hal_amdgpu_flag_int64_to_host_size(
    const char* flag_name, int64_t flag_value, iree_host_size_t* out_value) {
  if (flag_value < 0) {
    return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                            "%s must be non-negative (got %" PRIi64 ")",
                            flag_name, flag_value);
  }
  *out_value = (iree_host_size_t)flag_value;
  return iree_ok_status();
}

static iree_status_t iree_hal_amdgpu_flag_int64_to_device_size(
    const char* flag_name, int64_t flag_value, iree_device_size_t* out_value) {
  if (flag_value < 0) {
    return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                            "%s must be non-negative (got %" PRIi64 ")",
                            flag_name, flag_value);
  }
  *out_value = (iree_device_size_t)flag_value;
  return iree_ok_status();
}

static iree_status_t iree_hal_amdgpu_driver_factory_enumerate(
    void* self, iree_host_size_t* out_driver_info_count,
    const iree_hal_driver_info_t** out_driver_infos) {
  static const iree_hal_driver_info_t default_driver_info = {
      .driver_name = IREE_SVL("amdgpu"),
      .full_name = IREE_SVL("AMD GPU Driver (HSA/ROCR)"),
  };
  *out_driver_info_count = 1;
  *out_driver_infos = &default_driver_info;
  return iree_ok_status();
}

static iree_status_t iree_hal_amdgpu_driver_factory_try_create(
    void* self, iree_string_view_t driver_name, iree_allocator_t host_allocator,
    iree_hal_driver_t** out_driver) {
  if (!iree_string_view_equal(driver_name, IREE_SV("amdgpu"))) {
    return iree_make_status(IREE_STATUS_UNAVAILABLE,
                            "no driver '%.*s' is provided by this factory",
                            (int)driver_name.size, driver_name.data);
  }

  iree_hal_amdgpu_driver_options_t options;
  iree_hal_amdgpu_driver_options_initialize(&options);
  iree_hal_amdgpu_logical_device_options_t* device_options =
      &options.default_device_options;

  options.libhsa_search_paths = FLAG_amdgpu_libhsa_search_path_list();

  if (FLAG_amdgpu_host_block_pool_small_size) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_host_size(
        "--amdgpu_host_block_pool_small_size",
        FLAG_amdgpu_host_block_pool_small_size,
        &device_options->host_block_pools.small.block_size));
  }
  if (FLAG_amdgpu_host_block_pool_large_size) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_host_size(
        "--amdgpu_host_block_pool_large_size",
        FLAG_amdgpu_host_block_pool_large_size,
        &device_options->host_block_pools.large.block_size));
  }
  if (FLAG_amdgpu_host_block_pool_command_buffer_size) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_host_size(
        "--amdgpu_host_block_pool_command_buffer_size",
        FLAG_amdgpu_host_block_pool_command_buffer_size,
        &device_options->host_block_pools.command_buffer.usable_block_size));
  }

  if (FLAG_amdgpu_device_block_pool_small_size) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_device_size(
        "--amdgpu_device_block_pool_small_size",
        FLAG_amdgpu_device_block_pool_small_size,
        &device_options->device_block_pools.small.block_size));
  }
  if (FLAG_amdgpu_device_block_pool_small_capacity) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_host_size(
        "--amdgpu_device_block_pool_small_capacity",
        FLAG_amdgpu_device_block_pool_small_capacity,
        &device_options->device_block_pools.small.initial_capacity));
  }
  if (FLAG_amdgpu_device_block_pool_large_size) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_device_size(
        "--amdgpu_device_block_pool_large_size",
        FLAG_amdgpu_device_block_pool_large_size,
        &device_options->device_block_pools.large.block_size));
  }
  if (FLAG_amdgpu_device_block_pool_large_capacity) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_host_size(
        "--amdgpu_device_block_pool_large_capacity",
        FLAG_amdgpu_device_block_pool_large_capacity,
        &device_options->device_block_pools.large.initial_capacity));
  }

  if (FLAG_amdgpu_default_pool_range_length) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_device_size(
        "--amdgpu_default_pool_range_length",
        FLAG_amdgpu_default_pool_range_length,
        &device_options->default_pool.range_length));
  }
  if (FLAG_amdgpu_default_pool_alignment) {
    IREE_RETURN_IF_ERROR(iree_hal_amdgpu_flag_int64_to_device_size(
        "--amdgpu_default_pool_alignment", FLAG_amdgpu_default_pool_alignment,
        &device_options->default_pool.alignment));
  }
  if (FLAG_amdgpu_default_pool_frontier_capacity) {
    if (FLAG_amdgpu_default_pool_frontier_capacity < 0) {
      return iree_make_status(
          IREE_STATUS_OUT_OF_RANGE,
          "default pool frontier capacity must be non-negative (got %d)",
          FLAG_amdgpu_default_pool_frontier_capacity);
    }
    if (FLAG_amdgpu_default_pool_frontier_capacity > UINT8_MAX) {
      return iree_make_status(
          IREE_STATUS_OUT_OF_RANGE,
          "default pool frontier capacity %d exceeds maximum %u",
          FLAG_amdgpu_default_pool_frontier_capacity, UINT8_MAX);
    }
    device_options->default_pool.frontier_capacity =
        (uint8_t)FLAG_amdgpu_default_pool_frontier_capacity;
  }

  if (strcmp(FLAG_amdgpu_queue_placement, "any") == 0) {
    device_options->queue_placement = IREE_HAL_AMDGPU_QUEUE_PLACEMENT_ANY;
  } else if (strcmp(FLAG_amdgpu_queue_placement, "host") == 0) {
    device_options->queue_placement = IREE_HAL_AMDGPU_QUEUE_PLACEMENT_HOST;
  } else if (strcmp(FLAG_amdgpu_queue_placement, "device") == 0) {
    device_options->queue_placement = IREE_HAL_AMDGPU_QUEUE_PLACEMENT_DEVICE;
  } else {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "unrecognized queue placement: '%s'",
                            FLAG_amdgpu_queue_placement);
  }

  if (strcmp(FLAG_amdgpu_command_buffer_mode, "aql") == 0) {
    device_options->command_buffer_mode =
        IREE_HAL_AMDGPU_COMMAND_BUFFER_MODE_AQL;
  } else if (strcmp(FLAG_amdgpu_command_buffer_mode, "pm4") == 0) {
    device_options->command_buffer_mode =
        IREE_HAL_AMDGPU_COMMAND_BUFFER_MODE_PM4;
  } else if (strcmp(FLAG_amdgpu_command_buffer_mode, "auto") == 0) {
    device_options->command_buffer_mode =
        IREE_HAL_AMDGPU_COMMAND_BUFFER_MODE_AUTO;
  } else {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "unrecognized command buffer mode: '%s'",
                            FLAG_amdgpu_command_buffer_mode);
  }

  if (strcmp(FLAG_amdgpu_pm4_command_buffer_publication_mode, "host-copy") ==
      0) {
    device_options->pm4_command_buffer_publication_mode =
        IREE_HAL_AMDGPU_PM4_COMMAND_BUFFER_PUBLICATION_MODE_HOST_COPY;
  } else if (strcmp(FLAG_amdgpu_pm4_command_buffer_publication_mode,
                    "host-async-copy") == 0) {
    device_options->pm4_command_buffer_publication_mode =
        IREE_HAL_AMDGPU_PM4_COMMAND_BUFFER_PUBLICATION_MODE_HOST_ASYNC_COPY;
  } else if (strcmp(FLAG_amdgpu_pm4_command_buffer_publication_mode,
                    "host-async-copy-nonblocking") == 0) {
    device_options->pm4_command_buffer_publication_mode =
        IREE_HAL_AMDGPU_PM4_COMMAND_BUFFER_PUBLICATION_MODE_HOST_ASYNC_COPY_NONBLOCKING;
  } else {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "unrecognized PM4 command-buffer publication mode: '%s'",
        FLAG_amdgpu_pm4_command_buffer_publication_mode);
  }

  device_options->preallocate_pools = FLAG_amdgpu_preallocate_pools;

  device_options->exclusive_execution = FLAG_amdgpu_exclusive_execution;

  device_options->force_wait_barrier_defer =
      FLAG_amdgpu_force_wait_barrier_defer;

  device_options->enable_experimental_pm4_command_buffers =
      FLAG_amdgpu_experimental_pm4_command_buffers;

  if (FLAG_amdgpu_wait_active_for_ns < 0) {
    return iree_make_status(
        IREE_STATUS_OUT_OF_RANGE,
        "--amdgpu_wait_active_for_ns must be non-negative (got %" PRIi64 ")",
        FLAG_amdgpu_wait_active_for_ns);
  }
  device_options->wait_active_for_ns = FLAG_amdgpu_wait_active_for_ns;

  iree_status_t status = iree_hal_amdgpu_driver_create(
      driver_name, &options, host_allocator, out_driver);

  return status;
}

IREE_API_EXPORT iree_status_t
iree_hal_amdgpu_driver_module_register(iree_hal_driver_registry_t* registry) {
  static const iree_hal_driver_factory_t factory = {
      .self = NULL,
      .enumerate = iree_hal_amdgpu_driver_factory_enumerate,
      .try_create = iree_hal_amdgpu_driver_factory_try_create,
  };
  return iree_hal_driver_registry_register_factory(registry, &factory);
}
