xref: /aosp_15_r20/external/pytorch/benchmarks/operator_benchmark/benchmark_runner.py (revision da0073e96a02ea20f0ac840b70461e3646d07c45) 
1import argparse
2
3import benchmark_core
4
5import benchmark_utils
6
7import torch
8
9
10"""Performance microbenchmarks's main binary.
11
12This is the main function for running performance microbenchmark tests.
13It also registers existing benchmark tests via Python module imports.
14"""
15parser = argparse.ArgumentParser(
16    description="Run microbenchmarks.",
17    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
18)
19
20
21def parse_args():
22    parser.add_argument(
23        "--tag-filter",
24        "--tag_filter",
25        help="tag_filter can be used to run the shapes which matches the tag. (all is used to run all the shapes)",
26        default="short",
27    )
28
29    # This option is used to filter test cases to run.
30    parser.add_argument(
31        "--operators",
32        help="Filter tests based on comma-delimited list of operators to test",
33        default=None,
34    )
35
36    parser.add_argument(
37        "--operator-range",
38        "--operator_range",
39        help="Filter tests based on operator_range(e.g. a-c or b,c-d)",
40        default=None,
41    )
42
43    parser.add_argument(
44        "--test-name",
45        "--test_name",
46        help="Run tests that have the provided test_name",
47        default=None,
48    )
49
50    parser.add_argument(
51        "--list-ops",
52        "--list_ops",
53        help="List operators without running them",
54        action="store_true",
55    )
56
57    parser.add_argument(
58        "--list-tests",
59        "--list_tests",
60        help="List all test cases without running them",
61        action="store_true",
62    )
63
64    parser.add_argument(
65        "--iterations",
66        help="Repeat each operator for the number of iterations",
67        type=int,
68    )
69
70    parser.add_argument(
71        "--num-runs",
72        "--num_runs",
73        help="Run each test for num_runs. Each run executes an operator for number of <--iterations>",
74        type=int,
75        default=1,
76    )
77
78    parser.add_argument(
79        "--min-time-per-test",
80        "--min_time_per_test",
81        help="Set the minimum time (unit: seconds) to run each test",
82        type=int,
83        default=0,
84    )
85
86    parser.add_argument(
87        "--warmup-iterations",
88        "--warmup_iterations",
89        help="Number of iterations to ignore before measuring performance",
90        default=100,
91        type=int,
92    )
93
94    parser.add_argument(
95        "--omp-num-threads",
96        "--omp_num_threads",
97        help="Number of OpenMP threads used in PyTorch runtime",
98        default=None,
99        type=int,
100    )
101
102    parser.add_argument(
103        "--mkl-num-threads",
104        "--mkl_num_threads",
105        help="Number of MKL threads used in PyTorch runtime",
106        default=None,
107        type=int,
108    )
109
110    parser.add_argument(
111        "--report-aibench",
112        "--report_aibench",
113        type=benchmark_utils.str2bool,
114        nargs="?",
115        const=True,
116        default=False,
117        help="Print result when running on AIBench",
118    )
119
120    parser.add_argument(
121        "--use-jit",
122        "--use_jit",
123        type=benchmark_utils.str2bool,
124        nargs="?",
125        const=True,
126        default=False,
127        help="Run operators with PyTorch JIT mode",
128    )
129
130    parser.add_argument(
131        "--forward-only",
132        "--forward_only",
133        type=benchmark_utils.str2bool,
134        nargs="?",
135        const=True,
136        default=False,
137        help="Only run the forward path of operators",
138    )
139
140    parser.add_argument(
141        "--device",
142        help="Run tests on the provided architecture (cpu, cuda)",
143        default="None",
144    )
145
146    args, _ = parser.parse_known_args()
147
148    if args.omp_num_threads:
149        # benchmark_utils.set_omp_threads sets the env variable OMP_NUM_THREADS
150        # which doesn't have any impact as C2 init logic has already been called
151        # before setting the env var.
152
153        # In general, OMP_NUM_THREADS (and other OMP env variables) needs to be set
154        # before the program is started.
155        # From Chapter 4 in OMP standard: https://www.openmp.org/wp-content/uploads/openmp-4.5.pdf
156        # "Modifications to the environment variables after the program has started,
157        # even if modified by the program itself, are ignored by the OpenMP implementation"
158        benchmark_utils.set_omp_threads(args.omp_num_threads)
159        torch.set_num_threads(args.omp_num_threads)
160    if args.mkl_num_threads:
161        benchmark_utils.set_mkl_threads(args.mkl_num_threads)
162
163    return args
164
165
166def main():
167    args = parse_args()
168    benchmark_core.BenchmarkRunner(args).run()
169
170
171if __name__ == "__main__":
172    main()
173