Optimizing autotvm task extraction speed

python/tvm/autotvm/task/relay_integration.py

@@ -31,23 +31,28 @@
 logger = logging.getLogger('autotvm')
 
 
-# TODO(moreau89) find a more elegant way to build for VTAs
-def _build(func,
+# TODO(moreau89) find a more elegant way to lower for VTAs
+def _lower(func,
            target,
-           target_host,
            params):
-    """ Helper to build VTA properly.
+    """ Helper to lower VTA properly.
     """
 
     from tvm import relay
+    from tvm.relay.backend import graph_runtime_codegen
 
     if hasattr(target, 'device_name') and target.device_name == "vta":
         with relay.build_config(opt_level=3, disabled_pass={"AlterOpLayout"}):
             import vta
             with vta.build_config():
-                return relay.build(func, target, target_host, params)
+                _, mod, _ = relay.optimize(func, target, params)
+                grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+                return grc.codegen(mod["main"])
     # default case
-    return relay.build(func, target, target_host, params)
+    _, mod, _ = relay.optimize(func, target, params)
+    grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
+    return grc.codegen(mod["main"])
+
 
 def extract_from_program(func, params, ops, target, target_host=None):
     """ Extract tuning tasks from a relay program.
@@ -133,8 +138,8 @@ def extract_from_multiple_program(funcs, params, ops, target, target_host=None):
             relay.backend.compile_engine.get().clear()
             # wrap build call in thread to avoid multiprocessing problems
             mod = relay.Module.from_expr(func)
-            build_thread = threading.Thread(target=_build,
-                                            args=(mod, target, target_host, param))
+            build_thread = threading.Thread(target=_lower,
+                                            args=(mod, target, param))
             build_thread.start()
             build_thread.join()
 

python/tvm/relay/__init__.py

@@ -28,7 +28,7 @@
 from . import adt
 from . import analysis
 from . import transform
-from .build_module import build, create_executor
+from .build_module import build, create_executor, optimize
 from .transform import build_config
 from . import prelude
 from . import parser

python/tvm/relay/backend/graph_runtime_codegen.py

@@ -36,15 +36,15 @@
 from __future__ import absolute_import
 
 from tvm.ndarray import empty
-from tvm.relay import build_module
+from tvm.relay import _build_module
 from tvm import target as _target
 from tvm import expr as _expr
 
 class GraphRuntimeCodegen(object):
     """The compiler from Relay to the TVM runtime system."""
 
     def __init__(self, mod, target):
-        self._mod = build_module._GraphRuntimeCodegen()
+        self._mod = _build_module._GraphRuntimeCodegen()
         self._init = self._mod["init"]
         self._codegen = self._mod["codegen"]
         self._get_graph_json = self._mod["get_graph_json"]

python/tvm/relay/build_module.py

@@ -60,6 +60,7 @@ def __init__(self):
         self._get_graph_json = self.mod["get_graph_json"]
         self._get_module = self.mod["get_module"]
         self._build = self.mod["build"]
+        self._optimize = self.mod["optimize"]
         self._set_params_func = self.mod["set_params"]
         self._get_params_func = self.mod["get_params"]
 
@@ -113,6 +114,46 @@ def build(self, func, target=None, target_host=None, params=None):
 
         return graph_json, mod, params
 
+    def optimize(self, func, target=None, params=None):
+        """
+        Parameters
+        ----------
+        func: relay.Function
+            The function to build.
+
+        target : str, :any:`tvm.target.Target`, or dict of str(i.e.
+        device/context name) to str/tvm.target.Target, optional
+            For heterogeneous compilation, it is a dictionary indicating context
+            to target mapping. For homogeneous compilation, it is a build target.
+
+        params : dict of str to NDArray
+            Input parameters to the graph that do not change
+            during inference time. Used for constant folding.
+
+        Returns
+        -------
+        graph_json : str
+            The json string that can be accepted by graph runtime.
+
+        mod : relay.Module
+            The optimized relay module.
+
+        params : dict
+            The parameters of the final graph.
+        """
+        target = _update_target(target)
+
+        # Setup the params.
+        if params:
+            self._set_params(params)
+        mod = self._optimize(func, target)
+        # Get artifacts
+        graph_json = self.get_json()
+        params = self.get_params()
+
+        return graph_json, mod, params
+
+
     def _set_params(self, params):
         inputs = {}
         for name, param in params.items():
@@ -208,6 +249,61 @@ def build(mod, target=None, target_host=None, params=None):
     return graph_json, mod, params
 
 
+def optimize(mod, target=None, params=None):
+    """Helper function that builds a Relay function to run on TVM graph
+    runtime.
+
+    Parameters
+    ----------
+    mod : relay.Module
+        The module to build. Using relay.Function is deprecated.
+
+    target : str, :any:`tvm.target.Target`, or dict of str(i.e. device/context
+    name) to str/tvm.target.Target, optional
+        For heterogeneous compilation, it is a dictionary indicating context to
+        target mapping. For homogeneous compilation, it is a build target.
+
+    params : dict of str to NDArray
+        Input parameters to the graph that do not change
+        during inference time. Used for constant folding.
+
+    Returns
+    -------
+    graph_json : str
+        The json string that can be accepted by graph runtime.
+
+    mod : relay.Module
+        The optimized relay module.
+
+    params : dict
+        The parameters of the final graph.
+    """
+    if isinstance(mod, _Module):
+        func = mod["main"]
+    elif isinstance(mod, _expr.Function):
+        func = mod
+        warnings.warn(
+            "Please use input parameter mod (tvm.relay.module.Module) "
+            "instead of deprecated parameter func (tvm.relay.expr.Function)",
+            DeprecationWarning)
+    else:
+        raise ValueError("Type of input parameter mod must be tvm.relay.module.Module")
+
+    target = _update_target(target)
+
+    # If current dispatch context is fallback context (the default root context),
+    # then load pre-tuned parameters from TopHub
+    if isinstance(autotvm.DispatchContext.current, autotvm.FallbackContext):
+        tophub_context = autotvm.tophub.context(list(target.values()))
+    else:
+        tophub_context = autotvm.util.EmptyContext()
+
+    with tophub_context:
+        bld_mod = BuildModule()
+        graph_json, mod, params = bld_mod.optimize(func, target, params)
+    return graph_json, mod, params
+
+
 class GraphExecutor(_interpreter.Executor):
     """Wrapper around Executor interface.
 

src/relay/backend/build_module.cc

@@ -148,6 +148,19 @@ class RelayBuildModule : public runtime::ModuleNode {
       return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
           *rv = this->graph_codegen_->GetLoweredFunc();
       });
+    } else if (name == "optimize") {
+      return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+        CHECK_EQ(args.num_args, 2);
+        Function func = args[0];
+        if (this->params_.size()) {
+          func = this->BindParamsByName(func, this->params_);
+        }
+        // Perform Module->Module optimizations.
+        relay::Module relay_module = relay::ModuleNode::FromExpr(func);
+        relay_module = Optimize(relay_module, args[1], this->params_);
+        CHECK(relay_module.defined());
+        *rv = relay_module;
+      });
     } else {
       LOG(FATAL) << "Unknown packed function: " << name;
       return PackedFunc([sptr_to_self, name](TVMArgs args, TVMRetValue* rv) {});