// Copyright 2022 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/tooling/context_util.h"

#include <memory.h>
#include <string.h>

#include "iree/async/frontier_tracker.h"
#include "iree/async/util/proactor_pool.h"
#include "iree/base/internal/path.h"
#include "iree/base/threading/numa.h"
#include "iree/base/tooling/flags.h"
#include "iree/hal/local/loaders/registration/init.h"
#include "iree/hal/local/plugins/registration/init.h"
#include "iree/hal/replay/recorder.h"
#include "iree/io/file_contents.h"
#include "iree/io/file_handle.h"
#include "iree/modules/hal/inline/module.h"
#include "iree/modules/hal/loader/module.h"
#include "iree/modules/hal/module.h"
#include "iree/tooling/device_util.h"
#include "iree/tooling/modules/resolver.h"
#include "iree/tooling/parameter_util.h"
#include "iree/vm/bytecode/module.h"
#include "iree/vm/dynamic/module.h"

//===----------------------------------------------------------------------===//
// Module loading
//===----------------------------------------------------------------------===//

IREE_FLAG_LIST(
    string, module,
    "A VM module to load; either a vmfb containing a compiled bytecode module\n"
    "or a native system library containing a dynamic native module. Modules\n"
    "are registered in the order defined by the flags with all dependencies\n"
    "for a module needing to have been registered prior to the dependent\n"
    "module. HAL modules are added automatically when required.");

IREE_FLAG(
    string, module_mode, "preload",
    "A module I/O mode of ['preload', 'mmap'].\n"
    "  preload: read entire module into wired memory on startup.\n"
    "  mmap: maps the module file into discardable memory - can increase\n"
    "        warm-up time and variance as mapped pages are swapped\n"
    "        by the OS.");

static iree_status_t iree_tooling_load_bytecode_module(
    iree_vm_instance_t* instance, iree_string_view_t path,
    iree_allocator_t host_allocator, iree_vm_module_t** out_module) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_TEXT(z0, path.data, path.size);

  // Fetch the file contents into memory.
  iree_io_file_contents_t* file_contents = NULL;
  if (iree_string_view_equal(path, IREE_SV("-"))) {
    IREE_RETURN_AND_END_ZONE_IF_ERROR(
        z0, iree_io_file_contents_read_stdin(host_allocator, &file_contents));
  } else {
    if (strcmp(FLAG_module_mode, "mmap") == 0) {
      IREE_RETURN_AND_END_ZONE_IF_ERROR(
          z0, iree_io_file_contents_map(path, IREE_IO_FILE_ACCESS_READ,
                                        host_allocator, &file_contents));
    } else if (strcmp(FLAG_module_mode, "preload") == 0) {
      IREE_RETURN_AND_END_ZONE_IF_ERROR(
          z0, iree_io_file_contents_read(path, host_allocator, &file_contents));
    } else {
      IREE_TRACE_ZONE_END(z0);
      return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                              "unrecognized --module_mode= value '%s'",
                              FLAG_module_mode);
    }
  }

  // Try to load the module as bytecode (all we have today that we can use).
  // We could sniff the file ID and switch off to other module types.
  // The module takes ownership of the file contents (when successful).
  iree_vm_module_t* module = NULL;
  iree_status_t status = iree_vm_bytecode_module_create(
      instance, IREE_VM_BYTECODE_MODULE_FLAG_NONE, file_contents->const_buffer,
      iree_io_file_contents_deallocator(file_contents), host_allocator,
      &module);

  if (iree_status_is_ok(status)) {
    *out_module = module;
  } else {
    iree_io_file_contents_free(file_contents);
  }

  IREE_TRACE_ZONE_END(z0);
  return status;
}

static iree_status_t iree_tooling_load_dynamic_module(
    iree_vm_instance_t* instance, iree_string_view_t path,
    iree_string_view_t export_name, iree_string_view_t params,
    iree_allocator_t host_allocator, iree_vm_module_t** out_module) {
  IREE_TRACE_ZONE_BEGIN(z0);
  IREE_TRACE_ZONE_APPEND_TEXT(z0, path.data, path.size);
  IREE_TRACE_ZONE_APPEND_TEXT(z0, params.data, params.size);

  // Split up params list - comes in `key=value&key=value` form.
  iree_host_size_t param_count = 0;
  iree_uri_split_params(params, 0, &param_count, NULL);
  iree_string_pair_t* param_list = NULL;
  if (param_count > 0) {
    param_list =
        (iree_string_pair_t*)iree_alloca(param_count * sizeof(*param_list));
    iree_uri_split_params(params, param_count, &param_count, param_list);
  }

  iree_status_t status = iree_vm_dynamic_module_load_from_file(
      instance, path, export_name, param_count, param_list, host_allocator,
      out_module);

  IREE_TRACE_ZONE_END(z0);
  return status;
}

