[Runtime][Python] Keep all results in benchmark_module (#24985) `benchmark_module` parsed the console table by splitting whitespace columns. It kept only the first user counter and failed on rows without counters. Read a separate Google Benchmark JSON output file and expose each complete row through `BenchmarkResult.report`, including metadata and every user counter, without maintaining a list of Google Benchmark fields. This replaces `user_counters`. The name, timing, and iteration fields retain their string representation. Percentage aggregates use `%`, and an unmatched benchmark filter returns an empty list. Subprocess failures and JSON rows marked `error_occurred` raise `BenchmarkToolError` with their diagnostics. The wrapper also accepts an explicit tool executable, benchmarks eligible exports of a module file when no entry function is given, supports repeated flags through list values, and serializes booleans as `true`/`false`. These support the lifecycle tests in the following changes. Assisted by AI. Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
diff --git a/runtime/bindings/python/iree/runtime/benchmark.py b/runtime/bindings/python/iree/runtime/benchmark.py index 3a2d9da..6b7b76a 100644 --- a/runtime/bindings/python/iree/runtime/benchmark.py +++ b/runtime/bindings/python/iree/runtime/benchmark.py
@@ -20,9 +20,11 @@ from os import PathLike from . import VmModule +import json import numpy import os import subprocess +import tempfile from .dtypes import DTYPE_TO_ABI_TYPE __all__ = [ @@ -31,12 +33,18 @@ ] BenchmarkResult = namedtuple( - "BenchmarkResult", "benchmark_name time cpu_time iterations user_counters" + "BenchmarkResult", "benchmark_name time cpu_time iterations report" ) +BenchmarkResult.__doc__ = """One benchmark or dispatch row reported by iree-benchmark-module. + +time and cpu_time are "<value> <unit>" strings, using % for percentage +aggregates, and iterations is a string. report is the complete Google Benchmark +JSON row, including metadata and user counters with their original values. +""" class BenchmarkToolError(Exception): - """Benchmark exception that preserves the command line and error output.""" + """Failure of the tool or a reported benchmark, with its diagnostics.""" def __init__(self, message): self.message = message @@ -59,9 +67,10 @@ module: Union[VmModule, PathLike], entry_function: str | None = None, inputs: list[Union[str, numpy.ndarray]] | None = None, + executable: Union[str, PathLike, None] = None, **kwargs, ) -> tuple[list[str], bytes | None]: - args = [benchmark_exe()] + args = [os.fspath(executable) if executable is not None else benchmark_exe()] if isinstance(module, VmModule): funcs = [a for a in module.function_names if a != "__init"] @@ -79,16 +88,20 @@ else: flatbuffer = None args.append(f"--module={module}") - if entry_function is None: - raise ValueError( - "When specifying the module as a filepath the entry function must be specified." - ) - args.append(f"--function={entry_function}") + # Without --function the tool benchmarks every exported function that takes + # no inputs. + if entry_function is not None: + args.append(f"--function={entry_function}") for k in kwargs: - v = kwargs[k] - args.append(f"--{k}={v}") + # A list or tuple repeats the flag once per value. + values = kwargs[k] if isinstance(kwargs[k], (list, tuple)) else [kwargs[k]] + for v in values: + # IREE flags only parse lowercase booleans. + if isinstance(v, bool): + v = "true" if v else "false" + args.append(f"--{k}={v}") for inp in inputs or []: if isinstance(inp, str): @@ -107,17 +120,14 @@ return args, flatbuffer -def benchmark_module( - module: Union[VmModule, PathLike], - entry_function: str | None = None, - inputs: list[Union[str, numpy.ndarray]] | None = None, - timeout: float | None = None, - **kwargs, -): - args, flatbuffer = _build_benchmark_args( - module=module, entry_function=entry_function, inputs=inputs, **kwargs - ) +def _run_benchmark( + args: list[str], flatbuffer: bytes | None = None, timeout: float | None = None +) -> tuple[str, str]: + """Runs a benchmark command line and returns its stdout and stderr. + Raises BenchmarkTimeoutError if the timeout expires and BenchmarkToolError + with both outputs if the tool fails. + """ try: benchmark_process = subprocess.run( args=args, @@ -133,42 +143,71 @@ if benchmark_process.returncode != 0: raise BenchmarkToolError(f"stderr:\n{err}\nstdout:\n{out}") + return out, err - if "INVALID_ARGUMENT;" in err: - raise ValueError("Invalid inputs specified for benchmarking") - # In the event benchmarking runs but encounters an internal error, - # return the internal error instead of benchmark results. - if "INTERNAL; CUDA driver error" in out: - raise BenchmarkToolError(out) +def benchmark_module( + module: Union[VmModule, PathLike], + entry_function: str | None = None, + inputs: list[Union[str, numpy.ndarray]] | None = None, + timeout: float | None = None, + executable: Union[str, PathLike, None] = None, + **kwargs, +) -> list[BenchmarkResult]: + """Benchmarks a module with iree-benchmark-module. - # Grab individual results by line (skip header lines) - bench_lines = out.split("\n")[3:] - benchmark_results = [] - for line in bench_lines: - split = line.split() - if len(split) == 0: - continue - elif len(split) <= 5: - raise ValueError( - f"Benchmark standard output is not expected. Line: '{line}' should have at least 6 columns.\nDetails: {args=},\n{out=}" - ) + Keyword arguments are passed to the tool as --key=value flags, once per + value of a list or tuple. Without an + entry function a module file runs every exported function that takes no + inputs. The tool defaults to the one in this package unless an executable + is given. - benchmark_name = split[0] - time = " ".join(split[1:3]) - cpu_time = " ".join(split[3:5]) - iterations = split[5] - user_counters = None - if len(split) > 5: - user_counters = split[6] - benchmark_results.append( + Returns a BenchmarkResult with the complete JSON report of every row, + including aggregates and dispatch rows. Raises BenchmarkToolError if the + tool exits unsuccessfully or a reported benchmark fails. + """ + if "benchmark_out" in kwargs or "benchmark_out_format" in kwargs: + raise ValueError("benchmark_module reads results from its own JSON file") + with tempfile.TemporaryDirectory() as directory: + # A results file keeps console and output processing text on stdout + # out of the parsed report. + report = os.path.join(directory, "results.json") + args, flatbuffer = _build_benchmark_args( + module=module, + entry_function=entry_function, + inputs=inputs, + executable=executable, + benchmark_out=report, + benchmark_out_format="json", + **kwargs, + ) + _run_benchmark(args, flatbuffer, timeout) + with open(report) as file: + return _parse_benchmark_results(file.read()) + + +def _parse_benchmark_results(report: str) -> list[BenchmarkResult]: + """Parses a Google Benchmark JSON report into benchmark results.""" + # Google Benchmark leaves the report empty when no benchmarks match. + if not report.strip(): + return [] + results = [] + for run in json.loads(report)["benchmarks"]: + if run.get("error_occurred"): + raise BenchmarkToolError(f"{run['name']}: {run['error_message']}") + unit = run["time_unit"] + real_time, cpu_time = run["real_time"], run["cpu_time"] + if run.get("aggregate_unit") == "percentage": + unit = "%" + real_time *= 100 + cpu_time *= 100 + results.append( BenchmarkResult( - benchmark_name=benchmark_name, - time=time, - cpu_time=cpu_time, - iterations=iterations, - user_counters=user_counters, + benchmark_name=run["name"], + time=f"{real_time} {unit}", + cpu_time=f"{cpu_time} {unit}", + iterations=str(run["iterations"]), + report=run, ) ) - - return benchmark_results + return results
diff --git a/runtime/bindings/python/tests/benchmark_test.py b/runtime/bindings/python/tests/benchmark_test.py index 959edd9..65c6813 100644 --- a/runtime/bindings/python/tests/benchmark_test.py +++ b/runtime/bindings/python/tests/benchmark_test.py
@@ -4,7 +4,9 @@ # See https://llvm.org/LICENSE.txt for license information. # SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception +import json import numpy as np +import sys import unittest import tempfile from pathlib import Path @@ -13,7 +15,11 @@ import iree.runtime from iree.runtime.benchmark import ( _build_benchmark_args as build_benchmark_args, + _parse_benchmark_results as parse_benchmark_results, + _run_benchmark as run_benchmark, benchmark_module, + BenchmarkResult, + BenchmarkToolError, BenchmarkTimeoutError, ) @@ -95,6 +101,36 @@ self.assertEqual(args, ref_args) self.assertIsNone(flatbuffer) + def testBuildBenchmarkArgsExecutableWithoutFunction(self): + args, _ = build_benchmark_args( + module="test_module.vmfb", + executable=Path("tools") / "iree-benchmark-module", + ) + # Without a function the tool benchmarks all exported functions. + ref_args = [ + str(Path("tools") / "iree-benchmark-module"), + "--module=test_module.vmfb", + ] + self.assertEqual(args, ref_args) + + def testBuildBenchmarkArgsFlagValues(self): + args, _ = build_benchmark_args( + module="test_module.vmfb", + entry_function="test_func", + dispatch_statistics=True, + print_statistics=False, + device=["local-task", "local-sync"], + ) + self.assertEqual( + args[3:], + [ + "--dispatch_statistics=true", + "--print_statistics=false", + "--device=local-task", + "--device=local-sync", + ], + ) + def testBuildBenchmarkArgsInputs(self): all_inputs = [ # empty @@ -124,6 +160,103 @@ self.assertEqual(args[3:], ref) self.assertIsNone(flatbuffer) + def testParseBenchmarkResults(self): + report = """{ + "context": {}, + "benchmarks": [ + {"name": "BM_main/process_time/real_time", "run_type": "iteration", + "repetitions": 1, "repetition_index": 0, "threads": 1, + "iterations": 4, "real_time": 2.5, "cpu_time": 3.0, + "time_unit": "ms", "items_per_second": 400.0}, + {"name": "BM_main/main_dispatch_0", "iterations": 4, + "real_time": 1.5, "cpu_time": 2.0, "time_unit": "ms", + "calls": 2.0, "percent": 100.0} + ] + }""" + # New Google Benchmark metadata must survive without needing to classify + # it as a standard field or a user counter. + report_data = json.loads(report) + report_data["benchmarks"][0]["new_metadata"] = {"version": 2} + report = json.dumps(report_data) + self.assertEqual( + parse_benchmark_results(report), + [ + BenchmarkResult( + benchmark_name="BM_main/process_time/real_time", + time="2.5 ms", + cpu_time="3.0 ms", + iterations="4", + report=report_data["benchmarks"][0], + ), + BenchmarkResult( + benchmark_name="BM_main/main_dispatch_0", + time="1.5 ms", + cpu_time="2.0 ms", + iterations="4", + report=report_data["benchmarks"][1], + ), + ], + ) + + def testParseBenchmarkPercentageAggregate(self): + report = """{ + "benchmarks": [ + {"name": "BM_main/process_time/real_time_cv", + "run_type": "aggregate", "aggregate_name": "cv", + "aggregate_unit": "percentage", "iterations": 2, + "real_time": 0.125, "cpu_time": 0.25, "time_unit": "ms", + "items_per_second": 0.125} + ] + }""" + (result,) = parse_benchmark_results(report) + self.assertEqual(result.time, "12.5 %") + self.assertEqual(result.cpu_time, "25.0 %") + self.assertEqual(result.report["aggregate_unit"], "percentage") + self.assertEqual(result.report["real_time"], 0.125) + self.assertEqual(result.report["items_per_second"], 0.125) + + def testParseBenchmarkError(self): + report = json.dumps( + { + "benchmarks": [ + { + "name": "BM_main", + "error_occurred": True, + "error_message": "RESOURCE_EXHAUSTED; device allocation failed", + } + ] + } + ) + with self.assertRaisesRegex( + BenchmarkToolError, + "BM_main: RESOURCE_EXHAUSTED; device allocation failed", + ): + parse_benchmark_results(report) + + def testBenchmarkToolFailurePreservesDiagnostics(self): + with self.assertRaises(BenchmarkToolError) as caught: + run_benchmark( + [ + sys.executable, + "-c", + "import sys; print('output diagnostic'); " + "print('error diagnostic', file=sys.stderr); sys.exit(1)", + ] + ) + self.assertIn("stdout:\noutput diagnostic", str(caught.exception)) + self.assertIn("stderr:\nerror diagnostic", str(caught.exception)) + + def testBenchmarkModuleUnmatchedFilter(self): + with tempfile.TemporaryDirectory() as directory: + module = Path(directory) / "module.vmfb" + iree.compiler.compile_str( + "func.func @main() { return }", output_file=str(module) + ) + results = benchmark_module( + module, entry_function="main", benchmark_filter="does_not_exist" + ) + self.assertEqual(results, []) + def testBenchmarkModule(self): ctx = iree.runtime.SystemContext() vm_module = create_simple_mul_module(ctx.instance) @@ -139,6 +272,7 @@ self.assertEqual(len(benchmark_results), 1) benchmark_time = float(benchmark_results[0].time.split(" ")[0]) self.assertGreater(benchmark_time, 0) + self.assertGreater(benchmark_results[0].report["items_per_second"], 0) def testBenchmarkModuleFromFilePath(self): ctx = iree.runtime.SystemContext()