[GitHub] [tvm] areusch commented on a change in pull request #7746: Add support for using the VM across the RPC boundary.

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

areusch commented on a change in pull request #7746:
URL: https://github.com/apache/tvm/pull/7746#discussion_r602604080



##########
File path: include/tvm/runtime/vm/executable.h
##########
@@ -63,6 +63,14 @@ class Executable : public ModuleNode {
    */
   PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
 
+  /*!
+   * \brief Save the entire executable to a binary stream.
+   * \param stream The binary stream to save to.
+   */
+  void SaveToBinary(dmlc::Stream* stream) final;
+
+  void SaveToFile(const std::string& path, const std::string& format) final;

Review comment:
       docstring

##########
File path: python/tvm/runtime/module.py
##########
@@ -269,24 +269,31 @@ def _collect_dso_modules(self):
         return self._collect_from_import_tree(is_dso_exportable)
 
     def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=None, **kwargs):
-        """Export the module and its imported device code one library.
+        """
+        Export the module and all imported modules into a single device library.
 
-        This function only works on host llvm modules.
-        It will pack all the imported modules
+        This function only works on hos LLVM modules, other runtime::Module
+        subclasses will work with this API but they must support implement
+        the save and load mechanisms of modules completely including saving
+        from streams and files. This will pack your non-shared library module
+        into a single shared library which can later be loaded by TVM.
 
         Parameters
         ----------
         file_name : str
             The name of the shared library.
 
         fcompile : function(target, file_list, kwargs), optional
-            Compilation function to use create dynamic library.
+            The compilation function to use create the final library object during
+            export. For example this is used to link together all produced artifacts

Review comment:
       For example, when fcompile=_cc.create_shared, or when it is not supplied but module is "llvm," 

##########
File path: include/tvm/runtime/vm/executable.h
##########
@@ -125,9 +133,17 @@ class Executable : public ModuleNode {
    * \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.
+   * \return The runtime module that contains the hardware dependent code.
    */
-  runtime::Module GetLib() const { return lib; }
+  runtime::Module GetLib() const { return this->imports_[0]; }
+
+  void SetLib(const runtime::Module& lib) {

Review comment:
       why do you do this rather than just calling Import() from a factory function?

##########
File path: python/tvm/runtime/module.py
##########
@@ -269,24 +269,31 @@ def _collect_dso_modules(self):
         return self._collect_from_import_tree(is_dso_exportable)
 
     def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=None, **kwargs):

Review comment:
       want to add type annotations?

##########
File path: src/runtime/vm/executable.cc
##########
@@ -74,6 +76,12 @@ PackedFunc Executable::GetFunction(const std::string& name, const ObjectPtr<Obje
       int index = args[1];
       *rv = this->GetFunctionParameterName(func_name, index);
     });
