Add DE-Surrogate hybrid autotuner algorithm + early stopping option for DE and DE-Surrogate (#1096)

Author: FranciscoThiesen

.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]'
+          SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,de-surrogate]'
           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]'
+          SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0" uv pip install -e .'[dev,de-surrogate]'
           python -c "import helion; print(helion.__name__)"
 
       - name: Run Tests

helion/autotuner/__init__.py

@@ -6,6 +6,7 @@
 from .config_fragment import ListOf as ListOf
 from .config_fragment import 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,
 )
@@ -20,6 +21,7 @@
 from .random_search import RandomSearch as RandomSearch
 
 search_algorithms = {
+    "DESurrogateHybrid": DESurrogateHybrid,
     "DifferentialEvolutionSearch": DifferentialEvolutionSearch,
     "FiniteSearch": FiniteSearch,
     "PatternSearch": PatternSearch,

helion/autotuner/base_search.py

@@ -624,6 +624,7 @@ class PopulationMember:
         perfs (list[float]): The performance of the configuration, accumulated over multiple benchmarks.
         flat_values (FlatConfig): The flat representation of the configuration values.
         config (Config): The full configuration object.
+        compile_time (float | None): The compilation time for this configuration.
     """
 
     fn: Callable[..., object]

helion/autotuner/config_fragment.py

@@ -57,6 +57,26 @@ def get_minimum(self) -> int:
         """
         raise NotImplementedError
 
+    def encode_scalar(self, value: object) -> float:
+        """
+        Encode a configuration value into a float for ML models.
+
+        This is used by surrogate-assisted algorithms to convert configurations
+        into numerical vectors for prediction models.
+
+        Args:
+            value: The configuration value to encode.
+
+        Returns:
+            A float 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)
+
 
 @dataclasses.dataclass
 class PermutationFragment(ConfigSpecFragment):
@@ -121,6 +141,14 @@ 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)
+
 
 class PowerOfTwoFragment(BaseIntegerFragment):
     def random(self) -> int:
@@ -152,6 +180,20 @@ 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:
+        """Encode power-of-2 values using log2 transformation."""
+        import math
+
+        if not isinstance(value, (int, float)):
+            raise TypeError(
+                f"Expected int/float for PowerOfTwoFragment, got {type(value).__name__}: {value!r}"
+            )
+        if value <= 0:
+            raise ValueError(
+                f"Expected positive value for PowerOfTwoFragment, got {value}"
+            )
+        return math.log2(float(value))
+
 
 class IntegerFragment(BaseIntegerFragment):
     def random(self) -> int:
@@ -193,6 +235,17 @@ 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:
+        """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 BooleanFragment(ConfigSpecFragment):
     def default(self) -> bool:

helion/autotuner/config_generation.py

@@ -181,3 +181,24 @@ def differential_mutation(
         # TODO(jansel): can this be larger? (too large and Triton compile times blow up)
         self.shrink_config(result, 8192)
         return result
+
+    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))
+
+        return encoded