Make to_triton_code config arg optional (#291)

Author: jansel

helion/runtime/kernel.py

@@ -274,6 +274,7 @@ def __init__(self, kernel: Kernel[_R], args: tuple[object, ...]) -> None:
         super().__init__()
         self.kernel = kernel
         self._run: Callable[..., _R] | None = None
+        self._config: Config | None = None
         self._compile_cache: dict[Config, CompiledConfig] = {}
         self.env = CompileEnvironment(_find_device(args), self.kernel.settings)
         with self.env:
@@ -338,7 +339,7 @@ def configs(self) -> list[Config]:
         """
         return self.kernel.configs
 
-    def to_triton_code(self, config: ConfigLike) -> str:
+    def to_triton_code(self, config: ConfigLike | None = None) -> str:
         """
         Generate Triton code for the kernel based on the given configuration.
 
@@ -348,6 +349,8 @@ def to_triton_code(self, config: ConfigLike) -> str:
         Returns:
             str: The generated Triton code as a string.
         """
+        if config is None:
+            config = self._require_implicit_config()
         with self.env:
             if not isinstance(config, Config):
                 config = Config(**config)  # pyright: ignore[reportArgumentType]
@@ -356,7 +359,7 @@ def to_triton_code(self, config: ConfigLike) -> str:
             return get_needed_imports(root) + unparse(root)
 
     def compile_config(
-        self, config: ConfigLike, *, allow_print: bool = True
+        self, config: ConfigLike | None = None, *, allow_print: bool = True
     ) -> CompiledConfig:
         """
         Compile the kernel for a specific configuration.
@@ -368,6 +371,8 @@ def compile_config(
         Returns:
             CompiledConfig: A callable object representing the compiled kernel.
         """
+        if config is None:
+            config = self._require_implicit_config()
         if not isinstance(config, Config):
             config = Config(
                 **config  # pyright: ignore[reportArgumentType]
@@ -458,6 +463,7 @@ def set_config(self, config: ConfigLike) -> None:
                 **config  # pyright: ignore[reportArgumentType]
             )
         self._run = self.compile_config(config)
+        self._config = config
 
     def _specialize_extra(self) -> list[Callable[[Sequence[object]], Hashable]]:
         """
@@ -492,6 +498,27 @@ def make_extractor(v: Source) -> Callable[[Sequence[object]], Hashable]:
             extractors.append(make_extractor(source))
         return extractors
 
+    def _implicit_config(self) -> Config | None:
+        """
+        Returns a single config that is implicitly used by this kernel, if any.
+        """
+        configs = self.kernel.configs
+        if self._config is not None:
+            return self._config
+        if len(configs) == 1:
+            return configs[0]
+        if len(configs) == 0 and self.kernel.settings.use_default_config:
+            return self.config_spec.default_config()
+        return None
+
+    def _require_implicit_config(self) -> Config:
+        """
+        Returns the implicit config for this kernel, or raises an error if no implicit config is available.
+        """
+        if (config := self._implicit_config()) is None:
+            raise RuntimeError("no config provided and no implicit config available")
+        return config
+
     def __call__(self, *args: object) -> _R:
         """
         Execute the kernel with the given arguments.
@@ -503,8 +530,8 @@ def __call__(self, *args: object) -> _R:
             _R: The result of the kernel execution.
         """
         if self._run is None:
-            if not self.configs and self.settings.use_default_config:
-                self.set_config(self.config_spec.default_config())
+            if (config := self._implicit_config()) is not None:
+                self.set_config(config)
             else:
                 self.autotune(args)
             assert self._run is not None

test/test_misc.py

@@ -236,6 +236,57 @@ def test_tile_block_size_usage(x: torch.Tensor) -> torch.Tensor:
         # The result should have 1s at positions that are last in their tile
         self.assertTrue(result.sum().item() > 0)
 
+    def test_to_triton_code_optional_config(self):
+        """Test that to_triton_code() works without explicit config argument."""
+
+        # Test 1: Kernel with single config - should use that config
+        @helion.kernel(config={"block_sizes": [64]})
+        def kernel_single_config(x: torch.Tensor) -> torch.Tensor:
+            result = torch.empty_like(x)
+            for tile in hl.tile(x.shape):
+                result[tile] = x[tile] * 2
+            return result
+
+        x = torch.randn([32], device=DEVICE)
+        bound_kernel = kernel_single_config.bind((x,))
+
+        # Should work without config argument
+        code_without_config = bound_kernel.to_triton_code()
+        code_with_config = bound_kernel.to_triton_code({"block_sizes": [64]})
+        self.assertEqual(code_without_config, code_with_config)
+
+        # Test 2: Kernel with use_default_config - should use default config
+        @helion.kernel(use_default_config=True)
+        def kernel_default_config(x: torch.Tensor) -> torch.Tensor:
+            result = torch.empty_like(x)
+            for tile in hl.tile(x.shape):
+                result[tile] = x[tile] * 3
+            return result
+
+        bound_kernel_default = kernel_default_config.bind((x,))
+
+        # Should work without config argument using default config
+        code_default = bound_kernel_default.to_triton_code()
+        self.assertIsInstance(code_default, str)
+        self.assertIn("def", code_default)  # Basic sanity check
+
+        # Test 3: Kernel with no configs and no default - should raise error
+        @helion.kernel
+        def kernel_no_config(x: torch.Tensor) -> torch.Tensor:
+            result = torch.empty_like(x)
+            for tile in hl.tile(x.shape):
+                result[tile] = x[tile] * 4
+            return result
+
+        bound_kernel_no_config = kernel_no_config.bind((x,))
+
+        # Should raise RuntimeError when no implicit config available
+        with self.assertRaises(RuntimeError) as cm:
+            bound_kernel_no_config.to_triton_code()
+        self.assertIn(
+            "no config provided and no implicit config available", str(cm.exception)
+        )
+
 
 if __name__ == "__main__":
     unittest.main()