[iree-benchmark-module] support processing outputs after benchmarking (#24789)
Add --enable_output_processing option to iree-benchmark-module to make
it process the outputs of the last benchmark iteration executed. For
example, this processing can be writing the outputs to disk. The output
processing is the same as done by iree-run-module.
In order to achieve this:
- Move output processing from run_module.c into a new process_results.c.
- Slightly change the benchmark loop in iree-benchmark-module-main.cc:
- Use a output VM list passed by the caller and only create a
local/temporary one if none is passed.
- Move clearing of the output list from the end of the benchmark loop
body to the beginning of the benchmark loop body, so that the results of
the last benchmark iteartion are available at the end of the loop.
- Clear the results/output list after the benchmark loop if a local list
has been used. Otherwise, keep the results in the output list passed by
the caller.
Also change the main function of iree-benchmark-module:
- Allocate an output list if output processing is enabled.
- Pass the output list (if allocated) to benchmarking loop (for now only
for synachonous execution of modules).
- Process outputs in the output list at the end (if list is allocated).
---------
Signed-off-by: Stefan Schuermans <schuermans@roofline.ai>
diff --git a/runtime/src/iree/tooling/BUILD.bazel b/runtime/src/iree/tooling/BUILD.bazel
index 46236ef..a2b4fce 100644
--- a/runtime/src/iree/tooling/BUILD.bazel
+++ b/runtime/src/iree/tooling/BUILD.bazel
@@ -218,21 +218,35 @@
)
iree_runtime_cc_library(
+ name = "process_results",
+ srcs = ["process_results.c"],
+ hdrs = ["process_results.h"],
+ deps = [
+ ":comparison",
+ ":function_io",
+ "//runtime/src/iree/base",
+ "//runtime/src/iree/base/tooling:flags",
+ "//runtime/src/iree/hal",
+ "//runtime/src/iree/io:stream",
+ "//runtime/src/iree/vm",
+ ],
+)
+
+iree_runtime_cc_library(
name = "run_module",
srcs = ["run_module.c"],
hdrs = ["run_module.h"],
deps = [
- ":comparison",
":context_util",
":device_util",
":function_io",
":function_util",
":instrument_util",
+ ":process_results",
"//runtime/src/iree/base",
"//runtime/src/iree/base/tooling:flags",
"//runtime/src/iree/hal",
"//runtime/src/iree/hal/replay:recorder",
- "//runtime/src/iree/io:stream",
"//runtime/src/iree/modules/hal:types",
"//runtime/src/iree/vm",
"//runtime/src/iree/vm/bytecode:module",
diff --git a/runtime/src/iree/tooling/CMakeLists.txt b/runtime/src/iree/tooling/CMakeLists.txt
index 130734d..e0420a0 100644
--- a/runtime/src/iree/tooling/CMakeLists.txt
+++ b/runtime/src/iree/tooling/CMakeLists.txt
@@ -249,23 +249,40 @@
iree_cc_library(
NAME
+ process_results
+ HDRS
+ "process_results.h"
+ SRCS
+ "process_results.c"
+ DEPS
+ ::comparison
+ ::function_io
+ iree::base
+ iree::base::tooling::flags
+ iree::hal
+ iree::io::stream
+ iree::vm
+ PUBLIC
+)
+
+iree_cc_library(
+ NAME
run_module
HDRS
"run_module.h"
SRCS
"run_module.c"
DEPS
- ::comparison
::context_util
::device_util
::function_io
::function_util
::instrument_util
+ ::process_results
iree::base
iree::base::tooling::flags
iree::hal
iree::hal::replay::recorder
- iree::io::stream
iree::modules::hal::types
iree::vm
iree::vm::bytecode::module
diff --git a/runtime/src/iree/tooling/process_results.c b/runtime/src/iree/tooling/process_results.c
new file mode 100644
index 0000000..2acf9fc
--- /dev/null
+++ b/runtime/src/iree/tooling/process_results.c
@@ -0,0 +1,152 @@
+// Copyright 2026 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/process_results.h"
+
+#include <stdio.h>
+#include <stdlib.h>
+
+#include "iree/base/api.h"
+#include "iree/base/tooling/flags.h"
+#include "iree/hal/api.h"
+#include "iree/io/stdio_stream.h"
+#include "iree/tooling/comparison.h"
+#include "iree/tooling/function_io.h"
+#include "iree/vm/api.h"
+
+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 an ndarray\n"
+ " e.g.: --output=@file.npy\n"
+ " `+file.npy`: create/append a numpy npy file and write an ndarray\n"
+ " e.g.: --output=+file.npy\n"
+ " `@file.bin`: create/overwrite a binary file and write value contents\n"
+ " e.g.: --output=@file.bin\n"
+ " `+file.bin`: create/append a binary file and write value contents\n"
+ " e.g.: --output=+file.bin\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"
+ "Primitive values are written as shape=() ndarrays and buffers are\n"
+ "written as i8 arrays with the length of the buffer.\n"
+ "\n"
+ "Binary files contain only the contents of the values/buffers provided\n"
+ "without metadata; users must know the shape/type of the output.\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=`.\n"
+ "When present the results of the invocation will be compared against\n"
+ "these values and the tool will return non-zero if any differ. If the\n"
+ "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 and elides\n"
+ "the remainder.");
+
+iree_status_t iree_tooling_process_results(iree_hal_device_t* device,
+ iree_string_view_t results_cconv,
+ iree_vm_list_t* results,
+ iree_io_stream_t* stream,
+ iree_allocator_t host_allocator,
+ int* out_exit_code) {
+ *out_exit_code = EXIT_SUCCESS;
+
+ // Basic output handling to route to the console or files.
+ if (FLAG_expected_output_list().count == 0) {
+ if (FLAG_output_list().count == 0) {
+ // Print all outputs.
+ return iree_status_annotate_f(
+ iree_tooling_print_variants(
+ IREE_SV("result"), results,
+ (iree_host_size_t)FLAG_output_max_element_count, stream,
+ host_allocator),
+ "printing results");
+ } else {
+ // Write (or ignore) all outputs.
+ return iree_status_annotate_f(
+ iree_tooling_write_variants(
+ results, FLAG_output_list(),
+ (iree_host_size_t)FLAG_output_max_element_count, stream,
+ host_allocator),
+ "outputting results");
+ }
+ }
+
+ // Compare against contents in host-local memory. This avoids polluting
+ // device memory statistics.
+ iree_hal_allocator_t* heap_allocator = NULL;
+ IREE_RETURN_IF_ERROR(iree_hal_allocator_create_heap(
+ IREE_SV("heap"), host_allocator, host_allocator, &heap_allocator));
+
+ // Parse expected list into host-local memory that we can easily access.
+ iree_vm_list_t* expected_list = NULL;
+ iree_status_t status = iree_status_annotate_f(
+ iree_tooling_parse_variants(results_cconv, FLAG_expected_output_list(),
+ device, heap_allocator, host_allocator,
+ &expected_list),
+ "parsing expected function outputs");
+
+ // Compare expected vs actual lists and output diffs.
+ if (iree_status_is_ok(status)) {
+ bool did_match = iree_tooling_compare_variant_lists(expected_list, results,
+ host_allocator, stdout);
+ if (did_match) {
+ fprintf(
+ stdout,
+ "[SUCCESS] all function outputs matched their expected values.\n");
+ }
+
+ // Exit code 0 if all results matched the expected values.
+ *out_exit_code = did_match ? EXIT_SUCCESS : EXIT_FAILURE;
+ }
+
+ iree_vm_list_release(expected_list);
+ iree_hal_allocator_release(heap_allocator);
+ return status;
+}
+
+iree_status_t iree_tooling_process_results_and_print(
+ iree_hal_device_t* device, iree_string_view_t results_cconv,
+ iree_vm_list_t* results, iree_allocator_t host_allocator,
+ int* out_exit_code) {
+ // Wrap stdout for printing results.
+ iree_io_stream_t* stdout_stream = NULL;
+ iree_status_t status = iree_status_annotate_f(
+ iree_io_stdio_stream_wrap(IREE_IO_STREAM_MODE_WRITABLE, stdout,
+ /*owns_handle=*/false, host_allocator,
+ &stdout_stream),
+ "opening stdout stream");
+
+ // Handle either printing/writing the outputs or checking them against
+ // expected values (basic pass/fail testing).
+ if (iree_status_is_ok(status)) {
+ status = iree_status_annotate_f(
+ iree_tooling_process_results(device, results_cconv, results,
+ stdout_stream, host_allocator,
+ out_exit_code),
+ "processing function outputs");
+ }
+
+ iree_io_stream_release(stdout_stream);
+ fflush(stdout);
+
+ return status;
+}
diff --git a/runtime/src/iree/tooling/process_results.h b/runtime/src/iree/tooling/process_results.h
new file mode 100644
index 0000000..12599ea
--- /dev/null
+++ b/runtime/src/iree/tooling/process_results.h
@@ -0,0 +1,41 @@
+// Copyright 2026 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
+
+#ifndef IREE_TOOLING_PROCESS_RESULTS_H_
+#define IREE_TOOLING_PROCESS_RESULTS_H_
+
+#include "iree/base/api.h"
+#include "iree/hal/api.h"
+#include "iree/io/stream.h"
+#include "iree/vm/api.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif // __cplusplus
+
+// Handles either printing/writing the |results| of an invocation or checking
+// them against expected values (basic pass/fail testing) as specified by the
+// --output=/--expected_output= flags. Textual output is written to |stream|.
+// Returns the process result code in |out_exit_code| (0 for success).
+iree_status_t iree_tooling_process_results(iree_hal_device_t* device,
+ iree_string_view_t results_cconv,
+ iree_vm_list_t* results,
+ iree_io_stream_t* stream,
+ iree_allocator_t host_allocator,
+ int* out_exit_code);
+
+// Processes the |results| of an invocation with stdout as the output stream.
+// Refer to iree_tooling_process_results for details.
+iree_status_t iree_tooling_process_results_and_print(
+ iree_hal_device_t* device, iree_string_view_t results_cconv,
+ iree_vm_list_t* results, iree_allocator_t host_allocator,
+ int* out_exit_code);
+
+#ifdef __cplusplus
+} // extern "C"
+#endif // __cplusplus
+
+#endif // IREE_TOOLING_PROCESS_RESULTS_H_
diff --git a/runtime/src/iree/tooling/run_module.c b/runtime/src/iree/tooling/run_module.c
index 04a461b..20ddbc2 100644
--- a/runtime/src/iree/tooling/run_module.c
+++ b/runtime/src/iree/tooling/run_module.c
@@ -10,14 +10,13 @@
#include "iree/base/tooling/flags.h"
#include "iree/hal/api.h"
#include "iree/hal/replay/recorder.h"
-#include "iree/io/stdio_stream.h"
#include "iree/modules/hal/types.h"
-#include "iree/tooling/comparison.h"
#include "iree/tooling/context_util.h"
#include "iree/tooling/device_util.h"
#include "iree/tooling/function_io.h"
#include "iree/tooling/function_util.h"
#include "iree/tooling/instrument_util.h"
+#include "iree/tooling/process_results.h"
#include "iree/vm/api.h"
#include "iree/vm/bytecode/module.h"
@@ -45,58 +44,9 @@
"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 an ndarray\n"
- " e.g.: --output=@file.npy\n"
- " `+file.npy`: create/append a numpy npy file and write an ndarray\n"
- " e.g.: --output=+file.npy\n"
- " `@file.bin`: create/overwrite a binary file and write value contents\n"
- " e.g.: --output=@file.bin\n"
- " `+file.bin`: create/append a binary file and write value contents\n"
- " e.g.: --output=+file.bin\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"
- "Primitive values are written as shape=() ndarrays and buffers are\n"
- "written as i8 arrays with the length of the buffer.\n"
- "\n"
- "Binary files contain only the contents of the values/buffers provided\n"
- "without metadata; users must know the shape/type of the output.\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=`.\n"
- "When present the results of the invocation will be compared against\n"
- "these values and the tool will return non-zero if any differ. If the\n"
- "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 and elides\n"
- "the remainder.");
-
IREE_FLAG(bool, print_statistics, false,
"Prints runtime statistics to stderr on exit.");
-static iree_status_t iree_tooling_process_results(
- iree_hal_device_t* device, iree_string_view_t results_cconv,
- iree_vm_list_t* results, iree_io_stream_t* stream,
- iree_allocator_t host_allocator, int* out_exit_code);
-
static iree_status_t iree_tooling_create_run_context(
iree_vm_instance_t* instance, iree_string_view_t default_device_uri,
iree_const_byte_span_t module_contents, iree_allocator_t host_allocator,
@@ -319,30 +269,13 @@
/*wait_fence=*/NULL, /*signal_fence=*/NULL);
}
- // Wrap stdout for printing results.
- iree_io_stream_t* stdout_stream = NULL;
+ // Print/write the outputs or check them against expected values.
if (iree_status_is_ok(status)) {
- status = iree_status_annotate_f(
- iree_io_stdio_stream_wrap(IREE_IO_STREAM_MODE_WRITABLE, stdout,
- /*owns_handle=*/false, host_allocator,
- &stdout_stream),
- "opening stdout stream");
- }
-
- // Handle either printing/writing the outputs or checking them against
- // expected values (basic pass/fail testing).
- if (iree_status_is_ok(status)) {
- status = iree_status_annotate_f(
- iree_tooling_process_results(device, results_cconv, outputs,
- stdout_stream, host_allocator,
- out_exit_code),
- "processing function outputs");
+ status = iree_tooling_process_results_and_print(
+ device, results_cconv, outputs, host_allocator, out_exit_code);
}
iree_vm_list_release(outputs);
- iree_io_stream_release(stdout_stream);
- fflush(stdout);
-
if (replay_execute_scope_open) {
status = iree_status_join(status, iree_hal_replay_recorder_scope_end(
replay_recorder, IREE_SV("execute")));
@@ -350,66 +283,6 @@
return status;
}
-static iree_status_t iree_tooling_process_results(
- iree_hal_device_t* device, iree_string_view_t results_cconv,
- iree_vm_list_t* results, iree_io_stream_t* stream,
- iree_allocator_t host_allocator, int* out_exit_code) {
- *out_exit_code = EXIT_SUCCESS;
-
- // Basic output handling to route to the console or files.
- if (FLAG_expected_output_list().count == 0) {
- if (FLAG_output_list().count == 0) {
- // Print all outputs.
- return iree_status_annotate_f(
- iree_tooling_print_variants(
- IREE_SV("result"), results,
- (iree_host_size_t)FLAG_output_max_element_count, stream,
- host_allocator),
- "printing results");
- } else {
- // Write (or ignore) all outputs.
- return iree_status_annotate_f(
- iree_tooling_write_variants(
- results, FLAG_output_list(),
- (iree_host_size_t)FLAG_output_max_element_count, stream,
- host_allocator),
- "outputting results");
- }
- }
-
- // Compare against contents in host-local memory. This avoids polluting
- // device memory statistics.
- iree_hal_allocator_t* heap_allocator = NULL;
- IREE_RETURN_IF_ERROR(iree_hal_allocator_create_heap(
- IREE_SV("heap"), host_allocator, host_allocator, &heap_allocator));
-
- // Parse expected list into host-local memory that we can easily access.
- iree_vm_list_t* expected_list = NULL;
- iree_status_t status = iree_status_annotate_f(
- iree_tooling_parse_variants(results_cconv, FLAG_expected_output_list(),
- device, heap_allocator, host_allocator,
- &expected_list),
- "parsing expected function outputs");
-
- // Compare expected vs actual lists and output diffs.
- if (iree_status_is_ok(status)) {
- bool did_match = iree_tooling_compare_variant_lists(expected_list, results,
- host_allocator, stdout);
- if (did_match) {
- fprintf(
- stdout,
- "[SUCCESS] all function outputs matched their expected values.\n");
- }
-
- // Exit code 0 if all results matched the expected values.
- *out_exit_code = did_match ? EXIT_SUCCESS : EXIT_FAILURE;
- }
-
- iree_vm_list_release(expected_list);
- iree_hal_allocator_release(heap_allocator);
- return status;
-}
-
iree_status_t iree_tooling_run_module_from_flags(
iree_vm_instance_t* instance, iree_allocator_t host_allocator,
int* out_exit_code) {
diff --git a/tools/BUILD.bazel b/tools/BUILD.bazel
index 2c94abb..3277371 100644
--- a/tools/BUILD.bazel
+++ b/tools/BUILD.bazel
@@ -70,6 +70,7 @@
"//runtime/src/iree/tooling:context_util",
"//runtime/src/iree/tooling:device_util",
"//runtime/src/iree/tooling:function_io",
+ "//runtime/src/iree/tooling:process_results",
"//runtime/src/iree/vm",
"@com_google_benchmark//:benchmark",
],
diff --git a/tools/CMakeLists.txt b/tools/CMakeLists.txt
index 16e2c03..e1598ba 100644
--- a/tools/CMakeLists.txt
+++ b/tools/CMakeLists.txt
@@ -101,6 +101,7 @@
iree::tooling::context_util
iree::tooling::device_util
iree::tooling::function_io
+ iree::tooling::process_results
iree::vm
COVERAGE ${IREE_ENABLE_RUNTIME_COVERAGE}
INSTALL_COMPONENT IREETools-Runtime
diff --git a/tools/iree-benchmark-module-main.cc b/tools/iree-benchmark-module-main.cc
index f7bda24..6c75474 100644
--- a/tools/iree-benchmark-module-main.cc
+++ b/tools/iree-benchmark-module-main.cc
@@ -69,6 +69,7 @@
#include "iree/tooling/context_util.h"
#include "iree/tooling/device_util.h"
#include "iree/tooling/function_io.h"
+#include "iree/tooling/process_results.h"
#include "iree/vm/api.h"
constexpr char kNanosecondsUnitString[] = "ns";
@@ -157,6 +158,11 @@
parse_time_unit, print_time_unit, &FLAG_time_unit, time_unit,
"The time unit to be printed in the results. Can be 'ms', 'us', or 'ns'.");
+IREE_FLAG(
+ bool, enable_output_processing, false,
+ "Enable keeping outputs of last benchmark iteration and processing "
+ "those. This needs to be enabled for --output* options to be effective.");
+
static iree_hal_profiling_from_flags_t* g_profiling = nullptr;
namespace iree {
@@ -233,24 +239,31 @@
const std::string& benchmark_name, int32_t batch_size,
iree_hal_replay_recorder_t* recorder, iree_hal_device_t* device,
iree_vm_context_t* context, iree_vm_function_t function,
- iree_vm_list_t* inputs, benchmark::State& state) {
+ iree_vm_list_t* inputs, iree_vm_list_t* outputs, benchmark::State& state) {
IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, benchmark_name.data(),
benchmark_name.size());
IREE_TRACE_FRAME_MARK();
- vm::ref<iree_vm_list_t> outputs;
- IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
- iree_allocator_system(), &outputs));
-
+ // Use output list passed by caller if available. Create local VM list
+ // otherwise.
+ vm::ref<iree_vm_list_t> local_outputs;
+ if (outputs) {
+ local_outputs = vm::retain_ref(outputs);
+ } else {
+ IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
+ iree_allocator_system(), &local_outputs));
+ }
// Benchmarking loop.
while (state.KeepRunningBatch(batch_size)) {
IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration");
IREE_TRACE_FRAME_MARK_NAMED("Iteration");
BeginReplayExecuteScope(recorder);
+ // Clear the output list at the beginning of loop, so we can keep the
+ // outputs of the last loop iteration.
+ IREE_CHECK_OK(iree_vm_list_resize(local_outputs.get(), 0));
IREE_CHECK_OK(iree_vm_invoke(
context, function, IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/nullptr,
- inputs, outputs.get(), iree_allocator_system()));
- IREE_CHECK_OK(iree_vm_list_resize(outputs.get(), 0));
+ inputs, local_outputs.get(), iree_allocator_system()));
EndReplayExecuteScope(recorder);
IREE_TRACE_ZONE_END(z1);
if (device) {
@@ -261,6 +274,12 @@
}
state.SetItemsProcessed(state.iterations());
+ // Clear the outputs if we used a local list. Keep the outputs in the list
+ // if the caller provided the outputs list.
+ if (!outputs) {
+ IREE_CHECK_OK(iree_vm_list_resize(local_outputs.get(), 0));
+ }
+
IREE_TRACE_ZONE_END(z0);
}
@@ -269,15 +288,15 @@
iree_hal_device_t* device,
iree_vm_context_t* context,
iree_vm_function_t function,
- iree_vm_list_t* inputs) {
+ iree_vm_list_t* inputs, iree_vm_list_t* outputs) {
auto benchmark_name = "BM_" + function_name;
int32_t batch_size = FLAG_batch_size;
- benchmark::RegisterBenchmark(benchmark_name.c_str(),
- [=](benchmark::State& state) -> void {
- BenchmarkGenericFunction(
- benchmark_name, batch_size, recorder,
- device, context, function, inputs, state);
- })
+ benchmark::RegisterBenchmark(
+ benchmark_name.c_str(),
+ [=](benchmark::State& state) -> void {
+ BenchmarkGenericFunction(benchmark_name, batch_size, recorder, device,
+ context, function, inputs, outputs, state);
+ })
// By default only the main thread is included in CPU time. Include all
// the threads instead.
->MeasureProcessCPUTime()
@@ -525,6 +544,7 @@
iree_status_t Shutdown() {
// Order matters. Tear down modules first to release resources.
inputs_.reset();
+ outputs_.reset();
context_.reset();
iree_status_t status = CloseReplayCapture();
iree_tooling_module_list_reset(&module_list_);
@@ -575,6 +595,32 @@
return iree_ok_status();
}
+ // Turn on keeping results of last benchmark iteration by setting up
+ // persistent output list.
+ iree_status_t EnableKeepResults() {
+ return iree_vm_list_create(iree_vm_make_undefined_type_def(), 16,
+ iree_allocator_system(), &outputs_);
+ }
+
+ // Handle printing/writing/checking the outputs kept from running the last
+ // iteration of the module/function.
+ iree_status_t ProcessResults(int* out_exit_code) {
+ IREE_TRACE_SCOPE_NAMED("IREEBenchmark::ProcessResults");
+
+ if (!outputs_) {
+ return iree_make_status(IREE_STATUS_FAILED_PRECONDITION,
+ "no output list to process");
+ }
+
+ IREE_RETURN_IF_ERROR(iree_tooling_process_results_and_print(
+ device_,
+ iree_make_string_view(results_cconv_.data(), results_cconv_.size()),
+ outputs_.get(), iree_vm_instance_allocator(instance_.get()),
+ out_exit_code));
+
+ return iree_ok_status();
+ }
+
private:
iree_status_t Init() {
IREE_TRACE_SCOPE_NAMED("IREEBenchmark::Init");
@@ -613,6 +659,7 @@
iree_string_view_t arguments_cconv, results_cconv;
IREE_RETURN_IF_ERROR(iree_vm_function_call_get_cconv_fragments(
&signature, &arguments_cconv, &results_cconv));
+ results_cconv_.assign(results_cconv.data, results_cconv.size);
IREE_CHECK_OK(iree_tooling_parse_variants(
arguments_cconv, FLAG_input_list(), device_, device_allocator_.get(),
@@ -627,7 +674,8 @@
} else {
// Synchronous invocation.
iree::RegisterGenericBenchmark(function_name, replay_recorder_, device_,
- context_.get(), function, inputs_.get());
+ context_.get(), function, inputs_.get(),
+ outputs_.get());
}
return iree_ok_status();
}
@@ -656,7 +704,7 @@
iree::RegisterGenericBenchmark(
std::string(function_name.data, function_name.size),
replay_recorder_, device_, context_.get(), function,
- /*inputs=*/nullptr);
+ /*inputs=*/nullptr, /*outputs=*/nullptr);
} else {
// Pick up generic () -> () functions.
if (iree_string_view_starts_with(function_name,
@@ -695,7 +743,7 @@
iree::RegisterGenericBenchmark(
std::string(function_name.data, function_name.size),
replay_recorder_, device_, context_.get(), function,
- /*inputs=*/nullptr);
+ /*inputs=*/nullptr, /*outputs=*/nullptr);
}
}
}
@@ -711,6 +759,8 @@
iree_hal_replay_recorder_t* replay_recorder_ = nullptr;
iree_tooling_module_list_t module_list_;
iree::vm::ref<iree_vm_list_t> inputs_;
+ iree::vm::ref<iree_vm_list_t> outputs_;
+ std::string results_cconv_;
};
} // namespace
} // namespace iree
@@ -733,6 +783,10 @@
::benchmark::Initialize(&argc, argv);
iree::IREEBenchmark iree_benchmark;
+ if (FLAG_enable_output_processing) {
+ IREE_CHECK_OK(iree_benchmark.EnableKeepResults());
+ }
+
iree_status_t status = iree_benchmark.Register();
if (!iree_status_is_ok(status)) {
status = iree_status_join(status, iree_benchmark.Shutdown());
@@ -743,13 +797,20 @@
}
IREE_CHECK_OK(iree_hal_begin_profiling_from_flags(
iree_benchmark.device(), iree_allocator_system(), &g_profiling));
+
::benchmark::RunSpecifiedBenchmarks();
IREE_CHECK_OK(iree_hal_end_profiling_from_flags(g_profiling));
g_profiling = nullptr;
+
+ int exit_code = EXIT_SUCCESS;
+ if (FLAG_enable_output_processing) {
+ IREE_CHECK_OK(iree_benchmark.ProcessResults(&exit_code));
+ }
+
IREE_CHECK_OK(iree_benchmark.Shutdown());
IREE_TRACE_ZONE_END(z0);
- return EXIT_SUCCESS;
+ return exit_code;
}
int main(int argc, char** argv) {
diff --git a/tools/test/BUILD.bazel b/tools/test/BUILD.bazel
index 2b1dd59..d71a773 100644
--- a/tools/test/BUILD.bazel
+++ b/tools/test/BUILD.bazel
@@ -37,6 +37,7 @@
"executable_sources.mlir",
"iree-benchmark-executable.mlir",
"iree-compile-help.txt",
+ "iree-benchmark-module-outputs.mlir",
"iree-benchmark-module.mlir",
"iree-convert-parameters.txt",
"iree-dump-module-control-flow.mlir",
diff --git a/tools/test/CMakeLists.txt b/tools/test/CMakeLists.txt
index 108f961..04e11b4 100644
--- a/tools/test/CMakeLists.txt
+++ b/tools/test/CMakeLists.txt
@@ -24,6 +24,7 @@
"executable_configurations.mlir"
"executable_sources.mlir"
"iree-benchmark-executable.mlir"
+ "iree-benchmark-module-outputs.mlir"
"iree-benchmark-module.mlir"
"iree-compile-help.txt"
"iree-convert-parameters.txt"
diff --git a/tools/test/iree-benchmark-module-outputs.mlir b/tools/test/iree-benchmark-module-outputs.mlir
new file mode 100644
index 0000000..c26d578
--- /dev/null
+++ b/tools/test/iree-benchmark-module-outputs.mlir
@@ -0,0 +1,56 @@
+// Tests that output processing is performed by iree-benchmark-module when
+// --enable_output_processing is passed (and only if it is passed).
+// Test that --output= options are working in general, by testing a single one.
+// See iree-run-module-outputs.mlir for tests of the --output= options
+// themselves, which are implemented by code shared between iree-run-module
+// and iree-benchmark-module.
+
+// RUN: (iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx %s | \
+// RUN: iree-benchmark-module --device=local-sync --module=- \
+// RUN: --function=default) | \
+// RUN: FileCheck --check-prefix=DISABLED %s
+// DISABLED-LABEL: BM_default
+// DISABLED-NOT: result[
+
+// RUN: (iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx %s | \
+// RUN: iree-benchmark-module --device=local-sync --module=- \
+// RUN: --function=default \
+// RUN: --enable_output_processing) | \
+// RUN: FileCheck --check-prefix=ENABLED %s
+// ENABLED-LABEL: BM_default
+// ENABLED: result[0]: i32=123
+
+// RUN: (iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx %s | \
+// RUN: iree-benchmark-module --device=local-sync --module=- \
+// RUN: --function=default \
+// RUN: --enable_output_processing \
+// RUN: --output=@%t.npy) && \
+// RUN: "%PYTHON" %S/echo_npy.py %t.npy | \
+// RUN: FileCheck --check-prefix=OUTPUT-OPTION %s
+// OUTPUT-OPTION{LITERAL}: 123
+func.func @default() -> (i32) {
+ %0 = arith.constant 123 : i32
+ return %0 : i32
+}
+
+// Tests that a failing --expected_output= check makes the tool exit with a
+// failure exit code, same as iree-run-module. See
+// iree-run-module-expected.mlir for full coverage of --expected_output=
+// comparison behavior itself.
+
+// RUN: (iree-compile --iree-hal-target-device=local \
+// RUN: --iree-hal-local-target-device-backends=vmvx %s | \
+// RUN: not iree-benchmark-module --device=local-sync --module=- \
+// RUN: --function=abs --input=f32=-2 \
+// RUN: --enable_output_processing \
+// RUN: --expected_output=f32=3) | \
+// RUN: FileCheck --check-prefix=EXPECTED-MISMATCH %s
+// EXPECTED-MISMATCH-LABEL: BM_abs
+// EXPECTED-MISMATCH: [FAILED]
+func.func @abs(%input: tensor<f32>) -> (tensor<f32>) {
+ %result = math.absf %input : tensor<f32>
+ return %result : tensor<f32>
+}
