[Refactor] move vm.py under runtime and adt to runtime.container.py

python/tvm/autotvm/task/relay_integration.py

@@ -39,7 +39,6 @@ def _lower(mod,
     """
     # pylint: disable=import-outside-toplevel
     from tvm import relay
-    from tvm import runtime
     from tvm.relay.backend import graph_runtime_codegen
 
     if hasattr(target, 'device_name') and target.device_name == "vta":
@@ -50,7 +49,7 @@ def _lower(mod,
                 grc = graph_runtime_codegen.GraphRuntimeCodegen(None, target)
                 grc.codegen(mod["main"])
     # default case
-    compiler = runtime.vm.VMCompiler()
+    compiler = relay.vm.VMCompiler()
     if params:
         compiler.set_params(params)
     compiler.lower(mod, target=target)

python/tvm/relay/__init__.py

@@ -36,6 +36,7 @@
 from . import debug
 from . import param_dict
 from . import feature
+from .backend import vm
 
 # Root operators
 from .op import Op

python/tvm/relay/backend/_vm.py

@@ -0,0 +1,21 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""The Relay virtual machine FFI namespace.
+"""
+import tvm._ffi
+
+tvm._ffi._init_api("relay._vm", __name__)

python/tvm/relay/backend/vm.py

@@ -0,0 +1,258 @@
+# License .to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=no-else-return, unidiomatic-typecheck, undefined-variable, invalid-name, redefined-builtin
+"""
+The Relay Virtual Machine.
+
+Implements a Python interface to compiling and executing on the Relay VM.
+"""
+import numpy as np
+
+import tvm
+import tvm.runtime.ndarray as _nd
+import tvm.runtime.vm as vm_rt
+from tvm import autotvm
+from tvm.relay import expr as _expr
+from tvm.relay.backend.interpreter import Executor
+from . import _vm
+
+
+def compile(mod, target=None, target_host=None, params=None):
+    """Compile the module to VM executable. A helper function for VMCompiler.
+
+    Parameters
+    ----------
+    mod : relay.Module
+        The Relay module 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.
+
+    target_host : str or :any:`tvm.target.Target`, optional
+        Host compilation target, if target is device.
+        When TVM compiles device specific program such as CUDA,
+        we also need host(CPU) side code to interact with the driver
+        to setup the dimensions and parameters correctly.
+        target_host is used to specify the host side codegen target.
+        By default, llvm is used if it is enabled,
+        otherwise a stackvm intepreter is used.
+
+    params : dict of str to NDArray
+        Input parameters to the graph that do not change
+        during inference time. Used for constant folding.
+
+    Returns
+    -------
+    exec : tvm.runtime.vmf.Executable
+        The VM executable that contains both library code and bytecode.
+    """
+    compiler = VMCompiler()
+    if params:
+        compiler.set_params(params)
+    compiler.lower(mod, target, target_host)
+    compiler.codegen()
+    return compiler.get_exec()
+
+
+class VMCompiler(object):
+    """Compiler that compiles Relay module to VM executable."""
+    def __init__(self):
+        self.mod = _vm._VMCompiler()
+        self._lower = self.mod["lower"]
+        self._codegen = self.mod["codegen"]
+        self._get_exec = self.mod["get_executable"]
+        self._set_params_func = self.mod["set_params"]
+        self._get_params_func = self.mod["get_params"]
+        self._optimize = self.mod["optimize"]
+
+    def set_params(self, params):
+        """Set constant parameters for the model.
+
+        Parameters
+        ----------
+        params : dict of str to NDArray
+            Input parameters to the graph that do not change
+            during inference time. Used for constant folding.
+        """
+        inputs = {}
+        for name, param in params.items():
+            if isinstance(param, np.ndarray):
+                param = _nd.array(param)
+            inputs[name] = _expr.const(param)
+        self._set_params_func(inputs)
+
+    def get_params(self):
+        """Return the updated weights."""
+        params = self._get_params_func()
+        ret = {}
+        for key, value in params.items():
+            ret[key] = value.data
+        return ret
+
+    def lower(self, mod, target=None, target_host=None):
+        """Lower the module to VM bytecode.
+
+        Parameters
+        ----------
+        mod : relay.Module
+            The Relay module 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.
+
+        target_host : str or :any:`tvm.target.Target`, optional
+            Host compilation target, if target is device.
+            When TVM compiles device specific program such as CUDA,
+            we also need host(CPU) side code to interact with the driver
+            to setup the dimensions and parameters correctly.
+            target_host is used to specify the host side codegen target.
+            By default, llvm is used if it is enabled,
+            otherwise a stackvm intepreter is used.
+        """
+        target = self._update_target(target)
+        target_host = self._update_target_host(target, target_host)
+        tophub_context = self._tophub_context(target)
+        with tophub_context:
+            self._lower(mod, target, target_host)
+
+    def codegen(self):
+        """Generate the kernel library."""
+        self._codegen()
+
+    def optimize(self, mod, target=None, params=None):
+        """Helper method that optimizes a Relay module via VM.
+
+        Parameters
+        ----------
+        mod : relay.Module
+
+        target : str, :any:`tvm.target.Target`, or dict of str (i.e.
+            device/context name) to str/tvm.target.Target, optional
+
+        params : dict of str to NDArray
+            Input parameters to the graph that do not change
+            during inference time. Used for constant folding.
+
+        Returns
+        -------
+        mod : relay.Module
+            The optimized relay module.
+
+        params : dict
+            The parameters of the final module.
+        """
+        target = self._update_target(target)
+        if params:
+            self.set_params(params)
+        return self._optimize(mod, target), self.get_params()
+
+    def get_exec(self):
+        """Get the VM executable.
+
+        Returns
+        -------
+        exec : tvm.runtime.vm.Executable
+            The VM executable that contains both library code and bytecode.
+        """
+        return vm_rt.Executable(self._get_exec())
+
+    def _update_target(self, target):
+        """Update target."""
+        target = target if target else tvm.target.current_target()
+        if target is None:
+            raise ValueError("Target is not set in env or passed as argument.")
+        tgts = {}
+        if isinstance(target, (str, tvm.target.Target)):
+            dev_type = tvm.expr.IntImm("int32", tvm.nd.context(str(target)).device_type)
+            tgts[dev_type] = tvm.target.create(target)
+        elif isinstance(target, dict):
+            for dev, tgt in target.items():
+                dev_type = tvm.expr.IntImm("int32", tvm.nd.context(dev).device_type)
+                tgts[dev_type] = tvm.target.create(tgt)
+        else:
+            raise TypeError("target is expected to be str, tvm.target.Target, " +
+                            "or dict of str to str/tvm.target.Target, but received " +
+                            "{}".format(type(target)))
+        return tgts
+
+    def _update_target_host(self, target, target_host):
+        """Update target host."""
+        target_host = None if target_host == "" else target_host
+        if not target_host:
+            for device_type, tgt in target.items():
+                if device_type.value == tvm.nd.cpu(0).device_type:
+                    target_host = tgt
+                    break
+        if not target_host:
+            target_host = "llvm" if tvm.runtime.enabled("llvm") else "stackvm"
+        if isinstance(target_host, str):
+            target_host = tvm.target.create(target_host)
+        return target_host
+
+    def _tophub_context(self, target):
+        """Get the autotvm context."""
+        # 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()
+        return tophub_context
+
+
+class VMExecutor(Executor):
+    """
+    An implementation of the executor interface for
+    the Relay VM.
+
+    Useful interface for experimentation and debugging
+    the VM can also be used directly from the API.
+    supported by `tvm.runtime.vm`.
+
+    Parameters
+    ----------
+    mod : :py:class:`~tvm.relay.module.Module`
+        The module to support the execution.
+
+    ctx : :py:class:`~tvm.TVMContext`
+        The runtime context to run the code on.
+
+    target : :py:class:`Target`
+        The target option to build the function.
+    """
+    def __init__(self, mod, ctx, target):
+        if mod is None:
+            raise RuntimeError("Must provide module to get VM executor.")
+        self.mod = mod
+        self.ctx = ctx
+        self.target = target
+        self.executable = compile(mod, target)
+        self.vm = vm_rt.VirtualMachine(self.executable)
+        self.vm.init(ctx)
+
+    def _make_executor(self, expr=None):
+        main = self.mod["main"]
+
+        def _vm_wrapper(*args, **kwargs):
+            args = self._convert_args(main, args, kwargs)
+            return self.vm.run(*args)
+
+        return _vm_wrapper

python/tvm/relay/build_module.py

@@ -22,14 +22,14 @@
 import numpy as np
 
 from tvm import expr as tvm_expr
-from tvm.runtime import vm
 from .. import nd as _nd, target as _target, autotvm
 from ..contrib import graph_runtime as _graph_rt
 from . import _build_module
 from . import ty as _ty
 from . import expr as _expr
 from .module import Module as _Module
 from .backend import interpreter as _interpreter
+from .backend.vm import VMExecutor
 
 def _update_target(target):
     target = target if target else _target.current_target()
@@ -408,5 +408,5 @@ def create_executor(kind="debug",
     if kind == "graph":
         return GraphExecutor(mod, ctx, target)
     if kind == "vm":
-        return vm.VMExecutor(mod, ctx, target)
+        return VMExecutor(mod, ctx, target)
     raise RuntimeError("unknown execution strategy: {0}".format(kind))

python/tvm/runtime/_vm.py

@@ -18,4 +18,4 @@
 """
 import tvm._ffi
 
-tvm._ffi._init_api("relay._vm", __name__)
+tvm._ffi._init_api("runtime._vm", __name__)