Print Triton code when error for easier debugging (#874)

Author: yf225

helion/autotuner/base_search.py

@@ -134,9 +134,11 @@ 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)
             raise exc.InvalidConfig(
                 "Default config failed while computing baseline.\n"
                 f"Default config: {decorator}\n"
+                f"\nGenerated Triton code:\n{triton_code}\n"
             ) from e
         original_args_flat, _ = tree_flatten(self._original_args)
         new_args_flat, _ = tree_flatten(new_args)
@@ -235,9 +237,10 @@ def benchmark_function(self, config: Config, fn: CompiledConfig) -> float:
                 raise exc.TritonError(
                     f"{type(e).__qualname__}: {e}",
                     self.kernel.format_kernel_decorator(config, self.settings),
+                    self.kernel.to_triton_code(config),
                 ) from e
             if action == "warn":
-                self.log.warning(format_triton_compile_failure(config, e))
+                self.log.warning(format_triton_compile_failure(config, e, self.kernel))
             else:
                 self.log.debug(f"Benchmarking failed: {type(e).__name__}: {e}")
             return inf
@@ -277,13 +280,18 @@ def extract_launcher(
             # Should not reach here
             raise RuntimeError("Expected _ExtractedLaunchArgs exception")
         except _ExtractedLaunchArgs as e:
-            precompiler = make_precompiler(e.kernel, config)(*e.args, **e.kwargs)
+            precompiler = make_precompiler(
+                e.kernel,
+                config,
+                self.kernel,
+            )(*e.args, **e.kwargs)
             if precompiler is already_compiled:
                 return PrecompileFuture.skip(self, config, True)
         except Exception:
             log.warning(
-                "Helion autotuner precompile error for %s",
+                "Helion autotuner precompile error for %s\n\nGenerated Triton code:\n%s",
                 self.kernel.format_kernel_decorator(config, self.settings),
+                self.kernel.to_triton_code(config),
                 exc_info=True,
             )
             raise

helion/autotuner/logger.py

@@ -14,6 +14,7 @@
 
 if TYPE_CHECKING:
     from ..runtime.config import Config
+    from ..runtime.kernel import BoundKernel
 
 
 class LambdaLogger:
@@ -92,12 +93,19 @@ def _maybe_call(fn: Callable[[], str] | str) -> str:
     return fn
 
 
-def format_triton_compile_failure(config: Config, err: BaseException) -> str:
+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: {config!r}\n"
-        f"Error: {type(err).__name__}: {err}"
+        f"Config: {kernel_decorator}\n"
+        f"Error: {type(err).__name__}: {err}\n\n"
+        f"Generated Triton code:\n{triton_code}"
     )
 
 

helion/exc.py

@@ -316,7 +316,7 @@ class TorchOpTracingError(_WrapException):
 
 
 class TritonError(BaseError):
-    message = "Error running generated Triton program:\n{1}\n{0}"
+    message = "Error running generated Triton program:\n{1}\n{0}\n\nGenerated Triton code:\n{2}"
 
 
 class BaseWarning(_FixedMessage):

helion/runtime/precompile_shim.py

@@ -13,10 +13,13 @@
     from triton.runtime.jit import JITFunction
 
     from .config import Config
+    from .kernel import BoundKernel
 
 
 def make_precompiler(
-    fn: JITFunction[object], config: Config
+    fn: JITFunction[object],
+    config: Config,
+    bound_kernel: BoundKernel,
 ) -> Callable[..., Callable[[], None]]:
     from triton.runtime.jit import find_paths_if
     from triton.runtime.jit import get_iterable_path
@@ -64,7 +67,10 @@ def finish_it() -> None:
             except Exception as e:
                 action = classify_triton_exception(e)
                 if action != "debug":
-                    print(format_triton_compile_failure(config, e), file=sys.stderr)
+                    print(
+                        format_triton_compile_failure(config, e, bound_kernel),
+                        file=sys.stderr,
+                    )
                 sys.exit(1)
 
         return finish_it