+  } else if (name == "vm_load_executable") {

Review comment:
       why is this a member function of Executable rather than a factory function for VM?

##########
File path: include/tvm/runtime/vm/executable.h
##########
@@ -125,9 +133,17 @@ class Executable : public ModuleNode {
    * \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.
+   * \return The runtime module that contains the hardware dependent code.
    */
-  runtime::Module GetLib() const { return lib; }
+  runtime::Module GetLib() const { return this->imports_[0]; }

Review comment:
       explain why it is always imports_[0]

##########
File path: src/relay/backend/vm/compiler.cc
##########
@@ -1155,18 +1155,20 @@ void VMCompiler::Codegen() {
 
   auto compile_engine = CompileEngine::Global();
   auto ext_mods = compile_engine->LowerExternalFunctions();
+  runtime::Module lib;
   if (funcs.size() > 0) {
     Map<String, IRModule> build_funcs;
     for (const auto& i : funcs) {
       build_funcs.Set(i.first, i.second);
     }
-    exec_->lib = tvm::build(build_funcs, target_host_);
+    lib = tvm::build(build_funcs, target_host_);
   } else {
     // There is no function handled by TVM. We create a virtual main module
     // to make sure a DSO module will be also available.
-    exec_->lib = codegen::CSourceModuleCreate(";", "", Array<String>{});
+    lib = codegen::CSourceModuleCreate(";", "", Array<String>{});
   }
-  exec_->lib = codegen::CreateMetadataModule(params_, exec_->lib, ext_mods, target_host_);
+  lib = codegen::CreateMetadataModule(params_, lib, ext_mods, target_host_);
+  exec_->SetLib(lib);

Review comment:
       here I think you can just `exec_->Import(lib)`?

##########
File path: python/tvm/runtime/vm.py
##########
@@ -299,12 +300,16 @@ class VirtualMachine(object):
     POOLED_ALLOCATOR = 2
 
     def __init__(self, exe, device, memory_cfg=None):
-        if not isinstance(exe, Executable):
+        if not isinstance(exe, Executable) and not isinstance(exe, Module):

Review comment:
       you don't need the first condition here

##########
File path: src/runtime/vm/executable.cc
##########
@@ -476,8 +484,19 @@ void LoadHeader(dmlc::Stream* strm) {
 }
 
 runtime::Module Executable::Load(const std::string& code, const runtime::Module lib) {
+  std::cout << "code: " << code.size() << std::endl;

Review comment:
       rm

##########
File path: include/tvm/runtime/vm/executable.h
##########
@@ -63,6 +63,14 @@ class Executable : public ModuleNode {
    */
   PackedFunc GetFunction(const std::string& name, const ObjectPtr<Object>& sptr_to_self) final;
 
+  /*!
+   * \brief Save the entire executable to a binary stream.
+   * \param stream The binary stream to save to.
+   */
+  void SaveToBinary(dmlc::Stream* stream) final;
+
+  void SaveToFile(const std::string& path, const std::string& format) final;

Review comment:
       docstring

##########
File path: src/runtime/library_module.cc
##########
@@ -99,6 +99,28 @@ void InitContextFunctions(std::function<void*(const char*)> fgetsymbol) {
 #undef TVM_INIT_CONTEXT_FUNC
 }
 
+Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {

Review comment:
       add a docstring and explain when this is used

##########
File path: python/tvm/runtime/vm.py
##########
@@ -299,12 +300,16 @@ class VirtualMachine(object):
     POOLED_ALLOCATOR = 2
 
     def __init__(self, exe, device, memory_cfg=None):

Review comment:
       add a docstring, maybe type annotations, and explain in what case exe would be a Module

##########
File path: src/runtime/vm/executable.cc
##########
@@ -476,8 +484,19 @@ void LoadHeader(dmlc::Stream* strm) {
 }
 
 runtime::Module Executable::Load(const std::string& code, const runtime::Module lib) {
+  std::cout << "code: " << code.size() << std::endl;
   auto exec = make_object<Executable>();
-  exec->lib = lib;
+
+  // Support null-initialization of lib, to enable initialization during
+  // deserialization before we have we have deserialized the imports.
+  if (lib.defined()) {

Review comment:
       perhaps it's worth either creating a const bool is_initialized() { ICHECK_LE(this->imports_.size(), 1); return this->imports_.size() == 1; } and adding it everywhere we access lib, or make the change I suggested above in GetLib and audit all uses of lib to ensure we GetLib always.

##########
File path: tests/python/relay/test_vm.py
##########
@@ -799,5 +802,30 @@ def test_constant_shape_with_external_codegen():
     assert "shape_func" in opt_mod.astext(False)
 
 
+def test_vm_rpc():
+    target = "llvm"
+    target_host = "llvm"
+
+    x = relay.var("x", shape=(10, 1))
+    f = relay.Function([x], x + x)
+    mod = IRModule.from_expr(f)
+    vm_exec = vm.compile(mod, target=target, target_host=target_host)
+
+    temp = utils.tempdir()

Review comment:
       can you just add a couple comments explaining what it is you're doing and what you're testing?

##########
File path: python/tvm/runtime/module.py
##########
@@ -269,24 +269,31 @@ def _collect_dso_modules(self):
         return self._collect_from_import_tree(is_dso_exportable)
 
     def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=None, **kwargs):
-        """Export the module and its imported device code one library.
+        """
+        Export the module and all imported modules into a single device library.
 
-        This function only works on host llvm modules.
-        It will pack all the imported modules
+        This function only works on hos LLVM modules, other runtime::Module

Review comment:
       nit: host

##########
File path: include/tvm/runtime/vm/executable.h
##########
@@ -125,9 +133,17 @@ class Executable : public ModuleNode {
    * \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.
+   * \return The runtime module that contains the hardware dependent code.
    */
-  runtime::Module GetLib() const { return lib; }
+  runtime::Module GetLib() const { return this->imports_[0]; }

Review comment:
       also add an ICHECK_EQ(this.imports_.size(), 1) and an explanatory error describing the programming error that causes 0 imports

##########
File path: src/runtime/vm/executable.cc
##########
@@ -476,8 +484,19 @@ void LoadHeader(dmlc::Stream* strm) {
 }
 
 runtime::Module Executable::Load(const std::string& code, const runtime::Module lib) {
+  std::cout << "code: " << code.size() << std::endl;
   auto exec = make_object<Executable>();
-  exec->lib = lib;
+
+  // Support null-initialization of lib, to enable initialization during
+  // deserialization before we have we have deserialized the imports.
+  if (lib.defined()) {
+    ICHECK_EQ(exec->imports_.size(), 0)

Review comment:
       maybe this should move to GetLib. Since you just made exec, it should not have imports?




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