Adding `!input`/`!output`/`!blackboard` support to iree-run-trace. This allows for trace-wide input and output tensors that are handled just like iree-run-module. Traces can reference inputs/outputs in place of any call argument and can set/push outputs in place of any result. A blackboard is provided for storing temporary values used within the trace. Fixes #12525. Fixes #12526.
diff --git a/build_tools/scripts/run_yamllint.sh b/build_tools/scripts/run_yamllint.sh index a465af8..979fe0d 100755 --- a/build_tools/scripts/run_yamllint.sh +++ b/build_tools/scripts/run_yamllint.sh
@@ -21,6 +21,8 @@ declare -a excluded_files_patterns=( "/third_party/" "^third_party/" + "/tools/test/" + "^tools/test/" ) # Join on |
diff --git a/runtime/src/iree/tooling/trace_replay.c b/runtime/src/iree/tooling/trace_replay.c index 64fe7a3..93eadba 100644 --- a/runtime/src/iree/tooling/trace_replay.c +++ b/runtime/src/iree/tooling/trace_replay.c
@@ -41,11 +41,31 @@ out_replay->driver_registry = driver_registry; - return iree_ok_status(); + iree_status_t status = iree_ok_status(); + if (iree_status_is_ok(status)) { + status = iree_vm_list_create(NULL, 8u, host_allocator, &out_replay->inputs); + } + if (iree_status_is_ok(status)) { + status = + iree_vm_list_create(NULL, 8u, host_allocator, &out_replay->outputs); + } + if (iree_status_is_ok(status)) { + status = + iree_vm_list_create(NULL, 8u, host_allocator, &out_replay->blackboard); + } + + if (!iree_status_is_ok(status)) { + iree_trace_replay_deinitialize(out_replay, + IREE_TRACE_REPLAY_SHUTDOWN_QUIET); + } + return status; } void iree_trace_replay_deinitialize(iree_trace_replay_t* replay, iree_trace_replay_shutdown_flags_t flags) { + iree_vm_list_release(replay->inputs); + iree_vm_list_release(replay->outputs); + iree_vm_list_release(replay->blackboard); iree_vm_context_release(replay->context); iree_vm_instance_release(replay->instance); @@ -117,11 +137,11 @@ IREE_RETURN_IF_ERROR(iree_yaml_mapping_find(document, module_node, IREE_SV("name"), &name_node)); if (iree_yaml_string_equal(name_node, IREE_SV("hal"))) { - yaml_node_t* driver_node = NULL; + yaml_node_t* device_node = NULL; IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( - document, module_node, IREE_SV("driver"), &driver_node)); + document, module_node, IREE_SV("device"), &device_node)); IREE_RETURN_IF_ERROR(iree_trace_replay_create_device( - replay, driver_node, replay->host_allocator, &replay->device)); + replay, device_node, replay->host_allocator, &replay->device)); IREE_RETURN_IF_ERROR(iree_hal_module_create( replay->instance, replay->device, IREE_HAL_MODULE_FLAG_NONE, replay->host_allocator, &module)); @@ -354,13 +374,123 @@ } //===----------------------------------------------------------------------===// -// Input +// List I/O macros //===----------------------------------------------------------------------===// -static iree_status_t iree_trace_replay_parse_item(iree_trace_replay_t* replay, - yaml_document_t* document, - yaml_node_t* value_node, - iree_vm_list_t* target_list); +// Parses an I/O macro referencing the replay-global inputs/outputs. +// If |is_move| is true then the value is consumed from |list| and the original +// value in the list is reset to NULL. +// +// ```yaml +// !input.get 0 +// !input.take 1 +// !output.get 2 +// ``` +static iree_status_t iree_trace_replay_parse_list_get_macro( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_list_t* list, bool is_move, + iree_vm_variant_t* out_result) { + iree_string_view_t value_str = iree_yaml_node_as_string(value_node); + int32_t ordinal = 0; + if (!iree_string_view_atoi_int32(value_str, &ordinal)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "failed to parse I/O ordinal from `%.*s`", + (int)value_str.size, value_str.data); + } + iree_vm_variant_t variant = iree_vm_variant_empty(); + if (is_move) { + IREE_RETURN_IF_ERROR( + iree_vm_list_get_variant_move(list, ordinal, &variant)); + } else { + IREE_RETURN_IF_ERROR( + iree_vm_list_get_variant_retain(list, ordinal, &variant)); + } + *out_result = variant; + return iree_ok_status(); +} + +// Parses a list load macro referencing a replay-global |list|. +// +// ```yaml +// # gets |variant| at index 2, leaving it in the list for future use +// [!input.]get 2 +// # takes |variant| at index 2, clearing the entry in the list +// [!input.]take 2 +// ``` +static iree_status_t iree_trace_replay_parse_list_load_macro( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_string_view_t name, iree_vm_list_t* list, + iree_vm_variant_t* out_result) { + if (iree_string_view_equal(name, IREE_SV("get"))) { + return iree_trace_replay_parse_list_get_macro( + replay, document, value_node, list, + /*is_move=*/false, out_result); + } else if (iree_string_view_equal(name, IREE_SV("take"))) { + return iree_trace_replay_parse_list_get_macro(replay, document, value_node, + list, + /*is_move=*/true, out_result); + } else if (iree_string_view_equal(name, IREE_SV("pop"))) { + iree_host_size_t i = iree_vm_list_size(list) - 1; + IREE_RETURN_IF_ERROR(iree_vm_list_get_variant_move(list, i, out_result)); + return iree_vm_list_resize(list, i); + } + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "unsupported list load macro: `%.*s`", (int)name.size, + name.data); +} + +// Parses an output-set macro referencing the replay-global outputs. +// The provided |variant| value is set at the specified index. +// +// ```yaml +// !output.set 2 +// ``` +static iree_status_t iree_trace_replay_parse_list_set_macro( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_list_t* list, iree_vm_variant_t variant) { + iree_string_view_t value_str = iree_yaml_node_as_string(value_node); + int32_t ordinal = 0; + if (!iree_string_view_atoi_int32(value_str, &ordinal)) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "failed to parse I/O ordinal from `%.*s`", + (int)value_str.size, value_str.data); + } + if (iree_vm_list_size(list) <= ordinal) { + IREE_RETURN_IF_ERROR(iree_vm_list_resize(list, ordinal + 1)); + } + return iree_vm_list_set_variant_retain(list, ordinal, &variant); +} + +// Parses a list store macro referencing a replay-global |list|. +// +// ```yaml +// # sets |variant| at index 2 in the output list +// [!output.]set 2 +// # pushes |variant| to the end of the output list +// [!output.]push +// ``` +static iree_status_t iree_trace_replay_parse_list_store_macro( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_string_view_t name, iree_vm_list_t* list, + iree_vm_variant_t variant) { + if (iree_string_view_equal(name, IREE_SV("set"))) { + return iree_trace_replay_parse_list_set_macro(replay, document, value_node, + list, variant); + } else if (iree_string_view_equal(name, IREE_SV("push"))) { + return iree_vm_list_push_variant_retain(list, &variant); + } + return iree_make_status(IREE_STATUS_UNIMPLEMENTED, + "unsupported list store macro: `%.*s`", + (int)name.size, name.data); +} + +//===----------------------------------------------------------------------===// +// YAML value parsing +//===----------------------------------------------------------------------===// + +static iree_status_t iree_trace_replay_parse_item( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_variant_t* out_result); static iree_status_t iree_trace_replay_parse_item_sequence( iree_trace_replay_t* replay, yaml_document_t* document, yaml_node_t* sequence_node, iree_vm_list_t* target_list); @@ -372,7 +502,7 @@ // ``` static iree_status_t iree_trace_replay_parse_scalar( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* value_node, iree_vm_variant_t* out_result) { yaml_node_t* data_node = NULL; IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find(document, value_node, IREE_SV("i8"), &data_node)); @@ -384,10 +514,9 @@ IREE_STATUS_INVALID_ARGUMENT, "failed to parse i8 value: '%.*s'", (int)data_node->data.scalar.length, data_node->data.scalar.value); } - iree_vm_variant_t variant = iree_vm_variant_empty(); - variant.type.value_type = IREE_VM_VALUE_TYPE_I8; - variant.i8 = (int8_t)value; - return iree_vm_list_push_variant_move(target_list, &variant); + *out_result = + iree_vm_make_variant_value(iree_vm_value_make_i8((int8_t)value)); + return iree_ok_status(); } IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find(document, value_node, IREE_SV("i16"), &data_node)); @@ -399,69 +528,66 @@ IREE_STATUS_INVALID_ARGUMENT, "failed to parse i16 value: '%.*s'", (int)data_node->data.scalar.length, data_node->data.scalar.value); } - iree_vm_variant_t variant = iree_vm_variant_empty(); - variant.type.value_type = IREE_VM_VALUE_TYPE_I16; - variant.i16 = (int16_t)value; - return iree_vm_list_push_variant_move(target_list, &variant); + *out_result = + iree_vm_make_variant_value(iree_vm_value_make_i16((int16_t)value)); + return iree_ok_status(); } IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find(document, value_node, IREE_SV("i32"), &data_node)); if (data_node) { - iree_vm_variant_t variant = iree_vm_variant_empty(); - variant.type.value_type = IREE_VM_VALUE_TYPE_I32; + int32_t value = 0; if (!iree_string_view_atoi_int32(iree_yaml_node_as_string(data_node), - &variant.i32)) { + &value)) { return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "failed to parse i32 value: '%.*s'", (int)data_node->data.scalar.length, data_node->data.scalar.value); } - return iree_vm_list_push_variant_move(target_list, &variant); + *out_result = iree_vm_make_variant_value(iree_vm_value_make_i32(value)); + return iree_ok_status(); } IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find(document, value_node, IREE_SV("i64"), &data_node)); if (data_node) { - iree_vm_variant_t variant = iree_vm_variant_empty(); - variant.type.value_type = IREE_VM_VALUE_TYPE_I64; + int64_t value = 0; if (!iree_string_view_atoi_int64(iree_yaml_node_as_string(data_node), - &variant.i64)) { + &value)) { return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "failed to parse i64 value: '%.*s'", (int)data_node->data.scalar.length, data_node->data.scalar.value); } - return iree_vm_list_push_variant_move(target_list, &variant); + *out_result = iree_vm_make_variant_value(iree_vm_value_make_i64(value)); + return iree_ok_status(); } IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find(document, value_node, IREE_SV("f32"), &data_node)); if (data_node) { - iree_vm_variant_t variant = iree_vm_variant_empty(); - variant.type.value_type = IREE_VM_VALUE_TYPE_F32; - if (!iree_string_view_atof(iree_yaml_node_as_string(data_node), - &variant.f32)) { + float value = 0.0f; + if (!iree_string_view_atof(iree_yaml_node_as_string(data_node), &value)) { return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "failed to parse f32 value: '%.*s'", (int)data_node->data.scalar.length, data_node->data.scalar.value); } - return iree_vm_list_push_variant_move(target_list, &variant); + *out_result = iree_vm_make_variant_value(iree_vm_value_make_f32(value)); + return iree_ok_status(); } IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find(document, value_node, IREE_SV("f64"), &data_node)); if (data_node) { - iree_vm_variant_t variant = iree_vm_variant_empty(); - variant.type.value_type = IREE_VM_VALUE_TYPE_F64; - if (!iree_string_view_atod(iree_yaml_node_as_string(data_node), - &variant.f64)) { + double value = 0.0; + if (!iree_string_view_atod(iree_yaml_node_as_string(data_node), &value)) { return iree_make_status( IREE_STATUS_INVALID_ARGUMENT, "failed to parse f64 value: '%.*s'", (int)data_node->data.scalar.length, data_node->data.scalar.value); } - return iree_vm_list_push_variant_move(target_list, &variant); + *out_result = iree_vm_make_variant_value(iree_vm_value_make_f64(value)); + return iree_ok_status(); } return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "(%zu): unimplemented scalar type parser", value_node->start_mark.line); } -// Parses a !vm.list and appends it to |target_list|. +// Parses a !vm.list into |out_result|. // // ```yaml // items: @@ -472,7 +598,7 @@ // ``` static iree_status_t iree_trace_replay_parse_vm_list( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* value_node, iree_vm_variant_t* out_result) { if (value_node->type != YAML_MAPPING_NODE) { return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, "(%zu): expected sequence node for type", @@ -494,10 +620,8 @@ } if (iree_status_is_ok(status)) { - iree_vm_ref_t list_ref = iree_vm_list_move_ref(list); - status = iree_vm_list_push_ref_move(target_list, &list_ref); - } - if (!iree_status_is_ok(status)) { + *out_result = iree_vm_make_variant_ref_assign(iree_vm_list_move_ref(list)); + } else { iree_vm_list_release(list); } return status; @@ -733,7 +857,7 @@ } } -// Parses a !hal.buffer and appends it to |target_list|. +// Parses a !hal.buffer into |out_result|. // // ```yaml // shape: @@ -742,7 +866,7 @@ // ``` static iree_status_t iree_trace_replay_parse_hal_buffer( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* value_node, iree_vm_variant_t* out_result) { yaml_node_t* shape_node = NULL; IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( document, value_node, IREE_SV("shape"), &shape_node)); @@ -779,13 +903,12 @@ }, allocation_size, iree_const_byte_span_empty(), &buffer)); - iree_vm_ref_t buffer_ref = iree_hal_buffer_move_ref(buffer); - iree_status_t status = iree_vm_list_push_ref_move(target_list, &buffer_ref); - iree_vm_ref_release(&buffer_ref); - return status; + *out_result = + iree_vm_make_variant_ref_assign(iree_hal_buffer_move_ref(buffer)); + return iree_ok_status(); } -// Parses a !hal.buffer_view and appends it to |target_list|. +// Parses a !hal.buffer_view into |out_result|. // // ```yaml // shape: @@ -796,7 +919,7 @@ // ``` static iree_status_t iree_trace_replay_parse_hal_buffer_view( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* value_node, iree_vm_variant_t* out_result) { yaml_node_t* shape_node = NULL; IREE_RETURN_IF_ERROR(iree_yaml_mapping_try_find( document, value_node, IREE_SV("shape"), &shape_node)); @@ -863,14 +986,12 @@ iree_const_byte_span_empty(), &buffer_view)); } - iree_vm_ref_t buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - iree_status_t status = - iree_vm_list_push_ref_move(target_list, &buffer_view_ref); - iree_vm_ref_release(&buffer_view_ref); - return status; + *out_result = iree_vm_make_variant_ref_assign( + iree_hal_buffer_view_move_ref(buffer_view)); + return iree_ok_status(); } -// Parses a !hal.buffer in tensor form and appends it to |target_list|. +// Parses a !hal.buffer in tensor form into |out_result|. // The tensor form is used to size and initialize the buffer but then the // metadata is thrown away. // @@ -879,35 +1000,35 @@ // ``` static iree_status_t iree_trace_replay_parse_inline_hal_buffer( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* value_node, iree_vm_variant_t* out_result) { iree_hal_buffer_view_t* buffer_view = NULL; IREE_RETURN_IF_ERROR(iree_hal_buffer_view_parse( iree_yaml_node_as_string(value_node), iree_hal_device_allocator(replay->device), &buffer_view)); - iree_vm_ref_t buffer_ref = - iree_hal_buffer_retain_ref(iree_hal_buffer_view_buffer(buffer_view)); - iree_status_t status = iree_vm_list_push_ref_move(target_list, &buffer_ref); + *out_result = iree_vm_make_variant_ref_assign( + iree_hal_buffer_retain_ref(iree_hal_buffer_view_buffer(buffer_view))); iree_hal_buffer_view_release(buffer_view); - return status; + return iree_ok_status(); } -// Parses a !hal.buffer_view in tensor form and appends it to |target_list|. +// Parses a !hal.buffer_view in tensor form into |out_result|. // // ```yaml // !hal.buffer_view 4xf32=[0 1 2 3] // ``` static iree_status_t iree_trace_replay_parse_inline_hal_buffer_view( iree_trace_replay_t* replay, yaml_document_t* document, - yaml_node_t* value_node, iree_vm_list_t* target_list) { + yaml_node_t* value_node, iree_vm_variant_t* out_result) { iree_hal_buffer_view_t* buffer_view = NULL; IREE_RETURN_IF_ERROR(iree_hal_buffer_view_parse( iree_yaml_node_as_string(value_node), iree_hal_device_allocator(replay->device), &buffer_view)); - iree_vm_ref_t buffer_view_ref = iree_hal_buffer_view_move_ref(buffer_view); - return iree_vm_list_push_ref_move(target_list, &buffer_view_ref); + *out_result = iree_vm_make_variant_ref_assign( + iree_hal_buffer_view_move_ref(buffer_view)); + return iree_ok_status(); } -// Parses a typed item from |value_node| and appends it to |target_list|. +// Parses a typed item from |value_node| into |out_result|. // // ```yaml // type: vm.list @@ -919,16 +1040,25 @@ // ```yaml // !hal.buffer_view 4xf32=[0 1 2 3] // ``` -static iree_status_t iree_trace_replay_parse_item(iree_trace_replay_t* replay, - yaml_document_t* document, - yaml_node_t* value_node, - iree_vm_list_t* target_list) { - if (strcmp(value_node->tag, "!hal.buffer") == 0) { +static iree_status_t iree_trace_replay_parse_item( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* value_node, iree_vm_variant_t* out_result) { + iree_string_view_t tag = iree_make_cstring_view(value_node->tag); + if (iree_string_view_consume_prefix(&tag, IREE_SV("!input."))) { + return iree_trace_replay_parse_list_load_macro( + replay, document, value_node, tag, replay->inputs, out_result); + } else if (iree_string_view_consume_prefix(&tag, IREE_SV("!output."))) { + return iree_trace_replay_parse_list_load_macro( + replay, document, value_node, tag, replay->outputs, out_result); + } else if (iree_string_view_consume_prefix(&tag, IREE_SV("!blackboard."))) { + return iree_trace_replay_parse_list_load_macro( + replay, document, value_node, tag, replay->blackboard, out_result); + } else if (strcmp(value_node->tag, "!hal.buffer") == 0) { return iree_trace_replay_parse_inline_hal_buffer(replay, document, - value_node, target_list); + value_node, out_result); } else if (strcmp(value_node->tag, "!hal.buffer_view") == 0) { return iree_trace_replay_parse_inline_hal_buffer_view( - replay, document, value_node, target_list); + replay, document, value_node, out_result); } yaml_node_t* type_node = NULL; @@ -936,20 +1066,20 @@ IREE_SV("type"), &type_node)); iree_string_view_t type = iree_yaml_node_as_string(type_node); if (iree_string_view_equal(type, IREE_SV("null"))) { - iree_vm_variant_t null_value = iree_vm_variant_empty(); - return iree_vm_list_push_variant_move(target_list, &null_value); + *out_result = iree_vm_variant_empty(); + return iree_ok_status(); } else if (iree_string_view_equal(type, IREE_SV("value"))) { return iree_trace_replay_parse_scalar(replay, document, value_node, - target_list); + out_result); } else if (iree_string_view_equal(type, IREE_SV("vm.list"))) { return iree_trace_replay_parse_vm_list(replay, document, value_node, - target_list); + out_result); } else if (iree_string_view_equal(type, IREE_SV("hal.buffer"))) { return iree_trace_replay_parse_hal_buffer(replay, document, value_node, - target_list); + out_result); } else if (iree_string_view_equal(type, IREE_SV("hal.buffer_view"))) { return iree_trace_replay_parse_hal_buffer_view(replay, document, value_node, - target_list); + out_result); } return iree_make_status(IREE_STATUS_UNIMPLEMENTED, "unimplemented type parser: '%.*s'", (int)type.size, @@ -963,8 +1093,13 @@ for (yaml_node_item_t* item = sequence_node->data.sequence.items.start; item != sequence_node->data.sequence.items.top; ++item) { yaml_node_t* item_node = yaml_document_get_node(document, *item); + iree_vm_variant_t variant = iree_vm_variant_empty(); IREE_RETURN_IF_ERROR( - iree_trace_replay_parse_item(replay, document, item_node, target_list)); + iree_trace_replay_parse_item(replay, document, item_node, &variant)); + iree_status_t status = + iree_vm_list_push_variant_move(target_list, &variant); + iree_vm_variant_reset(&variant); + IREE_RETURN_IF_ERROR(status); } return iree_ok_status(); } @@ -973,70 +1108,36 @@ // Output //===----------------------------------------------------------------------===// -static iree_status_t iree_trace_replay_print_item( - iree_vm_variant_t* value, iree_allocator_t host_allocator); - -static iree_status_t iree_trace_replay_print_scalar(iree_vm_variant_t* value) { - switch (value->type.value_type) { - case IREE_VM_VALUE_TYPE_I8: - fprintf(stdout, "i8=%" PRIi8, value->i8); - break; - case IREE_VM_VALUE_TYPE_I16: - fprintf(stdout, "i16=%" PRIi16, value->i16); - break; - case IREE_VM_VALUE_TYPE_I32: - fprintf(stdout, "i32=%" PRIi32, value->i32); - break; - case IREE_VM_VALUE_TYPE_I64: - fprintf(stdout, "i64=%" PRIi64, value->i64); - break; - case IREE_VM_VALUE_TYPE_F32: - fprintf(stdout, "f32=%G", value->f32); - break; - case IREE_VM_VALUE_TYPE_F64: - fprintf(stdout, "f64=%G", value->f64); - break; - default: - fprintf(stdout, "?"); - break; +// Parses a single item. +static iree_status_t iree_trace_replay_parse_result_item( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* item_node, iree_vm_variant_t variant) { + iree_string_view_t tag = iree_make_cstring_view(item_node->tag); + if (iree_string_view_consume_prefix(&tag, IREE_SV("!output."))) { + return iree_trace_replay_parse_list_store_macro( + replay, document, item_node, tag, replay->outputs, variant); + } else if (iree_string_view_consume_prefix(&tag, IREE_SV("!blackboard."))) { + return iree_trace_replay_parse_list_store_macro( + replay, document, item_node, tag, replay->blackboard, variant); } + // NOTE: we ignore other types currently; we could parse them and compare + // against the |source_list| values or something. return iree_ok_status(); } -static iree_status_t iree_trace_replay_print_vm_list( - iree_vm_list_t* list, iree_allocator_t host_allocator) { - for (iree_host_size_t i = 0; i < iree_vm_list_size(list); ++i) { +// Parses a sequence of items and checks each against |source_list|. +static iree_status_t iree_trace_replay_parse_result_item_sequence( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* sequence_node, iree_vm_list_t* source_list) { + iree_host_size_t i = 0; + for (yaml_node_item_t* item = sequence_node->data.sequence.items.start; + item != sequence_node->data.sequence.items.top; ++item, ++i) { iree_vm_variant_t variant = iree_vm_variant_empty(); - IREE_RETURN_IF_ERROR(iree_vm_list_get_variant_assign(list, i, &variant), - "variant %zu not present", i); IREE_RETURN_IF_ERROR( - iree_trace_replay_print_item(&variant, host_allocator)); - fprintf(stdout, "\n"); - } - return iree_ok_status(); -} - -static iree_status_t iree_trace_replay_print_item( - iree_vm_variant_t* value, iree_allocator_t host_allocator) { - if (iree_vm_variant_is_value(*value)) { - IREE_RETURN_IF_ERROR(iree_trace_replay_print_scalar(value)); - } else if (iree_vm_variant_is_ref(*value)) { - if (iree_hal_buffer_view_isa(value->ref)) { - iree_hal_buffer_view_t* buffer_view = - iree_hal_buffer_view_deref(value->ref); - IREE_RETURN_IF_ERROR(iree_hal_buffer_view_fprint( - stdout, buffer_view, - /*max_element_count=*/1024, host_allocator)); - } else if (iree_vm_list_isa(value->ref)) { - iree_vm_list_t* list = iree_vm_list_deref(value->ref); - IREE_RETURN_IF_ERROR( - iree_trace_replay_print_vm_list(list, host_allocator)); - } else { - // TODO(benvanik): a way for ref types to describe themselves. - fprintf(stdout, "(no printer)"); - } - } else { - fprintf(stdout, "(null)"); + iree_vm_list_get_variant_assign(source_list, i, &variant)); + yaml_node_t* item_node = yaml_document_get_node(document, *item); + IREE_RETURN_IF_ERROR(iree_trace_replay_parse_result_item( + replay, document, item_node, variant)); } return iree_ok_status(); } @@ -1085,12 +1186,31 @@ return status; } -iree_status_t iree_trace_replay_event_call(iree_trace_replay_t* replay, - yaml_document_t* document, - yaml_node_t* event_node, - iree_vm_list_t** out_output_list) { +iree_status_t iree_trace_replay_event_call_finish(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) { IREE_TRACE_ZONE_BEGIN(z0); - if (out_output_list) *out_output_list = NULL; + + yaml_node_t* results_node = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_yaml_mapping_try_find(document, event_node, IREE_SV("results"), + &results_node)); + if (results_node) { + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_trace_replay_parse_result_item_sequence( + replay, document, results_node, output_list)); + } + + IREE_TRACE_ZONE_END(z0); + return iree_ok_status(); +} + +iree_status_t iree_trace_replay_event_call( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* event_node, const iree_trace_replay_call_hooks_t* hooks) { + IREE_TRACE_ZONE_BEGIN(z0); iree_vm_function_t function; iree_vm_list_t* input_list = NULL; @@ -1098,62 +1218,91 @@ z0, iree_trace_replay_event_call_prepare(replay, document, event_node, &function, &input_list)); - // Invoke the function to produce outputs. + iree_status_t status = iree_ok_status(); + if (hooks && hooks->before) { + status = hooks->before(hooks->user_data, replay, document, event_node, + function, input_list); + } + iree_vm_list_t* output_list = NULL; - iree_status_t status = - iree_vm_list_create(/*element_type=*/NULL, /*initial_capacity=*/8, - replay->host_allocator, &output_list); if (iree_status_is_ok(status)) { - status = iree_vm_invoke(replay->context, function, - IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/NULL, - input_list, output_list, replay->host_allocator); + status = iree_vm_list_create(/*element_type=*/NULL, /*initial_capacity=*/8, + replay->host_allocator, &output_list); + } + + // Invoke the function to produce outputs. + iree_status_t call_status = iree_ok_status(); + if (iree_status_is_ok(status)) { + call_status = iree_vm_invoke( + replay->context, function, IREE_VM_INVOCATION_FLAG_NONE, + /*policy=*/NULL, input_list, output_list, replay->host_allocator); } iree_vm_list_release(input_list); - if (iree_status_is_ok(status) && out_output_list) { - *out_output_list = output_list; - } else { - iree_vm_list_release(output_list); + if (!iree_status_is_ok(call_status)) { + if (hooks && hooks->error) { + status = hooks->error(hooks->user_data, replay, document, event_node, + function, call_status); + } else { + status = call_status; + } + } else if (hooks && hooks->after) { + status = hooks->after(hooks->user_data, replay, document, event_node, + function, output_list); } + + if (iree_status_is_ok(status)) { + status = iree_trace_replay_event_call_finish(replay, document, event_node, + function, output_list); + } + iree_vm_list_release(output_list); + IREE_TRACE_ZONE_END(z0); return status; } -static iree_status_t iree_trace_replay_event_call_stdout( +//===----------------------------------------------------------------------===// +// Blackboard management +//===----------------------------------------------------------------------===// + +static iree_status_t iree_trace_replay_event_blackboard_clear( iree_trace_replay_t* replay, yaml_document_t* document, yaml_node_t* event_node) { - yaml_node_t* function_node = NULL; - IREE_RETURN_IF_ERROR(iree_yaml_mapping_find( - document, event_node, IREE_SV("function"), &function_node)); - iree_string_view_t function_name = iree_yaml_node_as_string(function_node); - fprintf(stdout, "--- CALL[%.*s] ---\n", (int)function_name.size, - function_name.data); + IREE_TRACE_ZONE_BEGIN(z0); + iree_vm_list_clear(replay->blackboard); + IREE_TRACE_ZONE_END(z0); + return iree_ok_status(); +} - // Prepare to call the function. - iree_vm_function_t function; - iree_vm_list_t* input_list = NULL; - IREE_RETURN_IF_ERROR(iree_trace_replay_event_call_prepare( - replay, document, event_node, &function, &input_list)); +static iree_status_t iree_trace_replay_event_blackboard_assign( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* event_node) { + IREE_TRACE_ZONE_BEGIN(z0); - // Invoke the function to produce outputs. - iree_vm_list_t* output_list = NULL; + yaml_node_t* from_node = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_yaml_mapping_find(document, event_node, IREE_SV("from"), + &from_node)); + yaml_node_t* to_node = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, + iree_yaml_mapping_find(document, event_node, IREE_SV("to"), &to_node)); + + iree_vm_list_t* list = NULL; + IREE_RETURN_AND_END_ZONE_IF_ERROR( + z0, iree_vm_list_create(/*element_type=*/NULL, 8u, replay->host_allocator, + &list)); + iree_status_t status = - iree_vm_list_create(/*element_type=*/NULL, /*initial_capacity=*/8, - replay->host_allocator, &output_list); + iree_trace_replay_parse_item_sequence(replay, document, from_node, list); if (iree_status_is_ok(status)) { - status = iree_vm_invoke(replay->context, function, - IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/NULL, - input_list, output_list, replay->host_allocator); + status = iree_trace_replay_parse_result_item_sequence(replay, document, + to_node, list); } - iree_vm_list_release(input_list); - // Print the outputs. - if (iree_status_is_ok(status)) { - status = - iree_trace_replay_print_vm_list(output_list, replay->host_allocator); - } - iree_vm_list_release(output_list); + iree_vm_list_release(list); + IREE_TRACE_ZONE_END(z0); return status; } @@ -1176,8 +1325,15 @@ return iree_trace_replay_event_context_load(replay, document, event_node); } else if (iree_yaml_string_equal(type_node, IREE_SV("module_load"))) { return iree_trace_replay_event_module_load(replay, document, event_node); + } else if (iree_yaml_string_equal(type_node, IREE_SV("blackboard_clear"))) { + return iree_trace_replay_event_blackboard_clear(replay, document, + event_node); + } else if (iree_yaml_string_equal(type_node, IREE_SV("assign"))) { + return iree_trace_replay_event_blackboard_assign(replay, document, + event_node); } else if (iree_yaml_string_equal(type_node, IREE_SV("call"))) { - return iree_trace_replay_event_call_stdout(replay, document, event_node); + return iree_trace_replay_event_call(replay, document, event_node, + &replay->call_hooks); } return iree_make_status( IREE_STATUS_UNIMPLEMENTED, "(%zu): unhandled type '%.*s'",
diff --git a/runtime/src/iree/tooling/trace_replay.h b/runtime/src/iree/tooling/trace_replay.h index db2de14..ddc240c 100644 --- a/runtime/src/iree/tooling/trace_replay.h +++ b/runtime/src/iree/tooling/trace_replay.h
@@ -16,12 +16,37 @@ extern "C" { #endif // __cplusplus +typedef struct iree_trace_replay_t iree_trace_replay_t; + enum iree_trace_replay_shutdown_flag_bits_e { IREE_TRACE_REPLAY_SHUTDOWN_QUIET = 0u, IREE_TRACE_REPLAY_SHUTDOWN_PRINT_STATISTICS = 1 << 0u, }; typedef uint32_t iree_trace_replay_shutdown_flags_t; +// Optional set of callbacks around a replay event function call. +// Functions not required by the caller may be omitted. +typedef struct iree_trace_replay_call_hooks_t { + // User context passed to each callback. + void* user_data; + // Issued before the call begins with the call inputs. + iree_status_t (*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); + // Issued after the call completes successfully with the call outputs. + iree_status_t (*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); + // Issued only when the call fails and not the replay operation itself. + // |status| is as returned from the call and ownership is transferred to the + // hook. + iree_status_t (*error)(void* user_data, iree_trace_replay_t* replay, + yaml_document_t* document, yaml_node_t* event_node, + iree_vm_function_t function, iree_status_t status); +} iree_trace_replay_call_hooks_t; + typedef struct iree_trace_replay_t { iree_allocator_t host_allocator; iree_string_view_t root_path; @@ -33,8 +58,21 @@ iree_host_size_t device_uri_count; const iree_string_view_t* device_uris; + // Context used within the replay, modules registered on-demand. iree_vm_context_t* context; + + // Active HAL device if any. Will be initialized on the first HAL module load. iree_hal_device_t* device; + + // Optional inputs available via `!input.get`/`!input.take`. + iree_vm_list_t* inputs; + // Optional outputs populated via `!output.set`/`!output.push`. + iree_vm_list_t* outputs; + // Blackboard used to track state within the trace. + iree_vm_list_t* blackboard; + + // Optional call hooks allowing reflection of calls and their I/O. + iree_trace_replay_call_hooks_t call_hooks; } iree_trace_replay_t; // Initializes a trace replay context. @@ -82,12 +120,11 @@ iree_vm_list_t** out_input_list); // Replays a `call` event against the replay context. -// Optionally |out_output_list| can be populated with a caller-owned set of -// outputs from the call. -iree_status_t iree_trace_replay_event_call(iree_trace_replay_t* replay, - yaml_document_t* document, - yaml_node_t* event_node, - iree_vm_list_t** out_output_list); +// Optionally |hooks| may be specified to inspect the inputs and outputs of the +// call operation. +iree_status_t iree_trace_replay_event_call( + iree_trace_replay_t* replay, yaml_document_t* document, + yaml_node_t* event_node, const iree_trace_replay_call_hooks_t* hooks); #ifdef __cplusplus } // extern "C"
diff --git a/tools/BUILD b/tools/BUILD index a6bdb90..fff940c 100644 --- a/tools/BUILD +++ b/tools/BUILD
@@ -183,6 +183,7 @@ "//runtime/src/iree/hal", "//runtime/src/iree/tooling:device_util", "//runtime/src/iree/tooling:trace_replay", + "//runtime/src/iree/tooling:vm_util", "//runtime/src/iree/tooling:yaml_util", "//runtime/src/iree/vm", "@com_github_yaml_libyaml//:yaml",
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt index 7d37b6c..ea8d47c 100644 --- a/tools/CMakeLists.txt +++ b/tools/CMakeLists.txt
@@ -161,6 +161,7 @@ iree::modules::hal iree::tooling::device_util iree::tooling::trace_replay + iree::tooling::vm_util iree::tooling::yaml_util iree::vm yaml
diff --git a/tools/iree-run-trace-main.c b/tools/iree-run-trace-main.c index 57a2053..56c66ca 100644 --- a/tools/iree-run-trace-main.c +++ b/tools/iree-run-trace-main.c
@@ -14,6 +14,7 @@ #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" @@ -22,6 +23,102 @@ 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, @@ -34,6 +131,11 @@ : IREE_VM_CONTEXT_FLAG_NONE, 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. @@ -51,24 +153,59 @@ } 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, FLAG_print_statistics ? IREE_TRACE_REPLAY_SHUTDOWN_PRINT_STATISTICS @@ -97,6 +234,114 @@ } 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: 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,
diff --git a/tools/test/BUILD b/tools/test/BUILD index f7bf4f8..af0a150 100644 --- a/tools/test/BUILD +++ b/tools/test/BUILD
@@ -28,6 +28,7 @@ "iree-run-module.mlir", "iree-run-module-expected.mlir", "iree-run-module-outputs.mlir", + "iree-run-trace.mlir", "multiple_args.mlir", "multiple_exported_functions.mlir", "null_values.mlir", @@ -39,6 +40,7 @@ cfg = "//tools:lit.cfg.py", data = [ "echo_npy.py", + "iree-run-trace.yml", ], tags = [ "driver=local-task", @@ -51,6 +53,7 @@ "//tools:iree-opt", "//tools:iree-run-mlir", "//tools:iree-run-module", + "//tools:iree-run-trace", "@llvm-project//lld", "@llvm-project//llvm:FileCheck", "@llvm-project//llvm:not",
diff --git a/tools/test/CMakeLists.txt b/tools/test/CMakeLists.txt index 7c92d9d..305345a 100644 --- a/tools/test/CMakeLists.txt +++ b/tools/test/CMakeLists.txt
@@ -24,6 +24,7 @@ "iree-run-module-expected.mlir" "iree-run-module-outputs.mlir" "iree-run-module.mlir" + "iree-run-trace.mlir" "multiple_args.mlir" "multiple_exported_functions.mlir" "null_values.mlir" @@ -37,9 +38,11 @@ iree-opt iree-run-mlir iree-run-module + iree-run-trace not DATA echo_npy.py + iree-run-trace.yml LABELS "driver=local-task" "driver=vulkan"
diff --git a/tools/test/iree-run-trace.mlir b/tools/test/iree-run-trace.mlir new file mode 100644 index 0000000..ad898e7 --- /dev/null +++ b/tools/test/iree-run-trace.mlir
@@ -0,0 +1,22 @@ +// Tests iree-run-trace usage by running two calls of @mul and passing the +// result between them. The outputs of both calls are produced as outputs from +// the trace and both are written to a .npy file for processing. Inputs can +// also come from an .npy file. See iree-run-module usage for more information +// on the `--input=` and `--output=` flags. + +// RUN: (iree-compile --iree-hal-target-backends=vmvx %s | \ +// RUN: iree-run-trace %S/iree-run-trace.yml \ +// RUN: --device=local-sync \ +// RUN: --input=4xf32=4,4,4,4 \ +// RUN: --output=@%t \ +// RUN: --output=+%t) && \ +// RUN: python3 %S/echo_npy.py %t | \ +// RUN: FileCheck %s + +// CHECK{LITERAL}: [ 0. 4. 8. 12.] +// CHECK-NEXT{LITERAL}: [ 0. 12. 24. 36.] + +func.func @mul(%arg0: tensor<4xf32>, %arg1: tensor<4xf32>) -> tensor<4xf32> { + %0 = arith.mulf %arg0, %arg1 : tensor<4xf32> + return %0 : tensor<4xf32> +}
diff --git a/tools/test/iree-run-trace.yml b/tools/test/iree-run-trace.yml new file mode 100644 index 0000000..b267dc2 --- /dev/null +++ b/tools/test/iree-run-trace.yml
@@ -0,0 +1,80 @@ +# Copyright 2023 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 + +# Tests loading and executing a bytecode module and issuing a few calls showing +# how to take input, produce output, and support temporary values within the +# trace session. See iree-run-trace.mlir for how to compile the module and +# invoke the iree-run-trace tool. + +# Prepare the VM context for use; effectively a reset. +# API: iree_vm_context_create +type: context_load + +--- + +# Load the builtin HAL module used to execute the program. +# API: iree_hal_module_create +type: module_load +module: + type: builtin + name: hal + +--- + +# Load the compiled bytecode module. +# API: iree_vm_bytecode_module_create +type: module_load +module: + type: bytecode + name: module + # The test pulls the .vmfb from stdin but you can also reference relative or + # absolute file paths: + # path: ../iree-tmp/iree-run-trace.vmfb + path: <stdin> + +--- + +# Call #0 of @mul. +# API: iree_vm_invoke +type: call +function: module.mul +args: +# arg[0]: the first `--input=` buffer. !input.get would retain the input for +# other calls to use but otherwise prefer taking ownership. +- !input.take 0 +# arg[1]: constant value defined inline. +- !hal.buffer_view 4xf32=0,1,2,3 +results: +# result[0]: store in blackboard slot 4 for later use. +- !blackboard.set 4 + +--- + +# Assigns one or more source values to a set of target values. +# Effectively: outputs.push(retain(blackboard[4])) +type: assign +from: +# from[0]: retain blackboard slot 4, leaving it for later use. +- !blackboard.get 4 +to: +# to[0]: push on to the trace output list. --output= can save off the results +# and otherwise they are printed to stdout. +- !output.push + +--- + +# Call #1 of @mul. +# API: iree_vm_invoke +type: call +function: module.mul +args: +# arg[0]: take the previously-stored value in blackboard slot 4. +- !blackboard.take 4 +# arg[1]: another constant. +- !hal.buffer_view 4xf32=3,3,3,3 +results: +# result[0]: push on to the trace output list. +- !output.push