iree_status_t iree_tooling_load_modules_from_flags(
    iree_vm_instance_t* instance, iree_allocator_t host_allocator,
    iree_tooling_module_list_t* list) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(list);
  iree_host_size_t new_count = list->count + FLAG_module_list().count;
  if (new_count > list->capacity) {
    return iree_make_status(IREE_STATUS_RESOURCE_EXHAUSTED,
                            "too many modules; currently only %" PRIhsz
                            " are supported but at least %" PRIhsz
                            " are requested",
                            list->capacity, new_count);
  }
  IREE_TRACE_ZONE_BEGIN(z0);

  for (iree_host_size_t i = 0; i < FLAG_module_list().count; ++i) {
    // We support `file.so@export_name?params` syntax to allow loading multiple
    // modules from the same shared library and passing in parameters. When
    // omitted we'll use the default export name.
    iree_string_view_t path, export_name, params;
    iree_string_view_split(FLAG_module_list().values[i], '@', &path,
                           &export_name);
    iree_string_view_split(export_name, '?', &export_name, &params);

    // Load the module based on its (guessed) type.
    iree_vm_module_t* module = NULL;
    if (iree_file_path_is_dynamic_library(path)) {
      IREE_RETURN_AND_END_ZONE_IF_ERROR(
          z0,
          iree_tooling_load_dynamic_module(instance, path, export_name, params,
                                           host_allocator, &module),
          "loading dynamic module at '%.*s'", (int)path.size, path.data);
    } else {
      IREE_RETURN_AND_END_ZONE_IF_ERROR(
          z0,
          iree_tooling_load_bytecode_module(instance, path, host_allocator,
                                            &module),
          "loading bytecode module at '%.*s'", (int)path.size, path.data);
    }

    // Store loaded module in the list. It'll be the caller's responsibility to
    // clean it up even if we fail while loading more.
    list->values[list->count++] = module;
  }

  IREE_TRACE_ZONE_END(z0);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// HAL module device selection policy
//===----------------------------------------------------------------------===//

IREE_FLAG(int32_t, device_lead_allocator, 0,
          "Device ordinal of the lead device that will be used for allocations "
          "when more than one device is available. Only functions when there "
          "are selection requests including all devices and otherwise the "
          "first device in the list will be selected.");

static iree_status_t iree_hal_module_device_allocator_select_specific(
    void* user_data, iree_host_size_t device_count,
    const iree_hal_device_queue_affinity_pair_t* devices,
    iree_hal_memory_type_t memory_types, iree_hal_buffer_usage_t buffer_usage,
    iree_hal_module_device_allocator_select_flags_t flags,
    iree_host_size_t* out_selection) {
  for (iree_host_size_t i = 0; i < device_count; ++i) {
    if (devices[i].device == user_data) {
      *out_selection = i;
      return iree_ok_status();
    }
  }
  *out_selection = 0;
  return iree_ok_status();
}

static iree_status_t iree_tooling_device_lead_allocator_index_from_flags(
    iree_hal_device_group_t* device_group, iree_host_size_t* out_index) {
  IREE_ASSERT_ARGUMENT(device_group);
  IREE_ASSERT_ARGUMENT(out_index);

  const iree_host_size_t device_count =
      iree_hal_device_group_device_count(device_group);
  if (FLAG_device_lead_allocator < 0 ||
      (iree_host_size_t)FLAG_device_lead_allocator >= device_count) {
    return iree_make_status(
        IREE_STATUS_OUT_OF_RANGE,
        "--device_lead_allocator=%d outside the device group range [0, %" PRIhsz
        ")",
        FLAG_device_lead_allocator, device_count);
  }

  *out_index = (iree_host_size_t)FLAG_device_lead_allocator;
  return iree_ok_status();
}

