Adding `numpy_load`/`numpy_save` trace events.
These allow for easily loading/storing from/to the blackboard/outputs
within the trace as opposed to the command line flags.
diff --git a/runtime/src/iree/tooling/BUILD b/runtime/src/iree/tooling/BUILD
index 18182a8..84e1a12 100644
--- a/runtime/src/iree/tooling/BUILD
+++ b/runtime/src/iree/tooling/BUILD
@@ -193,6 +193,7 @@
srcs = ["trace_replay.c"],
hdrs = ["trace_replay.h"],
deps = [
+ ":numpy_io",
":yaml_util",
"//runtime/src/iree/base",
"//runtime/src/iree/base:tracing",
diff --git a/runtime/src/iree/tooling/CMakeLists.txt b/runtime/src/iree/tooling/CMakeLists.txt
index 5e41789..b3f7284 100644
--- a/runtime/src/iree/tooling/CMakeLists.txt
+++ b/runtime/src/iree/tooling/CMakeLists.txt
@@ -214,6 +214,7 @@
SRCS
"trace_replay.c"
DEPS
+ ::numpy_io
::yaml_util
iree::base
iree::base::internal
diff --git a/runtime/src/iree/tooling/trace_replay.c b/runtime/src/iree/tooling/trace_replay.c
index 93eadba..ea5a921 100644
--- a/runtime/src/iree/tooling/trace_replay.c
+++ b/runtime/src/iree/tooling/trace_replay.c
@@ -17,6 +17,7 @@
#include "iree/base/internal/path.h"
#include "iree/base/tracing.h"
#include "iree/modules/hal/module.h"
+#include "iree/tooling/numpy_io.h"
#include "iree/vm/bytecode/module.h"
//===----------------------------------------------------------------------===//
@@ -1307,6 +1308,177 @@
}
//===----------------------------------------------------------------------===//
+// Numpy ndarray management
+//===----------------------------------------------------------------------===//
+
+// Loads one or more ndarrays from a .npy file.
+//
+// Example:
+// ```yaml
+// type: numpy_load
+// path: three_ndarrays.npy
+// arrays:
+// - !blackboard.set 2
+// - !blackboard.set 3
+// - !output.set 4
+// ```
+//
+// NOTE: this currently reads things into new buffers; we could try mapping from
+// disk and other fancy things where possible.
+static iree_status_t iree_trace_replay_event_numpy_load(
+ iree_trace_replay_t* replay, yaml_document_t* document,
+ yaml_node_t* event_node) {
+ if (!replay->device) {
+ return iree_make_status(
+ IREE_STATUS_FAILED_PRECONDITION,
+ "HAL module must be loaded before loading numpy arrays");
+ }
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ yaml_node_t* path_node = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_yaml_mapping_find(document, event_node, IREE_SV("path"),
+ &path_node));
+ iree_string_view_t path_str = iree_yaml_node_as_string(path_node);
+
+ yaml_node_t* arrays_node = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_yaml_mapping_find(document, event_node, IREE_SV("arrays"),
+ &arrays_node));
+
+ char* full_path = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_file_path_join(replay->root_path, path_str,
+ replay->host_allocator, &full_path));
+ FILE* file = fopen(full_path, "rb");
+ iree_allocator_free(replay->host_allocator, full_path);
+ if (!file) {
+ IREE_TRACE_ZONE_END(z0);
+ return iree_make_status(iree_status_code_from_errno(errno),
+ "failed to open file `%.*s` for read",
+ (int)path_str.size, path_str.data);
+ }
+
+ uint64_t file_length = 0;
+ iree_status_t status = iree_file_query_length(file, &file_length);
+
+ iree_hal_buffer_params_t buffer_params = {0};
+ buffer_params.usage = IREE_HAL_BUFFER_USAGE_DEFAULT;
+ buffer_params.access = IREE_HAL_MEMORY_ACCESS_READ;
+ buffer_params.type = IREE_HAL_MEMORY_TYPE_DEVICE_LOCAL;
+ iree_hal_allocator_t* device_allocator =
+ iree_hal_device_allocator(replay->device);
+
+ if (iree_status_is_ok(status)) {
+ for (yaml_node_item_t* item = arrays_node->data.sequence.items.start;
+ item != arrays_node->data.sequence.items.top; ++item) {
+ // Parse the next array in the file. Note that we may be at the end!
+ if (iree_file_is_at(file, file_length)) {
+ status = iree_make_status(IREE_STATUS_OUT_OF_RANGE,
+ "file ended before all arrays were decoded");
+ break;
+ }
+ iree_hal_buffer_view_t* buffer_view = NULL;
+ status = iree_numpy_npy_load_ndarray(
+ file, IREE_NUMPY_NPY_LOAD_OPTION_DEFAULT, buffer_params,
+ device_allocator, &buffer_view);
+ if (!iree_status_is_ok(status)) break;
+
+ // Route the loaded value to its destination.
+ iree_vm_variant_t variant = iree_vm_make_variant_ref_assign(
+ iree_hal_buffer_view_move_ref(buffer_view));
+ yaml_node_t* item_node = yaml_document_get_node(document, *item);
+ status = iree_trace_replay_parse_result_item(replay, document, item_node,
+ variant);
+ iree_vm_variant_reset(&variant);
+ if (!iree_status_is_ok(status)) break;
+ }
+ }
+
+ fclose(file);
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+// Saves one or more ndarrays to a .npy file.
+//
+// Example:
+// ```yaml
+// type: numpy_save
+// path: three_ndarrays.npy
+// append: true
+// arrays:
+// - !blackboard.get 2
+// - !blackboard.take 3
+// - !input.get 4
+// - !hal.buffer_view 4xf32=0,1,2,3
+// ```
+static iree_status_t iree_trace_replay_event_numpy_save(
+ iree_trace_replay_t* replay, yaml_document_t* document,
+ yaml_node_t* event_node) {
+ IREE_TRACE_ZONE_BEGIN(z0);
+
+ yaml_node_t* path_node = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_yaml_mapping_find(document, event_node, IREE_SV("path"),
+ &path_node));
+ iree_string_view_t path_str = iree_yaml_node_as_string(path_node);
+
+ yaml_node_t* append_node = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_yaml_mapping_try_find(document, event_node, IREE_SV("append"),
+ &append_node));
+
+ yaml_node_t* arrays_node = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_yaml_mapping_find(document, event_node, IREE_SV("arrays"),
+ &arrays_node));
+
+ char* full_path = NULL;
+ IREE_RETURN_AND_END_ZONE_IF_ERROR(
+ z0, iree_file_path_join(replay->root_path, path_str,
+ replay->host_allocator, &full_path));
+ const char* mode = append_node ? "ab" : "wb";
+ FILE* file = fopen(full_path, mode);
+ iree_allocator_free(replay->host_allocator, full_path);
+ if (!file) {
+ IREE_TRACE_ZONE_END(z0);
+ return iree_make_status(iree_status_code_from_errno(errno),
+ "failed to open file `%.*s` for write",
+ (int)path_str.size, path_str.data);
+ }
+
+ iree_status_t status = iree_ok_status();
+ for (yaml_node_item_t* item = arrays_node->data.sequence.items.start;
+ item != arrays_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();
+ status =
+ iree_trace_replay_parse_item(replay, document, item_node, &variant);
+ if (!iree_status_is_ok(status)) break;
+ if (!iree_vm_variant_is_ref(variant) ||
+ !iree_hal_buffer_view_isa(variant.ref)) {
+ status =
+ iree_make_status(IREE_STATUS_INVALID_ARGUMENT,
+ "only buffer views can be saved to numpy files");
+ break;
+ }
+ iree_hal_buffer_view_t* buffer_view =
+ iree_hal_buffer_view_deref(variant.ref);
+ status =
+ iree_numpy_npy_save_ndarray(file, IREE_NUMPY_NPY_SAVE_OPTION_DEFAULT,
+ buffer_view, replay->host_allocator);
+ iree_vm_variant_reset(&variant);
+ }
+
+ fclose(file);
+
+ IREE_TRACE_ZONE_END(z0);
+ return status;
+}
+
+//===----------------------------------------------------------------------===//
// Event dispatch
//===----------------------------------------------------------------------===//
@@ -1331,6 +1503,10 @@
} 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("numpy_load"))) {
+ return iree_trace_replay_event_numpy_load(replay, document, event_node);
+ } else if (iree_yaml_string_equal(type_node, IREE_SV("numpy_save"))) {
+ return iree_trace_replay_event_numpy_save(replay, document, event_node);
} else if (iree_yaml_string_equal(type_node, IREE_SV("call"))) {
return iree_trace_replay_event_call(replay, document, event_node,
&replay->call_hooks);
diff --git a/tools/iree-run-trace-main.c b/tools/iree-run-trace-main.c
index 56c66ca..9c01f06 100644
--- a/tools/iree-run-trace-main.c
+++ b/tools/iree-run-trace-main.c
@@ -295,6 +295,14 @@
"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"
