From 64ea15b01e6538037d6fc0b1d07bf93460a7e3ed Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Thu, 25 Mar 2021 16:45:06 -0700
Subject: [PATCH 01/10] Get basic verison of VM RPC working

---
 include/tvm/runtime/vm/executable.h | 10 +++-
 python/tvm/runtime/module.py        |  9 ++-
 src/relay/backend/vm/compiler.cc    |  1 +
 src/runtime/library_module.cc       | 43 ++++++++-------
 src/runtime/library_module.h        |  3 +
 src/runtime/vm/executable.cc        | 86 +++++++++++++++++++++++++++++
 src/runtime/vm/vm.cc                | 19 ++++++-
 7 files changed, 145 insertions(+), 26 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 8d3f651758d1..1b300b5a4bac 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/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;
+
   /*!
    * \brief Serialize the executable into global section, constant section, and
    * code section.
@@ -125,7 +133,7 @@ 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; }
 
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 41d60683aa3b..8c321ae33ecc 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -269,10 +269,13 @@ 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 host LLVM modules, other runtime::Module
+        subclasses DO NOT work with this API. If you do in fact have an LLVM
+        module, this API will pack the module with all imported modules into
+        a single binary library which can be used with TVM.
 
         Parameters
         ----------
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index dafaed111c03..8cd6f642aff1 100644
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -1167,6 +1167,7 @@ void VMCompiler::Codegen() {
     exec_->lib = codegen::CSourceModuleCreate(";", "", Array<String>{});
   }
   exec_->lib = codegen::CreateMetadataModule(params_, exec_->lib, ext_mods, target_host_);
+  exec_->Import(exec_->lib);
 }
 
 ExprDeviceMap VMCompiler::AnalyzeContext() const {
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index 30ef2141c508..9699a941bd57 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -99,6 +99,29 @@ void InitContextFunctions(std::function<void*(const char*)> fgetsymbol) {
 #undef TVM_INIT_CONTEXT_FUNC
 }
 
+Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {
+  std::string loadkey = "runtime.module.loadbinary_";
+  std::string fkey = loadkey + type_key;
+  const PackedFunc* f = Registry::Get(fkey);
+  if (f == nullptr) {
+    std::string loaders = "";
+    for (auto name : Registry::ListNames()) {
+      if (name.rfind(loadkey, 0) == 0) {
+        if (loaders.size() > 0) {
+          loaders += ", ";
+        }
+        loaders += name.substr(loadkey.size());
+      }
+    }
+    ICHECK(f != nullptr)
+        << "Binary was created using " << type_key
+        << " but a loader of that name is not registered. Available loaders are " << loaders
+        << ". Perhaps you need to recompile with this runtime enabled.";
+  }
+
+  return (*f)(static_cast<void*>(stream));
+}
+
 /*!
  * \brief Load and append module blob to module list
  * \param mblob The module blob.
@@ -133,25 +156,7 @@ runtime::Module ProcessModuleBlob(const char* mblob, ObjectPtr<Library> lib) {
       ICHECK(stream->Read(&import_tree_row_ptr));
       ICHECK(stream->Read(&import_tree_child_indices));
     } else {
-      std::string loadkey = "runtime.module.loadbinary_";
-      std::string fkey = loadkey + tkey;
-      const PackedFunc* f = Registry::Get(fkey);
-      if (f == nullptr) {
-        std::string loaders = "";
-        for (auto name : Registry::ListNames()) {
-          if (name.rfind(loadkey, 0) == 0) {
-            if (loaders.size() > 0) {
-              loaders += ", ";
-            }
-            loaders += name.substr(loadkey.size());
-          }
-        }
-        ICHECK(f != nullptr)
-            << "Binary was created using " << tkey
-            << " but a loader of that name is not registered. Available loaders are " << loaders
-            << ". Perhaps you need to recompile with this runtime enabled.";
-      }
-      Module m = (*f)(static_cast<void*>(stream));
+      auto m = LoadModuleFromBinary(tkey, stream);
       modules.emplace_back(m);
     }
   }
diff --git a/src/runtime/library_module.h b/src/runtime/library_module.h
index 91918c1ccaa3..75bd287ed49d 100644
--- a/src/runtime/library_module.h
+++ b/src/runtime/library_module.h
@@ -32,6 +32,9 @@
 
 namespace tvm {
 namespace runtime {
+
+Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream);
+
 /*!
  * \brief Library is the common interface
  *  for storing data in the form of shared libaries.
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index 6992097e8d69..3246f197f8af 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -38,6 +38,8 @@
 #include <vector>
 
 #include "serialize_utils.h"
+#include "../library_module.h"
+#include "../file_utils.h"
 
 namespace tvm {
 namespace runtime {
@@ -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 == "create_vm") {
+    return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
+      auto vm = make_object<VirtualMachine>();
+      vm->LoadExecutable(this);
+      *rv = Module(vm);
+    });
   } else {
     LOG(FATAL) << "Unknown packed function: " << name;
     return PackedFunc(nullptr);
@@ -476,11 +484,17 @@ 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;
   exec->code_ = code;
   dmlc::MemoryStringStream strm(&exec->code_);
 
+  if (lib.defined()) {
+    std::cout << "Importing: " << std::endl;
+    exec->Import(lib);
+  }
+
   // Load header.
   LoadHeader(&strm);
 
@@ -765,6 +779,78 @@ void Executable::LoadCodeSection(dmlc::Stream* strm) {
   }
 }
 
+void Executable::SaveToBinary(dmlc::Stream* stream) {
+  auto code_bytes = this->Save();
+  std::string code(code_bytes.data, code_bytes.size);
+  stream->Write(code);
+
+  CHECK(this->lib.defined())
+    << "the library must be defined before serialization";
+
+  // this->lib->SaveToBinary(stream);
+  // std::vector<std::string> names;
+  // std::vector<DLTensor*> arrays;
+  // for (const auto& v : params_) {
+  //   names.emplace_back(v.first);
+  //   arrays.emplace_back(const_cast<DLTensor*>(v.second.operator->()));
+  // }
+  // uint64_t sz = arrays.size();
+  // ICHECK(sz == names.size());
+  // stream->Write(sz);
+  // stream->Write(names);
+  // for (size_t i = 0; i < sz; ++i) {
+  //   tvm::runtime::SaveDLTensor(stream, arrays[i]);
+  // }
+}
+
+Module ExecutableLoadBinary(void* strm) {
+  dmlc::Stream* stream = static_cast<dmlc::Stream*>(strm);
+  std::string code;
+  stream->Read(&code);
+  auto exec = Executable::Load(code, Module());
+  auto exec_node = exec.as<Executable>();
+  std::cout << exec_node->primitive_map.size() << std::endl;
+
+  // // std::unordered_map<std::string, tvm::runtime::NDArray> params;
+  // std::string module_name;
+  // ICHECK(stream->Read(&graph_json));
+  // uint64_t sz;
+  // ICHECK(stream->Read(&sz));
+  // std::vector<std::string> names;
+  // ICHECK(stream->Read(&names));
+  // ICHECK(sz == names.size());
+  // for (size_t i = 0; i < sz; ++i) {
+  //   tvm::runtime::NDArray temp;
+  //   temp.Load(stream);
+  //   params[names[i]] = temp;
+  // }
+
+  return exec;
+}
+
+void Executable::SaveToFile(const std::string& path, const std::string& format) {
+  std::string data;
+  dmlc::MemoryStringStream writer(&data);
+  dmlc::SeekStream* strm = &writer;
+  SaveToBinary(strm);
+  SaveBinaryToFile(path, data);
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_VMExecutable")
+    .set_body_typed(ExecutableLoadBinary);
+
+  // Load module from module.
+Module ExecutableLoadFile(const std::string& file_name, const std::string& format) {
+  std::string data;
+  LoadBinaryFromFile(file_name, &data);
+  dmlc::MemoryStringStream reader(&data);
+  dmlc::Stream* strm = &reader;
+  auto exec = ExecutableLoadBinary(reinterpret_cast<void*>(strm));
+  return exec;
+}
+
+TVM_REGISTER_GLOBAL("runtime.module.loadfile_VMExecutable").set_body_typed(ExecutableLoadFile);
+
 TVM_REGISTER_GLOBAL("runtime.GetNumOfGlobals").set_body([](TVMArgs args, TVMRetValue* rv) {
   runtime::Module mod = args[0];
   const auto* exec = dynamic_cast<Executable*>(mod.operator->());
diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc
index ee06da83bd92..04dd2eb47601 100644
--- a/src/runtime/vm/vm.cc
+++ b/src/runtime/vm/vm.cc
@@ -281,11 +281,24 @@ void VirtualMachine::LoadExecutable(const Executable* exec) {
   ICHECK(exec) << "The executable is not created yet.";
   exec_ = exec;
 
-  runtime::Module lib = exec_->lib;
+  runtime::Module lib;
+  if (exec_->lib.defined()) {
+    lib = exec_->lib;
+  } else {
+    ICHECK(exec_->imports().size() > 0)
+      << "fix";
+    lib = exec_->imports()[0];
+  }
+
   // Get the list of packed functions.
-  ICHECK(exec->primitive_map.empty() || lib.operator->())
-      << "runtime module should have been built for primitive functions"
+  ICHECK(!exec->primitive_map.empty())
+      << "runtime module primitive map is empty"
       << "\n";
+
+  ICHECK(lib.operator->())
+      << "library is null"
+      << "\n";
+
   for (const auto& it : exec_->primitive_map) {
     const auto& packed_name = it.first;
     auto packed_index = static_cast<size_t>(it.second);

From 9896d619df889a7bfd40570fe51ad43e5d3b106f Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Thu, 25 Mar 2021 17:44:30 -0700
Subject: [PATCH 02/10] Test case passes

---
 python/tvm/runtime/vm.py      |  9 +++++++--
 tests/python/relay/test_vm.py | 28 +++++++++++++++++++++++++++-
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index a503da53c465..31bb3662e8db 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -23,6 +23,7 @@
 import numpy as np
 
 import tvm
+from tvm.runtime import Module
 from tvm._ffi.runtime_ctypes import TVMByteArray
 from tvm._ffi import base as _base
 from .object import Object
@@ -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):
             raise TypeError(
                 "exe is expected to be the type of Executable, "
                 + "but received {}".format(type(exe))
             )
-        self.module = _ffi_api._VirtualMachine(exe.module)
+
+        if not isinstance(exe, Executable):
+            exe = Executable(exe)
+
+        self.module = exe.mod["create_vm"]()
         self._exec = exe
         self._init = self.module["init"]
         self._invoke = self.module["invoke"]
diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py
index 4ecd0d9189ea..f782e4323071 100644
--- a/tests/python/relay/test_vm.py
+++ b/tests/python/relay/test_vm.py
@@ -19,11 +19,14 @@
 
 import tvm
 from tvm import runtime
-from tvm import relay
+from tvm import relay, IRModule
+from tvm.relay.backend import vm
 from tvm.relay.scope_builder import ScopeBuilder
 from tvm.relay.prelude import Prelude
 from tvm.relay.loops import while_loop
 from tvm.relay import testing
+from tvm.contrib import utils
+from tvm import rpc
 import tvm.testing
 
 
@@ -798,6 +801,29 @@ def test_constant_shape_with_external_codegen():
     opt_mod, _ = comp.optimize(mod, target="llvm")
     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()
+    path = temp.relpath("vm_library.so")
+    vm_exec.mod.export_library(path)
+
+    remote = rpc.LocalSession()
+    remote.upload(path)
+    rexec = remote.load_module("vm_library.so")
+
+    ctx = remote.cpu()
+    vm_factory = runtime.vm.VirtualMachine(rexec, ctx)
+    np_input = np.random.uniform(size=(10, 1)).astype("float32")
+    input_tensor = tvm.nd.array(np_input, ctx)
+    out = vm_factory.invoke("main", [input_tensor])
+    np.testing.assert_allclose(out.asnumpy(), np_input + np_input)
 
 if __name__ == "__main__":
     pytest.main([__file__])

From 94f42019bbd7cda5c65b38551f1370269e8c0e85 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Fri, 26 Mar 2021 14:10:18 -0700
Subject: [PATCH 03/10] Clean up PR

---
 include/tvm/runtime/vm/executable.h | 15 ++++++---
 python/tvm/runtime/module.py        | 18 ++++++----
 python/tvm/runtime/vm.py            |  2 +-
 src/relay/backend/vm/compiler.cc    |  9 ++---
 src/runtime/library_module.h        |  8 +++++
 src/runtime/vm/executable.cc        | 51 +++++++----------------------
 src/runtime/vm/vm.cc                | 20 +++--------
 7 files changed, 52 insertions(+), 71 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 1b300b5a4bac..2e8fd27e620e 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -135,7 +135,17 @@ class Executable : public ModuleNode {
    *
    * \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) {
+    ICHECK(lib.defined())
+      << "library can not be null";
+
+    ICHECK_EQ(this->imports().size(), 0)
+      << "can only import the library once";
+
+    this->Import(lib);
+  }
 
   /*!
    * \brief Get the arity of the VM Fucntion.
@@ -156,9 +166,6 @@ class Executable : public ModuleNode {
 
   const char* type_key() const final { return "VMExecutable"; }
 
-  /*! \brief The runtime module/library that contains both the host and also the device
-   * code when executing on non-CPU devices. */
-  runtime::Module lib;
   /*! \brief The global constant pool. */
   std::vector<ObjectRef> constants;
   /*! \brief A map from globals (as strings) to their index in the function map. */
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 8c321ae33ecc..0f339d23f729 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -272,10 +272,11 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
         """
         Export the module and all imported modules into a single device library.
 
-        This function only works on host LLVM modules, other runtime::Module
-        subclasses DO NOT work with this API. If you do in fact have an LLVM
-        module, this API will pack the module with all imported modules into
-        a single binary library which can be used with TVM.
+        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
         ----------
@@ -283,13 +284,16 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
             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
+            into a final dynamic library.
+            This behavior is controlled by the type of object exported.
             If fcompile has attribute object_format, will compile host library
             to that format. Otherwise, will use default format "o".
 
         workspace_dir : str, optional
-            the path to a directory used to create intermediary
-            artifacts for the process exporting of the library.
+            The path of the directory used to create the intermediate
+            artifacts when exporting the module.
             If this is not provided a temporary dir will be created.
 
         kwargs : dict, optional
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index 31bb3662e8db..1fdb0fa0f777 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -309,7 +309,7 @@ def __init__(self, exe, device, memory_cfg=None):
         if not isinstance(exe, Executable):
             exe = Executable(exe)
 
-        self.module = exe.mod["create_vm"]()
+        self.module = exe.mod["vm_load_executable"]()
         self._exec = exe
         self._init = self.module["init"]
         self._invoke = self.module["invoke"]
diff --git a/src/relay/backend/vm/compiler.cc b/src/relay/backend/vm/compiler.cc
index 8cd6f642aff1..906250c1bb0d 100644
--- a/src/relay/backend/vm/compiler.cc
+++ b/src/relay/backend/vm/compiler.cc
@@ -1155,19 +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_);
-  exec_->Import(exec_->lib);
+  lib = codegen::CreateMetadataModule(params_, lib, ext_mods, target_host_);
+  exec_->SetLib(lib);
 }
 
 ExprDeviceMap VMCompiler::AnalyzeContext() const {
diff --git a/src/runtime/library_module.h b/src/runtime/library_module.h
index 75bd287ed49d..7b226a16dc95 100644
--- a/src/runtime/library_module.h
+++ b/src/runtime/library_module.h
@@ -33,6 +33,14 @@
 namespace tvm {
 namespace runtime {
 
+/*! \brief Load a module with the given type key directly from the stream.
+ *  This function wraps the registry mechanism used to store type based deserializers
+ *  for each runtime::Module sub-class.
+ *
+ * \param type_key The type key of the serialized module.
+ * \param stream A pointer to the stream containing the serialized module.
+ * \return module The deserialized module.
+*/
 Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream);
 
 /*!
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index 3246f197f8af..f4f914043862 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -76,7 +76,7 @@ PackedFunc Executable::GetFunction(const std::string& name, const ObjectPtr<Obje
       int index = args[1];
       *rv = this->GetFunctionParameterName(func_name, index);
     });
-  } else if (name == "create_vm") {
+  } else if (name == "vm_load_executable") {
     return PackedFunc([sptr_to_self, this](TVMArgs args, TVMRetValue* rv) {
       auto vm = make_object<VirtualMachine>();
       vm->LoadExecutable(this);
@@ -486,15 +486,20 @@ 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;
-  exec->code_ = code;
-  dmlc::MemoryStringStream strm(&exec->code_);
 
+  // Support null-initialization of lib, to enable initialization during
+  // deserialization before we have we have deserialized the imports.
   if (lib.defined()) {
-    std::cout << "Importing: " << std::endl;
+    ICHECK_EQ(exec->imports_.size(), 0)
+      << "A VMExecutable should never have more than one import inside an the executable, \n"
+      << "the first import should *always* be the library containing"
+      << "the platform specific kernel code";
     exec->Import(lib);
   }
 
+  exec->code_ = code;
+  dmlc::MemoryStringStream strm(&exec->code_);
+
   // Load header.
   LoadHeader(&strm);
 
@@ -784,23 +789,8 @@ void Executable::SaveToBinary(dmlc::Stream* stream) {
   std::string code(code_bytes.data, code_bytes.size);
   stream->Write(code);
 
-  CHECK(this->lib.defined())
-    << "the library must be defined before serialization";
-
-  // this->lib->SaveToBinary(stream);
-  // std::vector<std::string> names;
-  // std::vector<DLTensor*> arrays;
-  // for (const auto& v : params_) {
-  //   names.emplace_back(v.first);
-  //   arrays.emplace_back(const_cast<DLTensor*>(v.second.operator->()));
-  // }
-  // uint64_t sz = arrays.size();
-  // ICHECK(sz == names.size());
-  // stream->Write(sz);
-  // stream->Write(names);
-  // for (size_t i = 0; i < sz; ++i) {
-  //   tvm::runtime::SaveDLTensor(stream, arrays[i]);
-  // }
+  ICHECK(this->imports()[0].defined())
+    << "the library must be imported before serialization";
 }
 
 Module ExecutableLoadBinary(void* strm) {
@@ -808,23 +798,6 @@ Module ExecutableLoadBinary(void* strm) {
   std::string code;
   stream->Read(&code);
   auto exec = Executable::Load(code, Module());
-  auto exec_node = exec.as<Executable>();
-  std::cout << exec_node->primitive_map.size() << std::endl;
-
-  // // std::unordered_map<std::string, tvm::runtime::NDArray> params;
-  // std::string module_name;
-  // ICHECK(stream->Read(&graph_json));
-  // uint64_t sz;
-  // ICHECK(stream->Read(&sz));
-  // std::vector<std::string> names;
-  // ICHECK(stream->Read(&names));
-  // ICHECK(sz == names.size());
-  // for (size_t i = 0; i < sz; ++i) {
-  //   tvm::runtime::NDArray temp;
-  //   temp.Load(stream);
-  //   params[names[i]] = temp;
-  // }
-
   return exec;
 }
 
diff --git a/src/runtime/vm/vm.cc b/src/runtime/vm/vm.cc
index 04dd2eb47601..76ca009bc741 100644
--- a/src/runtime/vm/vm.cc
+++ b/src/runtime/vm/vm.cc
@@ -281,23 +281,11 @@ void VirtualMachine::LoadExecutable(const Executable* exec) {
   ICHECK(exec) << "The executable is not created yet.";
   exec_ = exec;
 
-  runtime::Module lib;
-  if (exec_->lib.defined()) {
-    lib = exec_->lib;
-  } else {
-    ICHECK(exec_->imports().size() > 0)
-      << "fix";
-    lib = exec_->imports()[0];
-  }
-
-  // Get the list of packed functions.
-  ICHECK(!exec->primitive_map.empty())
-      << "runtime module primitive map is empty"
-      << "\n";
+  runtime::Module lib = exec_->GetLib();
 
-  ICHECK(lib.operator->())
-      << "library is null"
-      << "\n";
+  ICHECK(exec->primitive_map.empty() || lib.operator->())
+      << "If the executable has declared primitive functions, the"
+      << "generated kernel library must non-be null.";
 
   for (const auto& it : exec_->primitive_map) {
     const auto& packed_name = it.first;

From e17520f9321157abea5a455deb8d1d666fd79924 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Fri, 26 Mar 2021 14:25:48 -0700
Subject: [PATCH 04/10] Lint

---
 src/runtime/library_module.h | 1 +
 1 file changed, 1 insertion(+)

diff --git a/src/runtime/library_module.h b/src/runtime/library_module.h
index 7b226a16dc95..fbb78852ae9c 100644
--- a/src/runtime/library_module.h
+++ b/src/runtime/library_module.h
@@ -29,6 +29,7 @@
 #include <tvm/runtime/module.h>
 
 #include <functional>
+#include <string>
 
 namespace tvm {
 namespace runtime {

From 02939dc8ca9503cfa7b6869de124a64c00cb65c1 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Fri, 26 Mar 2021 14:26:25 -0700
Subject: [PATCH 05/10] Format

---
 include/tvm/runtime/vm/executable.h |  6 ++----
 src/runtime/library_module.cc       |  7 +++----
 src/runtime/library_module.h        |  2 +-
 src/runtime/vm/executable.cc        | 18 ++++++++----------
 tests/python/relay/test_vm.py       |  2 ++
 5 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 2e8fd27e620e..88941a991e47 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -138,11 +138,9 @@ class Executable : public ModuleNode {
   runtime::Module GetLib() const { return this->imports_[0]; }
 
   void SetLib(const runtime::Module& lib) {
-    ICHECK(lib.defined())
-      << "library can not be null";
+    ICHECK(lib.defined()) << "library can not be null";
 
-    ICHECK_EQ(this->imports().size(), 0)
-      << "can only import the library once";
+    ICHECK_EQ(this->imports().size(), 0) << "can only import the library once";
 
     this->Import(lib);
   }
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index 9699a941bd57..d91ac5d9aae3 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -113,10 +113,9 @@ Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {
         loaders += name.substr(loadkey.size());
       }
     }
-    ICHECK(f != nullptr)
-        << "Binary was created using " << type_key
-        << " but a loader of that name is not registered. Available loaders are " << loaders
-        << ". Perhaps you need to recompile with this runtime enabled.";
+    ICHECK(f != nullptr) << "Binary was created using " << type_key
+                         << " but a loader of that name is not registered. Available loaders are "
+                         << loaders << ". Perhaps you need to recompile with this runtime enabled.";
   }
 
   return (*f)(static_cast<void*>(stream));
diff --git a/src/runtime/library_module.h b/src/runtime/library_module.h
index fbb78852ae9c..00c79e8248f4 100644
--- a/src/runtime/library_module.h
+++ b/src/runtime/library_module.h
@@ -41,7 +41,7 @@ namespace runtime {
  * \param type_key The type key of the serialized module.
  * \param stream A pointer to the stream containing the serialized module.
  * \return module The deserialized module.
-*/
+ */
 Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream);
 
 /*!
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index f4f914043862..d7256aa91940 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -37,9 +37,9 @@
 #include <utility>
 #include <vector>
 
-#include "serialize_utils.h"
-#include "../library_module.h"
 #include "../file_utils.h"
+#include "../library_module.h"
+#include "serialize_utils.h"
 
 namespace tvm {
 namespace runtime {
@@ -491,9 +491,9 @@ runtime::Module Executable::Load(const std::string& code, const runtime::Module
   // deserialization before we have we have deserialized the imports.
   if (lib.defined()) {
     ICHECK_EQ(exec->imports_.size(), 0)
-      << "A VMExecutable should never have more than one import inside an the executable, \n"
-      << "the first import should *always* be the library containing"
-      << "the platform specific kernel code";
+        << "A VMExecutable should never have more than one import inside an the executable, \n"
+        << "the first import should *always* be the library containing"
+        << "the platform specific kernel code";
     exec->Import(lib);
   }
 
@@ -789,8 +789,7 @@ void Executable::SaveToBinary(dmlc::Stream* stream) {
   std::string code(code_bytes.data, code_bytes.size);
   stream->Write(code);
 
-  ICHECK(this->imports()[0].defined())
-    << "the library must be imported before serialization";
+  ICHECK(this->imports()[0].defined()) << "the library must be imported before serialization";
 }
 
 Module ExecutableLoadBinary(void* strm) {
@@ -809,10 +808,9 @@ void Executable::SaveToFile(const std::string& path, const std::string& format)
   SaveBinaryToFile(path, data);
 }
 
-TVM_REGISTER_GLOBAL("runtime.module.loadbinary_VMExecutable")
-    .set_body_typed(ExecutableLoadBinary);
+TVM_REGISTER_GLOBAL("runtime.module.loadbinary_VMExecutable").set_body_typed(ExecutableLoadBinary);
 
-  // Load module from module.
+// Load module from module.
 Module ExecutableLoadFile(const std::string& file_name, const std::string& format) {
   std::string data;
   LoadBinaryFromFile(file_name, &data);
diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py
index f782e4323071..fd54991f784d 100644
--- a/tests/python/relay/test_vm.py
+++ b/tests/python/relay/test_vm.py
@@ -801,6 +801,7 @@ def test_constant_shape_with_external_codegen():
     opt_mod, _ = comp.optimize(mod, target="llvm")
     assert "shape_func" in opt_mod.astext(False)
 
+
 def test_vm_rpc():
     target = "llvm"
     target_host = "llvm"
@@ -825,5 +826,6 @@ def test_vm_rpc():
     out = vm_factory.invoke("main", [input_tensor])
     np.testing.assert_allclose(out.asnumpy(), np_input + np_input)
 
+
 if __name__ == "__main__":
     pytest.main([__file__])

From 2fba1c389ac497bbd3568a9ca4af85cb9263719c Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Fri, 26 Mar 2021 15:49:28 -0700
Subject: [PATCH 06/10] Address Andrew R and TK feedback

---
 include/tvm/runtime/vm/executable.h | 34 ++++++++++++++++++++++++-----
 python/tvm/runtime/module.py        |  8 +++++--
 python/tvm/runtime/vm.py            | 20 +++++++++++++++++
 src/runtime/library_module.cc       |  8 +++----
 src/runtime/vm/executable.cc        |  1 -
 tests/python/relay/test_vm.py       | 16 ++++++++++++++
 6 files changed, 74 insertions(+), 13 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 88941a991e47..c65ce7a411f6 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -64,11 +64,16 @@ 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.
+   * \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) final;
 
+  /*!
+   * \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.
+   * \param format The format of the serialized blob.
+   */
   void SaveToFile(const std::string& path, const std::string& format) final;
 
   /*!
@@ -135,18 +140,35 @@ class Executable : public ModuleNode {
    *
    * \return The runtime module that contains the hardware dependent code.
    */
-  runtime::Module GetLib() const { return this->imports_[0]; }
+  runtime::Module GetLib() const {
+    ICHECK_EQ(this->imports_.size(), 1)
+      << "The kernel library must be imported as the only module in an Executable";
+
+    return this->imports_[0];
+  }
 
+  /*!
+   * \brief Set the `lib` module in an executable.
+   *
+   * This allows us to do partial initialization in the case of (de|ser)ialization cases.
+   * This method also ensures correct initialization of library ensuring we only Import a
+   * single library.
+   *
+   * NB: This also provides some abstraction over how libraries are stored as there are plans
+   * to iterate on the way runtime::Module works in the backend of the compiler.
+   */
   void SetLib(const runtime::Module& lib) {
-    ICHECK(lib.defined()) << "library can not be null";
+    ICHECK(lib.defined())
+      << "the provided library can not be null";
 
-    ICHECK_EQ(this->imports().size(), 0) << "can only import the library once";
+    ICHECK_EQ(this->imports().size(), 0)
+      << "you can only import one device specific library";
 
     this->Import(lib);
   }
 
   /*!
-   * \brief Get the arity of the VM Fucntion.
+   * \brief Get the arity of the VMFunction.
    * \param func Function name.
    * \return The number of parameters.
    */
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index 0f339d23f729..3bd43ceda86c 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -272,7 +272,7 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
         """
         Export the module and all imported modules into a single device library.
 
