[relay][vm] Separate VM runtime with executable

Author: zhiics

include/tvm/runtime/vm.h

@@ -27,6 +27,7 @@
 
 #include <tvm/runtime/object.h>
 #include <tvm/runtime/packed_func.h>
+#include <tvm/runtime/registry.h>
 #include <memory>
 #include <string>
 #include <unordered_map>
@@ -362,15 +363,101 @@ struct VMFrame {
         caller_return_register(0) {}
 };
 
+/*! \brief The executable emitted by the VM compiler.
+ *
+ * The executable contains information (e.g. data in different memory regions)
+ * to create a virtual machine.
+ */
+class Executable : public ModuleNode {
+ public:
+  /*!
+   * \brief Get a PackedFunc from an executable module.
+   *
+   * \param name the name of the function.
+   * \param sptr_to_self The shared_ptr that points to this module node.
+   *
+   * \return PackedFunc or nullptr when it is not available.
+   */
+  PackedFunc GetFunction(const std::string& name,
+                         const std::shared_ptr<ModuleNode>& sptr_to_self) final;
+
+  /*!
+   * \brief Get the serialized form of the `functions` in `vm_`. This is
+   * essentially bytecode serialization.
+   *
+   * \return The serialized vm bytecode.
+   *
+   * \note 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:
+   *   opcode num_fields field1 ... fieldX # The text format.
+   *
+   * The field starting from # is only used for debugging. The serialized code
+   * doesn't contain it, therefore the deserializer doens't need to handle it.
+   */
+  std::string GetBytecode() const;
+
+/*!
+   * \brief Print the detailed statistics of the given code, i.e. number of
+   * globls and constants, etc.
+   */
+  std::string Stats() const;
+
+  /*! \brief Get the `lib` module in an executable. Users have the flexibility to call
+   * `export_library` from the frontend to save the library to disk.
+   *
+   * \return The runtime module that contains the hardwre dependent code.
+   */
+  runtime::Module GetLib() const { return lib; }
+
+  /*!
+   * \brief Set the execution context for the executable.
+   *
+   * \param ctxs The list of TVMContext.
+   */
+  void SetContext(const std::vector<TVMContext>& ctxs);
+
+  /*! \brief Get device context for params.
+   */
+  TVMContext GetParamsContext() const;
+
+  virtual ~Executable() {}
+
+  const char* type_key() const final {
+    return "VMExecutable";
+  }
+
+  /*! \brief The runtime module/library that contains hardware dependent code. */
+  runtime::Module lib;
+  /*! \brief The global constant pool. */
+  std::vector<Object> constants;
+  /*! \brief A map from globals (as strings) to their index in the function map. */
+  std::unordered_map<std::string, Index> global_map;
+  /*! \brief A mapping from the packed function (as string) to the index that
+   * corresponds to the position of the `packed_funcs` list in a `VirtualMachine` object.
+   */
+  std::unordered_map<std::string, Index> primitive_map;
+  /*! \brief The virtual machine's function table. */
+  std::vector<VMFunction> functions;
+
+  /*! \brief The set of TVM contexts the VM is currently executing on. */
+  std::vector<TVMContext> ctxs;
+};
+
 /*! \brief The virtual machine.
  *
  * The virtual machine contains all the current execution state,
- * as well as the global view of functions, the global constant
- * table, the compiled operators.
+ * as well as the executable.
  *
  * The goal is to have a single self-contained object,
  * enabling one to easily pass around VMs, execute them on
- * multiple threads, or serialized them to disk or over the
+ * multiple threads, or serialize them to disk or over the
  * wire.
  */
 class VirtualMachine : public runtime::ModuleNode {
@@ -415,16 +502,10 @@ class VirtualMachine : public runtime::ModuleNode {
     return "VirtualMachine";
   }
 
-  /*! \brief The runtime module/library that contains generated code. */
-  runtime::Module lib;
   /*! \brief The virtual machine's packed function table. */
   std::vector<PackedFunc> packed_funcs;
-  /*! \brief The virtual machine's function table. */
-  std::vector<VMFunction> functions;
   /*! \brief The current stack of call frames. */
   std::vector<VMFrame> frames;
-  /*! \brief The global constant pool. */
-  std::vector<Object> constants;
   /*! \brief The fuction table index of the current function. */
   Index func_index;
   /*! \brief The current pointer to the code section. */
@@ -435,8 +516,8 @@ class VirtualMachine : public runtime::ModuleNode {
   /*! \brief The special return register. */
   Object return_register;
 
-  /*! \brief The set of TVM contexts the VM is currently executing on. */
-  std::vector<TVMContext> ctxs;
+  /*! \brief The executable the VM will operate on. */
+  const Executable* exec;
 
   /*! \brief Push a call frame on to the call stack. */
   void PushFrame(Index arg_count, Index ret_pc, const VMFunction& vm_func);
@@ -478,44 +559,24 @@ class VirtualMachine : public runtime::ModuleNode {
    */
   Object Invoke(const std::string& name, const std::vector<Object>& args);
 
-  VirtualMachine() : functions(), frames(), func_index(0), code(nullptr), pc(0) {}
+  VirtualMachine() : frames(), func_index(0), code(nullptr), pc(0), exec(nullptr) {}
 
-  /*! \brief Initialize the virtual machine for a set of contexts.
-   *  \param contexts The set of TVM contexts.
+  /*! \brief Initialize the virtual machine using an executable.
+   *  \param exec The executable.
    */
-  void Init(const std::vector<TVMContext>& contexts);
+  void Init(const Executable* exec);
 
   /*! \brief Run VM dispatch loop.
    */
   void RunLoop();
 
-  /*! \brief Get device context for params.
-   */
-  TVMContext GetParamsContext() const;
-
-  /*!
-   * \brief Load parameters from the parameter bytearray.
-   * \param params The binary file that contains parameters.
-   */
-  void LoadParams(const std::string& params);
-
-  /*! \brief A map from globals (as strings) to their index in the function map.
-   */
-  std::unordered_map<std::string, Index> global_map;
-
-  /*! \brief A mapping from the packed function (as string) to the index that
-   * corresponds to the position of the `packed_funcs` list.
-   */
-  std::unordered_map<std::string, Index> primitive_map;
-
  private:
   /*! \brief Invoke a global setting up the VM state to execute.
    *
    * This does not begin execution of the VM.
    */
   void InvokeGlobal(const VMFunction& func, const std::vector<Object>& args);
 
-
   /*! \brief The parameter name to data mapping. */
   std::unordered_map<std::string, Object> params_;
 };

python/tvm/relay/backend/deserializer.py

@@ -31,7 +31,7 @@ def _create_deserializer(code, lib):
     Parameters
     ----------
     code : bytearray
-        The serialized virtual machine code.
+        The serialized virtual machine bytecode.
 
     lib : :py:class:`~tvm.module.Module`
         The serialized runtime module/library that contains the hardware
@@ -40,7 +40,7 @@ def _create_deserializer(code, lib):
     Returns
     -------
     ret : Deserializer
-        The created virtual machine deserializer.
+        The created virtual machine executable deserializer.
     """
     if isinstance(code, (bytes, str)):
         code = bytearray(code)
@@ -55,12 +55,12 @@ def _create_deserializer(code, lib):
 
 
 class Deserializer:
-    """Relay VM deserializer.
+    """Relay VM executable deserializer.
 
     Parameters
     ----------
     code : bytearray
-        The serialized virtual machine code.
+        The serialized virtual machine bytecode.
 
     lib : :py:class:`~tvm.module.Module`
         The serialized runtime module/library that contains the hardware
@@ -71,11 +71,11 @@ def __init__(self, code, lib):
         self._deserialize = self.mod["deserialize"]
 
     def deserialize(self):
-        """Deserialize the serialized bytecode into a Relay VM.
+        """Deserialize the serialized bytecode into a Relay VM executable.
 
         Returns
         -------
-        ret : VirtualMachine
-            The deserialized Relay VM.
+        ret : Executable
+            The deserialized Relay VM executable.
         """
-        return rly_vm.VirtualMachine(self._deserialize())
+        return rly_vm.Executable(self._deserialize())

python/tvm/relay/backend/profiler_vm.py

@@ -49,8 +49,8 @@ def compile(mod, target=None, target_host=None, params=None):
 
     Returns
     -------
-    vm : VirtualMachineProfiler
-        The profile VM runtime.
+    exec : Executable
+        The executable with profiling code.
     """
     compiler = VMCompilerProfiler()
     target = compiler.update_target(target)
@@ -60,21 +60,24 @@ def compile(mod, target=None, target_host=None, params=None):
     tophub_context = compiler.tophub_context(target)
     with tophub_context:
         compiler._compile(mod, target, target_host)
-    return VirtualMachineProfiler(compiler._get_vm())
+    return vm.Executable(compiler._get_exec())
 
 class VMCompilerProfiler(vm.VMCompiler):
     """Build Relay module to run on VM runtime."""
     def __init__(self):
         super().__init__()
         self.mod = _vm._VMCompilerProfiler()
         self._compile = self.mod["compile"]
-        self._get_vm = self.mod["get_vm"]
+        self._get_exec = self.mod["get_executable"]
         self._set_params_func = self.mod["set_params"]
 
 class VirtualMachineProfiler(vm.VirtualMachine):
     """Relay profile VM runtime."""
     def __init__(self, mod):
         super().__init__(mod)
+        m = mod.module if isinstance(mod, vm.Executable) else mod
+        self.mod = _vm._VirtualMachineDebug(m)
+        self._invoke = self.mod["invoke"]
         self._get_stat = self.mod["get_stat"]
 
     def get_stat(self):