xref: /aosp_15_r20/external/pytorch/benchmarks/compare-fastrnn-results.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1import argparse
2import json
3from collections import namedtuple
4
5
6Result = namedtuple("Result", ["name", "base_time", "diff_time"])
7
8
9def construct_name(fwd_bwd, test_name):
10    bwd = "backward" in fwd_bwd
11    suite_name = fwd_bwd.replace("-backward", "")
12    return f"{suite_name}[{test_name}]:{'bwd' if bwd else 'fwd'}"
13
14
15def get_times(json_data):
16    r = {}
17    for fwd_bwd in json_data:
18        for test_name in json_data[fwd_bwd]:
19            name = construct_name(fwd_bwd, test_name)
20            r[name] = json_data[fwd_bwd][test_name]
21    return r
22
23
24parser = argparse.ArgumentParser("compare two pytest jsons")
25parser.add_argument("base", help="base json file")
26parser.add_argument("diff", help="diff json file")
27parser.add_argument(
28    "--format", default="md", type=str, help="output format (csv, md, json, table)"
29)
30args = parser.parse_args()
31
32with open(args.base) as base:
33    base_times = get_times(json.load(base))
34with open(args.diff) as diff:
35    diff_times = get_times(json.load(diff))
36
37all_keys = set(base_times.keys()).union(diff_times.keys())
38results = [
39    Result(name, base_times.get(name, float("nan")), diff_times.get(name, float("nan")))
40    for name in sorted(all_keys)
41]
42
43header_fmt = {
44    "table": "{:48s} {:>13s} {:>15s} {:>10s}",
45    "md": "| {:48s} | {:>13s} | {:>15s} | {:>10s} |",
46    "csv": "{:s}, {:s}, {:s}, {:s}",
47}
48data_fmt = {
49    "table": "{:48s} {:13.6f} {:15.6f} {:9.1f}%",
50    "md": "| {:48s} | {:13.6f} | {:15.6f} | {:9.1f}% |",
51    "csv": "{:s}, {:.6f}, {:.6f}, {:.2f}%",
52}
53
54if args.format in ["table", "md", "csv"]:
55    header_fmt_str = header_fmt[args.format]
56    data_fmt_str = data_fmt[args.format]
57    print(header_fmt_str.format("name", "base time (s)", "diff time (s)", "% change"))
58    if args.format == "md":
59        print(header_fmt_str.format(":---", "---:", "---:", "---:"))
60    for r in results:
61        print(
62            data_fmt_str.format(
63                r.name,
64                r.base_time,
65                r.diff_time,
66                (r.diff_time / r.base_time - 1.0) * 100.0,
67            )
68        )
69elif args.format == "json":
70    print(json.dumps(results))
71else:
72    raise ValueError("Unknown output format: " + args.format)
73