// Copyright 2021 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 <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "iree/base/api.h"
#include "iree/base/internal/flags.h"
#include "iree/base/internal/path.h"
#include "iree/hal/api.h"
#include "iree/tooling/device_util.h"
#include "iree/tooling/trace_replay.h"
#include "iree/tooling/vm_util.h"
#include "iree/tooling/yaml_util.h"
#include "iree/vm/api.h"

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

IREE_FLAG(bool, print_statistics, false,
          "Prints runtime statistics to stderr on exit.");

IREE_FLAG(bool, print_calls, false, "Prints all I/O for each call to stdout.");
IREE_FLAG(bool, print_call_inputs, false,
          "Prints all inputs for each call before they are made to stdout.");
IREE_FLAG(bool, print_call_outputs, false,
          "Prints all outputs for each call after they are made to stdout.");

IREE_FLAG_LIST(
    string, input,
    "An input (a) value or (b) buffer of the format:\n"
    "  (a) scalar value\n"
    "     value\n"
    "     e.g.: --input=\"3.14\"\n"
    "  (b) buffer:\n"
    "     [shape]xtype=[value]\n"
    "     e.g.: --input=\"2x2xi32=1 2 3 4\"\n"
    "Optionally, brackets may be used to separate the element values:\n"
    "  2x2xi32=[[1 2][3 4]]\n"
    "Raw binary files can be read to provide buffer contents:\n"
    "  2x2xi32=@some/file.bin\n"
    "\n"
    "Numpy npy files from numpy.save can be read to provide 1+ values:\n"
    "  @some.npy\n"
    "\n"
    "Each occurrence of the flag indicates an input in the order they were\n"
    "specified on the command line.");

IREE_FLAG_LIST(
    string, output,
    "Specifies how to handle an output from the invocation:\n"
    "  `` (empty): ignore output\n"
    "     e.g.: --output=\n"
    "  `-`: print textual form to stdout\n"
    "     e.g.: --output=-\n"
    "  `@file.npy`: create/overwrite a numpy npy file and write buffer view\n"
    "     e.g.: --output=@file.npy\n"
    "  `+file.npy`: create/append a numpy npy file and write buffer view\n"
    "     e.g.: --output=+file.npy\n"
    "\n"
    "Numpy npy files can be read in Python using numpy.load, for example an\n"
    "invocation producing two outputs can be concatenated as:\n"
    "    --output=@file.npy --output=+file.npy\n"
    "And then loaded in Python by reading from the same file:\n"
    "  with open('file.npy', 'rb') as f:\n"
    "    print(numpy.load(f))\n"
    "    print(numpy.load(f))\n"
    "\n"
    "Each occurrence of the flag indicates an output in the order they were\n"
    "specified on the command line.");

IREE_FLAG_LIST(string, expected_output,
               "An expected function output following the same format as "
               "--input. When present the results of the "
               "invocation will be compared against these values and the "
               "tool will return non-zero if any differ. If the value of a "
               "particular output is not of interest provide `(ignored)`.");

IREE_FLAG(int32_t, output_max_element_count, 1024,
          "Prints up to the maximum number of elements of output tensors, "
          "eliding the remainder.");

static iree_status_t iree_trace_replay_call_before(void* user_data,
                                                   iree_trace_replay_t* replay,
                                                   yaml_document_t* document,
                                                   yaml_node_t* event_node,
                                                   iree_vm_function_t function,
                                                   iree_vm_list_t* input_list) {
  if (FLAG_print_calls || FLAG_print_call_inputs) {
    iree_string_view_t function_name = iree_vm_function_name(&function);
    fprintf(stdout, "--- CALL[%.*s] ---\n", (int)function_name.size,
            function_name.data);
    IREE_RETURN_IF_ERROR(iree_tooling_variant_list_fprint(
        IREE_SV("arg"), input_list,
        (iree_host_size_t)FLAG_output_max_element_count, stdout));
  }
  return iree_ok_status();
}

