Completed merge with recent changes to searchspace

Author: fjwillemsen

kernel_tuner/interface.py

@@ -614,16 +614,24 @@ def tune_kernel(
     tuning_options = Options([(k, opts[k]) for k in _tuning_options.keys()])
     device_options = Options([(k, opts[k]) for k in _device_options.keys()])
     tuning_options["unique_results"] = {}
-    if strategy_options and "max_fevals" in strategy_options:
-        tuning_options["max_fevals"] = strategy_options["max_fevals"]
-    if strategy_options and "time_limit" in strategy_options:
-        tuning_options["time_limit"] = strategy_options["time_limit"]
 
+    # copy some values from strategy_options
+    searchspace_construction_options = {}
+    if strategy_options:
+        if "max_fevals" in strategy_options:
+            tuning_options["max_fevals"] = strategy_options["max_fevals"]
+        if "time_limit" in strategy_options:
+            tuning_options["time_limit"] = strategy_options["time_limit"] 
+        if "searchspace_construction_options" in strategy_options:
+            searchspace_construction_options = strategy_options["searchspace_construction_options"]         
+
+    # log the user inputs
     logging.debug("tune_kernel called")
     logging.debug("kernel_options: %s", util.get_config_string(kernel_options))
     logging.debug("tuning_options: %s", util.get_config_string(tuning_options))
     logging.debug("device_options: %s", util.get_config_string(device_options))
 
+    # check whether the selected strategy and options are valid
     strategy_string = strategy
     if strategy:
         if strategy in strategy_map:
@@ -669,7 +677,7 @@ def preprocess_cache(filepath):
 
     # create search space
     tuning_options.restrictions_unmodified = deepcopy(restrictions)
-    searchspace = Searchspace(tune_params, restrictions, runner.dev.max_threads)
+    searchspace = Searchspace(tune_params, restrictions, runner.dev.max_threads, **searchspace_construction_options)
     restrictions = searchspace._modified_restrictions
     tuning_options.restrictions = restrictions
     if verbose:

kernel_tuner/searchspace.py

@@ -31,6 +31,7 @@
 from kernel_tuner.util import check_restrictions as check_instance_restrictions
 from kernel_tuner.util import (
     compile_restrictions,
+    convert_constraint_lambdas,
     default_block_size_names,
     get_interval,
 )
@@ -114,15 +115,14 @@ def __init__(
                 )
             )
             and not (
-                framework_l == "pysmt" or framework_l == "bruteforce" or solver_method.lower() == "pc_parallelsolver"
+                framework_l == "pysmt" or framework_l == "bruteforce" or framework_l == "pythonconstraint" or solver_method.lower() == "pc_parallelsolver"
             )
         ):
             self.restrictions = compile_restrictions(
                 restrictions,
                 tune_params,
                 monolithic=False,
                 format=framework_l if framework_l == "pyatf" else None,
-                try_to_constraint=framework_l == "pythonconstraint",
             )
 
         # if an imported cache, skip building and set the values directly