-        This function only works on hos LLVM modules, other runtime::Module
+        This function only works on host 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
@@ -285,8 +285,12 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
 
         fcompile : function(target, file_list, kwargs), optional
             The compilation function to use create the final library object during
-            export. For example this is used to link together all produced artifacts
+            export.
+
+            For example, when fcompile=_cc.create_shared, or when it is not supplied but
+            module is "llvm," this is used to link all produced artifacts
             into a final dynamic library.
+
             This behavior is controlled by the type of object exported.
             If fcompile has attribute object_format, will compile host library
             to that format. Otherwise, will use default format "o".
diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index 1fdb0fa0f777..a5ecf2ba751e 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -300,6 +300,26 @@ class VirtualMachine(object):
     POOLED_ALLOCATOR = 2
 
     def __init__(self, exe, device, memory_cfg=None):
+        """
+        Construct a VirtualMachine wrapper class which provides a simple
+        interface over the raw C++ Module based API.
+
+        Parameters
+        ----------
+        exe: Union[Executable, Module]
+            The executable either with the wrapper Python type or the raw runtime.Module.
+
+        device: Union[Device, List[Device]]
+            The device, or devices on which to execute the VM code.
+
+        memory_cfg: Optional[str]
+            The allocator behavior to use for the VM.
+
+        Returns
+        -------
+        vm: VirtualMachine
+            A VM wrapper object.
+        """
         if not isinstance(exe, Executable) and not isinstance(exe, Module):
             raise TypeError(
                 "exe is expected to be the type of Executable, "
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index d91ac5d9aae3..28aad65935af 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -106,16 +106,16 @@ Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {
   if (f == nullptr) {
     std::string loaders = "";
     for (auto name : Registry::ListNames()) {
-      if (name.rfind(loadkey, 0) == 0) {
+      if (name.find(loadkey, 0) == 0) {
         if (loaders.size() > 0) {
           loaders += ", ";
         }
         loaders += name.substr(loadkey.size());
       }
     }
-    ICHECK(f != nullptr) << "Binary was created using " << type_key
-                         << " but a loader of that name is not registered. Available loaders are "
-                         << loaders << ". Perhaps you need to recompile with this runtime enabled.";
+    LOG(FATAL) << "Binary was created using " << type_key
+               << " but a loader of that name is not registered. Available loaders are "
+               << loaders << ". Perhaps you need to recompile with this runtime enabled.";
   }
 
   return (*f)(static_cast<void*>(stream));
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index d7256aa91940..32d0160c197f 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -484,7 +484,6 @@ 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>();
 
   // Support null-initialization of lib, to enable initialization during
diff --git a/tests/python/relay/test_vm.py b/tests/python/relay/test_vm.py
index fd54991f784d..c1bdc3ff9fd0 100644
--- a/tests/python/relay/test_vm.py
+++ b/tests/python/relay/test_vm.py
@@ -803,27 +803,43 @@ def test_constant_shape_with_external_codegen():
 
 
 def test_vm_rpc():
+    """
+    This test checks to make sure you can export a VMExecutable,
+    upload it to a remote machine using RPC and then execute it
+    on the other machine.
+    """
     target = "llvm"
     target_host = "llvm"
 
+    # Build a IRModule.
     x = relay.var("x", shape=(10, 1))
     f = relay.Function([x], x + x)
     mod = IRModule.from_expr(f)
+
+    # Compile to VMExecutable.
     vm_exec = vm.compile(mod, target=target, target_host=target_host)
 
+    # Export to Disk
     temp = utils.tempdir()
     path = temp.relpath("vm_library.so")
     vm_exec.mod.export_library(path)
 
+    # Use LocalRPC for testing.
     remote = rpc.LocalSession()
+
+    # Upload the serialized Executable.
     remote.upload(path)
+    # Get a handle to remote Executable.
     rexec = remote.load_module("vm_library.so")
 
     ctx = remote.cpu()
+    # Build a VM out of the executable and context.
     vm_factory = runtime.vm.VirtualMachine(rexec, ctx)
     np_input = np.random.uniform(size=(10, 1)).astype("float32")
     input_tensor = tvm.nd.array(np_input, ctx)
+    # Invoke its "main" function.
     out = vm_factory.invoke("main", [input_tensor])
+    # Check the result.
     np.testing.assert_allclose(out.asnumpy(), np_input + np_input)
 
 

From 77719962adee590c535f8b91192d7658811d77f6 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Fri, 26 Mar 2021 16:36:24 -0700
Subject: [PATCH 07/10] Add comment for Andrew

---
 python/tvm/runtime/vm.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/python/tvm/runtime/vm.py b/python/tvm/runtime/vm.py
index a5ecf2ba751e..d0de0520a674 100644
--- a/python/tvm/runtime/vm.py
+++ b/python/tvm/runtime/vm.py
@@ -309,6 +309,14 @@ def __init__(self, exe, device, memory_cfg=None):
         exe: Union[Executable, Module]
             The executable either with the wrapper Python type or the raw runtime.Module.
 
+            In most cases this will be the Python wrapper class tvm.runtime.vm.Executable but
+            if you instead get the underlying runtime.Module subclass (i.e `exe.mod`) you
+            can directly pass it to this method.
+
+            This case can occur when doing things such as RPC where TVM's module APIs
+            return the raw modules, not the wrapped modules. This constructor will
+            handle this internally.
+
         device: Union[Device, List[Device]]
             The device, or devices on which to execute the VM code.
 

From 7a11d187c8a2e93c27f0431665fb26922b00df50 Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 29 Mar 2021 12:00:07 -0700
Subject: [PATCH 08/10] Address Zhi's comment

---
 include/tvm/runtime/vm/executable.h | 10 +---------
 src/runtime/vm/executable.cc        | 19 ++++++++++++++-----
 2 files changed, 15 insertions(+), 14 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index c65ce7a411f6..8f33862e9b1e 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -157,15 +157,7 @@ class Executable : public ModuleNode {
    * NB: This also provides some abstraction over how libraries are stored as there are plans
    * to iterate on the way runtime::Module works in the backend of the compiler.
    */
-  void SetLib(const runtime::Module& lib) {
-    ICHECK(lib.defined())
-      << "the provided library can not be null";
-
-    ICHECK_EQ(this->imports().size(), 0)
-      << "you can only import one device specific library";
-
-    this->Import(lib);
-  }
+  void SetLib(const runtime::Module& lib);
 
   /*!
    * \brief Get the arity of the VMFunction.
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index 32d0160c197f..d8b6551e56bd 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -483,17 +483,26 @@ void LoadHeader(dmlc::Stream* strm) {
   STREAM_CHECK(version == TVM_VERSION, "version");
 }
 
+void Executable::SetLib(const runtime::Module& lib) {
+    ICHECK(lib.defined())
+      << "the provided library can not be null";
+
+    ICHECK_EQ(this->imports_.size(), 0)
+        << "A VMExecutable should never have more than one import inside an the executable, \n"
+        << "the first import should *always* be the library containing"
+        << "the platform specific kernel code";
+
+    this->Import(lib);
+  }
+
+
 runtime::Module Executable::Load(const std::string& code, const runtime::Module lib) {
   auto exec = make_object<Executable>();
 
   // 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)
-        << "A VMExecutable should never have more than one import inside an the executable, \n"
-        << "the first import should *always* be the library containing"
-        << "the platform specific kernel code";
-    exec->Import(lib);
+    exec->SetLib(lib);
   }
 
   exec->code_ = code;

From 6278e406e65c591ce7d47b2a2d2b7c733c4ec03d Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 29 Mar 2021 12:00:50 -0700
Subject: [PATCH 09/10] Format

---
 include/tvm/runtime/vm/executable.h |  2 +-
 src/runtime/library_module.cc       |  4 ++--
 src/runtime/vm/executable.cc        | 16 +++++++---------
 3 files changed, 10 insertions(+), 12 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index 8f33862e9b1e..f3b868aa7ecf 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -142,7 +142,7 @@ class Executable : public ModuleNode {
    */
   runtime::Module GetLib() const {
     ICHECK_EQ(this->imports_.size(), 1)
-      << "The kernel library must be imported as the only module in an Executable";
+        << "The kernel library must be imported as the only module in an Executable";
 
     return this->imports_[0];
   }
diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc
index 28aad65935af..370dc838839f 100644
--- a/src/runtime/library_module.cc
+++ b/src/runtime/library_module.cc
@@ -114,8 +114,8 @@ Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) {
       }
     }
     LOG(FATAL) << "Binary was created using " << type_key
-               << " but a loader of that name is not registered. Available loaders are "
-               << loaders << ". Perhaps you need to recompile with this runtime enabled.";
+               << " but a loader of that name is not registered. Available loaders are " << loaders
+               << ". Perhaps you need to recompile with this runtime enabled.";
   }
 
   return (*f)(static_cast<void*>(stream));
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index d8b6551e56bd..2a525bdae4ab 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -484,17 +484,15 @@ void LoadHeader(dmlc::Stream* strm) {
 }
 
 void Executable::SetLib(const runtime::Module& lib) {
-    ICHECK(lib.defined())
-      << "the provided library can not be null";
+  ICHECK(lib.defined()) << "the provided library can not be null";
 
-    ICHECK_EQ(this->imports_.size(), 0)
-        << "A VMExecutable should never have more than one import inside an the executable, \n"
-        << "the first import should *always* be the library containing"
-        << "the platform specific kernel code";
-
-    this->Import(lib);
-  }
+  ICHECK_EQ(this->imports_.size(), 0)
+      << "A VMExecutable should never have more than one import inside an the executable, \n"
+      << "the first import should *always* be the library containing"
+      << "the platform specific kernel code";
 
+  this->Import(lib);
+}
 
 runtime::Module Executable::Load(const std::string& code, const runtime::Module lib) {
   auto exec = make_object<Executable>();

From 99c80a8b2bcf47a3740d09fcd773f8e0a95094ce Mon Sep 17 00:00:00 2001
From: Jared Roesch <jroesch@octoml.ai>
Date: Mon, 29 Mar 2021 16:49:55 -0700
Subject: [PATCH 10/10] Fix broken test

---
 include/tvm/runtime/vm/executable.h |  7 +------
 src/runtime/vm/executable.cc        | 11 +++++++++++
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/include/tvm/runtime/vm/executable.h b/include/tvm/runtime/vm/executable.h
index f3b868aa7ecf..95c6d6f4ab47 100644
--- a/include/tvm/runtime/vm/executable.h
+++ b/include/tvm/runtime/vm/executable.h
@@ -140,12 +140,7 @@ class Executable : public ModuleNode {
    *
    * \return The runtime module that contains the hardware dependent code.
    */
-  runtime::Module GetLib() const {
-    ICHECK_EQ(this->imports_.size(), 1)
-        << "The kernel library must be imported as the only module in an Executable";
-
-    return this->imports_[0];
-  }
+  runtime::Module GetLib() const;
 
   /*!
    * \brief Set the `lib` module in an executable.
diff --git a/src/runtime/vm/executable.cc b/src/runtime/vm/executable.cc
index 2a525bdae4ab..e8b948d3d2ae 100644
--- a/src/runtime/vm/executable.cc
+++ b/src/runtime/vm/executable.cc
@@ -483,6 +483,17 @@ void LoadHeader(dmlc::Stream* strm) {
   STREAM_CHECK(version == TVM_VERSION, "version");
 }
 
+runtime::Module Executable::GetLib() const {
+  ICHECK_LE(this->imports_.size(), 1)
+      << "The kernel library must be imported as the only module in an Executable";
+
+  if (this->imports().size() == 0) {
+    return Module(nullptr);
+  } else {
+    return this->imports_[0];
+  }
+}
+
 void Executable::SetLib(const runtime::Module& lib) {
   ICHECK(lib.defined()) << "the provided library can not be null";