Adding new setting, autotune_effort=[none/quick/full] (#913)

Co-authored-by: Jason Ansel <jansel@meta.com>

Author: choijon5

docs/api/autotuner.md

@@ -2,6 +2,8 @@
 
 The `helion.autotuner` module provides automatic optimization of kernel configurations.
 
+Autotuning effort can be adjusted via :attr:`helion.Settings.autotune_effort`, which configures how much each algorithm explores (``"none"`` disables autotuning, ``"quick"`` runs a smaller search, ``"full"`` uses the full search budget). Users may still override individual autotuning parameters if they need finer control.
+
 ```{eval-rst}
 .. currentmodule:: helion.autotuner
 

docs/api/kernel.md

@@ -155,6 +155,7 @@ Settings control **how the kernel is compiled** and the development environment:
 @helion.kernel(
     # Settings parameters
     use_default_config=True,      # Skip autotuning for development
+    autotune_effort="quick",     # Smaller autotuning budget when search is enabled
     print_output_code=True,       # Debug: show generated Triton code
     static_shapes=True,           # Compilation optimization strategy
     autotune_log_level=logging.DEBUG  # Verbose autotuning output

docs/api/settings.md

@@ -160,6 +160,18 @@ with helion.set_default_settings(
 .. autoattribute:: Settings.autotune_config_overrides
 
    Dict of config key/value pairs to force during autotuning. Useful for disabling problematic candidates or pinning experimental options.
+
+.. autoattribute:: Settings.autotune_effort
+
+   Select the autotuning effort preset. Available values:
+
+   - ``"none"`` – skip autotuning and run the default configuration (equivalent to ``use_default_config=True``).
+   - ``"quick"`` – limited search for faster runs with decent performance.
+   - ``"full"`` – exhaustive autotuning (current default behavior).
+
+   Users can still override individual ``autotune_*`` settings; explicit values win over the preset. Controlled by ``HELION_AUTOTUNE_EFFORT``.
+
+
 ```
 
 ### Autotuning Cache

helion/autotuner/__init__.py

@@ -9,6 +9,10 @@
 from .differential_evolution import (
     DifferentialEvolutionSearch as DifferentialEvolutionSearch,
 )
+from .effort_profile import AutotuneEffortProfile as AutotuneEffortProfile
+from .effort_profile import DifferentialEvolutionConfig as DifferentialEvolutionConfig
+from .effort_profile import PatternSearchConfig as PatternSearchConfig
+from .effort_profile import RandomSearchConfig as RandomSearchConfig
 from .finite_search import FiniteSearch as FiniteSearch
 from .local_cache import LocalAutotuneCache as LocalAutotuneCache
 from .local_cache import StrictLocalAutotuneCache as StrictLocalAutotuneCache

helion/autotuner/base_search.py

@@ -550,10 +550,9 @@ def should_rebenchmark(self, member: PopulationMember) -> bool:
         Returns:
             True if the member should be re-benchmarked, False otherwise.
         """
-        return (
+        threshold = self.settings.get_rebenchmark_threshold()
+        return member.perf < threshold * self.best_perf_so_far and math.isfinite(
             member.perf
-            < self.settings.autotune_rebenchmark_threshold * self.best_perf_so_far
-            and math.isfinite(member.perf)
         )
 
     def rebenchmark(
@@ -568,7 +567,14 @@ def rebenchmark(
         """
         if len(members) < 2:
             return
-        repeat = min(1000, max(3, int(200 / self.best_perf_so_far)))
+
+        # Calculate repeat count based on best performance
+        base_repeat = (
+            int(200 / self.best_perf_so_far)
+            if math.isfinite(self.best_perf_so_far) and self.best_perf_so_far > 0
+            else 1000
+        )
+        repeat = min(1000, max(3, base_repeat))
         iterator = [functools.partial(m.fn, *self.args) for m in members]
         if self.settings.autotune_progress_bar:
             new_timings = interleaved_bench(iterator, repeat=repeat, desc=desc)

helion/autotuner/differential_evolution.py

@@ -8,6 +8,7 @@
 from .base_search import PopulationMember
 from .base_search import performance
 from .base_search import population_statistics
+from .effort_profile import DIFFERENTIAL_EVOLUTION_DEFAULTS
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -26,8 +27,8 @@ def __init__(
         self,
         kernel: BoundKernel,
         args: Sequence[object],
-        population_size: int = 40,
-        max_generations: int = 40,
+        population_size: int = DIFFERENTIAL_EVOLUTION_DEFAULTS.population_size,
+        max_generations: int = DIFFERENTIAL_EVOLUTION_DEFAULTS.max_generations,
         crossover_rate: float = 0.8,
         immediate_update: bool | None = None,
     ) -> None:

helion/autotuner/effort_profile.py

@@ -0,0 +1,82 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Literal
+
+AutotuneEffort = Literal["none", "quick", "full"]
+
+
+@dataclass(frozen=True)
+class PatternSearchConfig:
+    initial_population: int
+    copies: int
+    max_generations: int
+
+
+@dataclass(frozen=True)
+class DifferentialEvolutionConfig:
+    population_size: int
+    max_generations: int
+
+
+@dataclass(frozen=True)
+class RandomSearchConfig:
+    count: int
+
+
+# Default values for each algorithm (single source of truth)
+PATTERN_SEARCH_DEFAULTS = PatternSearchConfig(
+    initial_population=100,
+    copies=5,
+    max_generations=20,
+)
+
+DIFFERENTIAL_EVOLUTION_DEFAULTS = DifferentialEvolutionConfig(
+    population_size=40,
+    max_generations=40,
+)
+
+RANDOM_SEARCH_DEFAULTS = RandomSearchConfig(
+    count=1000,
+)
+
+
+@dataclass(frozen=True)
+class AutotuneEffortProfile:
+    pattern_search: PatternSearchConfig | None
+    differential_evolution: DifferentialEvolutionConfig | None
+    random_search: RandomSearchConfig | None
+    rebenchmark_threshold: float = 1.5
+
+
+_PROFILES: dict[AutotuneEffort, AutotuneEffortProfile] = {
+    "none": AutotuneEffortProfile(
+        pattern_search=None,
+        differential_evolution=None,
+        random_search=None,
+    ),
+    "quick": AutotuneEffortProfile(
+        pattern_search=PatternSearchConfig(
+            initial_population=30,
+            copies=2,
+            max_generations=5,
+        ),
+        differential_evolution=DifferentialEvolutionConfig(
+            population_size=20,
+            max_generations=8,
+        ),
+        random_search=RandomSearchConfig(
+            count=100,
+        ),
+        rebenchmark_threshold=0.9,  # <1.0 effectively disables rebenchmarking
+    ),
+    "full": AutotuneEffortProfile(
+        pattern_search=PATTERN_SEARCH_DEFAULTS,
+        differential_evolution=DIFFERENTIAL_EVOLUTION_DEFAULTS,
+        random_search=RANDOM_SEARCH_DEFAULTS,
+    ),
+}
+
+
+def get_effort_profile(effort: AutotuneEffort) -> AutotuneEffortProfile:
+    return _PROFILES[effort]

helion/autotuner/pattern_search.py

@@ -8,6 +8,7 @@
 from .base_search import PopulationBasedSearch
 from .base_search import PopulationMember
 from .base_search import performance
+from .effort_profile import PATTERN_SEARCH_DEFAULTS
 
 if TYPE_CHECKING:
     from collections.abc import Iterator
@@ -25,9 +26,9 @@ def __init__(
         kernel: BoundKernel,
         args: Sequence[object],
         *,
-        initial_population: int = 100,
-        copies: int = 5,
-        max_generations: int = 20,
+        initial_population: int = PATTERN_SEARCH_DEFAULTS.initial_population,
+        copies: int = PATTERN_SEARCH_DEFAULTS.copies,
+        max_generations: int = PATTERN_SEARCH_DEFAULTS.max_generations,
         min_improvement_delta: float = 0.001,
     ) -> None:
         """

helion/autotuner/random_search.py

@@ -3,6 +3,7 @@
 from typing import TYPE_CHECKING
 
 from .config_generation import ConfigGeneration
+from .effort_profile import RANDOM_SEARCH_DEFAULTS
 from .finite_search import FiniteSearch
 
 if TYPE_CHECKING:
@@ -31,7 +32,7 @@ def __init__(
         self,
         kernel: BoundKernel,
         args: Sequence[object],
-        count: int = 1000,
+        count: int = RANDOM_SEARCH_DEFAULTS.count,
     ) -> None:
         super().__init__(
             kernel,

helion/runtime/kernel.py

@@ -571,7 +571,7 @@ def _implicit_config(self) -> Config | None:
             return self._config
         if len(configs) == 1:
             return configs[0]
-        if len(configs) == 0 and self.kernel.settings.use_default_config:
+        if len(configs) == 0 and self.kernel.settings.autotune_effort == "none":
             config = self.config_spec.default_config()
             if not is_ref_mode_enabled(self.kernel.settings):
                 kernel_decorator = self.format_kernel_decorator(config, self.settings)