141 lines · python
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 at6#7# http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# 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 and13# limitations under the License.14"""Python benchmarking utilities.15 16Example usage:17 import google_benchmark as benchmark18 19 @benchmark.register20 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 atexit30 31from absl import app32 33from google_benchmark import _benchmark34from 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 = func65 self.builder_calls = []66 67 @classmethod68 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_options72 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 function80 # or the previous Options81 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 decorator85 # and needs a final call to @register86 return options87 88 return __decorator89 90 return __builder_method91 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.range96option = __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 decorator103 return lambda f: register(f, name=name)104 105 # We have either the function to benchmark (simple case) or an instance of Options106 # (@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 the113 # benchmark builder pattern114 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 it119 return options.func120 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.Initialize139run_benchmarks = _benchmark.RunSpecifiedBenchmarks140atexit.register(_benchmark.ClearRegisteredBenchmarks)141