Do not print output code durring autotuning

Printing 1500+ triton kernels during autotuning is not useful.

stack-info: PR: #130, branch: jansel/stack/26

Author: jansel

helion/autotuner/base_search.py

@@ -10,6 +10,7 @@
 from math import inf
 from multiprocessing import connection
 import re
+import sys
 import time
 from typing import TYPE_CHECKING
 from typing import NamedTuple
@@ -82,7 +83,7 @@ def benchmark(self, config: Config) -> float:
         :return: The performance of the configuration in seconds.
         :rtype: float
         """
-        fn = self.kernel.compile_config(config)
+        fn = self.kernel.compile_config(config, allow_print=False)
         if self.start_precompile_and_check_for_hangs(config, fn)():
             return self.benchmark_function(config, fn)
         return inf
@@ -160,7 +161,7 @@ def parallel_benchmark(self, configs: list[Config]) -> list[tuple[Config, float]
         :return: A list of tuples containing configurations and their performance.
         :rtype: list[tuple[Config, float]]
         """
-        fns = [self.kernel.compile_config(c) for c in configs]
+        fns = [self.kernel.compile_config(c, allow_print=False) for c in configs]
         if self.settings.autotune_precompile:
             is_workings = PrecompileFuture.wait_for_all(
                 [
@@ -200,6 +201,9 @@ def autotune(self) -> Config:
             f"    @helion.kernel(config={best!r})\n",
             level=logging.INFO + 5,
         )
+        if self.settings.print_output_code:
+            triton_code = self.kernel.to_triton_code(best)
+            print(triton_code, file=sys.stderr)
         return best
 
     def _autotune(self) -> Config:

helion/runtime/kernel.py

@@ -328,12 +328,16 @@ def to_triton_code(self, config: ConfigLike) -> str:
             root = generate_ast(self.host_function, config)
             return get_needed_imports(root) + unparse(root)
 
-    def compile_config(self, config: ConfigLike) -> CompiledConfig:
+    def compile_config(
+        self, config: ConfigLike, *, allow_print: bool = True
+    ) -> CompiledConfig:
         """
         Compile the kernel for a specific configuration.
 
         :param config: The configuration to compile the kernel with.
         :type config: Config or dict[str, object]
+        :param allow_print: Set to suppress printing the output code when autotuning.
+        :type allow_print: bool
         :return: A callable object representing the compiled kernel.
         :rtype: Callable[..., object]
         """
@@ -342,10 +346,11 @@ def compile_config(self, config: ConfigLike) -> CompiledConfig:
         if (rv := self._compile_cache.get(config)) is not None:
             return rv
         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)
+        if allow_print:
+            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

@@ -54,9 +54,11 @@ class _Settings:
     index_dtype: torch.dtype = torch.int32
     dot_precision: Literal["tf32", "tf32x3", "ieee"] = "tf32"
     static_shapes: bool = False
-    use_default_config: bool = False
+    use_default_config: bool = os.environ.get("HELION_USE_DEFAULT_CONFIG", "0") == "1"
     autotune_log_level: int = logging.INFO
-    autotune_compile_timeout: int = 60
+    autotune_compile_timeout: int = int(
+        os.environ.get("HELION_AUTOTUNE_COMPILE_TIMEOUT", "60")
+    )
     autotune_precompile: bool = sys.platform != "win32"
     print_output_code: bool = os.environ.get("HELION_PRINT_OUTPUT_CODE", "0") == "1"
 
@@ -91,8 +93,6 @@ def __init__(self, **settings: object) -> None:
             settings = {**defaults.to_dict(), **settings}
         # pyre-ignore[6]
         super().__init__(**settings)
-        if os.getenv("HELION_USE_DEFAULT_CONFIG") == "1":
-            self.use_default_config: bool = True
 
     def to_dict(self) -> dict[str, object]:
         """