Generating intervals for pyATF where possible

Author: fjwillemsen

kernel_tuner/searchspace.py

@@ -18,7 +18,7 @@
 )
 
 from kernel_tuner.util import check_restrictions as check_instance_restrictions
-from kernel_tuner.util import compile_restrictions, default_block_size_names
+from kernel_tuner.util import compile_restrictions, default_block_size_names, get_interval
 
 supported_neighbor_methods = ["strictly-adjacent", "adjacent", "Hamming"]
 
@@ -262,7 +262,7 @@ def all_smt(formula, keys) -> list:
 
     def __build_searchspace_pyATF(self, block_size_names: list, max_threads: int, solver: Solver):
         """Builds the searchspace using pyATF."""
-        from pyatf import TP, Set, Tuner
+        from pyatf import TP, Set, Interval, Tuner
         from pyatf.cost_functions.generic import CostFunction
         from pyatf.search_techniques import Exhaustive
 
@@ -289,7 +289,8 @@ def get_params():
             params = list()
             print("get_params")
             for index, (key, values) in enumerate(self.tune_params.items()):
-                vals = Set(*np.array(values).flatten())  # TODO check if can be interval
+                vi = get_interval(values)
+                vals = Interval(vi[0], vi[1], vi[2]) if vi is not None else Set(*np.array(values).flatten())
                 constraint = res_dict.get(key, None)
                 if len(res_dict) == 0 and index == len(self.tune_params) - 1 and constraint is None:
                     res = self.restrictions[0][0]

kernel_tuner/util.py

@@ -439,6 +439,23 @@ def get_instance_string(params):
     return "_".join([str(i) for i in params.values()])
 
 
+def get_interval(a: list):
+    """Checks if an array can be an interval. Returns (start, end, step) if interval, otherwise None."""
+    if not all(isinstance(e, (int, float)) for e in a):
+        return None
+    a_min = min(a)
+    a_max = max(a)
+    if len(a) <= 2:
+        return (a_min, a_max, a_max-a_min)
+    # determine the first step size
+    step = a[1]-a_min
+    # for each element, the step size should be equal to the first step
+    for i, e in enumerate(a):
+        if e-a[i-1] != step:
+            return None 
+    return (a_min, a_max, step)
+
+
 def get_kernel_string(kernel_source, params=None):
     """Retrieve the kernel source and return as a string.