Print autotune benchmark timings (triton-lang#1392)

Co-authored-by: zhanwei33 <22679270+zhanwei33@users.noreply.github.com>

Author: zhanwei33

third_party/ascend/backend/runtime/autotuner.py

@@ -80,6 +80,14 @@ def _inject_default_simt_stack_limit(options: Dict[str, object], stack_limit: in
         options["simt_stack_limit"] = stack_limit
 
 
+def _format_autotune_timing(timing) -> str:
+    """Format the timing returned by the active autotune benchmarker."""
+    if isinstance(timing, (tuple, list)):
+        labels = ("p50", "p20", "p80")
+        return ", ".join(f"{label}={value:.4f} ms" for label, value in zip(labels, timing))
+    return f"mean={timing:.4f} ms"
+
+
 def _get_constexpr_candidates_from_fn(fn) -> List[str]:
     """
     Returns all constexpr parameter names from the kernel function definition.
@@ -2131,6 +2139,7 @@ def benchmark():
         if self.print_autotuning and did_benchmark:
             print(f"Triton autotuning for function {self.base_fn.__name__} finished after "
                   f"{self.bench_time:.2f}s; best config selected: {self.best_config};")
+            self._print_benchmark_results(self.configs_timings)
 
         if did_benchmark and self.auto_profile_dir is not None:
             self._profile(*args, config=self.best_config, **kwargs)
@@ -2173,6 +2182,15 @@ def _try_ubtuner(self, *args, config, excp, run_fns, **kwargs):
             if self.print_autotuning:
                 print(f"[WARN] encounter exception when try ubtune, Details: {e}")
 
+    def _print_benchmark_results(self, timings) -> None:
+        if not self.print_autotuning:
+            return
+
+        print(f"Triton autotuning benchmark results for function {self.base_fn.__name__}:")
+        for config, timing in timings.items():
+            selected = " [selected]" if config == self.best_config else ""
+            print(f"  config={config}; {_format_autotune_timing(timing)}{selected}")
+
     def _batch_bench(self, *args, configs, **kwargs):
         from triton.compiler.errors import CompileTimeAssertionFailure, MLIRCompilationError
         from triton.runtime.errors import OutOfResources
@@ -2607,7 +2625,8 @@ def kernel(x_ptr, x_size, **META):
 
     If the environment variable :code:`TRITON_PRINT_AUTOTUNING` is set to
     :code:`"1"`, Triton will print a message to stdout after autotuning each
-    kernel, including the time spent autotuning and the best configuration.
+    kernel, including the benchmark timing for each valid configuration, the
+    time spent autotuning, and the best configuration.
 
     :param configs: a list of :code:`triton.Config` objects
     :type configs: list[triton.Config]

third_party/ascend/unittest/autotune_ut/test_do_bench_compat.py

@@ -176,6 +176,62 @@ def _dummy_kernel():
     assert marker["called"] is False
 
 
+def test_print_benchmark_results_formats_quantiles_and_selected_config(capsys):
+
+    def _dummy_kernel():
+        return None
+
+    selected = Config({"BLOCK_SIZE": 128}, num_warps=4)
+    other = Config({"BLOCK_SIZE": 256}, num_warps=8)
+    tuner = object.__new__(AutoTilingTuner)
+    tuner.print_autotuning = True
+    tuner.base_fn = _dummy_kernel
+    tuner.best_config = selected
+
+    tuner._print_benchmark_results({
+        selected: (1.23456, 1.0, 1.8),
+        other: 2.34567,
+    })
+
+    output = capsys.readouterr().out
+    assert "Triton autotuning benchmark results for function _dummy_kernel:" in output
+    assert "config=BLOCK_SIZE: 128, num_warps: 4" in output
+    assert "p50=1.2346 ms, p20=1.0000 ms, p80=1.8000 ms [selected]" in output
+    assert "config=BLOCK_SIZE: 256, num_warps: 8" in output
+    assert "mean=2.3457 ms" in output
+
+
+def test_print_benchmark_results_is_disabled_without_debug_flag(capsys):
+    tuner = object.__new__(AutoTilingTuner)
+    tuner.print_autotuning = False
+
+    tuner._print_benchmark_results({Config({"BLOCK_SIZE": 128}): 1.0})
+
+    assert capsys.readouterr().out == ""
+
+
+def test_run_prints_benchmark_results_after_tuning(capsys):
+
+    def _dummy_kernel():
+        return None
+
+    selected = Config({"BLOCK_SIZE": 128}, num_warps=4)
+    other = Config({"BLOCK_SIZE": 256}, num_warps=8)
+    tuner, _ = _make_run_tuner([selected, other])
+    tuner.cache_results = False
+    tuner.print_autotuning = True
+    tuner.base_fn = _dummy_kernel
+    tuner._batch_bench = lambda *args, configs, **kwargs: {
+        selected: (1.0, 0.9, 1.1),
+        other: (2.0, 1.8, 2.2),
+    }
+
+    assert tuner.run() == "kernel-result"
+    output = capsys.readouterr().out
+    assert "Triton autotuning benchmark results for function _dummy_kernel:" in output
+    assert "p50=1.0000 ms, p20=0.9000 ms, p80=1.1000 ms [selected]" in output
+
+
 def test_ascend_autotune_decorator_forwards_do_bench(monkeypatch):
     import triton.backends.ascend.runtime.autotuner as ascend_autotuner