Use autotuner's BoundKernel in caching (#388)

Author: oulgen

helion/autotuner/base_cache.py

@@ -8,7 +8,6 @@
 import os
 from typing import TYPE_CHECKING
 from typing import Hashable
-from typing import Sequence
 
 from torch._inductor.codecache import build_code_hash
 from torch._inductor.codecache import torch_key
@@ -18,7 +17,6 @@
 
 if TYPE_CHECKING:
     from ..runtime.config import Config
-    from ..runtime.kernel import BoundKernel
     from .base_search import BaseSearch
 
 log: logging.Logger = logging.getLogger(__name__)
@@ -114,12 +112,10 @@ class AutotuneCacheBase(abc.ABC):
     provide implementations for get and put methods.
     """
 
-    def __init__(
-        self, kernel: BoundKernel, args: Sequence[object], autotuner: BaseSearch
-    ) -> None:
+    def __init__(self, autotuner: BaseSearch) -> None:
         self.autotuner = autotuner
-        self.kernel = kernel
-        self.args = args
+        self.kernel = self.autotuner.kernel
+        self.args = self.autotuner.args
 
     @abc.abstractmethod
     def get(self) -> Config | None:

helion/autotuner/local_cache.py

@@ -7,7 +7,6 @@
 from pathlib import Path
 import textwrap
 from typing import TYPE_CHECKING
-from typing import Sequence
 
 import torch
 from torch._inductor.runtime.cache_dir_utils import (
@@ -20,7 +19,6 @@
 from .base_cache import StrictAutotuneCacheKey
 
 if TYPE_CHECKING:
-    from ..runtime.kernel import BoundKernel
     from .base_search import BaseSearch
 
 log: logging.Logger = logging.getLogger(__name__)
@@ -38,10 +36,8 @@ class LocalAutotuneCache(AutotuneCacheBase):
     PyTorch. Use StrictLocalAutotuneCache to consider these properties.
     """
 
-    def __init__(
-        self, kernel: BoundKernel, args: Sequence[object], autotuner: BaseSearch
-    ) -> None:
-        super().__init__(kernel, args, autotuner)
+    def __init__(self, autotuner: BaseSearch) -> None:
+        super().__init__(autotuner)
         self.key = self._generate_key()
 
     def _generate_key(self) -> LooseAutotuneCacheKey:

helion/runtime/kernel.py

@@ -460,8 +460,6 @@ def autotune(
             from ..autotuner import LocalAutotuneCache
 
             config = LocalAutotuneCache(
-                self,
-                args,
                 DifferentialEvolutionSearch(
                     self,
                     args,

test/test_cache.py

@@ -28,11 +28,8 @@ def test_basic(self):
         b = torch.randn(16, device=DEVICE, dtype=torch.float16)
         args_b = (b, b)
 
-        # TODO(oulgen): Using a custom autotuner is very verbose, requires passing args 3 times etc
         bound_kernel = basic_kernels.add.bind(args_a)
-        config = StrictLocalAutotuneCache(
-            bound_kernel, args_a, BasicSearch(bound_kernel, args_a)
-        ).autotune()
+        config = StrictLocalAutotuneCache(BasicSearch(bound_kernel, args_a)).autotune()
         bound_kernel.set_config(config)
         result = bound_kernel(*args_a)
         torch.testing.assert_close(result, a + a)
@@ -44,9 +41,7 @@ def test_basic(self):
         basic_kernels.add.reset()
 
         bound_kernel = basic_kernels.add.bind(args_a)
-        config = StrictLocalAutotuneCache(
-            bound_kernel, args_a, BasicSearch(bound_kernel, args_a)
-        ).autotune()
+        config = StrictLocalAutotuneCache(BasicSearch(bound_kernel, args_a)).autotune()
         bound_kernel.set_config(config)
         result = bound_kernel(*args_a)
         torch.testing.assert_close(result, a + a)
@@ -58,9 +53,7 @@ def test_basic(self):
         basic_kernels.add.reset()
 
         bound_kernel = basic_kernels.add.bind(args_b)
-        config = StrictLocalAutotuneCache(
-            bound_kernel, args_b, BasicSearch(bound_kernel, args_b)
-        ).autotune()
+        config = StrictLocalAutotuneCache(BasicSearch(bound_kernel, args_b)).autotune()
         bound_kernel.set_config(config)
         result = bound_kernel(*args_b)
         torch.testing.assert_close(result, b + b)