static iree_status_t iree_hal_module_device_policy_from_flags(
    iree_hal_device_group_t* device_group,
    iree_hal_module_device_policy_t* out_policy) {
  IREE_ASSERT_ARGUMENT(out_policy);

  iree_host_size_t lead_index = 0;
  IREE_RETURN_IF_ERROR(iree_tooling_device_lead_allocator_index_from_flags(
      device_group, &lead_index));

  *out_policy = iree_hal_module_device_policy_default();
  out_policy->allocator_select.fn =
      iree_hal_module_device_allocator_select_specific;
  out_policy->allocator_select.user_data =
      iree_hal_device_group_device_at(device_group, lead_index);
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// HAL execution model management
//===----------------------------------------------------------------------===//

IREE_FLAG(string, device_replay_output, "",
          "Writes a HAL replay capture stream to the specified .ireereplay "
          "file. The HAL module device group is wrapped so operations across "
          "all devices are emitted into one ordered replay.");
IREE_FLAG(
    string, device_replay_file_policy, "reference",
    "Policy for imported fd-backed HAL files while recording a HAL replay. "
    "Options are `reference` to capture external file paths, "
    "`capture-ranges` to embed bytes read by queue_read operations, "
    "`capture-all` to embed every file byte inline, or `fail` to reject "
    "fd-backed files.");
IREE_FLAG(
    string, device_replay_file_validation, "identity",
    "Validation metadata captured for referenced fd-backed HAL files. Options "
    "are `identity` for cheap length/dev/inode/mtime validation or `digest` "
    "to read every file byte at capture and replay time. Use `digest` only "
    "when referenced files will be copied or staged to a different "
    "filesystem.");

static bool iree_tooling_device_replay_capture_requested(void) {
  return strlen(FLAG_device_replay_output) != 0;
}

static iree_status_t iree_tooling_parse_device_replay_file_policy(
    iree_string_view_t value,
    iree_hal_replay_recorder_external_file_policy_t* out_policy) {
  IREE_ASSERT_ARGUMENT(out_policy);
  if (iree_string_view_equal(value, IREE_SV("reference"))) {
    *out_policy = IREE_HAL_REPLAY_RECORDER_EXTERNAL_FILE_POLICY_REFERENCE;
    return iree_ok_status();
  } else if (iree_string_view_equal(value, IREE_SV("capture-ranges"))) {
    *out_policy = IREE_HAL_REPLAY_RECORDER_EXTERNAL_FILE_POLICY_CAPTURE_RANGES;
    return iree_ok_status();
  } else if (iree_string_view_equal(value, IREE_SV("capture-all"))) {
    *out_policy = IREE_HAL_REPLAY_RECORDER_EXTERNAL_FILE_POLICY_CAPTURE_ALL;
    return iree_ok_status();
  } else if (iree_string_view_equal(value, IREE_SV("fail"))) {
    *out_policy = IREE_HAL_REPLAY_RECORDER_EXTERNAL_FILE_POLICY_FAIL;
    return iree_ok_status();
  }
  return iree_make_status(
      IREE_STATUS_INVALID_ARGUMENT,
      "--device_replay_file_policy must be `reference`, `capture-ranges`, "
      "`capture-all`, or `fail`; got `%.*s`",
      (int)value.size, value.data);
}

static iree_status_t iree_tooling_parse_device_replay_file_validation(
    iree_string_view_t value,
    iree_hal_replay_recorder_external_file_validation_t* out_validation) {
  IREE_ASSERT_ARGUMENT(out_validation);
  if (iree_string_view_equal(value, IREE_SV("identity"))) {
    *out_validation =
        IREE_HAL_REPLAY_RECORDER_EXTERNAL_FILE_VALIDATION_IDENTITY;
    return iree_ok_status();
  } else if (iree_string_view_equal(value, IREE_SV("digest"))) {
    *out_validation =
        IREE_HAL_REPLAY_RECORDER_EXTERNAL_FILE_VALIDATION_CONTENT_DIGEST;
    return iree_ok_status();
  }
  return iree_make_status(
      IREE_STATUS_INVALID_ARGUMENT,
      "--device_replay_file_validation must be `identity` or `digest`; got "
      "`%.*s`",
      (int)value.size, value.data);
}

static iree_status_t iree_tooling_create_hal_replay_recorder_from_flags(
    iree_allocator_t host_allocator,
    iree_hal_replay_recorder_t** out_replay_recorder) {
  IREE_ASSERT_ARGUMENT(out_replay_recorder);
  *out_replay_recorder = NULL;

  if (!iree_tooling_device_replay_capture_requested()) return iree_ok_status();

  iree_hal_replay_recorder_options_t options =
      iree_hal_replay_recorder_options_default();
  IREE_RETURN_IF_ERROR(iree_tooling_parse_device_replay_file_policy(
      iree_make_cstring_view(FLAG_device_replay_file_policy),
      &options.external_file_policy));
  IREE_RETURN_IF_ERROR(iree_tooling_parse_device_replay_file_validation(
      iree_make_cstring_view(FLAG_device_replay_file_validation),
      &options.external_file_validation));

  iree_io_file_handle_t* file_handle = NULL;
  iree_status_t status = iree_io_file_handle_create(
      IREE_IO_FILE_MODE_WRITE | IREE_IO_FILE_MODE_SEQUENTIAL_SCAN |
          IREE_IO_FILE_MODE_SHARE_READ,
      iree_make_cstring_view(FLAG_device_replay_output),
      /*initial_size=*/0, host_allocator, &file_handle);
  if (iree_status_is_ok(status)) {
    status = iree_hal_replay_recorder_create(
        file_handle, &options, host_allocator, out_replay_recorder);
  }
  iree_io_file_handle_release(file_handle);
  return status;
}

static iree_status_t iree_tooling_close_hal_replay_recorder(
    iree_status_t base_status, iree_hal_replay_recorder_t* replay_recorder) {
  if (replay_recorder) {
    base_status = iree_status_join(
        base_status,
        iree_status_annotate_f(iree_hal_replay_recorder_close(replay_recorder),
                               "closing HAL replay capture"));
    iree_hal_replay_recorder_release(replay_recorder);
  }
  return base_status;
}

static iree_status_t iree_tooling_load_hal_async_module(
    iree_vm_instance_t* instance, iree_string_view_t default_device_uri,
    iree_allocator_t host_allocator, iree_vm_module_t** out_module,
    iree_hal_device_list_t** out_device_list,
    iree_hal_allocator_t** out_device_allocator,
    iree_hal_replay_recorder_t** out_replay_recorder) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(out_module);
  IREE_ASSERT_ARGUMENT(out_device_list);
  IREE_ASSERT_ARGUMENT(out_device_allocator);
  if (*out_device_list || *out_device_allocator) {
    return iree_make_status(
        IREE_STATUS_FAILED_PRECONDITION,
        "async HAL module can not be used with other primary HAL module types");
  }
  *out_module = NULL;
  *out_device_list = NULL;
  *out_device_allocator = NULL;
  if (out_replay_recorder) *out_replay_recorder = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  if (iree_tooling_device_replay_capture_requested() && !out_replay_recorder) {
    IREE_TRACE_ZONE_END(z0);
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "--device_replay_output requires a caller-owned "
                            "replay recorder output");
  }

  // Register required types before creating the module.
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_module_register_all_types(instance));

  // Create a proactor pool for async I/O on the device(s).
  iree_async_proactor_pool_t* proactor_pool = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_async_proactor_pool_create(
              iree_numa_node_count(), /*node_ids=*/NULL,
              iree_async_proactor_pool_options_default(), host_allocator,
              &proactor_pool));

  iree_async_frontier_tracker_t* frontier_tracker = NULL;
  iree_status_t status = iree_async_frontier_tracker_create(
      iree_async_frontier_tracker_options_default(), host_allocator,
      &frontier_tracker);

  // Create the device(s) to use.
  if (iree_string_view_is_empty(default_device_uri)) {
    default_device_uri = iree_hal_default_device_uri();
  }
  iree_hal_device_create_params_t create_params =
      iree_hal_device_create_params_default();
  create_params.proactor_pool = proactor_pool;
  iree_hal_device_list_t* device_list = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_create_devices_from_flags(
        iree_hal_available_driver_registry(), default_device_uri,
        &create_params, host_allocator, &device_list);
  }
  iree_async_proactor_pool_release(proactor_pool);
  if (!iree_status_is_ok(status)) {
    iree_async_frontier_tracker_release(frontier_tracker);
    IREE_TRACE_ZONE_END(z0);
    return status;
  }

  // Build a device group from all enumerated devices.
  iree_hal_device_group_builder_t group_builder;
  iree_hal_device_group_builder_initialize(&group_builder, frontier_tracker);
  iree_async_frontier_tracker_release(frontier_tracker);
  for (iree_host_size_t i = 0;
       i < device_list->count && iree_status_is_ok(status); ++i) {
    status = iree_hal_device_group_builder_add_device(&group_builder,
                                                      device_list->devices[i]);
  }
  iree_hal_device_group_t* device_group = NULL;
  if (iree_status_is_ok(status)) {
    status = iree_hal_device_group_builder_finalize(
        &group_builder, host_allocator, &device_group);
  } else {
    iree_hal_device_group_builder_deinitialize(&group_builder);
  }

  // Optionally replace each group device with a replay-recording wrapper. The
  // module and convenience outputs must all observe the same effective group.
  iree_hal_replay_recorder_t* replay_recorder = NULL;
  iree_hal_device_group_t* replay_device_group = NULL;
  bool replay_init_scope_open = false;
  if (iree_status_is_ok(status) &&
      iree_tooling_device_replay_capture_requested()) {
    status = iree_tooling_create_hal_replay_recorder_from_flags(
        host_allocator, &replay_recorder);
    if (iree_status_is_ok(status)) {
      status = iree_hal_replay_recorder_scope_begin(replay_recorder,
                                                    IREE_SV("init"));
      replay_init_scope_open = iree_status_is_ok(status);
    }
    if (iree_status_is_ok(status)) {
      status = iree_hal_replay_wrap_device_group(
          replay_recorder, device_group, host_allocator, &replay_device_group);
    }
  }

  iree_hal_device_group_t* module_device_group =
      replay_device_group ? replay_device_group : device_group;

  // Pick the lead device we'll use for bookkeeping.
  iree_hal_device_t* device = NULL;
  iree_hal_allocator_t* device_allocator = NULL;
  iree_host_size_t lead_index = 0;
  if (iree_status_is_ok(status)) {
    status = iree_tooling_device_lead_allocator_index_from_flags(
        module_device_group, &lead_index);
  }
  if (iree_status_is_ok(status)) {
    device = iree_hal_device_group_device_at(module_device_group, lead_index);
    device_allocator = iree_hal_device_allocator(device);
    iree_hal_allocator_retain(device_allocator);
  }

  // Create HAL module wrapping the device group.
  iree_hal_module_flags_t flags = IREE_HAL_MODULE_FLAG_NONE;
  if (iree_status_is_ok(status)) {
    bool retain_profile_metadata = false;
    status =
        iree_hal_profiling_from_flags_requires_retained_command_buffer_metadata(
            &retain_profile_metadata);
    if (iree_status_is_ok(status) && retain_profile_metadata) {
      flags |= IREE_HAL_MODULE_FLAG_RETAIN_COMMAND_BUFFER_PROFILE_METADATA;
    }
  }
  iree_vm_module_t* module = NULL;
  if (iree_status_is_ok(status)) {
    iree_hal_module_device_policy_t device_policy;
    status = iree_hal_module_device_policy_from_flags(module_device_group,
                                                      &device_policy);
    if (!iree_status_is_ok(status)) {
      iree_hal_allocator_release(device_allocator);
      device_allocator = NULL;
      device = NULL;
    } else {
      status = iree_hal_module_create(
          instance, device_policy, module_device_group, flags,
          iree_hal_module_debug_sink_stdio(stderr), host_allocator, &module);
    }
  }
  if (replay_init_scope_open) {
    status = iree_status_join(status, iree_hal_replay_recorder_scope_end(
                                          replay_recorder, IREE_SV("init")));
  }
  iree_hal_device_group_release(replay_device_group);
  iree_hal_device_group_release(device_group);

  if (iree_status_is_ok(status)) {
    *out_module = module;
    *out_device_list = device_list;
    *out_device_allocator = device_allocator;
    if (out_replay_recorder) {
      *out_replay_recorder = replay_recorder;
      replay_recorder = NULL;
    }
  } else {
    status = iree_tooling_close_hal_replay_recorder(status, replay_recorder);
    iree_hal_allocator_release(device_allocator);
    iree_hal_device_list_free(device_list);
    iree_vm_module_release(module);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

// Creates a HAL device allocator for host-local heap usage.
// TODO(benvanik): generalize to allocator wrapper flag (suballocator, etc).
static iree_status_t iree_tooling_create_inline_device_allocator_from_flags(
    iree_allocator_t host_allocator,
    iree_hal_allocator_t** out_device_allocator) {
  IREE_TRACE_ZONE_BEGIN(z0);
  iree_status_t status = iree_hal_allocator_create_heap(
      IREE_SV("heap"), host_allocator, host_allocator, out_device_allocator);
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static iree_status_t iree_tooling_load_hal_inline_module(
    iree_vm_instance_t* instance, iree_allocator_t host_allocator,
    iree_vm_module_t** out_module,
    iree_hal_allocator_t** out_device_allocator) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(out_module);
  IREE_ASSERT_ARGUMENT(out_device_allocator);
  if (*out_device_allocator) {
    return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
                            "inline HAL module cannot be used with other "
                            "primary HAL module types");
  }
  *out_module = NULL;
  *out_device_allocator = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  // Register required types before creating the module.
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_module_register_inline_types(instance));

  // Create default heap device allocator.
  iree_hal_allocator_t* device_allocator = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_tooling_create_inline_device_allocator_from_flags(
              host_allocator, &device_allocator));

  // Create the module; it's immutable and can be reused but we don't do that in
  // this tooling.
  iree_hal_inline_module_flags_t flags = IREE_HAL_INLINE_MODULE_FLAG_NONE;
  iree_vm_module_t* module = NULL;
  iree_status_t status = iree_hal_inline_module_create(
      instance, flags, iree_hal_module_debug_sink_stdio(stderr),
      device_allocator, host_allocator, &module);

  if (iree_status_is_ok(status)) {
    *out_module = module;
    *out_device_allocator = device_allocator;
  } else {
    iree_hal_allocator_release(device_allocator);
    iree_vm_module_release(module);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

static iree_status_t iree_tooling_load_hal_loader_module(
    iree_vm_instance_t* instance, iree_allocator_t host_allocator,
    iree_vm_module_t** out_module) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(out_module);
  *out_module = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  // Register required types before creating the module.
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_module_register_loader_types(instance));

  // Create plugin manager for executable imports.
  iree_hal_executable_plugin_manager_t* plugin_manager = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_hal_executable_plugin_manager_create_from_flags(
              host_allocator, &plugin_manager));

  // Create all executable loaders built into the binary.
  // We could allow users to choose the set with a flag.
  iree_host_size_t loader_count = 0;
  iree_hal_executable_loader_t* loaders[16];
  iree_status_t status = iree_hal_create_all_available_executable_loaders(
      plugin_manager, IREE_ARRAYSIZE(loaders), &loader_count, loaders,
      host_allocator);

  // Create the module; it retains the loaders for its lifetime.
  iree_vm_module_t* module = NULL;
  if (iree_status_is_ok(status)) {
    iree_hal_loader_module_flags_t flags = IREE_HAL_LOADER_MODULE_FLAG_NONE;
    status = iree_hal_loader_module_create(instance, flags, loader_count,
                                           loaders, host_allocator, &module);
  }

  // Always release loaders; loader module has retained them.
  for (iree_host_size_t i = 0; i < loader_count; ++i) {
    iree_hal_executable_loader_release(loaders[i]);
  }
  iree_hal_executable_plugin_manager_release(plugin_manager);

  if (iree_status_is_ok(status)) {
    *out_module = module;
  } else {
    iree_vm_module_release(module);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

//===----------------------------------------------------------------------===//
// Module management
//===----------------------------------------------------------------------===//

void iree_tooling_module_list_initialize(iree_tooling_module_list_t* out_list) {
  out_list->capacity = IREE_ARRAYSIZE(out_list->values);
  out_list->count = 0;
  memset(out_list->values, 0, sizeof(out_list->values));
}

void iree_tooling_module_list_clone(
    const iree_tooling_module_list_t* source_list,
    iree_tooling_module_list_t* out_list) {
  iree_tooling_module_list_initialize(out_list);
  for (iree_host_size_t i = 0; i < source_list->count; ++i) {
    iree_vm_module_t* module = source_list->values[i];
    iree_vm_module_retain(module);
    out_list->values[out_list->count++] = module;
  }
}

void iree_tooling_module_list_reset(iree_tooling_module_list_t* list) {
  for (iree_host_size_t i = 0; i < list->count; ++i) {
    iree_vm_module_release(list->values[i]);
  }
  list->count = 0;
}

// Returns true if |list| contains a module with the given |module_name|.
static bool iree_tooling_module_list_contains(
    const iree_tooling_module_list_t* list, iree_string_view_t module_name) {
  for (iree_host_size_t i = 0; i < list->count; ++i) {
    if (iree_string_view_equal(iree_vm_module_name(list->values[i]),
                               module_name)) {
      return true;
    }
  }
  return false;
}

iree_status_t iree_tooling_module_list_push_back(
    iree_tooling_module_list_t* list, iree_vm_module_t* module) {
  if (list->count + 1 > list->capacity) {
    return iree_make_status(IREE_STATUS_OUT_OF_RANGE,
                            "resolved module list capacity %" PRIhsz
                            " too small to fit all resolved modules",
                            list->capacity);
  }
  iree_vm_module_retain(module);
  list->values[list->count++] = module;
  return iree_ok_status();
}

iree_vm_module_t* iree_tooling_module_list_back(
    const iree_tooling_module_list_t* list) {
  return list->count ? list->values[list->count - 1] : NULL;
}

typedef struct {
  // VM instance used to register dependency modules.
  iree_vm_instance_t* instance;
  // Host allocator used for transient and retained resolved objects.
  iree_allocator_t host_allocator;
  // Resolved module list receiving discovered dependency modules.
  iree_tooling_module_list_t* resolved_list;
  // Default device URI used when --device= was not specified.
  iree_string_view_t default_device_uri;
  // HAL device list retained from the resolved HAL module, if any.
  iree_hal_device_list_t* device_list;
  // Lead HAL allocator retained from the resolved HAL module, if any.
  iree_hal_allocator_t* device_allocator;
  // Optional replay recorder retained when --device_replay_output= is
  // specified.
  iree_hal_replay_recorder_t* replay_recorder;
} iree_tooling_resolve_state_t;
static iree_status_t iree_tooling_resolve_module_dependency_callback(
    void* user_data_ptr, const iree_vm_module_dependency_t* dependency) {
  iree_tooling_resolve_state_t* state =
      (iree_tooling_resolve_state_t*)user_data_ptr;
  if (iree_tooling_module_list_contains(state->resolved_list,
                                        dependency->name)) {
    // Already registered (redundant system dep or another user module).
    return iree_ok_status();
  }

  // Register one of the known modules. Note that today this is not recursive
  // but it could be in the future.
  iree_vm_module_t* module = NULL;
  if (iree_string_view_equal(dependency->name, IREE_SV("hal"))) {
    IREE_RETURN_IF_ERROR(iree_tooling_load_hal_async_module(
        state->instance, state->default_device_uri, state->host_allocator,
        &module, &state->device_list, &state->device_allocator,
        &state->replay_recorder));
  } else if (iree_string_view_equal(dependency->name, IREE_SV("hal_inline"))) {
    IREE_RETURN_IF_ERROR(iree_tooling_load_hal_inline_module(
        state->instance, state->host_allocator, &module,
        &state->device_allocator));
  } else if (iree_string_view_equal(dependency->name, IREE_SV("hal_loader"))) {
    IREE_RETURN_IF_ERROR(iree_tooling_load_hal_loader_module(
        state->instance, state->host_allocator, &module));
  } else if (iree_string_view_equal(dependency->name,
                                    IREE_SV("io_parameters"))) {
    IREE_RETURN_IF_ERROR(iree_tooling_create_parameters_module_from_flags(
        state->instance, /*additional_provider_count=*/0,
        /*additional_providers=*/NULL, state->host_allocator, &module));
  } else {
    // Defer to the generic module resolver registry.
    IREE_RETURN_IF_ERROR(iree_tooling_resolve_module_dependency(
        state->instance, dependency, state->host_allocator, &module));
  }
  if (!module) return iree_ok_status();

  iree_status_t status =
      iree_tooling_module_list_push_back(state->resolved_list, module);
  iree_vm_module_release(module);
  return status;
}

// NOTE: today we don't have fancy resolution and only check for specific system
// modules. A dynamic resolution utility that made callbacks for each module
// would be really handy to have in the base API.
//
// With our current single-pass approach we just scan the user modules to find
// the system deps and then append the user modules. We could instead do a
// recursive scan to build a potential set list, coalesce to find the minimum
// required versions, and then call a user callback to load runtime-registered
// modules.
iree_status_t iree_tooling_resolve_modules(
    iree_vm_instance_t* instance, iree_host_size_t user_module_count,
    iree_vm_module_t** user_modules, iree_string_view_t default_device_uri,
    iree_allocator_t host_allocator, iree_tooling_module_list_t* resolved_list,
    iree_hal_device_list_t** out_device_list,
    iree_hal_allocator_t** out_device_allocator,
    iree_hal_replay_recorder_t** out_replay_recorder) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(!user_module_count || user_modules);
  IREE_ASSERT_ARGUMENT(resolved_list);
  if (out_device_list) *out_device_list = NULL;
  if (out_device_allocator) *out_device_allocator = NULL;
  if (out_replay_recorder) *out_replay_recorder = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  if (iree_tooling_device_replay_capture_requested() && !out_replay_recorder) {
    IREE_TRACE_ZONE_END(z0);
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "--device_replay_output requires a caller-owned "
                            "replay recorder output");
  }

  iree_tooling_resolve_state_t resolve_state = {
      .instance = instance,
      .host_allocator = host_allocator,
      .resolved_list = resolved_list,
      .default_device_uri = default_device_uri,
      .device_list = NULL,
      .device_allocator = NULL,
      .replay_recorder = NULL,
  };
  iree_status_t status = iree_ok_status();
  for (iree_host_size_t i = 0; i < user_module_count; ++i) {
    iree_vm_module_t* user_module = user_modules[i];
    status = iree_vm_module_enumerate_dependencies(
        user_module, iree_tooling_resolve_module_dependency_callback,
        &resolve_state);
    if (!iree_status_is_ok(status)) {
      iree_string_view_t module_name = iree_vm_module_name(user_module);
      (void)module_name;
      status =
          iree_status_annotate_f(status, "resolving dependencies for '%.*s'",
                                 (int)module_name.size, module_name.data);
      break;
    }
    status = iree_tooling_module_list_push_back(resolved_list, user_modules[i]);
    if (!iree_status_is_ok(status)) break;
  }

  if (iree_status_is_ok(status) &&
      iree_tooling_device_replay_capture_requested() &&
      !resolve_state.replay_recorder) {
    status = iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "--device_replay_output requires a module importing the async HAL "
        "module");
  }

  if (iree_status_is_ok(status)) {
    if (out_device_allocator) {
      *out_device_allocator = resolve_state.device_allocator;
    } else {
      iree_hal_allocator_release(resolve_state.device_allocator);
    }
    if (out_device_list) {
      *out_device_list = resolve_state.device_list;
    } else {
      iree_hal_device_list_free(resolve_state.device_list);
    }
    if (out_replay_recorder) {
      *out_replay_recorder = resolve_state.replay_recorder;
    } else {
      iree_hal_replay_recorder_release(resolve_state.replay_recorder);
    }
  } else {
    status = iree_tooling_close_hal_replay_recorder(
        status, resolve_state.replay_recorder);
    iree_hal_allocator_release(resolve_state.device_allocator);
    iree_hal_device_list_free(resolve_state.device_list);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

iree_status_t iree_tooling_find_single_exported_function(
    iree_vm_module_t* module, iree_vm_function_t* out_function) {
  memset(out_function, 0, sizeof(*out_function));
  iree_vm_module_signature_t module_signature =
      iree_vm_module_signature(module);
  iree_host_size_t exported_functions = 0;
  for (iree_host_size_t i = 0; i < module_signature.export_function_count;
       ++i) {
    iree_vm_function_t function = {0};
    IREE_RETURN_IF_ERROR(
        iree_vm_module_lookup_function_by_ordinal(
            module, IREE_VM_FUNCTION_LINKAGE_EXPORT, i, &function),
        "looking up function export %" PRIhsz, i);
    iree_string_view_t function_name = iree_vm_function_name(&function);
    if (iree_string_view_starts_with(function_name,
                                     iree_make_cstring_view("__")) ||
        iree_string_view_find_char(function_name, '$', 0) !=
            IREE_STRING_VIEW_NPOS) {
      // Function was either internal or special; we don't want to run these
      // as they have special ABI requirements or must only be called in
      // specific situations (module initializers, etc).
      continue;
    }
    if (exported_functions == 0) *out_function = function;
    ++exported_functions;
  }
  if (exported_functions == 0) {
    return iree_make_status(
        IREE_STATUS_INVALID_ARGUMENT,
        "no exported functions found in module; at least one must be present");
  } else if (exported_functions > 1) {
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "more than one exported function present; "
                            "--function= must be specified explicitly");
  }
  return iree_ok_status();
}

