Add LFBO Pattern Search (#1115)

Co-authored-by: Ethan Che <eche@fb.com>

Author: ethche

.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,surrogate]'
           python -c "import helion; print(helion.__name__)"
 
       - name: Install Benchmark Requirements

.github/workflows/lint.yml

@@ -41,7 +41,7 @@ jobs:
         run: |
           source .venv/bin/activate
           uv pip install pyrefly
-          uv pip install .'[dev]'
+          uv pip install .'[dev,surrogate]'
 
       - name: Run pre-commit
         run: |

.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,surrogate]'
           python -c "import helion; print(helion.__name__)"
 
       - name: Run Tests

helion/autotuner/__init__.py

@@ -19,9 +19,11 @@
 from .local_cache import StrictLocalAutotuneCache as StrictLocalAutotuneCache
 from .pattern_search import PatternSearch as PatternSearch
 from .random_search import RandomSearch as RandomSearch
+from .surrogate_pattern_search import LFBOPatternSearch
 
 search_algorithms = {
     "DESurrogateHybrid": DESurrogateHybrid,
+    "LFBOPatternSearch": LFBOPatternSearch,
     "DifferentialEvolutionSearch": DifferentialEvolutionSearch,
     "FiniteSearch": FiniteSearch,
     "PatternSearch": PatternSearch,

helion/autotuner/config_fragment.py

@@ -51,15 +51,15 @@ def differential_mutation(self, a: object, b: object, c: object) -> object:
     def is_block_size(self) -> bool:
         return False
 
-    def get_minimum(self) -> int:
+    def dim(self) -> int:
         """
-        Return the minimum allowed value for this fragment.
+        Returns the dimension of the output of encode
         """
         raise NotImplementedError
 
-    def encode_scalar(self, value: object) -> float:
+    def encode(self, value: object) -> list[float]:
         """
-        Encode a configuration value into a float for ML models.
+        Encode a configuration value into a list of floats for ML models.
 
         This is used by surrogate-assisted algorithms to convert configurations
         into numerical vectors for prediction models.
@@ -68,14 +68,15 @@ def encode_scalar(self, value: object) -> float:
             value: The configuration value to encode.
 
         Returns:
-            A float representing the encoded value.
+            A list of floats representing the encoded value.
         """
-        # Default: convert to float if possible
-        if not isinstance(value, (int, float, bool)):
-            raise TypeError(
-                f"Cannot encode {type(value).__name__} value {value!r} for ML"
-            )
-        return float(value)
+        raise NotImplementedError
+
+    def get_minimum(self) -> int:
+        """
+        Return the minimum allowed value for this fragment.
+        """
+        raise NotImplementedError
 
 
 @dataclasses.dataclass
@@ -106,6 +107,17 @@ def pattern_neighbors(self, current: object) -> list[object]:
                 neighbors.append(swapped)
         return neighbors
 
+    def dim(self) -> int:
+        return self.length
+
+    def encode(self, value: object) -> list[float]:
+        assert isinstance(value, list)
+        encoded = []
+        for val in value:
+            assert isinstance(val, int)
+            encoded.append(float(val))
+        return value
+
 
 @dataclasses.dataclass
 class BaseIntegerFragment(ConfigSpecFragment):
@@ -129,6 +141,9 @@ def clamp(self, val: int) -> int:
     def get_minimum(self) -> int:
         return self.low
 
+    def dim(self) -> int:
+        return 1
+
     def pattern_neighbors(self, current: object) -> list[object]:
         if type(current) is not int:  # bool is not allowed
             raise TypeError(f"Expected int, got {type(current).__name__}")
@@ -141,13 +156,9 @@ 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(self, value: object) -> list[float]:
+        assert isinstance(value, int)
+        return [float(value)]
 
 
 class PowerOfTwoFragment(BaseIntegerFragment):
@@ -180,7 +191,7 @@ 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(self, value: object) -> list[float]:
         """Encode power-of-2 values using log2 transformation."""
         import math
 
@@ -192,7 +203,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):
@@ -235,7 +246,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 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 +258,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 +279,14 @@ def differential_mutation(self, a: object, b: object, c: object) -> bool:
             return a
         return not a
 
+    def 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:
@@ -320,3 +342,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 dim(self) -> int:
+        return self.length * self.inner.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

@@ -199,6 +199,8 @@ def encode_config(self, flat_config: FlatConfig) -> list[float]:
 
         for flat_idx, spec in enumerate(self.flat_spec):
             value = flat_config[flat_idx]
-            encoded.append(spec.encode_scalar(value))
+            encoded_value = spec.encode(value)
+            assert len(encoded_value) == spec.dim()
+            encoded.extend(encoded_value)
 
         return encoded

helion/autotuner/de_surrogate_hybrid.py

@@ -91,7 +91,7 @@ def __init__(
         if not HAS_ML_DEPS:
             raise ImportError(
                 "DESurrogateHybrid requires numpy and scikit-learn. "
-                "Install them with: pip install helion[de-surrogate]"
+                "Install them with: pip install helion[surrogate]"
             )
 
         # Initialize parent with early stopping parameters

helion/autotuner/pattern_search.py

@@ -134,20 +134,35 @@ def _pattern_search_from(
             if len(candidates) <= 1:
                 return  # no new candidates, stop searching
             yield candidates  # yield new population to benchmark in parallel
+            # update search copy and check early stopping criteria
             best = min(candidates, key=performance)
-            if best is current:
-                return  # no improvement, stop searching
-            # Stop if the relative improvement is smaller than a user-specified delta
-            if (
-                self.min_improvement_delta > 0.0
-                and math.isfinite(best.perf)
-                and math.isfinite(current.perf)
-                and current.perf != 0.0
-                and abs(best.perf / current.perf - 1.0) < self.min_improvement_delta
-            ):
+            if self._check_early_stopping(best, current):
                 return
             current = best
 
+    def _check_early_stopping(
+        self, best: PopulationMember, current: PopulationMember
+    ) -> bool:
+        """
+        Check if early stopping criteria are met for the search copy
+
+        Early stops if either the best config has not changed or if
+        the relative improvement is smaller than a user-specified delta
+
+        Returns:
+            True the search copy is terminated, False otherwise.
+        """
+        if best is current:
+            return True  # no improvement, stop searching
+        # Stop if the relative improvement is smaller than a user-specified delta
+        return bool(
+            self.min_improvement_delta > 0.0
+            and math.isfinite(best.perf)
+            and math.isfinite(current.perf)
+            and current.perf != 0.0
+            and abs(best.perf / current.perf - 1.0) < self.min_improvement_delta
+        )
+
     def _generate_neighbors(self, base: FlatConfig) -> list[FlatConfig]:
         """
         Generate neighboring configurations by changing one or two parameters at a time.