[VM] Enhance VM Executable as a Subclass of runtime::Module (#95)

Enhance VM Executable as a Subclass of runtime::Module

apps/relax_examples/mlp.py

@@ -48,8 +48,8 @@ def build_mlp(data, weight):
 
     # build and create vm executor
     target = tvm.target.Target("llvm", host="llvm")
-    ex, lib = relax.vm.build(mod, target)
-    vm = relax.VirtualMachine(ex, tvm.cpu(), mod=lib)
+    ex = relax.vm.build(mod, target)
+    vm = relax.VirtualMachine(ex, tvm.cpu())
 
     # run the mlp model on relax vm
     data = tvm.nd.array(np.random.rand(16, 32).astype(np.float32))

apps/relax_examples/nn_module.py

@@ -48,16 +48,16 @@
         data = nn.Placeholder((n, input_size), name="data")
         output = model(data)
         params = [data] + model.parameters()
-        builder.emit_func_output(output, params=params) 
+        builder.emit_func_output(output, params=params)
 
     # get and print the IRmodule being built
     mod = builder.get()
     print(R.parser.astext(mod))
 
     # build the IRModule and create relax vm
     target = tvm.target.Target("llvm", host="llvm")
-    ex, lib = relax.vm.build(mod, target)
-    vm = relax.VirtualMachine(ex, tvm.cpu(), mod=lib)
+    ex = relax.vm.build(mod, target)
+    vm = relax.VirtualMachine(ex, tvm.cpu())
 
     # init parameters
     params = nn.init_params(mod)

apps/relax_examples/resnet.py

@@ -36,8 +36,8 @@
 
     # build the IRModule and create relax vm
     target = tvm.target.Target("llvm", host="llvm")
-    ex, lib = relax.vm.build(relax_mod, target)
-    vm = relax.VirtualMachine(ex, tvm.cpu(), mod=lib)
+    ex = relax.vm.build(relax_mod, target)
+    vm = relax.VirtualMachine(ex, tvm.cpu())
 
     # init weights and run the model on relax vm
     shape = (1, 3, 224, 224)

include/tvm/relax/vm/exec_builder.h

@@ -19,7 +19,6 @@
 
 /*!
  * \file tvm/relax/vm/exec_builder.h
- * \brief
  */
 #ifndef TVM_RELAX_VM_EXEC_BUILDER_H_
 #define TVM_RELAX_VM_EXEC_BUILDER_H_
@@ -48,8 +47,8 @@ class ExecBuilder;
  */
 class ExecBuilderNode : public Object {
  public:
-  /*! \brief The mutable internal executable node. */
-  ObjectPtr<vm::ExecutableNode> exec;  // mutable
+  /*! \brief The mutable internal executable. */
+  ObjectPtr<vm::Executable> exec;  // mutable
   /*!
    * \brief To annotate the start of a vm function.
    * \param func The function name.
@@ -81,16 +80,17 @@ class ExecBuilderNode : public Object {
   void EmitIf(vm::RegName cond, vm::Index false_offset);
   /*!
    * \brief Emit a constant value to the constant pool.
+   * \param obj The constant value to be emitted
    * \return The index that represents the constant.
    */
   vm::Index EmitConstant(TVMRetValue obj);
   /*!
    * \brief Get the built executable.
    * \return The built executable.
    */
-  vm::Executable Get();
+  ObjectPtr<vm::Executable> Get();
   /*!
-   * \brief Create a ExecBuilder.
+   * \brief Create an ExecBuilder.
    * \return The ExecBuilder.
    */
   TVM_DLL static ExecBuilder Create();
@@ -102,6 +102,11 @@ class ExecBuilderNode : public Object {
   TVM_DECLARE_FINAL_OBJECT_INFO(ExecBuilderNode, Object);
 
  private:
+  /*!
+   * \brief A helper function to check if an executable is legal by checking if registers are used
+   * properly
+   */
+  void CheckExecutable();
   /*!
    * \brief Formalize the executable.
    */

include/tvm/relax/vm/executable.h

@@ -38,8 +38,6 @@ namespace tvm {
 namespace runtime {
 namespace relax_vm {
 
-class Executable;
-
 /*!
  * \brief A representation of a Relax function in the VM.
  *
@@ -63,50 +61,68 @@ struct VMFunction {
  * The executable contains information (e.g. data in different memory regions)
  * to run in a virtual machine.
  */
-class ExecutableNode : public Object {
+class Executable : public runtime::ModuleNode {
  public:
+  /*!
+   * \brief Get a PackedFunc from the 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 ObjectPtr<Object>& sptr_to_self) final;
   /*!
    * \brief Print the detailed statistics of the given code, i.e. number of
-   * globls and constants, etc.
+   * globals and constants, etc.
+   * \return The statistics represented by a string.
    */
   std::string Stats() const;
   /*!
    * \brief Get the i-th instruction from the executable.
+   * \param i The index of the instruction to be fetched.
    * \return The instruction.
    */
   Instruction GetInstruction(Index i) const;
   /*!
    * \brief Set j-th byte data of i-th instruction to val.
+   * \param i The index of the instruction to be updated.
+   * \param j The index of the byte data of the instruction to be updated.
+   * \param val The value to be set
    */
   void SetInstructionData(Index i, Index j, ExecWord val);
   /*!
    * \brief Print the instructions as text format.
+   * \return The text format of the instructions.
    */
   String AsText() const;
   /*!
    * \brief Print the instructions as python program.
+   * \return The python program of the instructions, represented by a string.
    */
   String AsPython() const;
   /*!
    * \brief Write the Executable to the binary stream in serialized form.
    * \param stream The binary stream to save the executable to.
    */
-  void SaveToBinary(dmlc::Stream* stream);
+  void SaveToBinary(dmlc::Stream* stream) final;
   /*!
    * \brief Load Executable from the binary stream in serialized form.
    * \param stream The binary stream that load the executable from.
+   * \return The loaded executable, in the form of a `runtime::Module`.
    */
-  static Executable LoadFromBinary(void* stream);
+  static Module LoadFromBinary(void* stream);
   /*!
-   * \brief Write the Executable to the provided path as a file contianing its serialized content.
-   * \param path The path to write the serialized data to.
+   * \brief Write the Executable to the provided path as a file containing its serialized content.
+   * \param file_name The name of the file to write the serialized data to.
+   * \param format The target format of the saved file.
    */
-  void SaveToFile(const std::string& path);
+  void SaveToFile(const std::string& file_name, const std::string& format) final;
   /*!
    * \brief Load Executable from the file.
-   * \param file_name The file that load the executable from.
+   * \param file_name The path of the file that load the executable from.
+   * \return The loaded executable, in the form of a `runtime::Module`.
    */
-  static Executable LoadFromFile(const std::string& file_name);
+  static Module LoadFromFile(const std::string& file_name);
+
   /*! \brief The virtual machine's function table. */
   std::vector<VMFunction> global_funcs;
   /*! \brief A map from globals (as strings) to their index in the function map. */
@@ -115,7 +131,8 @@ class ExecutableNode : public Object {
   std::vector<TVMRetValue> constants;
   /*! \brief The name of packed functions. */
   std::vector<std::string> func_names;
-  /*! \brief A mapping from the packed function (as string) to the index that
+  /*!
+   * \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> func2idx;
@@ -124,9 +141,9 @@ class ExecutableNode : public Object {
   /*! \brief The byte data of instruction. */
   std::vector<ExecWord> instr_data;
 
-  static constexpr const uint32_t _type_index = TypeIndex::kDynamic;
-  static constexpr const char* _type_key = "relax.Executable";
-  TVM_DECLARE_FINAL_OBJECT_INFO(ExecutableNode, Object);
+  virtual ~Executable() {}
+
+  const char* type_key() const final { return "relax.Executable"; }
 
  private:
   /*!
@@ -171,12 +188,6 @@ class ExecutableNode : public Object {
   void LoadPackedFuncNames(dmlc::Stream* strm);
 };
 
-/*! \brief Reference to Executable. */
-class Executable : public ObjectRef {
- public:
-  TVM_DEFINE_MUTABLE_OBJECT_REF_METHODS(Executable, ObjectRef, ExecutableNode);
-};
-
 }  // namespace relax_vm
 }  // namespace runtime
 }  // namespace tvm

include/tvm/relax/vm/vm.h

@@ -66,8 +66,8 @@ struct VMFrame {
 struct VMState {
   /*! \brief The memory allocators. */
   std::vector<Allocator*> allocators;
-  /*! \brief The loaded module. */
-  runtime::Module mod_;
+  /*! \brief The kernel library. */
+  Optional<runtime::Module> lib;
 };
 
 /*!
@@ -90,11 +90,10 @@ class VirtualMachine : public runtime::ModuleNode {
    */
   void Init(const std::vector<Device>& devices, const std::vector<AllocatorType>& alloc_types);
   /*!
-   * \brief load the executable and module for the virtual machine.
+   * \brief Load the executable for the virtual machine.
    * \param exec The executable.
-   * \param mod The library module.
    */
-  void Load(Executable exec, runtime::Module mod);
+  void LoadExecutable(ObjectPtr<Executable> exec);
   /*!
    * \brief Get a PackedFunc from module.
    *
@@ -117,17 +116,19 @@ class VirtualMachine : public runtime::ModuleNode {
   ~VirtualMachine() final {}
 
   const char* type_key() const final { return "relax.VirtualMachine"; }
-  /*! \brief The state of the virtual machine, which can be referred by
-   *  instructions.
-   */
+
+  /*! \brief The state of the virtual machine, which can be referred by instructions. */
   VMState state;
 
  protected:
-  /*! \brief Push a call frame on to the call stack. */
+  /*!
+   * \brief Push a call frame on to the call stack.
+   * \param ret_pc The program counter to return to.
+   * \param vm_func The function to be pushed to the call stack.
+   */
   void PushFrame(Index ret_pc, const VMFunction& vm_func);
   /*!
    * \brief Pop a frame off the call stack.
-   * \return The number of frames left.
    */
   void PopFrame();
   /*!
@@ -139,7 +140,7 @@ class VirtualMachine : public runtime::ModuleNode {
   /*!
    * \brief Read a VM register.
    * \param reg The register to read from.
-   * \return The read object.
+   * \return The value of the register.
    */
   inline RegType ReadRegister(RegName reg) const;
   /*!
@@ -154,7 +155,7 @@ class VirtualMachine : public runtime::ModuleNode {
 
  private:
   /*! \brief The loaded executable. */
-  Executable exec_;
+  ObjectPtr<Executable> exec_;
   /*! \brief The current stack of call frames. */
   std::vector<VMFrame> frames_;
   /*! \brief The virtual machine PC. */

python/tvm/relax/exec_builder.py

@@ -124,4 +124,4 @@ def emit_if(self, cond, false_offset):
 
     def get(self) -> Executable:
         """return the executable"""
-        return _ffi_api.ExecBuilderGet(self)
+        return Executable(_ffi_api.ExecBuilderGet(self))