@@ -342,9 +342,7 @@ def get_tune_params_pyatf(self, block_size_names: list = None, max_threads: int
             # adding the default blocksize restriction requires recompilation because pyATF requires combined restrictions for the same parameter
             max_block_size_product = f"{' * '.join(valid_block_size_names)} <= {max_threads}"
             restrictions = self._modified_restrictions.copy() + [max_block_size_product]
-            self.restrictions = compile_restrictions(
-                restrictions, self.tune_params, format="pyatf", try_to_constraint=False
-            )
+            self.restrictions = compile_restrictions(restrictions, self.tune_params, format="pyatf")
 
         # build a dictionary of the restrictions, combined based on last parameter
         res_dict = dict()
@@ -371,6 +369,7 @@ def get_tune_params_pyatf(self, block_size_names: list = None, max_threads: int
             vals = (
                 Interval(vi[0], vi[1], vi[2]) if vi is not None and vi[2] != 0 else Set(*np.array(values).flatten())
             )
+            assert vals is not None, f"Values for parameter {key} are None, this should not happen."
             constraint = res_dict.get(key, None)
             constraint_source = None
             if constraint is not None:
@@ -442,7 +441,7 @@ def __parameter_space_list_to_lookup_and_return_type(
             parameter_space_dict,
             size_list,
         )
-
+    
     def __build_searchspace(self, block_size_names: list, max_threads: int, solver: Solver):
         """Compute valid configurations in a search space based on restrictions and max_threads."""
         # instantiate the parameter space with all the variables
@@ -451,6 +450,9 @@ def __build_searchspace(self, block_size_names: list, max_threads: int, solver:
             parameter_space.addVariable(str(param_name), param_values)
 
         # add the user-specified restrictions as constraints on the parameter space
+        if not isinstance(self.restrictions, (list, tuple)):
+            self.restrictions = [self.restrictions]
+        self.restrictions = convert_constraint_lambdas(self.restrictions)
         parameter_space = self.__add_restrictions(parameter_space)
 
         # add the default blocksize threads restrictions last, because it is unlikely to reduce the parameter space by much
@@ -477,20 +479,29 @@ def __add_restrictions(self, parameter_space: Problem) -> Problem:
             for restriction in self.restrictions:
                 required_params = self.param_names
 
-                # convert to a Constraint type if necessary
-                if isinstance(restriction, tuple):
-                    restriction, required_params, _ = restriction
+                # (un)wrap where necessary
+                if isinstance(restriction, tuple) and len(restriction) >= 2:
+                    required_params = restriction[1]
+                    restriction = restriction[0]
                 if callable(restriction) and not isinstance(restriction, Constraint):
-                    restriction = FunctionConstraint(restriction)
-
-                # add the Constraint
+                    # def restrictions_wrapper(*args):
+                    #     return check_instance_restrictions(restriction, dict(zip(self.param_names, args)), False)
+                    # print(restriction, isinstance(restriction, Constraint))
+                    # restriction = FunctionConstraint(restrictions_wrapper)
+                    restriction = FunctionConstraint(restriction, required_params)
+
+                # add as a Constraint
+                all_params_required = all(param_name in required_params for param_name in self.param_names)
+                variables = None if all_params_required else required_params
                 if isinstance(restriction, FunctionConstraint):
-                    parameter_space.addConstraint(restriction, required_params)
+                    parameter_space.addConstraint(restriction, variables)
                 elif isinstance(restriction, Constraint):
-                    all_params_required = all(param_name in required_params for param_name in self.param_names)
-                    parameter_space.addConstraint(restriction, None if all_params_required else required_params)
-                elif isinstance(restriction, str) and self.solver_method.lower() == "pc_parallelsolver":
-                    parameter_space.addConstraint(restriction)
+                    parameter_space.addConstraint(restriction, variables)
+                elif isinstance(restriction, str):
+                    if self.solver_method.lower() == "pc_parallelsolver":
+                        parameter_space.addConstraint(restriction)
+                    else:
+                        parameter_space.addConstraint(restriction, variables)
                 else:
                     raise ValueError(f"Unrecognized restriction type {type(restriction)} ({restriction})")
 

kernel_tuner/strategies/common.py

@@ -51,10 +51,10 @@ def make_strategy_options_doc(strategy_options):
 
 def get_options(strategy_options, options):
     """Get the strategy-specific options or their defaults from user-supplied strategy_options."""
-    accepted = list(options.keys()) + ["max_fevals", "time_limit"]
+    accepted = list(options.keys()) + ["max_fevals", "time_limit", "searchspace_construction_options"]
     for key in strategy_options:
         if key not in accepted:
-            raise ValueError(f"Unrecognized option {key} in strategy_options")
+            raise ValueError(f"Unrecognized option {key} in strategy_options (allowed: {accepted})")
     assert isinstance(options, dict)
     return [strategy_options.get(opt, default) for opt, (_, default) in options.items()]