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

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

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



##########
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
+    POOLED_ALLOCATOR = 2
+
+    def __init__(self, exe, ctx, memory_cfg=None):
+        if not isinstance(exe, Executable):
+            raise TypeError("mod is expected to be the type of Executable, " +
+                            "but received {}".format(type(mod)))
+        self.module = _ffi_api._VirtualMachine(exe.module)
+        self._exec = exe
+        self._init = self.module["init"]
+        self._invoke = self.module["invoke"]
+        self._set_input = self.module["set_input"]
+        self._setup_ctx(ctx, memory_cfg)
+
+    def _setup_ctx(self, ctx, memory_cfg):
+        """Init context and allocators."""
+        if isinstance(ctx, tvm.runtime.TVMContext):
+            ctx = [ctx]
+        default_alloc_type = VirtualMachine.POOLED_ALLOCATOR
+        if memory_cfg is None:
+            memory_cfg = {}
+        elif isinstance(memory_cfg, str):
+            assert memory_cfg in ["naive", "pooled"]
+            if memory_cfg == "naive":
+                default_alloc_type = VirtualMachine.NAIVE_ALLOCATOR
+            memory_cfg = {}

Review comment:
       Just in case of wrong input type, we probably want to add `assert isinstance(memory_cfg, dict)` before init_args? 




----------------------------------------------------------------
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