diff --git a/python/tvm/container.py b/python/tvm/container.py index 3019ad2368fd..111257eec1fb 100644 --- a/python/tvm/container.py +++ b/python/tvm/container.py @@ -17,7 +17,7 @@ """Container data structures used in TVM DSL.""" import tvm._ffi -from tvm.runtime import Object, ObjectTypes +from tvm.runtime import Object from tvm.runtime.container import getitem_helper from tvm.runtime import _ffi_node_api from . import _api_internal @@ -104,56 +104,3 @@ class LoweredFunc(Object): MixedFunc = 0 HostFunc = 1 DeviceFunc = 2 - - -@tvm._ffi.register_object("vm.ADT") -class ADT(Object): - """Algebatic data type(ADT) object. - - Parameters - ---------- - tag : int - The tag of ADT. - - fields : list[Object] or tuple[Object] - The source tuple. - """ - def __init__(self, tag, fields): - for f in fields: - assert isinstance(f, ObjectTypes), "Expect object or " \ - "tvm NDArray type, but received : {0}".format(type(f)) - self.__init_handle_by_constructor__(_ADT, tag, *fields) - - @property - def tag(self): - return _GetADTTag(self) - - def __getitem__(self, idx): - return getitem_helper( - self, _GetADTFields, len(self), idx) - - def __len__(self): - return _GetADTSize(self) - - -def tuple_object(fields=None): - """Create a ADT object from source tuple. - - Parameters - ---------- - fields : list[Object] or tuple[Object] - The source tuple. - - Returns - ------- - ret : ADT - The created object. - """ - fields = fields if fields else [] - for f in fields: - assert isinstance(f, ObjectTypes), "Expect object or tvm " \ - "NDArray type, but received : {0}".format(type(f)) - return _Tuple(*fields) - - -tvm._ffi._init_api("tvm.container") diff --git a/python/tvm/relay/__init__.py b/python/tvm/relay/__init__.py index 2432ec31cfe5..25956b4656fd 100644 --- a/python/tvm/relay/__init__.py +++ b/python/tvm/relay/__init__.py @@ -37,7 +37,6 @@ from . import param_dict from . import feature from .backend import vm -from .backend import profiler_vm # Root operators from .op import Op diff --git a/python/tvm/relay/backend/interpreter.py b/python/tvm/relay/backend/interpreter.py index 3759bc9950af..58596ec4f247 100644 --- a/python/tvm/relay/backend/interpreter.py +++ b/python/tvm/relay/backend/interpreter.py @@ -20,7 +20,7 @@ import numpy as np -from tvm import container +from tvm.runtime import container from . import _backend from .. import _make, analysis, transform from .. import module diff --git a/python/tvm/relay/backend/vm.py b/python/tvm/relay/backend/vm.py index 67afe0945ccc..b61fafda1d1d 100644 --- a/python/tvm/relay/backend/vm.py +++ b/python/tvm/relay/backend/vm.py @@ -24,343 +24,11 @@ import tvm import tvm.runtime.ndarray as _nd -from tvm.runtime import Object -from tvm import autotvm, container +import tvm.runtime.vm as vm_rt +from tvm import autotvm from tvm.relay import expr as _expr -from tvm._ffi.runtime_ctypes import TVMByteArray -from tvm._ffi import base as _base +from tvm.relay.backend.interpreter import Executor from . import _vm -from .interpreter import Executor - -def _convert(arg, cargs): - if isinstance(arg, _expr.Constant): - cargs.append(arg.data) - elif isinstance(arg, Object): - cargs.append(arg) - elif isinstance(arg, np.ndarray): - nd_arr = tvm.nd.array(arg, ctx=tvm.cpu(0)) - cargs.append(nd_arr) - elif isinstance(arg, tvm.nd.NDArray): - cargs.append(arg) - elif isinstance(arg, (tuple, list)): - field_args = [] - for field in arg: - _convert(field, field_args) - cargs.append(container.tuple_object(field_args)) - elif isinstance(arg, (_base.numeric_types, bool)): - dtype = "int32" if isinstance(arg, (int, bool)) else "float32" - value = tvm.nd.array(np.array(arg, dtype=dtype), ctx=tvm.cpu(0)) - cargs.append(value) - else: - raise TypeError("Unsupported type: %s" % (type(arg))) - - -def convert(args): - cargs = [] - for arg in args: - _convert(arg, cargs) - - return cargs - - -class Executable(object): - """Relay VM executable""" - def __init__(self, mod): - self.mod = mod - self._function_params = {} - self._save = self.mod["save"] - self._get_lib = self.mod["get_lib"] - self._get_bytecode = self.mod["get_bytecode"] - self._get_stats = self.mod["get_stats"] - self._get_function_arity = self.mod["get_function_arity"] - self._get_function_param_name = self.mod["get_function_param_name"] - - def save(self): - """Save the Relay VM Executable. - - Returns - ------- - code : bytearray - The binary blob representing a serialized Relay VM executable. It - can then be saved to disk and later deserialized into a new - Executable. - - lib : :py:class:`~tvm.runtime.Module` - The runtime module that contains the generated code. It is - basically a library that is composed of hardware dependent code. - - Notes - ----- - The returned code is organized with the following sections in order. - - Global section. This section contains the globals used by the - virtual machine. - - Constant section. This section is used to store the constant pool of - a virtual machine. - - Primitive name section. This section is introduced to accommodate - the list of primitive operator names that will be invoked by the - virtual machine. - - Code section. The VM functions, including bytecode, are sitting in - this section. - - Examples - -------- - - .. code-block:: python - - import numpy as np - import tvm - from tvm import relay - # define a simple network. - x = relay.var('x', shape=(10, 10)) - f = relay.Function([x], x + x) - mod = relay.Module({"main": f}) - # create a Relay VM. - ctx = tvm.cpu() - target = "llvm" - executable = relay.vm.compile(mod, target) - code, lib = executable.save() - # save and load the code and lib file. - tmp = tvm.contrib.util.tempdir() - path_lib = tmp.relpath("lib.so") - lib.export_library(path_lib) - with open(tmp.relpath("code.ro"), "wb") as fo: - fo.write(code) - loaded_lib = tvm.runtime.load_module(path_lib) - loaded_code = bytearray(open(tmp.relpath("code.ro"), "rb").read()) - # deserialize. - des_exec = relay.vm.Executable.load_exec(loaded_code, loaded_code) - # execute the deserialized executable. - x_data = np.random.rand(10, 10).astype('float32') - des_vm = relay.vm.VirtualMachine(des_exec) - des_vm.init(ctx) - res = des_vm.run(x_data) - print(res.asnumpy()) - """ - return self._save(), self._get_lib() - - @staticmethod - def load_exec(bytecode, lib): - """Construct an executable from saved artifacts. - - Parameters - ---------- - bytecode : bytearray - The binary blob representing a the Relay VM bytecode. - - lib : :py:class:`~tvm.runtime.Module` - The runtime module that contains the generated code. - - Returns - ------- - exec: Executable - An executable constructed using the provided artifacts. - """ - if isinstance(bytecode, (bytes, str)): - code = bytearray(bytecode) - elif not isinstance(bytecode, (bytearray, TVMByteArray)): - raise TypeError("bytecode is expected to be the type of bytearray " + - "or TVMByteArray, but received {}".format(type(code))) - - if lib is not None and not isinstance(lib, tvm.runtime.Module): - raise TypeError("lib is expected to be the type of tvm.runtime.Module" + - ", but received {}".format(type(lib))) - - return Executable(_vm.Load_Executable(bytecode, lib)) - - @property - def lib(self): - """Get the library that contains hardware dependent code. - - Returns - ------- - ret : :py:class:`~tvm.Module` - The runtime module that contains hardware dependent code. - """ - return self._get_lib() - - @property - def stats(self): - """Get the statistics of the Relay VM executable. - - Returns - ------- - ret : String - The statistic information of the VM executable. - """ - return self._get_stats() - - @property - def primitive_ops(self): - """Get the name of the primitive ops contained in the executable. - - Returns - ------- - ret : List[String] - The list of primitive ops. - """ - ret = [] - num_primitives = _vm.GetNumOfPrimitives(self.module) - for i in range(num_primitives): - ret.append(_vm.GetPrimitiveFields(self.module, i)) - return ret - - @property - def bytecode(self): - """Get the bytecode of the Relay VM executable. - - Returns - ------- - ret : String - The bytecode of the executable. - - Notes - ----- - The bytecode is in the following format: - func_name reg_file_size num_instructions - param1 param2 ... paramM - instruction1 - instruction2 - ... - instructionN - - Each instruction is printed in the following format: - hash opcode field1 ... fieldX # The text format. - - The part starting from # is only used for visualization and debugging. - The real serialized code doesn't contain it, therefore the deserializer - doesn't need to deal with it as well. - """ - return self._get_bytecode() - - @property - def globals(self): - """Get the globals used by the Relay VM executable. - - Returns - ------- - ret : List[String] - The globals contained in the executable. - """ - ret = [] - num_globals = _vm.GetNumOfGlobals(self.module) - for i in range(num_globals): - ret.append(_vm.GetGlobalFields(self.module, i)) - return ret - - @property - def module(self): - """Return the runtime module contained in a virtual machine executable.""" - return self.mod - - def get_function_params(self, func_name): - """Get VM Function parameters""" - if func_name in self._function_params: - return self._function_params[func_name] - arity = self._get_function_arity(func_name) - assert arity >= 0 - params = [] - for i in range(arity): - p = self._get_function_param_name(func_name, i) - assert p - params.append(p) - self._function_params[func_name] = params - return params - - -class VirtualMachine(object): - """Relay VM runtime.""" - def __init__(self, mod): - if not isinstance(mod, (Executable, tvm.runtime.Module)): - raise TypeError("mod is expected to be the type of Executable or " + - "tvm.Module, but received {}".format(type(mod))) - m = mod.module if isinstance(mod, Executable) else mod - self.mod = _vm._VirtualMachine(m) - self._exec = mod - self._init = self.mod["init"] - self._invoke = self.mod["invoke"] - self._set_input = self.mod["set_input"] - - def init(self, ctx): - """Initialize the context in the VM. - - Parameters - ---------- - ctx : :py:class:`TVMContext` - The runtime context to run the code on. - """ - args = [ctx.device_type, ctx.device_id] - self._init(*args) - - def set_input(self, func_name, *args, **kwargs): - """Set the input to a function. - - Parameters - ---------- - func_name : str - The name of the function. - - args : list[NDArray] or list[np.ndarray] - The arguments to the function. - - kwargs: dict of str to NDArray or np.ndarray - Named arguments to the function. - """ - if kwargs: - func_params = self._exec.get_function_params(func_name) - new_args = [None] * len(func_params) - assert len(args) + len(kwargs) == len(func_params) - for k in kwargs: - idx = func_params.index(k) - new_args[idx] = kwargs[k] - idx = 0 - for i, arg in enumerate(new_args): - if arg is None: - new_args[i] = args[idx] - idx += 1 - args = new_args - cargs = convert(args) - self._set_input(func_name, *cargs) - - def invoke(self, func_name, *args, **kwargs): - """Invoke a function. - - Parameters - ---------- - func_name : str - The name of the function. - - args : list[NDArray] or list[np.ndarray] - The arguments to the function. - - kwargs: dict of str to NDArray or np.ndarray - Named arguments to the function. - - Returns - ------- - result : Object - The output. - """ - if args or kwargs: - self.set_input(func_name, *args, **kwargs) - return self._invoke(func_name) - - def run(self, *args, **kwargs): - """Run the main function. - - Parameters - ---------- - args : list[NDArray] or list[np.ndarray] - The arguments to the function. - - kwargs: dict of str to NDArray or np.ndarray - Named arguments to the function. - - Returns - ------- - result : Object - The output. - """ - return self.invoke("main", *args, **kwargs) def compile(mod, target=None, target_host=None, params=None): @@ -391,7 +59,7 @@ def compile(mod, target=None, target_host=None, params=None): Returns ------- - exec : Executable + exec : tvm.runtime.vm.Executable The VM executable that contains both library code and bytecode. """ compiler = VMCompiler() @@ -501,10 +169,10 @@ def get_exec(self): Returns ------- - exec : Executable + exec : tvm.runtime.vm.Executable The VM executable that contains both library code and bytecode. """ - return Executable(self._get_exec()) + return vm_rt.Executable(self._get_exec()) def _update_target(self, target): """Update target.""" @@ -549,6 +217,7 @@ def _tophub_context(self, target): tophub_context = autotvm.util.EmptyContext() return tophub_context + class VMExecutor(Executor): """ An implementation of the executor interface for @@ -556,7 +225,7 @@ class VMExecutor(Executor): Useful interface for experimentation and debugging the VM can also be used directly from the API. - supported by `tvm.relay.vm`. + supported by `tvm.runtime.vm`. Parameters ---------- @@ -576,7 +245,7 @@ def __init__(self, mod, ctx, target): self.ctx = ctx self.target = target self.executable = compile(mod, target) - self.vm = VirtualMachine(self.executable) + self.vm = vm_rt.VirtualMachine(self.executable) self.vm.init(ctx) def _make_executor(self, expr=None): diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index 72b835dddee7..e2825f815b67 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -32,15 +32,15 @@ # import numpy # import tvm # from tvm import relay -# from tvm import import container as _container # from tvm import nd +# from tvm.runtime import import container as _container # from tvm.relay.backend.interpreter import RefValue, ConstructorValue PROLOGUE = [ ast.Import([alias('numpy', None)]), ast.Import([alias('tvm', None)]), ast.ImportFrom('tvm', [alias('relay', None)], 0), ast.ImportFrom('tvm', [alias('nd', None)], 0), - ast.ImportFrom('tvm', [alias('container', '_container')], + ast.ImportFrom('tvm.runtime', [alias('container', '_container')], 0), ast.ImportFrom('tvm.relay.backend.interpreter', [alias('RefValue', None), diff --git a/python/tvm/runtime/container.py b/python/tvm/runtime/container.py index 196291fd325a..02a082affc77 100644 --- a/python/tvm/runtime/container.py +++ b/python/tvm/runtime/container.py @@ -15,6 +15,9 @@ # specific language governing permissions and limitations # under the License. """Runtime container structures.""" +import tvm._ffi + +from tvm.runtime import Object, ObjectTypes def getitem_helper(obj, elem_getter, length, idx): """Helper function to implement a pythonic getitem function. @@ -54,3 +57,56 @@ def getitem_helper(obj, elem_getter, length, idx): if idx < 0: idx += length return elem_getter(obj, idx) + + +@tvm._ffi.register_object("vm.ADT") +class ADT(Object): + """Algebatic data type(ADT) object. + + Parameters + ---------- + tag : int + The tag of ADT. + + fields : list[Object] or tuple[Object] + The source tuple. + """ + def __init__(self, tag, fields): + for f in fields: + assert isinstance(f, ObjectTypes), "Expect object or " \ + "tvm NDArray type, but received : {0}".format(type(f)) + self.__init_handle_by_constructor__(_ADT, tag, *fields) + + @property + def tag(self): + return _GetADTTag(self) + + def __getitem__(self, idx): + return getitem_helper( + self, _GetADTFields, len(self), idx) + + def __len__(self): + return _GetADTSize(self) + + +def tuple_object(fields=None): + """Create a ADT object from source tuple. + + Parameters + ---------- + fields : list[Object] or tuple[Object] + The source tuple. + + Returns + ------- + ret : ADT + The created object. + """ + fields = fields if fields else [] + for f in fields: + assert isinstance(f, ObjectTypes), "Expect object or tvm " \ + "NDArray type, but received : {0}".format(type(f)) + return _Tuple(*fields) + + +tvm._ffi._init_api("tvm.runtime.container") diff --git a/python/tvm/relay/backend/profiler_vm.py b/python/tvm/runtime/profiler_vm.py similarity index 92% rename from python/tvm/relay/backend/profiler_vm.py rename to python/tvm/runtime/profiler_vm.py index fa0326edb95b..11cedd5fa200 100644 --- a/python/tvm/relay/backend/profiler_vm.py +++ b/python/tvm/runtime/profiler_vm.py @@ -20,18 +20,19 @@ Provides extra APIs for profiling vm execution. """ -from . import vm, _vm +from tvm.runtime import _ffi_api +from . import vm def enabled(): """Whether vm profiler is enabled.""" - return hasattr(_vm, "_VirtualMachineDebug") + return hasattr(_ffi_api, "_VirtualMachineDebug") class VirtualMachineProfiler(vm.VirtualMachine): """Relay profile VM runtime.""" def __init__(self, mod): super(VirtualMachineProfiler, self).__init__(mod) m = mod.module if isinstance(mod, vm.Executable) else mod - self.mod = _vm._VirtualMachineDebug(m) + self.mod = _ffi_api._VirtualMachineDebug(m) self._init = self.mod["init"] self._invoke = self.mod["invoke"] self._get_stat = self.mod["get_stat"] diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py new file mode 100644 index 000000000000..211bee32ed3c --- /dev/null +++ b/python/tvm/runtime/vm.py @@ -0,0 +1,365 @@ +# 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 runtime. + +Implements a Python interface to executing the compiled VM object. +""" +import numpy as np + +import tvm +from tvm._ffi.runtime_ctypes import TVMByteArray +from tvm._ffi import base as _base +from .object import Object +from . import _ffi_api, container + +def _convert(arg, cargs): + if isinstance(arg, Object): + cargs.append(arg) + elif isinstance(arg, np.ndarray): + nd_arr = tvm.nd.array(arg, ctx=tvm.cpu(0)) + cargs.append(nd_arr) + elif isinstance(arg, tvm.runtime.NDArray): + cargs.append(arg) + elif isinstance(arg, (tuple, list)): + field_args = [] + for field in arg: + _convert(field, field_args) + cargs.append(container.tuple_object(field_args)) + elif isinstance(arg, (_base.numeric_types, bool)): + dtype = "int32" if isinstance(arg, (int, bool)) else "float32" + value = tvm.nd.array(np.array(arg, dtype=dtype), ctx=tvm.cpu(0)) + cargs.append(value) + else: + raise TypeError("Unsupported type: %s" % (type(arg))) + + +def convert(args): + cargs = [] + for arg in args: + _convert(arg, cargs) + + return cargs + + +class Executable(object): + """Relay VM executable""" + def __init__(self, mod): + self.mod = mod + self._function_params = {} + self._save = self.mod["save"] + self._get_lib = self.mod["get_lib"] + self._get_bytecode = self.mod["get_bytecode"] + self._get_stats = self.mod["get_stats"] + self._get_function_arity = self.mod["get_function_arity"] + self._get_function_param_name = self.mod["get_function_param_name"] + + def save(self): + """Save the Relay VM Executable. + + Returns + ------- + code : bytearray + The binary blob representing a serialized Relay VM executable. It + can then be saved to disk and later deserialized into a new + Executable. + + lib : :py:class:`~tvm.runtime.Module` + The runtime module that contains the generated code. It is + basically a library that is composed of hardware dependent code. + + Notes + ----- + The returned code is organized with the following sections in order. + - Global section. This section contains the globals used by the + virtual machine. + + - Constant section. This section is used to store the constant pool of + a virtual machine. + + - Primitive name section. This section is introduced to accommodate + the list of primitive operator names that will be invoked by the + virtual machine. + + - Code section. The VM functions, including bytecode, are sitting in + this section. + + Examples + -------- + + .. code-block:: python + + import numpy as np + import tvm + from tvm import relay + # define a simple network. + x = relay.var('x', shape=(10, 10)) + f = relay.Function([x], x + x) + mod = relay.Module({"main": f}) + # create a Relay VM. + ctx = tvm.cpu() + target = "llvm" + executable = relay.vm.compile(mod, target) + code, lib = executable.save() + # save and load the code and lib file. + tmp = tvm.contrib.util.tempdir() + path_lib = tmp.relpath("lib.so") + lib.export_library(path_lib) + with open(tmp.relpath("code.ro"), "wb") as fo: + fo.write(code) + loaded_lib = tvm.runtime.load_module(path_lib) + loaded_code = bytearray(open(tmp.relpath("code.ro"), "rb").read()) + # deserialize. + des_exec = tvm.runtime.vm.Executable.load_exec(loaded_code, loaded_code) + # execute the deserialized executable. + x_data = np.random.rand(10, 10).astype('float32') + des_vm = tvm.runtime.vm.VirtualMachine(des_exec) + des_vm.init(ctx) + res = des_vm.run(x_data) + print(res.asnumpy()) + """ + return self._save(), self._get_lib() + + @staticmethod + def load_exec(bytecode, lib): + """Construct an executable from saved artifacts. + + Parameters + ---------- + bytecode : bytearray + The binary blob representing a the Relay VM bytecode. + + lib : :py:class:`~tvm.runtime.Module` + The runtime module that contains the generated code. + + Returns + ------- + exec: Executable + An executable constructed using the provided artifacts. + """ + if isinstance(bytecode, (bytes, str)): + code = bytearray(bytecode) + elif not isinstance(bytecode, (bytearray, TVMByteArray)): + raise TypeError("bytecode is expected to be the type of bytearray " + + "or TVMByteArray, but received {}".format(type(code))) + + if lib is not None and not isinstance(lib, tvm.runtime.Module): + raise TypeError("lib is expected to be the type of tvm.runtime.Module" + + ", but received {}".format(type(lib))) + + return Executable(_ffi_api.Load_Executable(bytecode, lib)) + + @property + def lib(self): + """Get the library that contains hardware dependent code. + + Returns + ------- + ret : :py:class:`~tvm.runtime.Module` + The runtime module that contains hardware dependent code. + """ + return self._get_lib() + + @property + def stats(self): + """Get the statistics of the Relay VM executable. + + Returns + ------- + ret : String + The statistic information of the VM executable. + """ + return self._get_stats() + + @property + def primitive_ops(self): + """Get the name of the primitive ops contained in the executable. + + Returns + ------- + ret : List[String] + The list of primitive ops. + """ + ret = [] + num_primitives = _ffi_api.GetNumOfPrimitives(self.module) + for i in range(num_primitives): + ret.append(_ffi_api.GetPrimitiveFields(self.module, i)) + return ret + + @property + def bytecode(self): + """Get the bytecode of the Relay VM executable. + + Returns + ------- + ret : String + The bytecode of the executable. + + Notes + ----- + The bytecode is in the following format: + func_name reg_file_size num_instructions + + param1 param2 ... paramM + + instruction1 + + instruction2 + + ... + + instructionN + + Each instruction is printed in the following format: + hash opcode field1 ... fieldX # The text format. + + The part starting from # is only used for visualization and debugging. + The real serialized code doesn't contain it, therefore the deserializer + doesn't need to deal with it as well. + """ + return self._get_bytecode() + + @property + def globals(self): + """Get the globals used by the Relay VM executable. + + Returns + ------- + ret : List[String] + The globals contained in the executable. + """ + ret = [] + num_globals = _ffi_api.GetNumOfGlobals(self.module) + for i in range(num_globals): + ret.append(_ffi_api.GetGlobalFields(self.module, i)) + return ret + + @property + def module(self): + """Return the runtime module contained in a virtual machine executable.""" + return self.mod + + def get_function_params(self, func_name): + """Get VM Function parameters""" + if func_name in self._function_params: + return self._function_params[func_name] + arity = self._get_function_arity(func_name) + assert arity >= 0 + params = [] + for i in range(arity): + p = self._get_function_param_name(func_name, i) + assert p + params.append(p) + self._function_params[func_name] = params + return params + + +class VirtualMachine(object): + """Relay VM runtime.""" + def __init__(self, mod): + if not isinstance(mod, (Executable, tvm.runtime.Module)): + raise TypeError("mod is expected to be the type of Executable or " + + "tvm.runtime.Module, but received {}".format(type(mod))) + m = mod.module if isinstance(mod, Executable) else mod + self.mod = _ffi_api._VirtualMachine(m) + self._exec = mod + self._init = self.mod["init"] + self._invoke = self.mod["invoke"] + self._set_input = self.mod["set_input"] + + def init(self, ctx): + """Initialize the context in the VM. + + Parameters + ---------- + ctx : :py:class:`TVMContext` + The runtime context to run the code on. + """ + args = [ctx.device_type, ctx.device_id] + self._init(*args) + + def set_input(self, func_name, *args, **kwargs): + """Set the input to a function. + + Parameters + ---------- + func_name : str + The name of the function. + + args : list[tvm.runtime.NDArray] or list[np.ndarray] + The arguments to the function. + + kwargs: dict of str to tvm.runtime.NDArray or np.ndarray + Named arguments to the function. + """ + if kwargs: + func_params = self._exec.get_function_params(func_name) + new_args = [None] * len(func_params) + assert len(args) + len(kwargs) == len(func_params) + for k in kwargs: + idx = func_params.index(k) + new_args[idx] = kwargs[k] + idx = 0 + for i, arg in enumerate(new_args): + if arg is None: + new_args[i] = args[idx] + idx += 1 + args = new_args + cargs = convert(args) + self._set_input(func_name, *cargs) + + def invoke(self, func_name, *args, **kwargs): + """Invoke a function. + + Parameters + ---------- + func_name : str + The name of the function. + + args : list[tvm.runtime.NDArray] or list[np.ndarray] + The arguments to the function. + + kwargs: dict of str to tvm.runtime.NDArray or np.ndarray + Named arguments to the function. + + Returns + ------- + result : Object + The output. + """ + if args or kwargs: + self.set_input(func_name, *args, **kwargs) + return self._invoke(func_name) + + def run(self, *args, **kwargs): + """Run the main function. + + Parameters + ---------- + args : list[tvm.runtime.NDArray] or list[np.ndarray] + The arguments to the function. + + kwargs: dict of str to tvm.runtime.NDArray or np.ndarray + Named arguments to the function. + + Returns + ------- + result : Object + The output. + """ + return self.invoke("main", *args, **kwargs) diff --git a/src/runtime/container.cc b/src/runtime/container.cc index cd426482d285..f54ae6cf9c23 100644 --- a/src/runtime/container.cc +++ b/src/runtime/container.cc @@ -32,14 +32,14 @@ namespace runtime { using namespace vm; -TVM_REGISTER_GLOBAL("container._GetADTTag") +TVM_REGISTER_GLOBAL("runtime.container._GetADTTag") .set_body([](TVMArgs args, TVMRetValue* rv) { ObjectRef obj = args[0]; const auto& adt = Downcast(obj); *rv = static_cast(adt.tag()); }); -TVM_REGISTER_GLOBAL("container._GetADTSize") +TVM_REGISTER_GLOBAL("runtime.container._GetADTSize") .set_body([](TVMArgs args, TVMRetValue* rv) { ObjectRef obj = args[0]; const auto& adt = Downcast(obj); @@ -47,7 +47,7 @@ TVM_REGISTER_GLOBAL("container._GetADTSize") }); -TVM_REGISTER_GLOBAL("container._GetADTFields") +TVM_REGISTER_GLOBAL("runtime.container._GetADTFields") .set_body([](TVMArgs args, TVMRetValue* rv) { ObjectRef obj = args[0]; int idx = args[1]; @@ -56,7 +56,7 @@ TVM_REGISTER_GLOBAL("container._GetADTFields") *rv = adt[idx]; }); -TVM_REGISTER_GLOBAL("container._Tuple") +TVM_REGISTER_GLOBAL("runtime.container._Tuple") .set_body([](TVMArgs args, TVMRetValue* rv) { std::vector fields; for (auto i = 0; i < args.size(); ++i) { @@ -65,7 +65,7 @@ TVM_REGISTER_GLOBAL("container._Tuple") *rv = ADT::Tuple(fields); }); -TVM_REGISTER_GLOBAL("container._ADT") +TVM_REGISTER_GLOBAL("runtime.container._ADT") .set_body([](TVMArgs args, TVMRetValue* rv) { int itag = args[0]; size_t tag = static_cast(itag); diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc index bd650665e196..c2036da46e09 100644 --- a/src/runtime/vm/executable.cc +++ b/src/runtime/vm/executable.cc @@ -738,7 +738,7 @@ void Executable::LoadCodeSection(dmlc::Stream* strm) { } } -TVM_REGISTER_GLOBAL("relay._vm.GetNumOfGlobals") +TVM_REGISTER_GLOBAL("runtime.GetNumOfGlobals") .set_body([](TVMArgs args, TVMRetValue* rv) { runtime::Module mod = args[0]; const auto* exec = dynamic_cast(mod.operator->()); @@ -746,7 +746,7 @@ TVM_REGISTER_GLOBAL("relay._vm.GetNumOfGlobals") *rv = static_cast(exec->global_map.size()); }); -TVM_REGISTER_GLOBAL("relay._vm.GetGlobalFields") +TVM_REGISTER_GLOBAL("runtime.GetGlobalFields") .set_body([](TVMArgs args, TVMRetValue* rv) { runtime::Module mod = args[0]; const auto* exec = dynamic_cast(mod.operator->()); @@ -763,7 +763,7 @@ TVM_REGISTER_GLOBAL("relay._vm.GetGlobalFields") *rv = globals[idx].first; }); -TVM_REGISTER_GLOBAL("relay._vm.GetNumOfPrimitives") +TVM_REGISTER_GLOBAL("runtime.GetNumOfPrimitives") .set_body([](TVMArgs args, TVMRetValue* rv) { runtime::Module mod = args[0]; const auto* exec = dynamic_cast(mod.operator->()); @@ -772,7 +772,7 @@ TVM_REGISTER_GLOBAL("relay._vm.GetNumOfPrimitives") }); -TVM_REGISTER_GLOBAL("relay._vm.GetPrimitiveFields") +TVM_REGISTER_GLOBAL("runtime.GetPrimitiveFields") .set_body([](TVMArgs args, TVMRetValue* rv) { runtime::Module mod = args[0]; const auto* exec = dynamic_cast(mod.operator->()); @@ -789,7 +789,7 @@ TVM_REGISTER_GLOBAL("relay._vm.GetPrimitiveFields") } }); -TVM_REGISTER_GLOBAL("relay._vm.Load_Executable") +TVM_REGISTER_GLOBAL("runtime.Load_Executable") .set_body_typed([]( std::string code, runtime::Module lib) { diff --git a/src/runtime/vm/profiler/vm.cc b/src/runtime/vm/profiler/vm.cc index 3b7b7aa8e73e..4dac66e50a82 100644 --- a/src/runtime/vm/profiler/vm.cc +++ b/src/runtime/vm/profiler/vm.cc @@ -133,7 +133,7 @@ runtime::Module CreateVirtualMachineDebug(const Executable* exec) { return runtime::Module(vm); } -TVM_REGISTER_GLOBAL("relay._vm._VirtualMachineDebug") +TVM_REGISTER_GLOBAL("runtime._VirtualMachineDebug") .set_body([](TVMArgs args, TVMRetValue* rv) { runtime::Module mod = args[0]; const auto* exec = dynamic_cast(mod.operator->()); diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc index 84a3e26fb7f9..af3a86cac17f 100644 --- a/src/runtime/vm/vm.cc +++ b/src/runtime/vm/vm.cc @@ -1057,7 +1057,7 @@ runtime::Module CreateVirtualMachine(const Executable* exec) { return runtime::Module(vm); } -TVM_REGISTER_GLOBAL("relay._vm._VirtualMachine") +TVM_REGISTER_GLOBAL("runtime._VirtualMachine") .set_body([](TVMArgs args, TVMRetValue* rv) { runtime::Module mod = args[0]; const auto* exec = dynamic_cast(mod.operator->()); diff --git a/tests/python/frontend/tensorflow/test_forward.py b/tests/python/frontend/tensorflow/test_forward.py index 830c3f6ea995..2340bd4e6318 100644 --- a/tests/python/frontend/tensorflow/test_forward.py +++ b/tests/python/frontend/tensorflow/test_forward.py @@ -62,7 +62,7 @@ def convert_to_list(x): def vmobj_to_list(o): if isinstance(o, tvm.nd.NDArray): return [o.asnumpy().tolist()] - elif isinstance(o, tvm.container.ADT): + elif isinstance(o, tvm.runtime.container.ADT): result = [] for f in o: result.extend(vmobj_to_list(f)) diff --git a/tests/python/relay/benchmarking/benchmark_vm.py b/tests/python/relay/benchmarking/benchmark_vm.py index 3513832184ce..55d788756b5c 100644 --- a/tests/python/relay/benchmarking/benchmark_vm.py +++ b/tests/python/relay/benchmarking/benchmark_vm.py @@ -19,7 +19,9 @@ import tvm from tvm.contrib import graph_runtime -from tvm import relay, container +from tvm import relay +from tvm.runtime import container +from tvm.runtime import vm as vm_rt from tvm.relay import testing from tvm.relay import vm @@ -58,7 +60,7 @@ def get_vm_output(mod, data, params, target, ctx, dtype='float32', number=2, repeat=20): with relay.build_config(opt_level=3): exe = vm.compile(mod, target, params=params) - rly_vm = vm.VirtualMachine(exe) + rly_vm = vm_rt.VirtualMachine(exe) rly_vm.init(ctx) result = rly_vm.run(data) diff --git a/tests/python/relay/test_adt.py b/tests/python/relay/test_adt.py index 00c07b928b1a..81594c0d04ed 100644 --- a/tests/python/relay/test_adt.py +++ b/tests/python/relay/test_adt.py @@ -117,7 +117,7 @@ def tree_to_dict(t): def vmobj_to_list(o, dtype="float32"): if isinstance(o, tvm.nd.NDArray): return [o.asnumpy().tolist()] - elif isinstance(o, tvm.container.ADT): + elif isinstance(o, tvm.runtime.container.ADT): if len(o) == 0: tensor_nil = p.get_var("tensor_nil", dtype=dtype) if tensor_nil.tag == o.tag: diff --git a/tests/python/relay/test_backend_interpreter.py b/tests/python/relay/test_backend_interpreter.py index 11a9e05b3e7f..28906f19ea3a 100644 --- a/tests/python/relay/test_backend_interpreter.py +++ b/tests/python/relay/test_backend_interpreter.py @@ -18,7 +18,8 @@ import tvm import tvm.testing from tvm import nd -from tvm import relay, container +from tvm import relay +from tvm.runtime import container from tvm.relay.backend.interpreter import RefValue, ConstructorValue from tvm.relay.scope_builder import ScopeBuilder from tvm.relay import testing, create_executor diff --git a/tests/python/relay/test_external_codegen.py b/tests/python/relay/test_external_codegen.py index 13193fc87e07..17585835c178 100644 --- a/tests/python/relay/test_external_codegen.py +++ b/tests/python/relay/test_external_codegen.py @@ -18,12 +18,12 @@ import os import sys import numpy as np -import pytest import tvm import tvm.relay.testing import tvm.relay.transform from tvm import relay +from tvm import runtime from tvm.contrib import util def check_result(mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", @@ -52,8 +52,8 @@ def check_vm_result(): exe = relay.vm.compile(mod, target=target) code, lib = exe.save() lib = update_lib(lib) - exe = relay.vm.Executable.load_exec(code, lib) - vm = relay.vm.VirtualMachine(exe) + exe = runtime.vm.Executable.load_exec(code, lib) + vm = runtime.vm.VirtualMachine(exe) vm.init(ctx) out = vm.run(**map_inputs) tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol) diff --git a/tests/python/relay/test_pass_partition_graph.py b/tests/python/relay/test_pass_partition_graph.py index 75d3c932f05a..869dba82bea4 100644 --- a/tests/python/relay/test_pass_partition_graph.py +++ b/tests/python/relay/test_pass_partition_graph.py @@ -24,6 +24,7 @@ import tvm.relay.testing import tvm.relay.transform as transform from tvm import relay +from tvm import runtime from tvm.contrib import util from tvm.relay.annotation import compiler_begin, compiler_end from tvm.relay.expr_functor import ExprMutator @@ -182,7 +183,7 @@ def update_lib(lib): lib_name = 'lib.so' lib_path = tmp_path.relpath(lib_name) lib.export_library(lib_path, fcompile=False, **kwargs) - lib = tvm.runtime.load_module(lib_path) + lib = runtime.load_module(lib_path) return lib @@ -191,8 +192,8 @@ def check_vm_result(): exe = relay.vm.compile(mod, target=target, params=params) code, lib = exe.save() lib = update_lib(lib) - exe = relay.vm.Executable.load_exec(code, lib) - vm = relay.vm.VirtualMachine(exe) + exe = runtime.vm.Executable.load_exec(code, lib) + vm = runtime.vm.VirtualMachine(exe) vm.init(ctx) out = vm.run(**map_inputs) tvm.testing.assert_allclose(out.asnumpy(), result, rtol=tol, atol=tol) diff --git a/tests/python/relay/test_py_converter.py b/tests/python/relay/test_py_converter.py index 76aa697a2aab..e46b6d41eeb5 100644 --- a/tests/python/relay/test_py_converter.py +++ b/tests/python/relay/test_py_converter.py @@ -19,7 +19,7 @@ from tvm import relay from tvm.relay.testing import to_python, run_as_python from tvm.relay.prelude import Prelude -from tvm.container import ADT +from tvm.runtime.container import ADT from tvm.relay.backend.interpreter import RefValue, ConstructorValue # helper: uses a dummy let binding to sequence a list diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py index 9ea939ce9c83..cd8a875f5bf7 100644 --- a/tests/python/relay/test_vm.py +++ b/tests/python/relay/test_vm.py @@ -14,16 +14,16 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import os +import numpy as np +import pytest import tvm -import numpy as np +from tvm import runtime from tvm import relay from tvm.relay.scope_builder import ScopeBuilder from tvm.relay.testing.config import ctx_list from tvm.relay.prelude import Prelude from tvm.relay import testing -import pytest def check_result(args, expected_result, mod=None): """ @@ -52,14 +52,14 @@ def veval(f, *args, ctx=tvm.cpu(), target="llvm"): assert isinstance(f, relay.Module), "expected expression or module" mod = f exe = relay.vm.compile(mod, target) - vm = relay.vm.VirtualMachine(exe) + vm = runtime.vm.VirtualMachine(exe) vm.init(ctx) return vm.invoke("main", *args) def vmobj_to_list(o): if isinstance(o, tvm.nd.NDArray): return [o.asnumpy().tolist()] - elif isinstance(o, tvm.container.ADT): + elif isinstance(o, tvm.runtime.container.ADT): result = [] for f in o: result.extend(vmobj_to_list(f)) @@ -573,7 +573,7 @@ def test_add_op_broadcast(): def test_vm_optimize(): mod, params = testing.resnet.get_workload(batch_size=1, num_layers=18) - comp = relay.backend.vm.VMCompiler() + comp = relay.vm.VMCompiler() opt_mod, _ = comp.optimize(mod, "llvm", params) if __name__ == "__main__": diff --git a/tests/python/relay/test_vm_serialization.py b/tests/python/relay/test_vm_serialization.py index 6f4e09a393ff..515baa2ef6ce 100644 --- a/tests/python/relay/test_vm_serialization.py +++ b/tests/python/relay/test_vm_serialization.py @@ -19,9 +19,10 @@ import numpy as np import tvm +from tvm.runtime import vm as _vm +from tvm.relay import vm as rly_vm from tvm import relay from tvm.relay.module import Module as rly_module -from tvm.relay import vm as _vm from tvm.relay.scope_builder import ScopeBuilder from tvm.relay.prelude import Prelude from tvm.contrib import util @@ -31,11 +32,11 @@ def create_exec(f, target="llvm", params=None): if isinstance(f, relay.Expr): mod = relay.Module() mod["main"] = f - executable = _vm.compile(mod, target=target, params=params) + executable = rly_vm.compile(mod, target=target, params=params) return executable else: assert isinstance(f, relay.Module), "expected mod as relay.Module" - executable = _vm.compile(f, target=target, params=params) + executable = rly_vm.compile(f, target=target, params=params) return executable diff --git a/tests/python/unittest/test_container.py b/tests/python/unittest/test_container.py index 7bdab82d7a65..5ed6e04e0b45 100644 --- a/tests/python/unittest/test_container.py +++ b/tests/python/unittest/test_container.py @@ -18,7 +18,7 @@ import numpy as np import tvm from tvm import nd, relay -from tvm import container as _container +from tvm.runtime import container as _container def test_adt_constructor(): diff --git a/tests/python/unittest/test_runtime_vm_profiler.py b/tests/python/unittest/test_runtime_vm_profiler.py index b7bbe2f64c96..849a9ef3f823 100644 --- a/tests/python/unittest/test_runtime_vm_profiler.py +++ b/tests/python/unittest/test_runtime_vm_profiler.py @@ -14,11 +14,10 @@ # KIND, either express or implied. See the License for the # specific language governing permissions and limitations # under the License. -import os -import tvm import numpy as np -import pytest +import tvm +from tvm.runtime import profiler_vm from tvm import relay from tvm.relay.testing import resnet @@ -26,10 +25,10 @@ def test_basic(): mod, params = resnet.get_workload() target = 'llvm' ctx = tvm.cpu() - if not relay.profiler_vm.enabled(): + if not profiler_vm.enabled(): return exe = relay.vm.compile(mod, target, params=params) - vm = relay.profiler_vm.VirtualMachineProfiler(exe) + vm = profiler_vm.VirtualMachineProfiler(exe) vm.init(ctx) data = np.random.rand(1, 3, 224, 224).astype('float32')