From 7f5b84940abef6e99313cd7f89bee4a83dc06e4f Mon Sep 17 00:00:00 2001 From: Peter Sharpe Date: Sat, 15 Jun 2024 20:01:15 -0400 Subject: [PATCH] Adds a warm-up period when timing functions, to correctly handle JIT-compiled functions, those with internal imports, or caches --- aerosandbox/tools/code_benchmarking.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/aerosandbox/tools/code_benchmarking.py b/aerosandbox/tools/code_benchmarking.py index df51221e..77d0f631 100644 --- a/aerosandbox/tools/code_benchmarking.py +++ b/aerosandbox/tools/code_benchmarking.py @@ -97,6 +97,7 @@ def time_function( repeats: int = None, desired_runtime: float = None, runtime_reduction=np.min, + warmup: bool = True, ) -> Tuple[float, Any]: """ Runs a given callable and tells you how long it took to run, in seconds. Also returns the result of the function @@ -114,6 +115,9 @@ def time_function( runtime_reduction: A function that takes in a list of runtimes and returns a reduced value. For example, np.min will return the minimum runtime, np.mean will return the mean runtime, etc. Default is np.min. + warmup: If True, runs the function once before starting the timer. This can be useful for functions with JIT-compilation, + with internal imports, or caches. + Returns: A Tuple of (time_taken, result). - time_taken is a float of the time taken to run the function, in seconds. - result is the result of the function, if any. @@ -122,6 +126,9 @@ def time_function( if (repeats is not None) and (desired_runtime is not None): raise ValueError("You can't specify both repeats and desired_runtime!") + if warmup: + func() + def time_function_once(): start_ns = time.perf_counter_ns() result = func() @@ -176,4 +183,4 @@ def f(): t = Timer() t.tic() time.sleep(0.1) - print(t.toc()) \ No newline at end of file + print(t.toc())