Code refine for tune_test.py & Add a pre load callback (apache#20)

* Bug fix for tutorials

* Add PreLoadMeasuredStates

* Add search_callback support for task tuner

* Code refine for tune_test.py

* Update

* Update

* Update

* Update

* Bug fix

Author: jcf94

python/tvm/ansor/__init__.py

@@ -29,7 +29,7 @@
 
 # Shortcut
 from .compute_dag import ComputeDAG
-from .auto_schedule import SearchTask, MetaTileRewritePolicy, TuneOption, HardwareParams
+from .auto_schedule import SearchTask, MetaTileRewritePolicy, TuneOption, HardwareParams, PreLoadMeasuredStatesCallback
 from .auto_schedule import auto_schedule
 from .measure import MeasureInput, LocalBuilder, LocalRunner, RPCRunner, LocalRPCMeasureContext
 from .cost_model import RandomModel

python/tvm/ansor/auto_schedule.py

@@ -69,7 +69,15 @@ def __init__(self, dag, workload_key, target, target_host=None,
 class SearchPolicy(Object):
     def continue_search(self, task, num_measure, verbose, measurer):
         return _ffi_api.SearchPolicyContinueSearchOneRound(self, task, num_measure, verbose, measurer)
+    
+    def set_task(self, task):
+        _ffi_api.SearchPolicySetTask(self, task);
 
+    def set_verbose(self, verbose):
+        _ffi_api.SearchPolicySetVerbose(self, verbose);
+
+    def run_callbacks(self, callbacks):
+        _ffi_api.SearchPolicyRunCallbacks(self, callbacks)
 
 @tvm._ffi.register_object("ansor.MetaTileRewritePolicy")
 class MetaTileRewritePolicy(SearchPolicy):
@@ -117,6 +125,21 @@ def __init__(self,
             seed or random.randint(1, 1 << 30))
 
 
+@tvm._ffi.register_object("ansor.SearchCallback")
+class SearchCallback(Object):
+    pass
+
+
+@tvm._ffi.register_object("ansor.PreLoadMeasuredStatesCallback")
+class PreLoadMeasuredStatesCallback(SearchCallback):
+    """ A SearchCallback that used for search policy to load measured hash
+        from the log file.
+    """
+    def __init__(self, filename: str):
+        self.__init_handle_by_constructor__(
+            _ffi_api.PreLoadMeasuredStatesCallback, filename)
+
+
 @tvm._ffi.register_object("ansor.TuneOption")
 class TuneOption(Object):
     """ The options for tuning
@@ -135,11 +158,13 @@ class TuneOption(Object):
       Builder which builds the program
     runner: Runner
       Runner which runs the program and measure time costs
-    callbacks: List[MeasureCallback]
+    measure_callbacks: List[MeasureCallback]
       Callback functions
+    pre_search_callbacks: List[SearchCallback]
     """
     def __init__(self, n_trials=0, early_stopping=-1, num_measure_per_iter=64,
-                 verbose=1, builder='local', runner='local', callbacks=None):
+                 verbose=1, builder='local', runner='local', measure_callbacks=None,
+                 pre_search_callbacks=None):
         if isinstance(builder, str):
             if builder == 'local':
                 builder = LocalBuilder()
@@ -152,12 +177,15 @@ def __init__(self, n_trials=0, early_stopping=-1, num_measure_per_iter=64,
             else:
                 raise ValueError("Invalid builder: " + runner)
 
-        if callbacks is None:
-            callbacks = []
+        if measure_callbacks is None:
+            measure_callbacks = []
+
+        if pre_search_callbacks is None:
+            pre_search_callbacks = []
 
         self.__init_handle_by_constructor__(
             _ffi_api.TuneOption, n_trials, early_stopping, num_measure_per_iter,
-            verbose, builder, runner, callbacks)
+            verbose, builder, runner, measure_callbacks, pre_search_callbacks)
 
 
 def auto_schedule(workload, target=None,

python/tvm/ansor/measure.py

@@ -174,6 +174,16 @@ def __init__(self,
 
 @tvm._ffi.register_object("ansor.ProgramMeasurer")
 class ProgramMeasurer(Object):
+    """
+    Parameters
+    ----------
+    builder : Builder
+    runner : Runner
+    callbacks : List[MeasureCallback]
+    verbose : Int
+    max_continuous_error : Float
+    """
+
     def __init__(self, builder: Builder, runner: Runner,
                  callbacks: List[MeasureCallback],
                  verbose: int, max_continuous_error: int = -1):
@@ -182,6 +192,21 @@ def __init__(self, builder: Builder, runner: Runner,
 
 @tvm._ffi.register_object("ansor.RPCRunner")
 class RPCRunner(Runner):
+    """
+    Parameters
+    ----------
+    key : Str
+    host : Str
+    port : Int
+    priority : Int
+    n_parallel : Int
+    timeout : Int
+    number : Int
+    repeat : Int
+    min_repeat_ms : Int
+    cooldown_interval : Float
+    """
+
     def __init__(self, key, host, port, priority=1,
                  n_parallel=1,
                  timeout=10,
@@ -203,6 +228,19 @@ def __init__(self, key, host, port, priority=1,
 
 
 class LocalRPCMeasureContext:
+    """ A context wrapper for RPCRunner.
+
+    Parameters
+    ----------
+    priority : Int
+    n_parallel : Int
+    timeout : Int
+    number : Int
+    repeat : Int
+    min_repeat_ms : Int
+    cooldown_interval : Float
+    """
+
     def __init__(self,
                  priority=1,
                  n_parallel=1,
@@ -228,8 +266,8 @@ def __init__(self,
         time.sleep(0.5)
 
     def __del__(self):
-        self.tracker.terminate()
         self.server.terminate()
+        self.tracker.terminate()
 
 
 class MeasureErrorNo(object):

python/tvm/ansor/task_scheduler.py

@@ -153,7 +153,7 @@ def tune(self, tune_option: TuneOption, search_policy: Union[str, List[SearchPol
         self.tune_option = tune_option
         if self.use_debug_measurement_simulator is None:
             self.measurer = ProgramMeasurer(tune_option.builder, tune_option.runner,
-                                            tune_option.callbacks, tune_option.verbose)
+                                            tune_option.measure_callbacks, tune_option.verbose)
         self.ct = 0
         self.tic = time.time()
         # reset num_measure_per_iter to make sure every task is tuned at least once
@@ -167,6 +167,13 @@ def tune(self, tune_option: TuneOption, search_policy: Union[str, List[SearchPol
         self.sequential_now_task_idx = 0
         self.sequential_now_task_begin_ct = 0
 
+        for i in range(len(self.tasks)):
+            search_policy = self.search_policies[i]
+            task = self.tasks[i]
+            search_policy.set_task(task)
+            search_policy.set_verbose(tune_option.verbose)
+            search_policy.run_callbacks(tune_option.pre_search_callbacks)
+
         # do a round robin first
         if self.strategy != 'sequential':
             for i in range(len(self.tasks)):