Add LFBO Pattern Search

.github/workflows/benchmark.yml

@@ -100,7 +100,7 @@ jobs:
       - name: Install Helion
         run: |
           source .venv/bin/activate
-          SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,de-surrogate]'
+          SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,de-surrogate,bayesopt]'
           python -c "import helion; print(helion.__name__)"
 
       - name: Install Benchmark Requirements

.github/workflows/test.yml

@@ -146,7 +146,7 @@ jobs:
         run: |
           source .venv/bin/activate
           uv pip install setuptools ninja
-          SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,de-surrogate]'
+          SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,de-surrogate,bayesopt]'
           python -c "import helion; print(helion.__name__)"
 
       - name: Run Tests

helion/autotuner/__init__.py

@@ -1,27 +1,35 @@
 from __future__ import annotations
 
-from .config_fragment import BooleanFragment as BooleanFragment
-from .config_fragment import EnumFragment as EnumFragment
-from .config_fragment import IntegerFragment as IntegerFragment
-from .config_fragment import ListOf as ListOf
-from .config_fragment import PowerOfTwoFragment as PowerOfTwoFragment
+from .config_fragment import (
+    BooleanFragment as BooleanFragment,
+    EnumFragment as EnumFragment,
+    IntegerFragment as IntegerFragment,
+    ListOf as ListOf,
+    PowerOfTwoFragment as PowerOfTwoFragment,
+)
 from .config_spec import ConfigSpec as ConfigSpec
 from .de_surrogate_hybrid import DESurrogateHybrid as DESurrogateHybrid
 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 .effort_profile import (
+    AutotuneEffortProfile as AutotuneEffortProfile,
+    DifferentialEvolutionConfig as DifferentialEvolutionConfig,
+    PatternSearchConfig as PatternSearchConfig,
+    RandomSearchConfig as RandomSearchConfig,
+)
 from .finite_search import FiniteSearch as FiniteSearch
-from .local_cache import LocalAutotuneCache as LocalAutotuneCache
-from .local_cache import StrictLocalAutotuneCache as StrictLocalAutotuneCache
+from .local_cache import (
+    LocalAutotuneCache as LocalAutotuneCache,
+    StrictLocalAutotuneCache as StrictLocalAutotuneCache,
+)
 from .pattern_search import PatternSearch as PatternSearch
 from .random_search import RandomSearch as RandomSearch
