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""" 29import atexit 30 31from absl import app 32 33from google_benchmark import _benchmark 34from google_benchmark._benchmark import ( 35 Counter as Counter, 36 State as State, 37 kMicrosecond as kMicrosecond, 38 kMillisecond as kMillisecond, 39 kNanosecond as kNanosecond, 40 kSecond as kSecond, 41 o1 as o1, 42 oAuto as oAuto, 43 oLambda as oLambda, 44 oLogN as oLogN, 45 oN as oN, 46 oNCubed as oNCubed, 47 oNLogN as oNLogN, 48 oNone as oNone, 49 oNSquared as oNSquared, 50) 51from google_benchmark.version import __version__ as __version__ 52 53 54class __OptionMaker: 55 """A stateless class to collect benchmark options. 56 57 Collect all decorator calls like @option.range(start=0, limit=1<<5). 58 """ 59 60 class Options: 61 """Pure data class to store options calls, along with the benchmarked function.""" 62 63 def __init__(self, func): 64 self.func = func 65 self.builder_calls = [] 66 67 @classmethod 68 def make(cls, func_or_options): 69 """Make Options from Options or the benchmarked function.""" 70 if isinstance(func_or_options, cls.Options): 71 return func_or_options 72 return cls.Options(func_or_options) 73 74 def __getattr__(self, builder_name): 75 """Append option call in the Options.""" 76 77 # The function that get returned on @option.range(start=0, limit=1<<5). 78 def __builder_method(*args, **kwargs): 79 # The decorator that get called, either with the benchmared function 80 # or the previous Options 81 def __decorator(func_or_options): 82 options = self.make(func_or_options) 83 options.builder_calls.append((builder_name, args, kwargs)) 84 # The decorator returns Options so it is not technically a decorator 85 # and needs a final call to @register 86 return options 87 88 return __decorator 89 90 return __builder_method 91 92 93# Alias for nicer API. 94# We have to instantiate an object, even if stateless, to be able to use __getattr__ 95# on option.range 96option = __OptionMaker() 97 98 99def register(undefined=None, *, name=None): 100 """Register function for benchmarking.""" 101 if undefined is None: 102 # Decorator is called without parenthesis so we return a decorator 103 return lambda f: register(f, name=name) 104 105 # We have either the function to benchmark (simple case) or an instance of Options 106 # (@option._ case). 107 options = __OptionMaker.make(undefined) 108 109 if name is None: 110 name = options.func.__name__ 111 112 # We register the benchmark and reproduce all the @option._ calls onto the 113 # benchmark builder pattern 114 benchmark = _benchmark.RegisterBenchmark(name, options.func) 115 for name, args, kwargs in options.builder_calls[::-1]: 116 getattr(benchmark, name)(*args, **kwargs) 117 118 # return the benchmarked function because the decorator does not modify it 119 return options.func 120 121 122def _flags_parser(argv): 123 argv = _benchmark.Initialize(argv) 124 return app.parse_flags_with_usage(argv) 125 126 127def _run_benchmarks(argv): 128 if len(argv) > 1: 129 raise app.UsageError("Too many command-line arguments.") 130 return _benchmark.RunSpecifiedBenchmarks() 131 132 133def main(argv=None): 134 return app.run(_run_benchmarks, argv=argv, flags_parser=_flags_parser) 135 136 137# Methods for use with custom main function. 138initialize = _benchmark.Initialize 139run_benchmarks = _benchmark.RunSpecifiedBenchmarks 140atexit.register(_benchmark.ClearRegisteredBenchmarks) 141