Log generated triton code at the DEBUG level rather than INFO (#986)

Author: jansel

.github/workflows/benchmark.yml

@@ -37,6 +37,9 @@ jobs:
   benchmark:
     name: benchmark-${{ inputs.runtime-version }}-${{ inputs.kernels }}-py${{ inputs.python-version }}-${{ inputs.alias }}
 
+    env:
+      HELION_AUTOTUNE_LOG_LEVEL: DEBUG
+
     container:
       image: ${{ inputs.image }}
       options: ${{ inputs.container-options }}

helion/autotuner/base_search.py

@@ -46,9 +46,11 @@
 from .benchmarking import interleaved_bench
 from .config_generation import ConfigGeneration
 from .config_generation import FlatConfig
+from .logger import SUPPRESSED_TRITON_CODE_MSG
 from .logger import LambdaLogger
 from .logger import classify_triton_exception
 from .logger import format_triton_compile_failure
+from .logger import log_generated_triton_code_debug
 from .logger import match_unrecoverable_runtime_error
 from .progress_bar import iter_with_progress
 
@@ -163,11 +165,16 @@ def _compute_baseline(self) -> tuple[object, bool, Sequence[object] | None]:
             decorator = self.kernel.format_kernel_decorator(
                 baseline_config, self.settings
             )
-            triton_code = self.kernel.to_triton_code(baseline_config)
+            log_generated_triton_code_debug(
+                self.log,
+                self.kernel,
+                baseline_config,
+                prefix=f"Generated Triton code for {decorator}:",
+            )
             raise exc.InvalidConfig(
                 "Default config failed while computing baseline.\n"
                 f"Default config: {decorator}\n"
-                f"\nGenerated Triton code:\n{triton_code}\n"
+                f"{SUPPRESSED_TRITON_CODE_MSG}\n"
             ) from e
         original_args_flat, _ = tree_flatten(self._original_args)
         new_args_flat, _ = tree_flatten(new_args)
@@ -326,16 +333,35 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
             if self.settings.autotune_ignore_errors:
                 pass
             elif action == "raise":
+                decorator = self.kernel.format_kernel_decorator(config, self.settings)
+                log_generated_triton_code_debug(
+                    self.log,
+                    self.kernel,
+                    config,
+                    prefix=f"Generated Triton code for {decorator}:",
+                )
                 raise exc.TritonError(
                     error=f"{type(e).__qualname__}: {e}",
-                    decorator=self.kernel.format_kernel_decorator(
-                        config, self.settings
-                    ),
-                    code=self.kernel.to_triton_code(config),
+                    decorator=decorator,
+                    code=SUPPRESSED_TRITON_CODE_MSG,
                 ) from e
             elif action == "warn":
+                decorator = self.kernel.format_kernel_decorator(config, self.settings)
+                log_generated_triton_code_debug(
+                    self.log,
+                    self.kernel,
+                    config,
+                    prefix=f"Generated Triton code for {decorator}:",
+                )
                 self.log.warning(format_triton_compile_failure(config, e, self.kernel))
             else:
+                decorator = self.kernel.format_kernel_decorator(config, self.settings)
+                log_generated_triton_code_debug(
+                    self.log,
+                    self.kernel,
+                    config,
+                    prefix=f"Generated Triton code for {decorator}:",
+                )
                 self.log.debug(f"Benchmarking failed: {type(e).__name__}: {e}")
             return inf
 
@@ -1143,15 +1169,31 @@ def _consume_result(self, *, raise_on_raise: bool) -> None:
         if classification == "raise":
             if raise_on_raise:
                 self._remote_error_handled = True
+                decorator = self.search.kernel.format_kernel_decorator(
+                    self.config, self.search.settings
+                )
+                log_generated_triton_code_debug(
+                    self.search.log,
+                    self.search.kernel,
+                    self.config,
+                    prefix=f"Generated Triton code for {decorator}:",
+                )
                 raise exc.TritonError(
                     error=f"{type(exc_obj).__qualname__}: {exc_obj}",
-                    decorator=self.search.kernel.format_kernel_decorator(
-                        self.config, self.search.settings
-                    ),
-                    code=self.search.kernel.to_triton_code(self.config),
+                    decorator=decorator,
+                    code=SUPPRESSED_TRITON_CODE_MSG,
                 ) from exc_obj
             return
 
+        decorator = self.search.kernel.format_kernel_decorator(
+            self.config, self.search.settings
+        )
+        log_generated_triton_code_debug(
+            self.search.log,
+            self.search.kernel,
+            self.config,
+            prefix=f"Generated Triton code for {decorator}:",
+        )
         formatted = format_triton_compile_failure(
             self.config, exc_obj, self.search.kernel
         )
@@ -1265,10 +1307,16 @@ def extract_launcher(
             return None
         return precompiler
     except Exception:
+        log_generated_triton_code_debug(
+            log,
+            kernel,
+            config,
+            prefix=f"Generated Triton code for {decorator}:",
+        )
         log.warning(
-            "Helion autotuner precompile error for %s\n\nGenerated Triton code:\n%s",
+            "Helion autotuner precompile error for %s. %s",
             decorator,
-            kernel.to_triton_code(config),
+            SUPPRESSED_TRITON_CODE_MSG,
             exc_info=True,
         )
         raise

helion/autotuner/logger.py

@@ -93,19 +93,56 @@ def _maybe_call(fn: Callable[[], str] | str) -> str:
     return fn
 
 
+SUPPRESSED_TRITON_CODE_MSG = (
+    "Enable HELION_AUTOTUNE_LOG_LEVEL=DEBUG to log generated Triton code."
+)
+
+
+def log_generated_triton_code_debug(
+    logger: logging.Logger | LambdaLogger,
+    bound_kernel: BoundKernel,
+    config: Config,
+    *,
+    prefix: str | None = None,
+) -> None:
+    """
+    Emit the generated Triton code at debug level if the logger allows it.
+
+    Args:
+        logger: Logger that should receive the message.
+        bound_kernel: Kernel whose Triton code should be logged.
+        config: Config used to generate the Triton code.
+        prefix: Optional prefix for the log message.
+    """
+    message_prefix = prefix or "Generated Triton code:"
+    if isinstance(logger, LambdaLogger):
+        logger.debug(lambda: _format_triton_code(bound_kernel, config, message_prefix))
+        return
+    if logger.isEnabledFor(logging.DEBUG):
+        logger.debug(
+            "%s\n%s",
+            message_prefix,
+            bound_kernel.to_triton_code(config),
+        )
+
+
+def _format_triton_code(bound_kernel: BoundKernel, config: Config, prefix: str) -> str:
+    code = bound_kernel.to_triton_code(config)
+    return f"{prefix}\n{code}"
+
+
 def format_triton_compile_failure(
     config: Config, err: BaseException, bound_kernel: BoundKernel
 ) -> str:
     kernel_decorator = bound_kernel.format_kernel_decorator(
         config, bound_kernel.settings
     )
-    triton_code = bound_kernel.to_triton_code(config)
     return (
         "Triton compile failed. This likely indicates a bug in Triton. "
         "Skipping failing config.\n"
         f"Config: {kernel_decorator}\n"
-        f"Error: {type(err).__name__}: {err}\n\n"
-        f"Generated Triton code:\n{triton_code}"
+        f"Error: {type(err).__name__}: {err}\n"
+        f"{SUPPRESSED_TRITON_CODE_MSG}"
     )