+from .ucb_pattern_search import UCBPatternSearch
 
 search_algorithms = {
     "DESurrogateHybrid": DESurrogateHybrid,
+    "UCBPatternSearch": UCBPatternSearch,
     "DifferentialEvolutionSearch": DifferentialEvolutionSearch,
     "FiniteSearch": FiniteSearch,
     "PatternSearch": PatternSearch,

helion/autotuner/config_fragment.py

@@ -3,9 +3,7 @@
 import dataclasses
 import enum
 import random
-from typing import Iterable
-from typing import TypeGuard
-from typing import cast
+from typing import cast, Iterable, TypeGuard
 
 from ..exc import InvalidConfig
 
@@ -51,6 +49,21 @@ def differential_mutation(self, a: object, b: object, c: object) -> object:
     def is_block_size(self) -> bool:
         return False
 
+    def is_categorical(self) -> bool:
+        return True
+
+    def encode_dim(self) -> int:
+        """
+        Returns the dimension of the output of encode
+        """
+        raise NotImplementedError
+
+    def encode(self, value: object) -> list[float]:
+        """
+        Returns a list of floats that can be used to encode the value of this fragment.
+        """
+        raise NotImplementedError
+
     def get_minimum(self) -> int:
         """
         Return the minimum allowed value for this fragment.
@@ -106,6 +119,17 @@ def pattern_neighbors(self, current: object) -> list[object]:
                 neighbors.append(swapped)
         return neighbors
 
+    def encode_dim(self) -> int:
+        return self.length * self.length
+
+    def encode(self, value: object) -> list[float]:
+        encoded = []
+        for pos in range(self.length):
+            val = value[pos]
+            for v in range(self.length):
+                encoded.append(1.0 if v == val else 0.0)
+        return encoded
+
 
 @dataclasses.dataclass
 class BaseIntegerFragment(ConfigSpecFragment):
@@ -126,6 +150,9 @@ def default(self) -> int:
     def clamp(self, val: int) -> int:
         return max(min(val, self.high), self.low)
 
+    def is_categorical(self) -> bool:
+        return False
+
     def get_minimum(self) -> int:
         return self.low
 
@@ -141,13 +168,19 @@ def pattern_neighbors(self, current: object) -> list[object]:
             neighbors.append(upper)
         return neighbors
 
-    def encode_scalar(self, value: object) -> float:
-        """Encode integer values directly as floats."""
-        if not isinstance(value, (int, float)):
-            raise TypeError(
-                f"Expected int/float for BaseIntegerFragment, got {type(value).__name__}: {value!r}"
-            )
-        return float(value)
+    def encode_dim(self) -> int:
+        return 1
+
+    def encode(self, value: object) -> list[float]:
+        """Encode enum values as their index."""
+        try:
+            choice_idx = self.choices.index(value)
+        except ValueError:
+            raise ValueError(
+                f"Invalid enum value {value!r} for EnumFragment. "
+                f"Valid choices: {self.choices}"
+            ) from None
+        return [float(choice_idx)]
 
 
 class PowerOfTwoFragment(BaseIntegerFragment):
@@ -180,7 +213,10 @@ def differential_mutation(self, a: object, b: object, c: object) -> int:
             return self.clamp(ai * 2)
         return ai
 
-    def encode_scalar(self, value: object) -> float:
+    def encode_dim(self) -> int:
+        return 1
+
+    def encode(self, value: object) -> list[float]:
         """Encode power-of-2 values using log2 transformation."""
         import math
 
@@ -192,7 +228,7 @@ def encode_scalar(self, value: object) -> float:
             raise ValueError(
                 f"Expected positive value for PowerOfTwoFragment, got {value}"
             )
-        return math.log2(float(value))
+        return [math.log2(float(value))]
 
 
 class IntegerFragment(BaseIntegerFragment):
@@ -211,6 +247,13 @@ def differential_mutation(self, a: object, b: object, c: object) -> int:
             return self.clamp(a + 1)
         return a
 
+    def encode_dim(self) -> int:
+        return 1
+
+    def encode(self, value: object) -> list[float]:
+        assert isinstance(value, int)
+        return [float(value)]
+
 
 @dataclasses.dataclass
 class EnumFragment(ConfigSpecFragment):
@@ -235,7 +278,10 @@ def differential_mutation(self, a: object, b: object, c: object) -> object:
             choices.remove(a)
         return random.choice(choices)
 
-    def encode_scalar(self, value: object) -> float:
+    def encode_dim(self) -> int:
+        return len(self.choices)
+
+    def encode(self, value: object) -> list[float]:
         """Encode enum values as their index."""
         try:
             choice_idx = self.choices.index(value)
@@ -244,7 +290,7 @@ def encode_scalar(self, value: object) -> float:
                 f"Invalid enum value {value!r} for EnumFragment. "
                 f"Valid choices: {self.choices}"
             ) from None
-        return float(choice_idx)
+        return [1.0 if i == choice_idx else 0.0 for i in range(len(self.choices))]
 
 
 class BooleanFragment(ConfigSpecFragment):
@@ -265,6 +311,14 @@ def differential_mutation(self, a: object, b: object, c: object) -> bool:
             return a
         return not a
 
+    def encode_dim(self) -> int:
+        return 1
+
+    def encode(self, value: object) -> list[float]:
+        """Encode enum values as their index."""
+        assert isinstance(value, bool)
+        return [1.0] if value else [0.0]
+
 
 class BlockSizeFragment(PowerOfTwoFragment):
     def category(self) -> Category:
@@ -296,6 +350,9 @@ def random(self) -> list[object]:
         """Return a list of random values."""
         return [self.inner.random() for _ in range(self.length)]
 
+    def is_categorical(self) -> bool:
+        return self.inner.is_categorical()
+
     def pattern_neighbors(self, current: object) -> list[object]:
         """Return neighbors by changing one element at a time."""
         if not isinstance(current, list) or len(current) != self.length:
@@ -320,3 +377,13 @@ def differential_mutation(self, a: object, b: object, c: object) -> list[object]
             self.inner.differential_mutation(a[i], b[i], c[i])
             for i in range(self.length)
         ]
+
+    def encode_dim(self):
+        return self.length * self.inner.encode_dim()
+
+    def encode(self, value: object) -> list[float]:
+        assert isinstance(value, list)
+        encoded = []
+        for v in value:
+            encoded.extend(self.inner.encode(v))
+        return encoded

helion/autotuner/config_generation.py

@@ -5,13 +5,10 @@
 import itertools
 import operator
 import random
-from typing import TYPE_CHECKING
-from typing import cast
+from typing import cast, TYPE_CHECKING
 
 from .._compat import warps_to_threads
-from .config_fragment import Category
-from .config_fragment import ConfigSpecFragment
-from .config_fragment import PowerOfTwoFragment
+from .config_fragment import Category, ConfigSpecFragment, PowerOfTwoFragment
 
 if TYPE_CHECKING:
     from collections.abc import Mapping
@@ -185,20 +182,17 @@ def differential_mutation(
     def encode_config(self, flat_config: FlatConfig) -> list[float]:
         """
         Encode a flat configuration into a numerical vector for ML models.
-
         This is used by surrogate-assisted algorithms (e.g., DE-Surrogate) that need
         to represent configurations as continuous vectors for prediction models.
-
         Args:
             flat_config: The flat configuration values to encode.
-
         Returns:
             A list of floats representing the encoded configuration.
         """
         encoded: list[float] = []
 
         for flat_idx, spec in enumerate(self.flat_spec):
             value = flat_config[flat_idx]
-            encoded.append(spec.encode_scalar(value))
+            encoded.extend(spec.encode(value))
 
         return encoded