static iree_status_t iree_trace_replay_call_after(void* user_data,
                                                  iree_trace_replay_t* replay,
                                                  yaml_document_t* document,
                                                  yaml_node_t* event_node,
                                                  iree_vm_function_t function,
                                                  iree_vm_list_t* output_list) {
  if (FLAG_print_calls || FLAG_print_call_outputs) {
    if (!FLAG_print_calls && !FLAG_print_call_inputs) {
      iree_string_view_t function_name = iree_vm_function_name(&function);
      fprintf(stdout, "--- CALL[%.*s] ---\n", (int)function_name.size,
              function_name.data);
    }
    IREE_RETURN_IF_ERROR(iree_tooling_variant_list_fprint(
        IREE_SV("result"), output_list,
        (iree_host_size_t)FLAG_output_max_element_count, stdout));
  }
  return iree_ok_status();
}

// Runs the trace in |file| using |root_path| as the base for any path lookups
// required for external files referenced in |file|.
static iree_status_t iree_run_trace_file(iree_string_view_t root_path,
                                         FILE* file,
                                         iree_vm_instance_t* instance) {
  iree_trace_replay_flags_t replay_flags = IREE_TRACE_REPLAY_FLAG_NONE;
  if (FLAG_print_statistics) {
    replay_flags |= IREE_TRACE_REPLAY_FLAG_PRINT_STATISTICS;
  }

  iree_vm_context_flags_t context_flags = IREE_VM_CONTEXT_FLAG_NONE;
  if (FLAG_trace_execution) {
    context_flags |= IREE_VM_CONTEXT_FLAG_TRACE_EXECUTION;
  }

  iree_trace_replay_t replay;
  IREE_RETURN_IF_ERROR(iree_trace_replay_initialize(
      root_path, instance, replay_flags, context_flags,
      iree_hal_available_driver_registry(), iree_allocator_system(), &replay));

  // Hook into all calls processed during the trace.
  replay.call_hooks.user_data = NULL;
  replay.call_hooks.before = iree_trace_replay_call_before;
  replay.call_hooks.after = iree_trace_replay_call_after;

  // Query device overrides, if any. When omitted the devices from the trace
  // file will be used.
  // TODO(#5724): remove this and instead provide a device set on initialize.
  iree_host_size_t device_uri_count = 0;
  const iree_string_view_t* device_uris = NULL;
  iree_hal_get_devices_flag_list(&device_uri_count, &device_uris);
  iree_trace_replay_set_hal_devices_override(&replay, device_uri_count,
                                             device_uris);

  yaml_parser_t parser;
  if (!yaml_parser_initialize(&parser)) {
    iree_trace_replay_deinitialize(&replay);
    return iree_make_status(IREE_STATUS_INTERNAL,
                            "yaml_parser_initialize failed");
  }
  yaml_parser_set_input_file(&parser, file);

  bool have_parsed_inputs = false;
  iree_status_t status = iree_ok_status();
  for (bool document_eof = false; !document_eof;) {
    // Parse the subdocument event.
    yaml_document_t document;
    if (!yaml_parser_load(&parser, &document)) {
      status = iree_status_from_yaml_parser_error(&parser);
      break;
    }

    // Execute the event or handle EOF (empty document).
    yaml_node_t* event_node = yaml_document_get_root_node(&document);
    if (event_node) {
      status = iree_trace_replay_event(&replay, &document, event_node);
    } else {
      document_eof = true;
    }

    // Reclaim subdocument resources before moving on to the next.
    yaml_document_delete(&document);
    if (!iree_status_is_ok(status)) break;

    // If the event created a device and we haven't yet performed our input
    // loading we can do that now before processing subsequent events.
    if (!have_parsed_inputs && replay.device) {
      status = iree_tooling_parse_into_variant_list(
          iree_hal_device_allocator(replay.device), FLAG_input_list().values,
          FLAG_input_list().count, replay.host_allocator, replay.inputs);
      have_parsed_inputs = true;
    }
    if (!iree_status_is_ok(status)) break;
  }

  yaml_parser_delete(&parser);

  // Optionally process outputs from the replay session.
  if (iree_status_is_ok(status)) {
    if (FLAG_output_list().count == 0) {
      IREE_RETURN_IF_ERROR(
          iree_tooling_variant_list_fprint(
              IREE_SV("output"), replay.outputs,
              (iree_host_size_t)FLAG_output_max_element_count, stdout),
          "printing results");
    } else {
      IREE_RETURN_IF_ERROR(
          iree_tooling_output_variant_list(
              replay.outputs, FLAG_output_list().values,
              FLAG_output_list().count,
              (iree_host_size_t)FLAG_output_max_element_count, stdout),
          "outputting results");
    }
  }

  iree_trace_replay_deinitialize(&replay);
  return status;
}

