Ansor Relay Integration (without layout rewrite) (apache#22)

* relay integration

Author: minminsun

python/tvm/ansor/__init__.py

@@ -28,14 +28,20 @@
 from . import task_scheduler
 
 # Shortcut
-from .compute_dag import ComputeDAG
+from .compute_dag import ComputeDAG, LayoutRewriteLevel, gen_schedule
 from .auto_schedule import SearchTask, MetaTileRewritePolicy, TuneOption, HardwareParams, \
     PreLoadMeasuredStates, PreAddCustomRule
 from .auto_schedule import auto_schedule
 from .measure import MeasureInput, LocalBuilder, LocalRunner, RPCRunner, LocalRPCMeasureContext
 from .cost_model import RandomModel
 from .cost_model.xgb_model import XGBModel
-from .serialization import LogToFile, LogReader, best_measure_pair_in_file, write_measure_records_to_file
+from .serialization import LogToFile, LogReader, best_measure_pair_in_file, \
+    load_from_file, write_measure_records_to_file
 from .workload_registry import register_auto_scheduler_workload_func, \
     workload_key_to_dag, make_workload_key_func
 from .task_scheduler import TaskScheduler, SimpleTaskScheduler
+from .dispatcher import DispatchContext, ApplyConfig, ApplyHistoryBest as apply_history_best, \
+    FallbackContext, clear_fallback_cache, ApplyGraphBest, BlockingEmptyContext
+from .topi_integration import register_topi_schedule, TaskExtractEnv
+from .relay_integration import extract_from_program, extract_from_multiple_program, \
+    finish_layout_rewrite

python/tvm/ansor/compute_dag.py

@@ -19,6 +19,7 @@
 
 import tvm._ffi
 from tvm.runtime import Object
+from tvm import te
 from .loop_state import State
 from . import _ffi_api
 
@@ -88,3 +89,13 @@ def infer_bound_from_state(self, state):
         state : StateObject
         """
         return _ffi_api.ComputeDAGInferBoundFromState(self, state)
+
+def gen_schedule(state, bufs):
+    if not state or not state.complete:
+        return te.create_schedule([x.op for x in bufs])
+    else:
+        dag = ComputeDAG(bufs)
+        # only update compute body, layout_rewrite_level = LayoutRewriteLevel.COMPUTE_REWRITE,
+        # since kernel layout has already been rewritten in relay pass
+        schedule, _ = dag.apply_steps_from_state(state, layout_rewrite_level=LayoutRewriteLevel.COMPUTE_REWRITE)
+    return schedule