//===----------------------------------------------------------------------===//
// Context management
//===----------------------------------------------------------------------===//

IREE_FLAG(bool, trace_execution, false, "Traces VM execution to stderr.");

iree_status_t iree_tooling_create_instance(iree_allocator_t host_allocator,
                                           iree_vm_instance_t** out_instance) {
  IREE_ASSERT_ARGUMENT(out_instance);
  *out_instance = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  iree_vm_instance_t* instance = NULL;
  IREE_RETURN_AND_END_ZONE_IF_ERROR(
      z0, iree_vm_instance_create(IREE_VM_TYPE_CAPACITY_DEFAULT, host_allocator,
                                  &instance));

  // HACK: to load modules we need the types registered even though we don't
  // know if the types are used.
  iree_status_t status = iree_hal_module_register_all_types(instance);
  if (iree_status_is_ok(status)) {
    status = iree_tooling_register_all_module_types(instance);
  }

  if (iree_status_is_ok(status)) {
    *out_instance = instance;
  } else {
    iree_vm_instance_release(instance);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}

iree_status_t iree_tooling_create_context_from_flags(
    iree_vm_instance_t* instance, iree_host_size_t user_module_count,
    iree_vm_module_t** user_modules, iree_string_view_t default_device_uri,
    iree_allocator_t host_allocator, iree_vm_context_t** out_context,
    iree_hal_device_list_t** out_device_list,
    iree_hal_allocator_t** out_device_allocator,
    iree_hal_replay_recorder_t** out_replay_recorder) {
  IREE_ASSERT_ARGUMENT(instance);
  IREE_ASSERT_ARGUMENT(!user_module_count || user_modules);
  IREE_ASSERT_ARGUMENT(out_context);
  *out_context = NULL;
  if (out_device_list) *out_device_list = NULL;
  if (out_device_allocator) *out_device_allocator = NULL;
  if (out_replay_recorder) *out_replay_recorder = NULL;
  IREE_TRACE_ZONE_BEGIN(z0);

  if (iree_tooling_device_replay_capture_requested() && !out_replay_recorder) {
    IREE_TRACE_ZONE_END(z0);
    return iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
                            "--device_replay_output requires a caller-owned "
                            "replay recorder output");
  }

  // Resolve all module dependencies into an ordered list.
  // All modules are retained in the list.
  iree_tooling_module_list_t resolved_list;
  iree_tooling_module_list_initialize(&resolved_list);
  iree_hal_device_list_t* device_list = NULL;
  iree_hal_allocator_t* device_allocator = NULL;
  iree_hal_replay_recorder_t* replay_recorder = NULL;
  iree_status_t status = iree_tooling_resolve_modules(
      instance, user_module_count, user_modules, default_device_uri,
      host_allocator, &resolved_list, &device_list, &device_allocator,
      &replay_recorder);
  if (!iree_status_is_ok(status)) {
    iree_tooling_module_list_reset(&resolved_list);
    IREE_TRACE_ZONE_END(z0);
    return status;
  }

  iree_vm_context_flags_t flags = IREE_VM_CONTEXT_FLAG_NONE;
  if (FLAG_trace_execution) {
    // This enables tracing for all invocations but even if not set each
    // invocation can have the flag specified to trace.
    flags |= IREE_VM_CONTEXT_FLAG_TRACE_EXECUTION;
  }

  // Create the context with the full list of resolved modules.
  // The context retains the modules and we can release them afterward.
  iree_vm_context_t* context = NULL;
  status = iree_vm_context_create_with_modules(
      instance, flags, resolved_list.count, resolved_list.values,
      host_allocator, &context);
  iree_tooling_module_list_reset(&resolved_list);

  // If no device allocator was created we'll create a default one just so that
  // callers have something to create buffer views from. This isn't strictly
  // required but a lot of tests do things like pass in buffers even if no HAL
  // methods are used and the HAL module is not needed.
  if (iree_status_is_ok(status) && !device_allocator && out_device_allocator) {
    status = iree_tooling_create_inline_device_allocator_from_flags(
        host_allocator, &device_allocator);
  }

  if (iree_status_is_ok(status)) {
    *out_context = context;
    if (out_device_allocator) {
      *out_device_allocator = device_allocator;
    } else {
      iree_hal_allocator_release(device_allocator);
    }
    if (out_device_list) {
      *out_device_list = device_list;
    } else {
      iree_hal_device_list_free(device_list);
    }
    if (out_replay_recorder) {
      *out_replay_recorder = replay_recorder;
    } else {
      iree_hal_replay_recorder_release(replay_recorder);
    }
  } else {
    status = iree_tooling_close_hal_replay_recorder(status, replay_recorder);
    iree_hal_allocator_release(device_allocator);
    iree_hal_device_list_free(device_list);
    iree_vm_context_release(context);
  }
  IREE_TRACE_ZONE_END(z0);
  return status;
}
