Add env HELION_PRINT_OUTPUT_CODE=1

helion/runtime/kernel.py

@@ -8,6 +8,7 @@
 import logging
 import operator
 import re
+import sys
 import types
 from typing import TYPE_CHECKING
 from typing import overload
@@ -343,6 +344,8 @@ def compile_config(self, config: ConfigLike) -> CompiledConfig:
         triton_code = self.to_triton_code(config)
         log.info("Output code: \n%s", triton_code)
         log.debug("Debug string: \n%s", LazyString(lambda: self._debug_str()))
+        if self.settings.print_output_code:
+            print(triton_code, file=sys.stderr)
         module = PyCodeCache.load(triton_code)
         rv = getattr(module, self.kernel.name)
         rv.make_precompiler = getattr(module, f"_{self.kernel.name}_make_precompiler")

helion/runtime/settings.py

@@ -58,6 +58,7 @@ class _Settings:
     autotune_log_level: int = logging.INFO
     autotune_compile_timeout: int = 60
     autotune_precompile: bool = sys.platform != "win32"
+    print_output_code: bool = os.environ.get("HELION_PRINT_OUTPUT_CODE", "0") == "1"
 
 
 class Settings(_Settings):
@@ -75,6 +76,7 @@ class Settings(_Settings):
         "autotune_log_level": "Log level for autotuning. 0 = no logging, 1 = only final config, 2 = default, 3 = verbose.",
         "autotune_compile_timeout": "Timeout for Triton compilation in seconds used for autotuning. Default is 60 seconds.",
         "autotune_precompile": "If True, precompile the kernel before autotuning. Requires fork-safe environment.",
+        "print_output_code": "If True, print the output code of the kernel to stderr.",
     }
     assert __slots__.keys() == {field.name for field in dataclasses.fields(_Settings)}