| // 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; |
| } |
