pyATF restrictions needed to be recombined in cases where multiple restrictions are defined on the same tunable parameter

Author: fjwillemsen

kernel_tuner/searchspace.py

@@ -287,12 +287,12 @@ def __build_searchspace_pyATF(self, block_size_names: list, max_threads: int, so
         # define the Tunable Parameters
         def get_params():
             params = list()
-            print("get_params")
             for index, (key, values) in enumerate(self.tune_params.items()):
                 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:
+                # in case of a leftover monolithic restriction, append at the last parameter
+                if index == len(self.tune_params) - 1 and len(res_dict) == 0 and len(self.restrictions) == 1:
                     res = self.restrictions[0][0]
                     assert callable(res)
                     constraint = res

kernel_tuner/util.py

@@ -1052,8 +1052,35 @@ def to_equality_constraint(
                     finalized_constraint = to_equality_constraint(parsed_restriction, params_used)
             if finalized_constraint is None:
                 # we must turn it into a general function
-                finalized_constraint = f"def r({', '.join(params_used)}): return {parsed_restriction} \n"
+                if format.lower() == "pyatf":
+                    finalized_constraint = parsed_restriction
+                else:
+                    finalized_constraint = f"def r({', '.join(params_used)}): return {parsed_restriction} \n"
             parsed_restrictions.append((finalized_constraint, params_used))
+
+        # if pyATF, restrictions that are set on the same parameter must be combined into one
+        if format.lower() == "pyatf":
+            res_dict = dict()
+            registered_params = list()
+            registered_restrictions = list()
+            for param in tune_params.keys():
+                registered_params.append(param)
+                for index, (res, params) in enumerate(parsed_restrictions):
+                    if index in registered_restrictions:
+                        continue
+                    if all(p in registered_params for p in params):
+                        if param not in res_dict:
+                            res_dict[param] = (list(), list())
+                        res_dict[param][0].append(res)
+                        res_dict[param][1].extend(params)
+                        registered_restrictions.append(index)
+            # combine multiple restrictions into one
+            parsed_restrictions_pyatf = list()
+            for res_tuple in res_dict.values():
+                res, params_used = res_tuple
+                params_used = list(dict.fromkeys(params_used))   # param_used should only contain unique, dict preserves order
+                parsed_restrictions_pyatf.append((f"def r({', '.join(params_used)}): return {' and '.join(res)} \n", params_used))
+            return parsed_restrictions_pyatf
     else:
         # create one monolithic function
         parsed_restrictions = ") and (".join(