The Searchspace neighbor cache now caches for all neighbor methods separately

Author: fjwillemsen

kernel_tuner/searchspace.py

@@ -64,7 +64,7 @@ def __init__(
         self.params_values_indices = None
         self.build_neighbors_index = build_neighbors_index
         self.solver_method = solver_method
-        self.__neighbor_cache = dict()
+        self.__neighbor_cache = { method: dict() for method in supported_neighbor_methods }
         self.neighbor_method = neighbor_method
         if (neighbor_method is not None or build_neighbors_index) and neighbor_method not in supported_neighbor_methods:
             raise ValueError(f"Neighbor method is {neighbor_method}, must be one of {supported_neighbor_methods}")
@@ -758,24 +758,25 @@ def get_neighbors_indices_no_cache(self, param_config: tuple, neighbor_method=No
         raise ValueError(f"The neighbor method {neighbor_method} is not in {supported_neighbor_methods}")
 
     def get_neighbors_indices(self, param_config: tuple, neighbor_method=None) -> List[int]:
-        """Get the neighbors indices for a parameter configuration, possibly cached."""
-        neighbors = self.__neighbor_cache.get(param_config, None)
+        """Get the neighbors indices for a parameter configuration, cached if requested before."""
+        if neighbor_method is None:
+            neighbor_method = self.neighbor_method
+            if neighbor_method is None:
+                raise ValueError("Neither the neighbor_method argument nor self.neighbor_method was set")
+        neighbors = self.__neighbor_cache[neighbor_method].get(param_config, None)
         # if there are no cached neighbors, compute them
         if neighbors is None:
             neighbors = self.get_neighbors_indices_no_cache(param_config, neighbor_method)
-            self.__neighbor_cache[param_config] = neighbors
-        # if the neighbors were cached but the specified neighbor method was different than the one initially used to build the cache, throw an error
-        elif (
-            self.neighbor_method is not None and neighbor_method is not None and self.neighbor_method != neighbor_method
-        ):
-            raise ValueError(
-                f"The neighbor method {neighbor_method} differs from the intially set {self.neighbor_method}, can not use cached neighbors. Use 'get_neighbors_no_cache()' when mixing neighbor methods to avoid this."
-            )
+            self.__neighbor_cache[neighbor_method][param_config] = neighbors
         return neighbors
 
-    def are_neighbors_indices_cached(self, param_config: tuple) -> bool:
+    def are_neighbors_indices_cached(self, param_config: tuple, neighbor_method=None) -> bool:
         """Returns true if the neighbor indices are in the cache, false otherwise."""
-        return param_config in self.__neighbor_cache
+        if neighbor_method is None:
+            neighbor_method = self.neighbor_method
+            if neighbor_method is None:
+                raise ValueError("Neither the neighbor_method argument nor self.neighbor_method was set")
+        return param_config in self.__neighbor_cache[neighbor_method]
 
     def get_neighbors_no_cache(self, param_config: tuple, neighbor_method=None) -> List[tuple]:
         """Get the neighbors for a parameter configuration (does not check running cache, useful when mixing neighbor methods)."""

test/test_searchspace.py

@@ -316,6 +316,34 @@ def test_neighbors_cached():
         assert neighbors == neighbors_2
 
 
+def test_neighbors_cached_mixed_methods():
+    """Test whether retrieving a set of neighbors with one method after another yields the correct neighbors."""
+    simple_searchspace_duplicate = Searchspace(
+        simple_tuning_options.tune_params,
+        simple_tuning_options.restrictions,
+        max_threads,
+    )
+
+    test_configs = simple_searchspace_duplicate.get_random_sample(5)
+    for test_config in test_configs:
+        assert not simple_searchspace_duplicate.are_neighbors_indices_cached(test_config, "Hamming")
+        neighbors_hamming = simple_searchspace_duplicate.get_neighbors(test_config, "Hamming")
+        assert simple_searchspace_duplicate.are_neighbors_indices_cached(test_config, "Hamming")
+
+        # now switch to a different method
+        neighbors_strictlyadjacent = simple_searchspace_duplicate.get_neighbors(test_config, "strictly-adjacent")
+        neighbors_strictlyadjacent_no_cache = simple_searchspace_duplicate.get_neighbors_no_cache(test_config, "strictly-adjacent")
+
+        neighbors_adjacent = simple_searchspace_duplicate.get_neighbors(test_config, "adjacent")
+        neighbors_adjacent_no_cache = simple_searchspace_duplicate.get_neighbors_no_cache(test_config, "adjacent")
+
+        # check that the neighbors are as expected
+        assert neighbors_strictlyadjacent == neighbors_strictlyadjacent_no_cache
+        assert neighbors_adjacent == neighbors_adjacent_no_cache
+        assert neighbors_hamming != neighbors_strictlyadjacent
+        assert neighbors_hamming != neighbors_adjacent
+
+
 def test_param_neighbors():
     """Test whether for a given parameter configuration and index the correct neighboring parameters are returned."""
     test_config = tuple([1.5, 4, "string_1"])