[iree-benchmark-module] Unify sync, async and dispatch execution (#24986) `iree-benchmark-module` had separate synchronous, asynchronous, and generated-dispatch execution paths that applied output processing, replay capture, and profiling flushes inconsistently. Run them through one batch lifecycle so discovery resolves each function's ABI and iteration accounting once and every execution model shares the measured loop. The shared lifecycle fixes several inconsistencies: - Retain asynchronous outputs for `--enable_output_processing`. - Supply fence arguments for discovered asynchronous functions. - Use the generated repeat count for dispatch benchmarks selected with `--function`. - Handle non-power-of-two `--batch_concurrency` without reading past the invocation list. - Return recording devices from context setup so asynchronous fences participate in replay capture. Errors go to stderr with `EXIT_FAILURE`, matching `iree-run-module`. Output processing requires a single benchmarked function, and nonpositive batch size or concurrency is rejected. Device profiling ends before output processing. Assisted by AI. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
diff --git a/docs/website/docs/developers/performance/benchmarking.md b/docs/website/docs/developers/performance/benchmarking.md index 015fce2..ae78f69 100644 --- a/docs/website/docs/developers/performance/benchmarking.md +++ b/docs/website/docs/developers/performance/benchmarking.md
@@ -103,6 +103,19 @@ Remember to [restore CPU scaling](#cpu-configuration) when you're done. +`--batch_size` counts logical benchmark iterations. Its execution depends on +the exported function's ABI: + +| Function | One measured batch | +| --- | --- | +| Synchronous entry point | One VM invocation. Match `--batch_size` to the compiler's `--iree-hal-benchmark-dispatch-repeat-count`. | +| Asynchronous entry point | `--batch_size` VM invocations, rounded up to a multiple of `--batch_concurrency`, followed by one completion wait. Fence and argument setup and output cleanup are outside the measured interval. | +| Generated dispatch benchmark | One VM invocation with `--batch_size` as its repeat count. | + +`--enable_output_processing` retains the last measured invocation's outputs +for the usual `--output` and `--expected_output` flags, including asynchronous +entry points. Select a single function when processing outputs. + ## Executable Benchmarks We also benchmark the performance of individual parts of the IREE system in
diff --git a/runtime/bindings/python/CMakeLists.txt b/runtime/bindings/python/CMakeLists.txt index 580cfe8..1b0c834 100644 --- a/runtime/bindings/python/CMakeLists.txt +++ b/runtime/bindings/python/CMakeLists.txt
@@ -374,6 +374,15 @@ iree_py_test( NAME + benchmark_lifecycle_test + SRCS + "tests/benchmark_lifecycle_test.py" + TIMEOUT + 300 + ) + + iree_py_test( + NAME hal_test SRCS "tests/hal_test.py"
diff --git a/runtime/bindings/python/tests/benchmark_lifecycle_test.py b/runtime/bindings/python/tests/benchmark_lifecycle_test.py new file mode 100644 index 0000000..6b12b11 --- /dev/null +++ b/runtime/bindings/python/tests/benchmark_lifecycle_test.py
@@ -0,0 +1,376 @@ +# 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 + +"""Exercises iree-benchmark-module lifecycle contracts. + +Covers execution models, batching, discovery, output ownership, capture and +failures through iree.runtime.benchmark. +""" + +import os +from pathlib import Path +import re +import subprocess +import tempfile +import unittest +from unittest import mock + +import iree.compiler +from iree._runtime import libs +from iree.runtime.benchmark import ( + _build_benchmark_args as build_benchmark_args, + _run_benchmark as run_benchmark_command, + benchmark_module, + BenchmarkToolError, +) + +# Two 64x64 matmuls and one 64x32 matmul: two dispatch functions, one of them +# called twice per invocation. +MATMUL = """ +func.func @main(%lhs: tensor<64x64xf32>, %rhs: tensor<64x64xf32>, + %narrow: tensor<64x32xf32>) -> tensor<64x32xf32> { + %zero = arith.constant 0.0 : f32 + %empty = tensor.empty() : tensor<64x64xf32> + %fill = linalg.fill ins(%zero : f32) outs(%empty : tensor<64x64xf32>) -> tensor<64x64xf32> + %0 = linalg.matmul ins(%lhs, %rhs : tensor<64x64xf32>, tensor<64x64xf32>) + outs(%fill : tensor<64x64xf32>) -> tensor<64x64xf32> + %1 = linalg.matmul ins(%0, %rhs : tensor<64x64xf32>, tensor<64x64xf32>) + outs(%fill : tensor<64x64xf32>) -> tensor<64x64xf32> + %narrow_empty = tensor.empty() : tensor<64x32xf32> + %narrow_fill = linalg.fill ins(%zero : f32) outs(%narrow_empty : tensor<64x32xf32>) -> tensor<64x32xf32> + %2 = linalg.matmul ins(%1, %narrow : tensor<64x64xf32>, tensor<64x32xf32>) + outs(%narrow_fill : tensor<64x32xf32>) -> tensor<64x32xf32> + return %2 : tensor<64x32xf32> +} +""" + +# Two exported functions without inputs, for discovery without --function. +NO_INPUTS = """ +func.func @one() -> tensor<4xf32> { + %input = util.unfoldable_constant dense<-1.0> : tensor<4xf32> + %result = math.absf %input : tensor<4xf32> + return %result : tensor<4xf32> +} +func.func @two() -> tensor<4xf32> { + %input = util.unfoldable_constant dense<-1.0> : tensor<4xf32> + %result = math.absf %input : tensor<4xf32> + return %result : tensor<4xf32> +} +""" + +# Uses no HAL device. +PURE = """ +func.func @main() -> i32 { + %result = arith.constant 42 : i32 + return %result : i32 +} +""" + +# Adds the number of invocations so far to the absolute input. +STATEFUL = """ +util.global private mutable @count = 0 : i32 +func.func @main(%input: tensor<4xf32>) -> tensor<4xf32> { + %old = util.global.load @count : i32 + %one = arith.constant 1 : i32 + %next = arith.addi %old, %one : i32 + util.global.store %next, @count : i32 + %count = arith.sitofp %next : i32 to f32 + %empty = tensor.empty() : tensor<4xf32> + %counts = linalg.fill ins(%count : f32) outs(%empty : tensor<4xf32>) -> tensor<4xf32> + %absolute = math.absf %input : tensor<4xf32> + %result = arith.addf %absolute, %counts : tensor<4xf32> + return %result : tensor<4xf32> +} +""" + +# One dispatch, dumped as a generated dispatch benchmark. +ABSOLUTE = """ +func.func @abs(%input: tensor<4xf32>) -> tensor<4xf32> { + %result = math.absf %input : tensor<4xf32> + return %result : tensor<4xf32> +} +""" + +LOCAL_TARGET = [ + "--iree-hal-target-device=local", + "--iree-hal-local-target-device-backends=llvm-cpu", +] +MATMUL_INPUTS = ["64x64xf32=1", "64x64xf32=1", "64x32xf32=1"] +DEVICE = "local-task" +# Flags of every run: four iterations on a device with two workers. +BENCHMARK_FLAGS = { + "device": DEVICE, + "task_topology_group_count": 2, + "benchmark_min_time": "4x", +} +# The tool under test and the replay tools come from the same package. +TOOLS = Path(libs.library_path) +BENCHMARK_TOOL = TOOLS / "iree-benchmark-module" +TIMEOUT_SECONDS = 60 + + +def setUpModule(): + # The checks must be independent of ambient Google Benchmark defaults. + environment = { + key: value + for key, value in os.environ.items() + if not key.startswith("BENCHMARK_") + } + patch = mock.patch.dict(os.environ, environment, clear=True) + patch.start() + unittest.addModuleCleanup(patch.stop) + + +def compile_module(directory, name, source, *flags): + """Compiles MLIR source to directory/name.vmfb and returns its path.""" + output = directory / f"{name}.vmfb" + iree.compiler.compile_str( + source, extra_args=[*LOCAL_TARGET, *flags], output_file=str(output) + ) + return output + + +def benchmark_results(module, function="main", inputs=MATMUL_INPUTS, **flags): + """Benchmarks a module with BENCHMARK_FLAGS and returns all result rows. + + Keyword flags are passed to the tool. Without a function every exported + function that takes no inputs runs. + """ + return benchmark_module( + module, + entry_function=function, + inputs=inputs, + timeout=TIMEOUT_SECONDS, + executable=BENCHMARK_TOOL, + **{**BENCHMARK_FLAGS, **flags}, + ) + + +def run_benchmark(module, function="main", inputs=MATMUL_INPUTS, **flags): + """Runs a module like benchmark_results and returns stdout and stderr.""" + args, _ = build_benchmark_args( + module, + entry_function=function, + inputs=inputs, + executable=BENCHMARK_TOOL, + **{**BENCHMARK_FLAGS, **flags}, + ) + return run_benchmark_command(args, timeout=TIMEOUT_SECONDS) + + +def run_replay_tool(name, *args): + """Runs a replay tool from the package and fails with its output.""" + result = subprocess.run( + [TOOLS / name, *map(str, args)], + capture_output=True, + text=True, + timeout=TIMEOUT_SECONDS, + ) + if result.returncode != 0: + raise AssertionError(f"{name} failed:\n{result.stderr}\n{result.stdout}") + + +def compile_dispatch_benchmark(directory): + """Compiles the generated benchmark of ABSOLUTE's one dispatch. + + Returns the module and the name of its benchmark function. + """ + sources = directory / "dispatches" + compile_module( + directory, + "abs", + ABSOLUTE, + f"--iree-hal-dump-executable-benchmarks-to={sources}", + ) + (source,) = sources.glob("*_benchmark.mlir") + module = directory / "dispatch.vmfb" + iree.compiler.compile_file(str(source), output_file=str(module)) + stdout, _ = run_benchmark( + module, function=None, inputs=[], benchmark_list_tests=True + ) + (name,) = re.findall(r"^BM_([^/\s]+)", stdout, re.MULTILINE) + return module, name + + +class BenchmarkLifecycleTest(unittest.TestCase): + @classmethod + def setUpClass(cls): + modules = tempfile.TemporaryDirectory(prefix="benchmark-lifecycle-test-") + cls.addClassCleanup(modules.cleanup) + directory = Path(modules.name) + cls.sync = compile_module(directory, "sync", MATMUL) + cls.async_module = compile_module( + directory, "async", MATMUL, "--iree-execution-model=async-external" + ) + cls.repeated = compile_module( + directory, + "repeated", + MATMUL, + "--iree-hal-benchmark-dispatch-repeat-count=4", + ) + # Four iterations of each execution model: four sync invocations, one + # batch of four async invocations, or one invocation that repeats its + # dispatches four times. + cls.batched = [(cls.sync, 1), (cls.async_module, 4), (cls.repeated, 4)] + + def setUp(self): + temporary = tempfile.TemporaryDirectory(prefix="benchmark-lifecycle-test-") + self.addCleanup(temporary.cleanup) + self.directory = Path(temporary.name) + + def assertBenchmarkFails(self, message, module, **kwargs): + """Checks that run_benchmark fails with message in its output.""" + with self.assertRaisesRegex(BenchmarkToolError, re.escape(message)): + run_benchmark(module, **kwargs) + + def test_unregistered_flags_pass_through(self): + # As in the other benchmark tools. + run_benchmark(self.sync, not_a_registered_flag=1) + + def test_iterations_for_each_execution_model(self): + for module, batch_size in self.batched: + with self.subTest(module=module.stem): + results = benchmark_results(module, batch_size=batch_size) + self.assertEqual([result.iterations for result in results], ["4"]) + + def test_output_processing_for_each_execution_model(self): + for module, batch_size in self.batched: + with self.subTest(module=module.stem): + stdout, _ = run_benchmark( + module, + batch_size=batch_size, + enable_output_processing=True, + expected_output="64x32xf32=262144", + ) + self.assertIn("[SUCCESS]", stdout) + + def test_batch_rounded_up_to_concurrency(self): + # Three timelines run a batch of four as six invocations. + results = benchmark_results( + self.async_module, batch_size=4, batch_concurrency=3 + ) + self.assertEqual([result.iterations for result in results], ["6"]) + + def test_discovery_prepares_async_arguments(self): + # Discovery must also prepare empty arguments for async entry points. + for model in ("async-internal", "async-external"): + with self.subTest(model=model): + module = compile_module( + self.directory, + model, + NO_INPUTS, + f"--iree-execution-model={model}", + ) + results = benchmark_results(module, function=None, inputs=[]) + self.assertEqual(len(results), 2, results) + + def test_filtered_discovery_uses_time_unit(self): + for model in ("async-internal", "async-external"): + with self.subTest(model=model): + module = compile_module( + self.directory, + model, + NO_INPUTS, + f"--iree-execution-model={model}", + ) + results = benchmark_results( + module, + function=None, + inputs=[], + benchmark_filter="BM_two", + time_unit="ns", + ) + self.assertEqual(len(results), 1, results) + self.assertTrue(results[0].time.endswith(" ns"), results) + + def test_output_processing_requires_one_function(self): + # Without --function both @one and @two run. + module = compile_module(self.directory, "no-inputs", NO_INPUTS) + self.assertBenchmarkFails( + "output processing requires one function", + module, + function=None, + inputs=[], + enable_output_processing=True, + ) + + def test_output_processing_without_devices(self): + module = compile_module(self.directory, "pure", PURE) + stdout, _ = run_benchmark( + module, inputs=[], enable_output_processing=True, output="-" + ) + self.assertIn("i32=42", stdout) + + def test_outputs_of_last_measured_invocation(self): + # The stateful function counts its invocations, and the outputs are + # those of the fourth. + for model in ("async-internal", "async-external"): + with self.subTest(model=model): + module = compile_module( + self.directory, + model, + STATEFUL, + f"--iree-execution-model={model}", + ) + stdout, _ = run_benchmark( + module, + inputs=["4xf32=-1"], + enable_output_processing=True, + expected_output="4xf32=5", + ) + self.assertIn("[SUCCESS]", stdout) + + def test_dispatch_benchmark_discovery_and_selection(self): + # Automatic discovery and explicit selection use the same dispatch ABI: + # two batches of three per four iterations. + module, name = compile_dispatch_benchmark(self.directory) + for function in (None, name): + with self.subTest(function=function): + results = benchmark_results( + module, function=function, inputs=[], batch_size=3 + ) + self.assertEqual([result.iterations for result in results], ["6"]) + + def test_dispatch_benchmark_rejects_inputs(self): + module, name = compile_dispatch_benchmark(self.directory) + self.assertBenchmarkFails( + "dispatch wrappers take --batch_size, not --input", + module, + function=name, + inputs=["4xf32=1"], + ) + + def test_device_statistics(self): + _, stderr = run_benchmark(self.sync, print_device_statistics=True) + self.assertIn("main_dispatch_", stderr) + + def test_replay_capture(self): + for module in (self.sync, self.async_module): + with self.subTest(module=module.stem): + capture = self.directory / f"{module.stem}.ireereplay" + run_benchmark(module, device_replay_output=capture) + run_replay_tool("iree-run-replay", f"--device={DEVICE}", capture) + + def test_replay_capture_finalized_on_error(self): + capture = self.directory / "failed.ireereplay" + self.assertBenchmarkFails( + "input0 shape rank mismatch", + self.sync, + inputs=["1xf32=1", *MATMUL_INPUTS[1:]], + device_replay_output=capture, + ) + run_replay_tool("iree-dump-replay", capture) + + def test_batch_size_and_concurrency_must_be_positive(self): + for flags in ({"batch_size": 0}, {"batch_concurrency": 0}): + with self.subTest(**flags): + self.assertBenchmarkFails( + "batch size and concurrency must be positive", self.sync, **flags + ) + + +if __name__ == "__main__": + unittest.main()
diff --git a/runtime/src/iree/tooling/context_util.c b/runtime/src/iree/tooling/context_util.c index e7b5b19..4e10ab7 100644 --- a/runtime/src/iree/tooling/context_util.c +++ b/runtime/src/iree/tooling/context_util.c
@@ -495,6 +495,17 @@ status = iree_status_join(status, iree_hal_replay_recorder_scope_end( replay_recorder, IREE_SV("init"))); } + if (iree_status_is_ok(status) && replay_device_group) { + // Callers creating fences and other HAL resources must use the same + // recording devices as the module so those resources enter the capture. + for (iree_host_size_t i = 0; i < device_list->count; ++i) { + iree_hal_device_t* recording_device = + iree_hal_device_group_device_at(replay_device_group, i); + iree_hal_device_retain(recording_device); + iree_hal_device_release(device_list->devices[i]); + device_list->devices[i] = recording_device; + } + } iree_hal_device_group_release(replay_device_group); iree_hal_device_group_release(device_group);
diff --git a/runtime/src/iree/tooling/context_util.h b/runtime/src/iree/tooling/context_util.h index b080c00..6dd186d 100644 --- a/runtime/src/iree/tooling/context_util.h +++ b/runtime/src/iree/tooling/context_util.h
@@ -57,6 +57,9 @@ // |default_device_uri| can be specified to provide a default if a device flag // is not provided by the user. // |out_device_list| will contain the created devices if using the full HAL. +// With --device_replay_output= it holds the replay-recording devices the +// modules use, so fences and other HAL resources created on them enter the +// capture, as do profiling sessions begun on them (replay skips those). // |out_device_allocator| can be used to allocate buffers for use with the // context and is available in all execution models. // |out_replay_recorder| will contain a recorder that must be closed after all @@ -99,6 +102,9 @@ // |default_device_uri| can be specified to provide a default if a device flag // is not provided by the user. // |out_device_list| will contain the created devices if using the full HAL. +// With --device_replay_output= it holds the replay-recording devices the +// modules use, so fences and other HAL resources created on them enter the +// capture, as do profiling sessions begun on them (replay skips those). // |out_device_allocator| can be used to allocate buffers for use with the // context and is available in all execution models. // |out_replay_recorder| will contain a recorder that must be closed after all
diff --git a/tools/iree-benchmark-module-main.cc b/tools/iree-benchmark-module-main.cc index fbd3293..c34cd0e 100644 --- a/tools/iree-benchmark-module-main.cc +++ b/tools/iree-benchmark-module-main.cc
@@ -52,11 +52,11 @@ // an appropriate device-specific tool before trusting the more generic and // higher-level numbers from this tool. -#include <array> #include <cstdio> -#include <iterator> +#include <cstdlib> +#include <cstring> +#include <memory> #include <string> -#include <type_traits> #include <utility> #include <vector> @@ -164,8 +164,6 @@ "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 { namespace { @@ -224,615 +222,462 @@ file); } -static void BeginReplayExecuteScope(iree_hal_replay_recorder_t* recorder) { - if (!recorder) return; - IREE_CHECK_OK( - iree_hal_replay_recorder_scope_begin(recorder, IREE_SV("execute"))); -} +// How a benchmarked function is called, which decides what one measured batch +// of --batch_size logical iterations executes. +enum class InvocationKind { + // A synchronous function. One VM invocation per batch; any repetition is + // compiled into the function (--iree-hal-benchmark-dispatch-repeat-count). + kSync, + // A function with the coarse-fences ABI. A batch is --batch_size VM + // invocations spread over --batch_concurrency fence timelines, followed by + // one wait for all of them. + kAsync, + // A generated dispatch benchmark (iree.benchmark = "dispatch"). One VM + // invocation per batch that takes the batch size as its repeat count. + kDispatch, +}; -static void EndReplayExecuteScope(iree_hal_replay_recorder_t* recorder) { - if (!recorder) return; - IREE_CHECK_OK( - iree_hal_replay_recorder_scope_end(recorder, IREE_SV("execute"))); -} - -static void BenchmarkGenericFunction( - 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, 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(); - - // 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, local_outputs.get(), iree_allocator_system())); - EndReplayExecuteScope(recorder); - IREE_TRACE_ZONE_END(z1); - if (device) { - state.PauseTiming(); - IREE_CHECK_OK(iree_hal_flush_profiling_from_flags(g_profiling)); - state.ResumeTiming(); - } - } - 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); -} - -void RegisterGenericBenchmark(const std::string& function_name, - 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, 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, outputs, state); - }) - // By default only the main thread is included in CPU time. Include all - // the threads instead. - ->MeasureProcessCPUTime() - // To make single and multi-threaded benchmarks more comparable, use the - // wall time to determine how many iterations to run. See - // https://github.com/google/benchmark#cpu-timers, - ->UseRealTime() - ->Unit(FLAG_time_unit.first ? FLAG_time_unit.second - : benchmark::kMillisecond); -} - -// Runs up to |batch_size| pipelined invocations in sequence along with -// concurrency. Example: -// batch_size=1, concurrency=1: -// [invocation 0] -// batch_size=2, concurrency=1: -// [invocation 0] -> [invocation 1] -// batch_size=2, concurrency=2: -// [invocation 0] -// [invocation 1] -// batch_size=4, concurrency=2: -// [invocation 0] -> [invocation 2] -// [invocation 1] -> [invocation 3] -static void BenchmarkAsyncFunction( - const std::string& benchmark_name, int32_t batch_size, - int32_t batch_concurrency, 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* common_inputs, - benchmark::State& state) { - IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, benchmark_name.data(), - benchmark_name.size()); - IREE_TRACE_FRAME_MARK(); - iree_allocator_t host_allocator = iree_allocator_system(); - - // Round up batch size to some multiple of concurrency. - batch_size = (int32_t)iree_host_align(batch_size, batch_concurrency); - - // Benchmarking loop. - while (state.KeepRunningBatch(batch_size)) { - state.PauseTiming(); - IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration"); - IREE_TRACE_FRAME_MARK_NAMED("Iteration"); - - IREE_TRACE_ZONE_BEGIN_NAMED(z_begin, "PrepareBatch"); - - // Each concurrent track of execution gets its own semaphore. - std::vector<vm::ref<iree_hal_semaphore_t>> timeline_semaphores; - for (int32_t i = 0; i < batch_concurrency; ++i) { - vm::ref<iree_hal_semaphore_t> timeline_semaphore; - IREE_CHECK_OK(iree_hal_semaphore_create( - device, IREE_HAL_QUEUE_AFFINITY_ANY, 0ull, - IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &timeline_semaphore)); - timeline_semaphores.push_back(std::move(timeline_semaphore)); - } - - // Preallocate fences and I/O for each invocation. - // The same inputs are used for each but we need a unique list to hold the - // unique fences. Each fence represents when the invocation has completed. - std::vector<vm::ref<iree_hal_fence_t>> invocation_fences; - std::vector<vm::ref<iree_vm_list_t>> invocation_inputs; - std::vector<vm::ref<iree_vm_list_t>> invocation_outputs; - vm::ref<iree_hal_fence_t> completion_fence; - IREE_CHECK_OK(iree_hal_fence_create(batch_concurrency, host_allocator, - &completion_fence)); - for (int32_t i = 0; i < batch_size / batch_concurrency; ++i) { - for (int32_t j = 0; j < batch_concurrency; ++j) { - // Chain each concurrent minibatch to the previous. Note that to start - // we wait on nothing and begin executing immediately. - vm::ref<iree_hal_fence_t> wait_fence; - if (i > 0) { - wait_fence = vm::retain_ref( - invocation_fences[(i - 1) * batch_concurrency + j]); - } - uint64_t signal_value = i + 1; - vm::ref<iree_hal_fence_t> signal_fence; - IREE_CHECK_OK(iree_hal_fence_create_at(timeline_semaphores[j].get(), - signal_value, host_allocator, - &signal_fence)); - invocation_fences.push_back(vm::retain_ref(signal_fence)); - - // Join the final minibatch on the completion fence. - if (i == batch_size / batch_concurrency - 1) { - IREE_CHECK_OK(iree_hal_fence_insert(completion_fence.get(), - timeline_semaphores[j].get(), - signal_value)); - } - - // Clone common inputs and add the invocation-specific fences. - vm::ref<iree_vm_list_t> inputs; - IREE_CHECK_OK( - iree_vm_list_clone(common_inputs, host_allocator, &inputs)); - IREE_CHECK_OK(iree_vm_list_push_ref_move(inputs.get(), wait_fence)); - IREE_CHECK_OK(iree_vm_list_push_ref_move(inputs.get(), signal_fence)); - invocation_inputs.push_back(std::move(inputs)); - - // Setup empty outputs. - vm::ref<iree_vm_list_t> outputs; - IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16, - host_allocator, &outputs)); - invocation_outputs.push_back(std::move(outputs)); - } - } - - IREE_TRACE_ZONE_END(z_begin); - - state.ResumeTiming(); - BeginReplayExecuteScope(recorder); - { - // TODO(benvanik): replace with async invocations. Today if the invocation - // performs any waits this will block on the initial invoke instead of - // actually overlapping things. - for (int32_t i = 0; i < batch_size; ++i) { - IREE_CHECK_OK( - iree_vm_invoke(context, function, IREE_VM_INVOCATION_FLAG_NONE, - /*policy=*/nullptr, invocation_inputs[i].get(), - invocation_outputs[i].get(), host_allocator)); - } - IREE_CHECK_OK(iree_hal_fence_wait(completion_fence.get(), - iree_infinite_timeout(), - IREE_ASYNC_WAIT_FLAG_NONE)); - } - EndReplayExecuteScope(recorder); - state.PauseTiming(); - - IREE_TRACE_ZONE_BEGIN_NAMED(z_end, "CleanupBatch"); - for (int32_t i = 0; i < batch_size; ++i) { - iree_vm_list_clear(invocation_outputs[i].get()); - } - invocation_fences.clear(); - invocation_inputs.clear(); - invocation_outputs.clear(); - completion_fence.reset(); - timeline_semaphores.clear(); - IREE_TRACE_ZONE_END(z_end); - - IREE_TRACE_ZONE_END(z1); - if (device) { - IREE_CHECK_OK(iree_hal_flush_profiling_from_flags(g_profiling)); - } - state.ResumeTiming(); - } - state.SetItemsProcessed(state.iterations()); - - IREE_TRACE_ZONE_END(z0); -} - -void RegisterAsyncBenchmark(const std::string& function_name, - 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) { - auto benchmark_name = "BM_" + function_name; - int32_t batch_size = FLAG_batch_size; - int32_t batch_concurrency = FLAG_batch_concurrency; - benchmark::RegisterBenchmark(benchmark_name.c_str(), - [=](benchmark::State& state) -> void { - BenchmarkAsyncFunction( - benchmark_name, batch_size, - batch_concurrency, recorder, device, - context, function, inputs, state); - }) - // By default only the main thread is included in CPU time. Include all - // the threads instead. - ->MeasureProcessCPUTime() - // To make single and multi-threaded benchmarks more comparable, use the - // wall time to determine how many iterations to run. See - // https://github.com/google/benchmark#cpu-timers, - ->UseRealTime() - ->Unit(FLAG_time_unit.first ? FLAG_time_unit.second - : benchmark::kMillisecond); -} - -static void BenchmarkDispatchFunction(const std::string& benchmark_name, - iree_hal_replay_recorder_t* recorder, - iree_vm_context_t* context, - iree_vm_function_t function, - benchmark::State& state) { - IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, benchmark_name.data(), - benchmark_name.size()); - IREE_TRACE_FRAME_MARK(); - +// A benchmarked function with the ABI and iteration accounting resolved during +// discovery. +struct Workload { + // Benchmark name, "BM_" followed by the function name. + std::string name; + // The exported VM function to invoke. + iree_vm_function_t function = {}; + // Calling convention of the results, used to print or check outputs. + std::string results_cconv; + InvocationKind kind = InvocationKind::kSync; + // Logical iterations per measured batch: --batch_size rounded up to a + // multiple of |concurrency|. + int32_t batch_size = 1; + // Number of fence timelines that async invocations run on; 1 otherwise. + int32_t concurrency = 1; + // Time unit of the reported rows. + benchmark::TimeUnit unit = benchmark::kMillisecond; + // Arguments shared by every invocation. Async invocations clone them and + // append their wait and signal fences; dispatch benchmarks take the repeat + // count. vm::ref<iree_vm_list_t> inputs; - IREE_CHECK_OK(iree_vm_list_create(iree_vm_make_undefined_type_def(), 16, - iree_allocator_system(), &inputs)); - iree_vm_value_t batch_size = iree_vm_value_make_i32(FLAG_batch_size); - IREE_CHECK_OK(iree_vm_list_push_value(inputs.get(), &batch_size)); - + // Outputs of the last measured invocation, retained only with + // --enable_output_processing. NULL if the function did not run. 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)); +}; - // Benchmarking loop. - while (state.KeepRunningBatch(FLAG_batch_size)) { - IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration"); - IREE_TRACE_FRAME_MARK_NAMED("Iteration"); - BeginReplayExecuteScope(recorder); - IREE_CHECK_OK(iree_vm_invoke( - context, function, IREE_VM_INVOCATION_FLAG_NONE, /*policy=*/nullptr, - inputs.get(), outputs.get(), iree_allocator_system())); - IREE_CHECK_OK(iree_vm_list_resize(outputs.get(), 0)); - EndReplayExecuteScope(recorder); - IREE_TRACE_ZONE_END(z1); +// Arguments and results of one VM call within a batch. +struct Invocation { + vm::ref<iree_vm_list_t> inputs; + vm::ref<iree_vm_list_t> outputs; +}; + +// Owns one batch's arguments, outputs and synchronization. Only the async ABI +// requires host-side batching; dispatch wrappers repeat inside the VM call. +struct Batch { + // One entry per VM call in the batch: |batch_size| for async workloads and + // one otherwise. + std::vector<Invocation> invocations; + // Signaled when the last invocation on every timeline has completed. NULL + // unless the workload is async. + vm::ref<iree_hal_fence_t> completion; + + iree_status_t Prepare(const Workload& workload, iree_hal_device_t* device) { + iree_allocator_t allocator = iree_allocator_system(); + const bool async = workload.kind == InvocationKind::kAsync; + const int count = async ? workload.batch_size : 1; + invocations.resize(count); + std::vector<vm::ref<iree_hal_semaphore_t>> timelines; + std::vector<vm::ref<iree_hal_fence_t>> previous; + if (async) { + timelines.resize(workload.concurrency); + previous.resize(workload.concurrency); + IREE_RETURN_IF_ERROR( + iree_hal_fence_create(workload.concurrency, allocator, &completion)); + for (auto& timeline : timelines) { + IREE_RETURN_IF_ERROR(iree_hal_semaphore_create( + device, IREE_HAL_QUEUE_AFFINITY_ANY, 0, + IREE_HAL_SEMAPHORE_FLAG_DEFAULT, &timeline)); + IREE_RETURN_IF_ERROR(iree_hal_fence_insert( + completion.get(), timeline.get(), count / workload.concurrency)); + } + } + for (int i = 0; i < count; ++i) { + auto& invocation = invocations[i]; + IREE_RETURN_IF_ERROR( + iree_vm_list_create(iree_vm_make_undefined_type_def(), 16, allocator, + &invocation.outputs)); + if (!async) { + invocation.inputs = vm::retain_ref(workload.inputs.get()); + continue; + } + const int track = i % workload.concurrency; + vm::ref<iree_hal_fence_t> signal; + IREE_RETURN_IF_ERROR(iree_hal_fence_create_at( + timelines[track].get(), i / workload.concurrency + 1, allocator, + &signal)); + IREE_RETURN_IF_ERROR(iree_vm_list_clone(workload.inputs.get(), allocator, + &invocation.inputs)); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_ref_move(invocation.inputs.get(), previous[track])); + previous[track] = vm::retain_ref(signal); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_ref_move(invocation.inputs.get(), signal)); + } + return iree_ok_status(); } - state.SetItemsProcessed(state.iterations()); - IREE_TRACE_ZONE_END(z0); -} + iree_status_t Execute(const Workload& workload, iree_vm_context_t* context) { + for (auto& invocation : invocations) { + // Preserve the measured lifetime of outputs: ordinary functions release + // the preceding results before invoking; dispatch wrappers clear after. + if (workload.kind == InvocationKind::kSync) { + IREE_RETURN_IF_ERROR(iree_vm_list_resize(invocation.outputs.get(), 0)); + } + IREE_RETURN_IF_ERROR(iree_vm_invoke( + context, workload.function, IREE_VM_INVOCATION_FLAG_NONE, nullptr, + invocation.inputs.get(), invocation.outputs.get(), + iree_allocator_system())); + if (workload.kind == InvocationKind::kDispatch) { + IREE_RETURN_IF_ERROR(iree_vm_list_resize(invocation.outputs.get(), 0)); + } + } + return completion + ? iree_hal_fence_wait(completion.get(), iree_infinite_timeout(), + IREE_ASYNC_WAIT_FLAG_NONE) + : iree_ok_status(); + } +}; -void RegisterDispatchBenchmark(const std::string& function_name, - iree_hal_replay_recorder_t* recorder, - iree_vm_context_t* context, - iree_vm_function_t function) { - auto benchmark_name = "BM_" + function_name; - benchmark::RegisterBenchmark(benchmark_name.c_str(), - [benchmark_name, recorder, context, - function](benchmark::State& state) -> void { - BenchmarkDispatchFunction(benchmark_name, - recorder, context, - function, state); - }) - // By default only the main thread is included in CPU time. Include all - // the threads instead. - ->MeasureProcessCPUTime() - // To make single and multi-threaded benchmarks more comparable, use the - // wall time to determine how many iterations to run. See - // https://github.com/google/benchmark#cpu-timers, - ->UseRealTime() - ->Unit(FLAG_time_unit.first ? FLAG_time_unit.second - : benchmark::kMicrosecond); -} - -// The lifetime of IREEBenchmark should be as long as -// ::benchmark::RunSpecifiedBenchmarks() where the resources are used during -// benchmarking. -class IREEBenchmark { +// Owns the VM and HAL state of one tool run: loads the modules, discovers the +// workloads, runs them for Google Benchmark and processes their outputs. +class BenchmarkSession { public: - IREEBenchmark() { iree_tooling_module_list_initialize(&module_list_); } + BenchmarkSession() { iree_tooling_module_list_initialize(&modules_); } + ~BenchmarkSession() { iree_status_ignore(Close()); } - ~IREEBenchmark() { - IREE_TRACE_SCOPE_NAMED("IREEBenchmark::dtor"); - IREE_CHECK_OK(Shutdown()); - }; + iree_status_t Open() { + if (FLAG_batch_size <= 0 || FLAG_batch_concurrency <= 0) { + return iree_make_status(IREE_STATUS_INVALID_ARGUMENT, + "batch size and concurrency must be positive"); + } + iree_allocator_t allocator = iree_allocator_system(); + IREE_RETURN_IF_ERROR(iree_tooling_create_instance(allocator, &instance_)); + IREE_RETURN_IF_ERROR(iree_tooling_load_modules_from_flags( + instance_.get(), allocator, &modules_)); + if (!modules_.count) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "no user module specified; use --module=file.vmfb"); + } + IREE_RETURN_IF_ERROR(iree_tooling_create_context_from_flags( + instance_.get(), modules_.count, modules_.values, + iree_string_view_empty(), allocator, &context_, &devices_, &allocator_, + &recorder_)); + const bool selected = strlen(FLAG_function) != 0; + iree_vm_module_t* module = iree_tooling_module_list_back(&modules_); + if (selected) { + iree_vm_function_t function; + IREE_RETURN_IF_ERROR(iree_vm_module_lookup_function_by_name( + module, IREE_VM_FUNCTION_LINKAGE_EXPORT, + iree_make_cstring_view(FLAG_function), &function)); + IREE_RETURN_IF_ERROR(AddWorkload(function, true)); + } else { + iree_vm_module_signature_t signature = iree_vm_module_signature(module); + for (iree_host_size_t i = 0; i < signature.export_function_count; ++i) { + iree_vm_function_t function; + IREE_RETURN_IF_ERROR(iree_vm_module_lookup_function_by_ordinal( + module, IREE_VM_FUNCTION_LINKAGE_EXPORT, i, &function)); + IREE_RETURN_IF_ERROR(AddWorkload(function, false)); + } + } + if (FLAG_enable_output_processing && workloads_.size() != 1) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "output processing requires one function; use --function"); + } + return iree_hal_begin_profiling_from_flags(Device(), allocator, + &profiling_); + } - iree_status_t Shutdown() { - // Order matters. Tear down modules first to release resources. - inputs_.reset(); - outputs_.reset(); + void Register() { + for (auto& workload : workloads_) { + Workload* target = workload.get(); + benchmark::RegisterBenchmark( + target->name.c_str(), + [this, target](benchmark::State& state) { Run(*target, state); }) + ->MeasureProcessCPUTime() + ->UseRealTime() + ->Unit(target->unit); + } + } + + iree_status_t ProcessOutputs(int* exit_code) { + if (!FLAG_enable_output_processing) return iree_ok_status(); + for (auto& workload : workloads_) { + if (!workload->outputs) continue; // Filtered/listed without execution. + if (Device()) { + iree_hal_buffer_params_t params = {}; + params.usage = + IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING; + params.access = IREE_HAL_MEMORY_ACCESS_ALL; + params.type = IREE_HAL_MEMORY_TYPE_HOST_LOCAL | + IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE; + params.queue_affinity = IREE_HAL_QUEUE_AFFINITY_ANY; + IREE_RETURN_IF_ERROR(iree_tooling_transfer_variants( + workload->outputs.get(), Device(), allocator_.get(), params, + nullptr, nullptr)); + } + IREE_RETURN_IF_ERROR(iree_tooling_process_results_and_print( + Device(), + iree_make_string_view(workload->results_cconv.data(), + workload->results_cconv.size()), + workload->outputs.get(), iree_allocator_system(), exit_code)); + } + return iree_ok_status(); + } + + iree_status_t Close() { + iree_status_t status = + iree_hal_end_profiling_from_flags(std::exchange(profiling_, nullptr)); + workloads_.clear(); context_.reset(); - iree_status_t status = CloseReplayCapture(); - iree_tooling_module_list_reset(&module_list_); + iree_tooling_module_list_reset(&modules_); instance_.reset(); - - // Tear down device last in order to get accurate statistics. - if (FLAG_print_statistics && device_list_) { - for (iree_host_size_t i = 0; i < device_list_->count; ++i) { - iree_hal_device_t* dev = iree_hal_device_list_at(device_list_, i); - iree_string_view_t device_id = iree_hal_device_id(dev); + if (recorder_) { + status = + iree_status_join(status, iree_hal_replay_recorder_close(recorder_)); + iree_hal_replay_recorder_release(recorder_); + recorder_ = nullptr; + } + if (FLAG_print_statistics && devices_) { + for (iree_host_size_t i = 0; i < devices_->count; ++i) { + iree_hal_device_t* device = devices_->devices[i]; + iree_string_view_t device_id = iree_hal_device_id(device); fprintf(stderr, "Device: %.*s\n", (int)device_id.size, device_id.data); - iree_hal_allocator_t* allocator = iree_hal_device_allocator(dev); status = iree_status_join( - status, iree_hal_allocator_statistics_fprint(stderr, allocator)); + status, iree_hal_allocator_statistics_fprint( + stderr, iree_hal_device_allocator(device))); fprintf(stderr, "\n"); } } - device_allocator_.reset(); - device_ = nullptr; - iree_hal_device_list_free(device_list_); - device_list_ = nullptr; + allocator_.reset(); + iree_hal_device_list_free(devices_); + devices_ = nullptr; return status; } - iree_hal_device_t* device() const { return device_; } - - iree_status_t CloseReplayCapture() { - if (!replay_recorder_) return iree_ok_status(); - iree_status_t status = iree_hal_replay_recorder_close(replay_recorder_); - iree_hal_replay_recorder_release(replay_recorder_); - replay_recorder_ = nullptr; - return status; - } - - iree_status_t Register() { - IREE_TRACE_SCOPE_NAMED("IREEBenchmark::Register"); - - if (!instance_ || !device_allocator_ || !context_ || !module_list_.count) { - IREE_RETURN_IF_ERROR(Init()); - } - - auto function_name = std::string(FLAG_function); - if (!function_name.empty()) { - IREE_RETURN_IF_ERROR(RegisterSpecificFunction(function_name)); - } else { - IREE_RETURN_IF_ERROR(RegisterAllExportedFunctions()); - } - 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"); - } - - // Transfer outputs to the host so they can be processed. Only required when - // using full HAL device-based execution. - if (device_ != NULL) { - iree_hal_buffer_params_t target_params = {}; - target_params.usage = - IREE_HAL_BUFFER_USAGE_TRANSFER | IREE_HAL_BUFFER_USAGE_MAPPING; - target_params.access = IREE_HAL_MEMORY_ACCESS_ALL; - target_params.type = - IREE_HAL_MEMORY_TYPE_HOST_LOCAL | IREE_HAL_MEMORY_TYPE_DEVICE_VISIBLE; - target_params.queue_affinity = IREE_HAL_QUEUE_AFFINITY_ANY; - target_params.min_alignment = 0; - IREE_RETURN_IF_ERROR(iree_tooling_transfer_variants( - outputs_.get(), device_, device_allocator_.get(), target_params, - /*wait_fence=*/NULL, /*signal_fence=*/NULL)); - } - - 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(); + // Ends the session of the device profiling flags so that output processing + // is not profiled, and returns the failures of all benchmark runs. + iree_status_t FinishMeasurements() { + iree_status_t status = std::exchange(error_, iree_ok_status()); + return iree_status_join(status, iree_hal_end_profiling_from_flags( + std::exchange(profiling_, nullptr))); } private: - iree_status_t Init() { - IREE_TRACE_SCOPE_NAMED("IREEBenchmark::Init"); - IREE_TRACE_FRAME_MARK_BEGIN_NAMED("init"); - - iree_allocator_t host_allocator = iree_allocator_system(); - IREE_RETURN_IF_ERROR( - iree_tooling_create_instance(host_allocator, &instance_)); - - IREE_RETURN_IF_ERROR(iree_tooling_load_modules_from_flags( - instance_.get(), host_allocator, &module_list_)); - - IREE_RETURN_IF_ERROR(iree_tooling_create_context_from_flags( - instance_.get(), module_list_.count, module_list_.values, - /*default_device_uri=*/iree_string_view_empty(), host_allocator, - &context_, &device_list_, &device_allocator_, &replay_recorder_)); - device_ = device_list_ ? iree_hal_device_list_at(device_list_, 0) : nullptr; - - IREE_TRACE_FRAME_MARK_END_NAMED("init"); - return iree_ok_status(); + iree_hal_device_t* Device() const { + return devices_ && devices_->count ? devices_->devices[0] : nullptr; } - iree_status_t RegisterSpecificFunction(const std::string& function_name) { - IREE_TRACE_SCOPE_NAMED("IREEBenchmark::RegisterSpecificFunction"); - - iree_vm_module_t* main_module = - iree_tooling_module_list_back(&module_list_); - iree_vm_function_t function; - IREE_RETURN_IF_ERROR(iree_vm_module_lookup_function_by_name( - main_module, IREE_VM_FUNCTION_LINKAGE_EXPORT, - iree_string_view_t{function_name.data(), - (iree_host_size_t)function_name.size()}, - &function)); + iree_status_t AddWorkload(iree_vm_function_t function, bool selected) { + iree_string_view_t name = iree_vm_function_name(&function); + iree_string_view_t type = iree_vm_function_lookup_attr_by_name( + &function, IREE_SV("iree.benchmark")); + iree_string_view_t model = iree_vm_function_lookup_attr_by_name( + &function, IREE_SV("iree.abi.model")); + InvocationKind kind = InvocationKind::kSync; + if (iree_string_view_equal(type, IREE_SV("dispatch"))) { + kind = InvocationKind::kDispatch; + } else if (iree_string_view_equal(model, IREE_SV("coarse-fences"))) { + kind = InvocationKind::kAsync; + } + // Without --function, run benchmark wrappers and any other public function + // whose arguments the tool can supply without --input. + const bool discovered = !selected && kind != InvocationKind::kDispatch && + !iree_string_view_equal(type, IREE_SV("entry")); + if (discovered && + (iree_string_view_starts_with(name, IREE_SV("__")) || + iree_string_view_find_char(name, '$', 0) != IREE_STRING_VIEW_NPOS)) { + return iree_ok_status(); // Internal or special function. + } iree_vm_function_signature_t signature = iree_vm_function_signature(&function); - 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(), - iree_vm_instance_allocator(instance_.get()), &inputs_)); - - iree_string_view_t invocation_model = iree_vm_function_lookup_attr_by_name( - &function, IREE_SV("iree.abi.model")); - if (iree_string_view_equal(invocation_model, IREE_SV("coarse-fences"))) { - // Asynchronous invocation. - iree::RegisterAsyncBenchmark(function_name, replay_recorder_, device_, - context_.get(), function, inputs_.get()); - } else { - // Synchronous invocation. - iree::RegisterGenericBenchmark(function_name, replay_recorder_, device_, - context_.get(), function, inputs_.get(), - outputs_.get()); - } - return iree_ok_status(); - } - - iree_status_t RegisterAllExportedFunctions() { - IREE_TRACE_SCOPE_NAMED("IREEBenchmark::RegisterAllExportedFunctions"); - iree_vm_module_t* main_module = - iree_tooling_module_list_back(&module_list_); - iree_vm_module_signature_t signature = - iree_vm_module_signature(main_module); - for (iree_host_size_t i = 0; i < signature.export_function_count; ++i) { - iree_vm_function_t function; - IREE_RETURN_IF_ERROR(iree_vm_module_lookup_function_by_ordinal( - main_module, IREE_VM_FUNCTION_LINKAGE_EXPORT, i, &function)); - iree_string_view_t function_name = iree_vm_function_name(&function); - - // We run anything with the 'benchmark' attribute. - // If the attribute is not present we'll run anything that looks runnable. - iree_string_view_t benchmark_type = iree_vm_function_lookup_attr_by_name( - &function, IREE_SV("iree.benchmark")); - if (iree_string_view_equal(benchmark_type, IREE_SV("dispatch"))) { - iree::RegisterDispatchBenchmark( - std::string(function_name.data, function_name.size), - replay_recorder_, context_.get(), function); - } else if (iree_string_view_equal(benchmark_type, IREE_SV("entry"))) { - iree::RegisterGenericBenchmark( - std::string(function_name.data, function_name.size), - replay_recorder_, device_, context_.get(), function, - /*inputs=*/nullptr, /*outputs=*/nullptr); - } else { - // Pick up generic () -> () functions. - if (iree_string_view_starts_with(function_name, - iree_make_cstring_view("__")) || - iree_string_view_find_char(function_name, '$', 0) != - IREE_STRING_VIEW_NPOS) { - // Skip internal or special functions. - continue; - } - - // Query function information to determine how to run it. - iree_vm_function_signature_t signature = - iree_vm_function_signature(&function); - iree_host_size_t argument_count = 0; - iree_host_size_t result_count = 0; - IREE_RETURN_IF_ERROR(iree_vm_function_call_count_arguments_and_results( - &signature, &argument_count, &result_count)); - iree_string_view_t invocation_model = - iree_vm_function_lookup_attr_by_name(&function, - IREE_SV("iree.abi.model")); - if (iree_string_view_equal(invocation_model, - IREE_SV("coarse-fences"))) { - // Asynchronous invocation with coarse fences. Expect just those. - if (argument_count == 2) { - // Only functions taking a (wait, signal) fence pair are run. - iree::RegisterAsyncBenchmark( - std::string(function_name.data, function_name.size), - replay_recorder_, device_, context_.get(), function, - /*inputs=*/nullptr); - } - } else { - // Basic synchronous invocation. - if (argument_count == 0) { - // Only functions with no inputs are run (because we can't pass - // anything). - iree::RegisterGenericBenchmark( - std::string(function_name.data, function_name.size), - replay_recorder_, device_, context_.get(), function, - /*inputs=*/nullptr, /*outputs=*/nullptr); - } - } + if (discovered) { + iree_host_size_t argument_count = 0, result_count = 0; + IREE_RETURN_IF_ERROR(iree_vm_function_call_count_arguments_and_results( + &signature, &argument_count, &result_count)); + if (argument_count != (kind == InvocationKind::kAsync ? 2u : 0u)) { + return iree_ok_status(); } } + iree_string_view_t arguments, results; + IREE_RETURN_IF_ERROR(iree_vm_function_call_get_cconv_fragments( + &signature, &arguments, &results)); + auto workload = std::make_unique<Workload>(); + workload->name = "BM_" + std::string(name.data, name.size); + workload->function = function; + workload->kind = kind; + workload->results_cconv.assign(results.data, results.size); + workload->batch_size = FLAG_batch_size; + workload->concurrency = + kind == InvocationKind::kAsync ? FLAG_batch_concurrency : 1; + // Concurrency need not be a power of two. Keep the logical iteration count + // consistent with the actual number of invocations in the rounded batch. + int64_t batch = + ((int64_t)workload->batch_size + workload->concurrency - 1) / + workload->concurrency * workload->concurrency; + if (batch > INT32_MAX) { + return iree_make_status(IREE_STATUS_OUT_OF_RANGE, + "rounded batch size is too large"); + } + workload->batch_size = static_cast<int32_t>(batch); + workload->unit = FLAG_time_unit.first ? FLAG_time_unit.second + : kind == InvocationKind::kDispatch + ? benchmark::kMicrosecond + : benchmark::kMillisecond; + if (kind == InvocationKind::kDispatch) { + if (selected && FLAG_input_list().count) { + return iree_make_status( + IREE_STATUS_INVALID_ARGUMENT, + "dispatch wrappers take --batch_size, not --input"); + } + IREE_RETURN_IF_ERROR( + iree_vm_list_create(iree_vm_make_undefined_type_def(), 1, + iree_allocator_system(), &workload->inputs)); + iree_vm_value_t count = iree_vm_value_make_i32(workload->batch_size); + IREE_RETURN_IF_ERROR( + iree_vm_list_push_value(workload->inputs.get(), &count)); + } else if (selected) { + IREE_RETURN_IF_ERROR(iree_tooling_parse_variants( + arguments, FLAG_input_list(), Device(), allocator_.get(), + iree_allocator_system(), &workload->inputs)); + } else { + IREE_RETURN_IF_ERROR( + iree_vm_list_create(iree_vm_make_undefined_type_def(), 2, + iree_allocator_system(), &workload->inputs)); + } + workloads_.push_back(std::move(workload)); return iree_ok_status(); } - iree::vm::ref<iree_vm_instance_t> instance_; - iree::vm::ref<iree_vm_context_t> context_; - iree_hal_device_list_t* device_list_ = NULL; - iree_hal_device_t* device_ = NULL; - iree::vm::ref<iree_hal_allocator_t> device_allocator_; - 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_; + bool Check(iree_status_t status) { + if (iree_status_is_ok(status)) return true; + if (state_) { + state_->SkipWithError(iree::Status(iree_status_clone(status)).ToString()); + } + error_ = iree_status_join(error_, status); + return false; + } + + void Run(Workload& workload, benchmark::State& state) { + state_ = &state; + const bool async = workload.kind == InvocationKind::kAsync; + Batch batch; + if (!iree_status_is_ok(error_)) { + state.SkipWithError("an earlier benchmark or capture failed"); + } else if (!async) { + Check(batch.Prepare(workload, Device())); + } + IREE_TRACE_ZONE_BEGIN_NAMED_DYNAMIC(z0, workload.name.data(), + workload.name.size()); + IREE_TRACE_FRAME_MARK(); + while (state.KeepRunningBatch(workload.batch_size)) { + if (async) { + state.PauseTiming(); + bool prepared = Check(batch.Prepare(workload, Device())); + if (!prepared) continue; + state.ResumeTiming(); + } + IREE_TRACE_ZONE_BEGIN_NAMED(z1, "BenchmarkIteration"); + IREE_TRACE_FRAME_MARK_NAMED("Iteration"); + iree_status_t status = recorder_ ? iree_hal_replay_recorder_scope_begin( + recorder_, IREE_SV("execute")) + : iree_ok_status(); + if (iree_status_is_ok(status)) { + status = batch.Execute(workload, context_.get()); + if (recorder_) { + status = iree_status_join(status, iree_hal_replay_recorder_scope_end( + recorder_, IREE_SV("execute"))); + } + } + IREE_TRACE_ZONE_END(z1); + if (!Check(status)) continue; + const bool pause = async || profiling_ || + (workload.kind == InvocationKind::kSync && Device()); + if (pause) state.PauseTiming(); + if (async) { + if (FLAG_enable_output_processing) { + workload.outputs = vm::retain_ref(batch.invocations.back().outputs); + } + batch = {}; + } + if (pause && Check(iree_hal_flush_profiling_from_flags(profiling_))) { + state.ResumeTiming(); + } + } + state.SetItemsProcessed(state.iterations()); + if (!async && FLAG_enable_output_processing && !state.skipped()) { + workload.outputs = std::move(batch.invocations.front().outputs); + } + IREE_TRACE_ZONE_END(z0); + state_ = nullptr; + } + + vm::ref<iree_vm_instance_t> instance_; + vm::ref<iree_vm_context_t> context_; + // Allocator of the first device, used for inputs and output transfers. + vm::ref<iree_hal_allocator_t> allocator_; + // Modules loaded from --module; the last one is benchmarked. + iree_tooling_module_list_t modules_; + // Devices the context uses; NULL or empty if the program uses none. + iree_hal_device_list_t* devices_ = nullptr; + // Recorder of --device_replay_output, or NULL. + iree_hal_replay_recorder_t* recorder_ = nullptr; + // Session requested by the device profiling flags, if any. + iree_hal_profiling_from_flags_t* profiling_ = nullptr; + // Functions to benchmark, in registration order. + std::vector<std::unique_ptr<Workload>> workloads_; + // State of the running benchmark function, or NULL between functions. + benchmark::State* state_ = nullptr; + // Failures of the measured runs, reported after all benchmarks finish. + iree_status_t error_ = iree_ok_status(); }; -} // namespace -} // namespace iree -static int runMain(int argc, char** argv) { - IREE_TRACE_ZONE_BEGIN_NAMED(z0, "iree-benchmark-module"); - - // Pass through flags to benchmark (allowing --help to fall through). - iree_flags_set_usage("iree-benchmark-module", - iree::kIreeBenchmarkModuleUsage); +static int RunBenchmarkModule(int argc, char** argv) { + iree_flags_set_usage("iree-benchmark-module", kIreeBenchmarkModuleUsage); iree_flags_parse_checked(IREE_FLAGS_PARSE_MODE_UNDEFINED_OK | IREE_FLAGS_PARSE_MODE_CONTINUE_AFTER_HELP, &argc, &argv); if (FLAG_agents_md) { - iree::PrintBenchmarkModuleAgentMarkdown(stdout); - fflush(stdout); - IREE_TRACE_ZONE_END(z0); + PrintBenchmarkModuleAgentMarkdown(stdout); return EXIT_SUCCESS; } ::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()); - int exit_code = static_cast<int>(iree_status_code(status)); - printf("%s\n", iree::Status(std::move(status)).ToString().c_str()); - IREE_TRACE_ZONE_END(z0); - return exit_code; - } - 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; - + BenchmarkSession session; + iree_status_t status = session.Open(); int exit_code = EXIT_SUCCESS; - if (FLAG_enable_output_processing) { - IREE_CHECK_OK(iree_benchmark.ProcessResults(&exit_code)); + if (iree_status_is_ok(status)) { + session.Register(); + ::benchmark::RunSpecifiedBenchmarks(); + ::benchmark::ClearRegisteredBenchmarks(); + status = session.FinishMeasurements(); + if (iree_status_is_ok(status)) status = session.ProcessOutputs(&exit_code); } - - IREE_CHECK_OK(iree_benchmark.Shutdown()); - - IREE_TRACE_ZONE_END(z0); + status = iree_status_join(status, session.Close()); + if (!iree_status_is_ok(status)) { + iree_status_fprint(stderr, status); + iree_status_free(status); + return EXIT_FAILURE; + } return exit_code; } +} // namespace +} // namespace iree + int main(int argc, char** argv) { IREE_TRACE_APP_ENTER(); - int exit_code = runMain(argc, argv); + IREE_TRACE_ZONE_BEGIN_NAMED(z0, "iree-benchmark-module"); + int exit_code = iree::RunBenchmarkModule(argc, argv); + IREE_TRACE_ZONE_END(z0); IREE_TRACE_APP_EXIT(exit_code); return exit_code; }