[GitHub] [incubator-tvm] jroesch commented on a change in pull request #6105: [Relay][VM] Allow to config allocator type and refactor vm code structure

[GitBox <gi...@apache.org>]

jroesch commented on a change in pull request #6105:
URL: https://github.com/apache/incubator-tvm/pull/6105#discussion_r458515786



##########
File path: python/tvm/runtime/vm.py
##########
@@ -273,29 +272,58 @@ def get_function_params(self, func_name):
 
 
 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)
+    """Relay VM runtime.
+
+    Parameters
+    ----------
+    exe : Executable
+        The VM executable.
+
+    ctx : tvm.runtime.TVMContext or List[tvm.runtime.TVMContext]
+        The context to deploy the module
+
+    memory_cfg : str or Dict[tvm.runtime.TVMContext, str], optional
+        Config the type of memory allocator. The allocator type can be ["naive",
+        "pooled"]. If memory_cfg is None, all contexts will use pooled allocator
+        by default. If memory_cfg is string, all contexts will use the specified
+        allocator type. If memory_cfg is a dict, each context uses the allocator
+        type specified in the dict, or pooled allocator if not specified in the
+        dict.
+    """
+
+    NAIVE_ALLOCATOR = 1

Review comment:
       Should we make this a C++ enum? 




----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

For queries about this service, please contact Infrastructure at:
users@infra.apache.org