Support all exported metrics in plotting scripts (#981)

scripts/plot_benchmark_comparison.py

@@ -20,7 +20,10 @@ import pathlib
 import matplotlib.pyplot as plt
 import numpy as np
 
+from plot_utils import METRICS, add_metric_argument, validate_metric
+
 parser = argparse.ArgumentParser(description=__doc__)
+add_metric_argument(parser)
 parser.add_argument(
     "files", nargs="+", type=pathlib.Path, help="JSON files with benchmark results"
 )
@@ -31,27 +34,29 @@ parser.add_argument(
 parser.add_argument("-o", "--output", help="Save image to the given filename")
 
 args = parser.parse_args()
+metric_label, metric_scale = METRICS[args.metric]
 
 commands = None
 data = []
 inputs = []
 
 if args.benchmark_names:
-    assert len(args.files) == len(
-        args.benchmark_names
-    ), "Number of benchmark names must match the number of input files."
+    assert len(args.files) == len(args.benchmark_names), (
+        "Number of benchmark names must match the number of input files."
+    )
 
 for i, filename in enumerate(args.files):
     with open(filename) as f:
         results = json.load(f)["results"]
+    validate_metric(parser, results, args.metric)
     benchmark_commands = [b.get("name", b["command"]) for b in results]
     if commands is None:
         commands = benchmark_commands
     else:
-        assert (
-            commands == benchmark_commands
-        ), f"Unexpected commands in {filename}: {benchmark_commands}, expected: {commands}"
-    data.append([b["summary"]["time_wall_clock"]["mean"] for b in results])
+        assert commands == benchmark_commands, (
+            f"Unexpected commands in {filename}: {benchmark_commands}, expected: {commands}"
+        )
+    data.append([b["summary"][args.metric]["mean"] / metric_scale for b in results])
     if args.benchmark_names:
         inputs.append(args.benchmark_names[i])
     else:
@@ -74,7 +79,7 @@ ax.grid(visible=True, axis="y")
 if args.title:
     plt.title(args.title)
 plt.xlabel("Benchmark")
-plt.ylabel("Time [s]")
+plt.ylabel(metric_label)
 plt.legend(title="Command")
 
 if args.output: