xref: /aosp_15_r20/external/pytorch/benchmarks/instruction_counts/main.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1"""Basic runner for the instruction count microbenchmarks.
2
3The contents of this file are placeholders, and will be replaced by more
4expressive and robust components (e.g. better runner and result display
5components) in future iterations. However this allows us to excercise the
6underlying benchmark generation infrastructure in the mean time.
7"""
8# mypy: ignore-errors
9import argparse
10import sys
11from typing import List
12
13from applications import ci
14from core.expand import materialize
15from definitions.standard import BENCHMARKS
16from execution.runner import Runner
17from execution.work import WorkOrder
18
19
20def main(argv: List[str]) -> None:
21    work_orders = tuple(
22        WorkOrder(label, autolabels, timer_args, timeout=600, retries=2)
23        for label, autolabels, timer_args in materialize(BENCHMARKS)
24    )
25
26    results = Runner(work_orders).run()
27    for work_order in work_orders:
28        print(
29            work_order.label,
30            work_order.autolabels,
31            work_order.timer_args.num_threads,
32            results[work_order].instructions,
33        )
34
35
36if __name__ == "__main__":
37    modes = {
38        "debug": main,
39        "ci": ci.main,
40    }
41
42    parser = argparse.ArgumentParser()
43    parser.add_argument("--mode", type=str, choices=list(modes.keys()), default="debug")
44
45    args, remaining_args = parser.parse_known_args(sys.argv)
46    modes[args.mode](remaining_args[1:])
47