// Runs each of the given traces files sequentially in isolated contexts.
static iree_status_t iree_run_trace_files(int file_count, char** file_paths,
                                          iree_vm_instance_t* instance) {
  for (int i = 0; i < file_count; ++i) {
    iree_string_view_t file_path = iree_make_cstring_view(file_paths[i]);
    iree_string_view_t root_path = iree_file_path_dirname(file_path);
    FILE* file = fopen(file_paths[i], "rb");
    if (!file) {
      return iree_make_status(iree_status_code_from_errno(errno),
                              "failed to open trace file '%.*s'",
                              (int)file_path.size, file_path.data);
    }
    iree_status_t status = iree_run_trace_file(root_path, file, instance);
    fclose(file);
    IREE_RETURN_IF_ERROR(status, "replaying trace file '%.*s'",
                         (int)file_path.size, file_path.data);
  }
  return iree_ok_status();
}

int main(int argc, char** argv) {
  iree_flags_set_usage(
      "iree-run-trace",
      "Executes a YAML trace file containing a sequence of context operations\n"
      "and calls represented as subdocuments.\n"
      "\n"
      "Example loading a bytecode module and calling a function:\n"
      "\n"
      "```yaml\n"
      "type: context_load\n"
      "---\n"
      "type: module_load\n"
      "module:\n"
      "  type: buildin\n"
      "  name: hal\n"
      "---\n"
      "type: module_load\n"
      "module:\n"
      "  type: bytecode\n"
      "  path: ../build/some_module.vmfb\n"
      "---\n"
      "type: call\n"
      "function: module.mul\n"
      "args:\n"
      "- !input.take 0\n"
      "- !input.take 1\n"
      "results:\n"
      "- !output.push\n"
      "- !output.push\n"
      "```\n"
      "\n"
      "This can be invoked like iree-run-module specifying inputs/outputs:\n"
      "  iree-run-trace trace.yml    \\\n"
      "      --device=local-sync     \\\n"
      "      --input=4xf32=0,1,2,3,4 \\\n"
      "      --input=@input1.npy     \\\n"
      "      --output=@outputs.npy   \\\n"
      "      --output=+outputs.npy\n"
      "\n"
      "In addition to `--input=`/`--output=` flag access a user-defined\n"
      "blackboard exists for preserving temporary values used within the\n"
      "trace. Blackboard slots are defined by ordinal and they can be used\n"
      "in any context and input/output can be, `!blackboard.get` instead of\n"
      "`!input.get` and `!blackboard.set` instead of `!output.set`.\n"
      "\n"
      "--- Events ---\n"
      "\n"
      "`type: context_load`\n"
      "Loads an empty VM context with no modules registered.\n"
      "\n"
      "`type: module_load`\n"
      "Loads a module into the current context. Modules may either be\n"
      "`builtin` (compiled into the binary) or dynamically-loaded `bytecode`.\n"
      "\n"
      "`type: blackboard_clear`\n"
      "Clears the contents of the blackboard and resets it to 0 elements.\n"
      "\n"
      "`type: assign`\n"
      "Assigns sources from a `from` sequence to targets in a `to` sequence.\n"
      "Equivalent to an identity function call and can be used to move\n"
      "between inputs, outputs, and the blackboard.\n"
      "\n"
      "`type: numpy_load`\n"
      "Loads one or more ndarrays from a .npy value. Each array has a target\n"
      "where the array will be retained such as `!blackboard.set 2`.\n"
      "\n"
      "`type: numpy_save\n"
      "Saves one or more ndarrays to a .npy value. Each array has a source\n"
      "where the array will be taken from such as `!blackboard.get 2`.\n"
      "\n"
      "`type: call`\n"
      "Invokes a function in the context by fully-qualified `function` name.\n"
      "Uses arguments from an `args` sequence and produces results into a\n"
      "`results` sequence.\n"
      "\n"
      "--- Sources ---\n"
      "\n"
      "`type: null`\n"
      "A null ref value.\n"
      "\n"
      "`!hal.buffer_view 4xf32=0,1,2,3`\n"
      "A constant iree_hal_buffer_view_t/!hal.buffer_view value using the\n"
      "same formatting as iree-run-module's `--input=` flag.\n"
      "\n"
      "`!hal.buffer 4xf32=0,1,2,3`\n"
      "An initialized iree_hal_buffer_t/!hal.buffer without the wrapping view\n"
      "metadata.\n"
      "\n"
      "`!input.get ORDINAL` / `!input.take ORDINAL`\n"
      "Returns a reference to `--input=` flag at ORDINAL. Note that a single\n"
      "npy file may expand to multiple inputs. The `take` variant transfers\n"
      "ownership and clears the slot in the list and is recommended to avoid\n"
      "keeping unneeded inputs around for the duration of the trace.\n"
      "\n"
      "`!output.get ORDINAL` / `!output.take ORDINAL`\n"
      "Returns a reference to the `--output=` flag at ORDINAL. These are\n"
      "initially empty until assigned by the trace.\n"
      "\n"
      "`!blackboard.get ORDINAL` / `!blackboard.take ORDINAL`\n"
      "Returns a reference to the blackboard slot ORDINAL. The blackboard is\n"
      "initially empty and slots must be assigned in order to define them.\n"
      "The `take` variant transfers ownership and clears the slot in the\n"
      "blackboard and is recommended to avoid keeping large resources live\n"
      "in the blackboard longer than they need to be.\n"
      "\n"
      "--- Targets ---\n"
      "\n"
      "`!output.set ORDINAL` / `!output.push`\n"
      "Sets the `--output=` flag result value at ORDINAL or pushes it to the\n"
      "back of the output list. Outputs can either be dumped to files or by\n"
      "default printed to stdout.\n"
      "\n"
      "`!blackboard.set ORDINAL` / `blackboard.push`\n"
      "Sets the value of the blackboard slot ORDINAL or pushes it to the back\n"
      "of the blackboard list. Blackboard values will be retained until they\n"
      "are consumed via `!blackboard.take` or the blackboard is cleared.\n"
      "\n");
  iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_DEFAULT, &argc, &argv);
  if (argc <= 1) {
    fprintf(stderr,
            "no trace files provided; pass one or more yaml file paths");
    return 1;
  }

  iree_vm_instance_t* instance = NULL;
  iree_status_t status = iree_vm_instance_create(
      IREE_VM_TYPE_CAPACITY_DEFAULT, iree_allocator_system(), &instance);
  if (iree_status_is_ok(status)) {
    status = iree_run_trace_files(argc - 1, argv + 1, instance);
  }
  iree_vm_instance_release(instance);
  if (!iree_status_is_ok(status)) {
    iree_status_fprint(stderr, status);
    iree_status_free(status);
    return 1;
  }
  return 0;
}