1# Copyright 2020 Google Inc. All rights reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14"""Python benchmarking utilities. 15 16Example usage: 17 import google_benchmark as benchmark 18 19 @benchmark.register 20 def my_benchmark(state): 21 ... # Code executed outside `while` loop is not timed. 22 23 while state: 24 ... # Code executed within `while` loop is timed. 25 26 if __name__ == '__main__': 27 benchmark.main() 28""" 29 30import atexit 31 32from absl import app 33 34from google_benchmark import _benchmark 35from google_benchmark._benchmark import ( 36 Counter as Counter, 37 State as State, 38 kMicrosecond as kMicrosecond, 39 kMillisecond as kMillisecond, 40 kNanosecond as kNanosecond, 41 kSecond as kSecond, 42 o1 as o1, 43 oAuto as oAuto, 44 oLambda as oLambda, 45 oLogN as oLogN, 46 oN as oN, 47 oNCubed as oNCubed, 48 oNLogN as oNLogN, 49 oNone as oNone, 50 oNSquared as oNSquared, 51) 52from google_benchmark.version import __version__ as __version__ 53 54 55class __OptionMaker: 56 """A stateless class to collect benchmark options. 57 58 Collect all decorator calls like @option.range(start=0, limit=1<<5). 59 """ 60 61 class Options: 62 """Pure data class to store options calls, along with the benchmarked function.""" 63 64 def __init__(self, func): 65 self.func = func 66 self.builder_calls = [] 67 68 @classmethod 69 def make(cls, func_or_options): 70 """Make Options from Options or the benchmarked function.""" 71 if isinstance(func_or_options, cls.Options): 72 return func_or_options 73 return cls.Options(func_or_options) 74 75 def __getattr__(self, builder_name): 76 """Append option call in the Options.""" 77 78 # The function that get returned on @option.range(start=0, limit=1<<5). 79 def __builder_method(*args, **kwargs): 80 # The decorator that get called, either with the benchmared function 81 # or the previous Options 82 def __decorator(func_or_options): 83 options = self.make(func_or_options) 84 options.builder_calls.append((builder_name, args, kwargs)) 85 # The decorator returns Options so it is not technically a decorator 86 # and needs a final call to @register 87 return options 88 89 return __decorator 90 91 return __builder_method 92 93 94# Alias for nicer API. 95# We have to instantiate an object, even if stateless, to be able to use __getattr__ 96# on option.range 97option = __OptionMaker() 98 99 100def register(undefined=None, *, name=None): 101 """Register function for benchmarking.""" 102 if undefined is None: 103 # Decorator is called without parenthesis so we return a decorator 104 return lambda f: register(f, name=name) 105 106 # We have either the function to benchmark (simple case) or an instance of Options 107 # (@option._ case). 108 options = __OptionMaker.make(undefined) 109 110 if name is None: 111 name = options.func.__name__ 112 113 # We register the benchmark and reproduce all the @option._ calls onto the 114 # benchmark builder pattern 115 benchmark = _benchmark.RegisterBenchmark(name, options.func) 116 for name, args, kwargs in options.builder_calls[::-1]: 117 getattr(benchmark, name)(*args, **kwargs) 118 119 # return the benchmarked function because the decorator does not modify it 120 return options.func 121 122 123def _flags_parser(argv): 124 argv = _benchmark.Initialize(argv) 125 return app.parse_flags_with_usage(argv) 126 127 128def _run_benchmarks(argv): 129 if len(argv) > 1: 130 raise app.UsageError("Too many command-line arguments.") 131 return _benchmark.RunSpecifiedBenchmarks() 132 133 134def main(argv=None): 135 return app.run(_run_benchmarks, argv=argv, flags_parser=_flags_parser) 136 137 138# Methods for use with custom main function. 139initialize = _benchmark.Initialize 140run_benchmarks = _benchmark.RunSpecifiedBenchmarks 141atexit.register(_benchmark.ClearRegisteredBenchmarks) 142