Add Index simplification & API update (apache#26)

* Add vectorized cooperative_fetching test

* Update math simplify for vectorized CF

* File rename

* Update tune_network

* API update

python/tvm/ansor/auto_schedule.py

@@ -245,8 +245,6 @@ def auto_schedule(workload, target=None,
 
     Returns
     -------
-    state : State
-
     sch : tvm.Schedule
 
     tensors : List[Tensor]
@@ -270,4 +268,3 @@ def auto_schedule(workload, target=None,
     else:
         raise ValueError("Invalid workload: " + workload +
                          ". Expect a string or SearchTask")
-

python/tvm/ansor/compute_dag.py

@@ -20,7 +20,7 @@
 import tvm._ffi
 from tvm.runtime import Object
 from tvm import te
-from .loop_state import State
+from .loop_state import State, StateObject
 from . import _ffi_api
 
 
@@ -63,8 +63,12 @@ def apply_steps_from_state(self, state, layout_rewrite_level=None):
         sch : Schedule
         args : List[Tensor]
         """
-        sch, args = _ffi_api.ComputeDAGApplyStepsFromState(self, state)
-        return sch, args
+        if isinstance(state, State):
+            return _ffi_api.ComputeDAGApplyStepsFromState(self, state.state_object)
+        elif isinstance(state, StateObject):
+            return _ffi_api.ComputeDAGApplyStepsFromState(self, state)
+        else:
+            raise ValueError("The input must be a State or StateObject")
 
     def print_python_code_from_state(self, state):
         """
@@ -76,7 +80,12 @@ def print_python_code_from_state(self, state):
         -------
         str : Str
         """
-        return _ffi_api.ComputeDAGPrintPythonCodeFromState(self, state)
+        if isinstance(state, State):
+            return _ffi_api.ComputeDAGPrintPythonCodeFromState(self, state.state_object)
+        elif isinstance(state, StateObject):
+            return _ffi_api.ComputeDAGPrintPythonCodeFromState(self, state)
+        else:
+            raise ValueError("The input must be a State or StateObject")
 
     def infer_bound_from_state(self, state):
         """
@@ -88,7 +97,12 @@ def infer_bound_from_state(self, state):
         -------
         state : StateObject
         """
-        return _ffi_api.ComputeDAGInferBoundFromState(self, state)
+        if isinstance(state, State):
+            return State(_ffi_api.ComputeDAGInferBoundFromState(self, state.state_object))
+        elif isinstance(state, StateObject):
+            return State(_ffi_api.ComputeDAGInferBoundFromState(self, state))
+        else:
+            raise ValueError("The input must be a State or StateObject")
 
 def gen_schedule(state, bufs):
     if not state or not state.complete:

python/tvm/ansor/feature.py

@@ -23,7 +23,7 @@
 import struct
 import numpy as np
 
-from .loop_state import StateObject
+from .loop_state import State, StateObject
 from .measure import MeasureInput, MeasureResult
 from . import _ffi_api
 
@@ -131,12 +131,16 @@ def get_per_stmt_features_from_measure_pairs(inputs: List[MeasureInput],
     return unpack_feature(byte_arr)
 
 
-def get_per_stmt_features_from_states(states: List[StateObject],
+def get_per_stmt_features_from_states(states,
                                       task: "SearchTask",
                                       max_n_bufs: int = None) -> List[np.ndarray]:
     """Get per_stmt features from states"""
+    if isinstance(states[0], State):
+        state_objects = [s.state_object for s in states]
+    elif isinstance(states[0], StateObject):
+        state_objects = states
     byte_arr = _ffi_api.GetPerStmtFeaturesFromStates(
-        states, task, max_n_bufs or DEFAULT_MAX_N_BUFS)
+        state_objects, task, max_n_bufs or DEFAULT_MAX_N_BUFS)
     return unpack_feature(byte_arr)[0]
 
 

python/tvm/ansor/loop_state.py

@@ -60,18 +60,6 @@ def iters(self):
             setattr(self, "iterators_cache", _ffi_api.StageGetIterators(self))
         return getattr(self, "iterators_cache")
 
-    def iter(self, index):
-        """
-        Parameters
-        ----------
-        index : Int
-
-        Returns
-        -------
-        iter : Iterator
-        """
-        return _ffi_api.StageGetIterator(self, index)
-
 
 @tvm._ffi.register_object("ansor.State")
 class StateObject(Object):
@@ -302,7 +290,7 @@ def bind_thread(self, stage_id, it, thread_name):
         }
         thread_id = trans_table[thread_name]
 
-        self.state_object, res = _ffi_api.StateUnroll(self.state_object, stage_id, it, thread_id)
+        self.state_object, res = _ffi_api.StateBindThread(self.state_object, stage_id, it, thread_id)
         self.clear_cache()
         return res
 

python/tvm/ansor/measure.py

@@ -62,7 +62,7 @@ class MeasureInput(Object):
     """
 
     def __init__(self, task, state):
-        self.__init_handle_by_constructor__(_ffi_api.MeasureInput, task, state)
+        self.__init_handle_by_constructor__(_ffi_api.MeasureInput, task, state.state_object)
 
 
 @tvm._ffi.register_object("ansor.BuildResult")

python/tvm/ansor/serialization.py

@@ -22,6 +22,7 @@
 import tvm._ffi
 from tvm.runtime import Object
 from .measure import MeasureCallback, MeasureErrorNo
+from .loop_state import State
 from . import _ffi_api
 
 
@@ -74,7 +75,8 @@ def write_measure_records_to_file(filename, inputs, results):
 
 def get_states_from_measure_inputs(inputs, task):
     """Get states from measure inputs"""
-    return _ffi_api.GetStatesFromMeasureInputs(inputs, task)
+    state_objects = _ffi_api.GetStatesFromMeasureInputs(inputs, task)
+    return [State(s) for s in state_objects]
 
 
 def best_measure_pair_in_file(filename, workload_key=None, target=None):

scripts/tune_network.py

@@ -191,22 +191,14 @@ def create_module(data_shape, graph, lib, target, input_name, params, debug_prof
     return module, ctx
 
 
-def tune_and_evaluate(network_name, model_path, batch_size, target, target_host,
-                      local_measure, device_key, host, port, n_parallel, ndk_cc,
-                      build_timeout, run_timeout, num_threads, tune, check_correctness,
-                      debug_profile, tuning_parameters, record_file, layout_set):
-    task_scheduler, model_type, policy, log_file, load_log_file = (tuning_parameters['task_scheduler'],
-            tuning_parameters['model_type'], tuning_parameters['policy'],
-            tuning_parameters['log_file'], tuning_parameters['load_log_file'])
-
-    if layout_set:
-        layout = layout_set
-
+def tune_and_evaluate(target, target_host, log_n_lines, search_policy, tune,
+                      debug_profile, check_correctness, network_parameters,
+                      task_scheduler_parameters, tune_parameters, module_parameters):
     # Extract workloads from relay program
-    print("=============== Extract workloads ===============")
-    mod, params, input_name, data_shape, out_shape = get_network(network_name, model_path, batch_size, layout)
+    mod, params, input_name, data_shape, out_shape = get_network(**network_parameters)
 
     if tune:
+        print("=============== Extracting workloads ===============")
         workloads, wkl_weights = ansor.extract_from_program(mod, target=target,
                 params=params, ops=(relay.op.nn.dense, relay.op.nn.softmax,
                                     relay.op.nn.conv2d, relay.op.nn.conv2d_transpose,
@@ -215,7 +207,7 @@ def tune_and_evaluate(network_name, model_path, batch_size, target, target_host,
                                     relay.op.nn.conv3d, relay.op.nn.adaptive_avg_pool3d,
                                     relay.op.nn.batch_matmul, relay.op.mean,
                                     ))
-        print("Total workload number: %d" % (len(workloads)))
+        print("Totally %d workload extracted." % (len(workloads)))
 
         # Tune workloads with auto scheduler
         print("=============== Tuning ===============")
@@ -225,23 +217,13 @@ def tune_and_evaluate(network_name, model_path, batch_size, target, target_host,
             print("[========= Task %d =========]\n" % i, dag)
             tasks.append(ansor.SearchTask(dag, wkl_key, target, target_host))
 
-        def objective_func(costs):
-            return sum(c * w for c, w in zip(costs, wkl_weights))
-
-        tuner = ansor.SimpleTaskScheduler(tasks, objective_func, strategy=task_scheduler,
-                                          load_log_file=load_log_file,
-                                          load_model_file=tuning_parameters['load_model'])
+        tuner = ansor.SimpleTaskScheduler(tasks,
+                lambda costs: sum(c * w for c, w in zip(costs, wkl_weights)),
+                **task_scheduler_parameters)
+        tune_option, measure_ctx = create_tune_option(target, **tune_parameters)
 
-        tune_option, measure_ctx = create_tune_option(target, log_file,
-                tuning_parameters['n_trials'], tuning_parameters['num_measure_per_iter'],
-                tuning_parameters['verbose'], n_parallel, build_timeout,
-                local_measure, device_key, host, port, ndk_cc,
-                tuning_parameters['early_stopping'])
-        search_policy = "%s.%s" % (policy, model_type)
-
-        if local_measure and target.target_name != 'cuda':
+        if tune_parameters['local_measure'] and target.target_name != 'cuda':
             os.environ['TVM_BIND_MASTER_CORE_0'] = "1"
-
         tuner.tune(tune_option, search_policy)
 
         if measure_ctx:
@@ -251,15 +233,13 @@ def objective_func(costs):
 
     # Compile graph with best states found by auto-scheduler
     print("=============== Compile ===============")
-    with ansor.apply_history_best(log_file, args.log_n_lines):
-    #if True:
-    #with ansor.BlockingEmptyContext():
+    with ansor.apply_history_best(tune_parameters['log_file'], log_n_lines):
         os.environ['TVM_AUTO_CACHE_FLUSH'] = "0"
         os.environ['TVM_BIND_MASTER_CORE_0'] = "1"
         if kernel_layout_rewrite:
             ansor.prepare_layout_rewrite(mod, target=target,
-                                                  params=params,
-                                                  ops=(relay.op.nn.dense, relay.op.nn.conv2d, relay.op.nn.conv3d))
+                                         params=params,
+                                         ops=(relay.op.nn.dense, relay.op.nn.conv2d, relay.op.nn.conv3d))
         else:
             # disable layout rewrite
             ansor.LayoutRewriteLevel.BOTH_REWRITE = ansor.LayoutRewriteLevel.NO_REWRITE
@@ -268,22 +248,12 @@ def objective_func(costs):
         with relay.build_config(opt_level=3):
             graph, lib, opt_params = relay.build_module.build(
                 mod, target=target, params=params)
-        '''
-        from tvm.relay.backend import graph_runtime_codegen
-        with relay.build_config(opt_level=3):
-            opt_mod, _ = relay.optimize(mod, target, params)
-            grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
-            grc.codegen(opt_mod["main"])
-        with tvm.transform.PassContext(opt_level=3):
-            graph, lib, opt_params = relay.build_module.build(
-                mod, target=target, params=params)
-        '''
+
         ansor.finish_layout_rewrite()
         print("=============== Compile Finish ===============")
 
-        module, ctx = create_module(data_shape, graph, lib, target, input_name, opt_params,
-                                    debug_profile, local_measure, ndk_cc,
-                                    device_key, host, port, run_timeout, num_threads)
+        module, ctx = create_module(data_shape, graph, lib, target, input_name,
+                                    opt_params, debug_profile, **module_parameters)
 
         # Evaluate
         print("========== Evaluate ==========")
@@ -315,9 +285,8 @@ def objective_func(costs):
             graph, lib, opt_params = relay.build_module.build(
                 mod, target=target, params=params)
 
-        module, _ = create_module(data_shape, graph, lib, target, input_name, opt_params,
-                                  debug_profile, local_measure, ndk_cc,
-                                  device_key, host, port, run_timeout, num_threads)
+        module, _ = create_module(data_shape, graph, lib, target, input_name,
+                                  opt_params, debug_profile, **module_parameters)
         module.run()
 
         expected_output = module.get_output(0).asnumpy()
@@ -343,7 +312,7 @@ def objective_func(costs):
     # Strategy related options
     parser.add_argument("--seed", type=int, default=0, help='random seed')
     parser.add_argument("--policy", type=str, choices=['multi-stage', 'meta-rewrite'],
-            default='meta-rewrite')
+                        default='meta-rewrite')
     parser.add_argument("--model-type", type=str, choices=['xgb', 'random', 'no-share'], default='xgb')
     parser.add_argument("--task-scheduler", type=str, default='gradient',
                         choices=['no', 'gradient', 'round-robin'],
@@ -359,6 +328,7 @@ def objective_func(costs):
     # Detailed control options
     parser.add_argument("--build-timeout", type=int, default=10)
     parser.add_argument("--run-timeout", type=int, default=10)
+    parser.add_argument("--early-stopping", type=int, default=-1)
     parser.add_argument("--verbose", type=int, default=1)
     parser.add_argument("--local-measure", type=str2bool, nargs='?', const=True, default=True)
     parser.add_argument("--device-key", type=str, default=None)
@@ -375,23 +345,59 @@ def objective_func(costs):
     logging.getLogger('ansor').setLevel(logging.DEBUG)
 
     target = tvm.target.create(args.target)
+    log_file =  args.log_file or "%s-B%d-%s.json" % (args.network, args.batch_size,
+                                                     target.target_name)
+    load_log_file = args.load_log or log_file
+    search_policy = "%s.%s" % (args.policy, args.model_type)
+    if args.layout:
+        layout = args.layout
+    elif target.target_name == "cuda":
+        layout = "NCHW"
+    else:
+        layout = "NHWC"
+
+    network_parameters = {
+        'name': args.network,
+        'model_path': args.model_path,
+        'batch_size': args.batch_size,
+        'layout': layout
+    }
+
+    task_scheduler_parameters = {
+        'strategy': args.task_scheduler,
+        'load_log_file': load_log_file,
+        'load_model_file': args.load_model,
+        'verbose': args.verbose,
+    }
 
-    tuning_parameters = {
+    control_parameters = {
+        'local_measure': args.local_measure,
+        'device_key': args.device_key,
+        'host': args.host,
+        'port': args.port,
+        'ndk_cc': args.ndk_cc,
+    }
+
+    tune_parameters = {
+        'log_file': log_file,
         'n_trials': args.n_trials,
         'num_measure_per_iter': args.num_measure_per_iter,
-        'log_file': args.log_file or "%s-B%d.json" % (args.network, args.batch_size),
-        'load_model': args.load_model,
-        'model_type': args.model_type,
-        'task_scheduler': args.task_scheduler,
-        'policy': args.policy,
-        'early_stopping': -1,
-        'verbose': 1,
+        'verbose': args.verbose,
+        'n_parallel': args.n_parallel,
+        'build_timeout': args.build_timeout,
+        'run_timeout': args.run_timeout,
+        'early_stopping': args.early_stopping,
+        **control_parameters
+    }
+
+    module_parameters = {
+        'run_timeout': args.run_timeout,
+        'num_threads': args.num_threads,
+        **control_parameters
     }
-    tuning_parameters['load_log_file'] = args.load_log or tuning_parameters['log_file']
 
     os.environ["TOPHUB_LOCATION"] = "NONE"
-    tune_and_evaluate(args.network, args.model_path, args.batch_size, target, args.target_host,
-                      args.local_measure, args.device_key, args.host,
-                      args.port, args.n_parallel, args.ndk_cc, args.build_timeout,
-                      args.run_timeout, args.num_threads, args.tune, args.check_correctness,
-                      args.debug_profile, tuning_parameters, args.out_file, args.layout)
+    tune_and_evaluate(target, args.target_host, args.log_n_lines, search_policy,
+                      args.tune, args.debug_profile, args.check_correctness,
+                      network_parameters, task_scheduler_parameters, tune_parameters,
+                      module_parameters)

scripts/tune_test.py

@@ -14,7 +14,7 @@
 
 def create_tune_option(target, log_file, n_trials, num_measure_per_iter, verbose,
                       n_parallel, build_timeout, local_measure, device_key, host,
-                      port, ndk_cc, early_stopping=-1):
+                      port, ndk_cc, early_stopping=-1, run_timeout=10):
     builder = runner = measure_ctx = None
     if local_measure:
         builder = ansor.LocalBuilder(timeout=build_timeout)
@@ -26,7 +26,7 @@ def create_tune_option(target, log_file, n_trials, num_measure_per_iter, verbose
     else:
         os.environ['TVM_NDK_CC'] = ndk_cc
         builder = ansor.LocalBuilder(timeout=build_timeout, build_func='ndk')
-        runner = ansor.RPCRunner(key=device_key, host=host, port=port,
+        runner = ansor.RPCRunner(key=device_key, host=host, port=port, timeout=run_timeout,
                                  n_parallel=n_parallel, repeat=1, min_repeat_ms=400)
 
     tune_option = ansor.TuneOption(n_trials=n_trials, early_stopping=early_stopping,

src/ansor/loop_state.cc

@@ -1019,11 +1019,6 @@ TVM_STATIC_IR_FUNCTOR(ReprPrinter, vtable)
   PrintState(&p->stream, node, true);
 });
 
-
-TVM_REGISTER_GLOBAL("ansor.StageGetIterator").set_body_typed([](const Stage& stage, int index) {
-  return stage->iters[index];
-});
-
 TVM_REGISTER_GLOBAL("ansor.StageGetIterators").set_body_typed([](const Stage& stage) {
   return Array<Iterator>(stage->iters);
 });