use existing bench_time; rename flag; update docstring

Author: bringlein

python/triton/runtime/autotuner.py

@@ -22,7 +22,7 @@ def __init__(
         prune_configs_by: Dict = None,
         warmup=25,
         rep=100,
-        report=False,
+        print_autotune_stats=False,
     ):
         """
         :param prune_configs_by: a dict of functions that are used to prune configs, fields:
@@ -78,7 +78,7 @@ def _post_hook(args):
         self.fn = fn
         self.num_warmups = warmup
         self.num_reps = rep
-        self.report = report
+        self.print_autotune_stats = print_autotune_stats
 
     def _bench(self, *args, config, **meta):
         # check for conflicts, i.e. meta-parameters both provided
@@ -111,7 +111,6 @@ def kernel_call():
 
     def run(self, *args, **kwargs):
         self.nargs = dict(zip(self.arg_names, args))
-        autotune_start = time.time()
         used_cached_result = True
         if len(self.configs) > 1:
             all_args = {**self.nargs, **kwargs}
@@ -139,10 +138,9 @@ def run(self, *args, **kwargs):
         else:
             config = self.configs[0]
         self.best_config = config
-        if self.report and not used_cached_result:
-            autotune_stop = time.time()
+        if self.print_autotune_stats and not used_cached_result:
             print(
-                f"Autotuner for function {self.fn} finished after {autotune_stop-autotune_start:.2f}s; best config selected: {self.best_config};"
+                f"Autotuner for function {self.fn} finished after {self.bench_time:.2f}s; best config selected: {self.best_config};"
             )
         full_nargs = {**self.nargs, **kwargs, **self.best_config.kwargs}
         if config.pre_hook is not None:
@@ -235,7 +233,7 @@ def __str__(self):
 
 
 def autotune(configs, key, prune_configs_by=None, reset_to_zero=None, restore_value=None, warmup=25, rep=100,
-             report=False):
+             print_autotune_stats=False):
     """
     Decorator for auto-tuning a :code:`triton.jit`'d function.
 
@@ -272,13 +270,13 @@ def kernel(x_ptr, x_size, **META):
     :type warmup: int
     :param rep: Repetition time (in ms) to pass to benchmarking, defaults to 100.
     :type rep: int
-    :param report: Flag to enable printing the selected configuration
+    :param print_autotune_stats: If set to true, a log message will be printed after each autotune evaluation containing the benchmark time and the selected configuration.
     :type report: bool
     """
 
     def decorator(fn):
         return Autotuner(fn, fn.arg_names, configs, key, reset_to_zero, restore_value, prune_configs_by, warmup, rep,
-                         report)
+                         print_autotune_stats)
 
     return decorator