[Doc] Introduction to module serialization

docs/dev/index.rst

@@ -37,3 +37,4 @@ In this part of documentation, we share the rationale for the specific choices m
    nnvm_json_spec
    nnvm_overview
    benchmark
+   introduction_to_module_serialization

docs/dev/introduction_to_module_serialization.rst

@@ -0,0 +1,211 @@
+..  Licensed to the Apache Software Foundation (ASF) under one
+    or more contributor license agreements.  See the NOTICE file
+    distributed with this work for additional information
+    regarding copyright ownership.  The ASF licenses this file
+    to you under the Apache License, Version 2.0 (the
+    "License"); you may not use this file except in compliance
+    with the License.  You may obtain a copy of the License at
+
+..    http://www.apache.org/licenses/LICENSE-2.0
+
+..  Unless required by applicable law or agreed to in writing,
+    software distributed under the License is distributed on an
+    "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+    KIND, either express or implied.  See the License for the
+    specific language governing permissions and limitations
+    under the License.
+
+Introduction to Module Serialization
+====================================
+
+When to deploy TVM runtime module, no matter it is CPU or GPU, TVM only needs one single DLL.
+The key is our unified module serialization mechanism. This document will introduce TVM module
+serialization format standard and implementation details.
+
+***************************************
+Module Export Example
+***************************************
+
+Let us build one resnet18 workload for GPU as our example firstly.
+
+.. code:: python
+
+   from tvm import relay
+   from tvm.relay import testing
+   from tvm.contrib import util
+   import tvm
+
+   # Resnet18 workload
+   resnet18_mod, resnet18_params = relay.testing.resnet.get_workload(num_layers=18)
+
+   # build
+   with relay.build_config(opt_level=3):
+       _, resnet18_lib, _ = relay.build_module.build(resnet18_mod, "cuda", params=resnet18_params)
+
+   # create one tempory directory
+   temp = util.tempdir()
+
+   # path lib
+   file_name = "deploy.so"
+   path_lib = temp.relpath(file_name)
+
+   # export library
+   resnet18_lib.export_library(path_lib)
+
+   # load it back
+   loaded_lib = tvm.module.load(path_lib)
+   assert loaded_lib.type_key == "library"
+   assert loaded_lib.imported_modules[0].type_key == "cuda"
+
+
+**************
+Serialization
+**************
+
+The entrance API is ``export_library`` of ``tvm.module.Module``.
+Inside this function, we will do the following steps:
+
+1. Collect all DSO modules (LLVM module or C module)
+
+
+2. If we have DSO modules, we will call ``save`` function to save them into files.
+
+
+3. Next, we will check whether we have imported modules. Like CUDA,
+   OpenCL or anything else, we don't restrict the module type here.
+   If we have imported modules, we will create one file named as ``dev.cc``
+   (so that we could compile into one dll), then call one function
+   ``_PackImportsToC`` to do module serialization.
+
+
+4. Finally, we use ``fcompile`` to call ``_cc.create_shared`` to get
+   dll.
+
+***************************************************
+Under the Hood of Serialization and Format Standard
+***************************************************
+
+As said before, we will do the serialization work in ``_PackImportsToC``.
+Inside this function, we firstly construct one helper class ``ModuleSerializer``.
+It will take ``module`` to do some initialize work, like marking module index.
+Then we could use its ``SerializeModule`` to serialize module.
+
+For the sake of understanding better, let us dig the implementation of this class a little deeper.
+
+When we construct ``ModuleSerializer``, we could see this code:
+
+.. code:: c++
+
+   explicit ModuleSerializer(runtime::Module mod) : mod_(mod) {
+     Init();
+   }
+   private:
+   void Init() {
+     CreateModuleIndex();
+     CreateImportTree();
+   }
+
+In ``CreateModuleIndex()``, We will inspect module import relationship
+using DFS and create index for them. Here, we have one invariance: root
+module is always at location 0. In our example, we have module relationship:
+``LLVM Module <- CUDA Module.`` So LLVM module will have index 0, CUDA
+module will have index 1.
+
+After constructing module index, we will try to construct import tree (``CreateImportTree()``),
+which will be used for restore module import relationship when we load
+the exported library back. In our design, we use CSR format to store
+import tree, each row is parent index, the child indices is its children
+index. In code, we use ``import_tree_row_ptr_`` and
+``import_tree_child_indices_`` to represent them.
+
+After initialization, we could serialize module using ``SerializeModule`` function.
+In its function logic, we will assume the serialization format like this:
+
+.. code:: c++
+
+   binary_blob_size
+   binary_blob_type_key
+   binary_blob_logic
+   ...
+   _import_tree
+   _import_tree_logic
+
+``binary_blob_size`` is how many blobs we will have in this
+serialization step. In our example, the number will equal to 3. One for
+LLVM module, one for CUDA module, one for ``_import_tree``.
+
+Then we will write the ``binary_blob_type_key``, for LLVM module / C
+module, the blob type key is ``_lib``. For CUDA module, it is
+``cuda``, which could get using ``module->type_key()``.
+
+Next we will do the ``binary_blob_logic``. Normally, we will call
+``SaveToBinary`` function to serialize blob into binary. So, for most
+modules, they should implement ``SaveToBinary`` function. But like DSO
+module, we don't implement ``SaveToBinary``, we also don't some special
+handling, so the ``binary_blob_logic`` for DSO module currently
+is empty in fact.
+
+Finally, we will write one key ``_import_tree``. Unless our module only
+have one DSO module and it is in the root, otherwise we will always
+have this key. It is used for reconstruct the module import relationship
+when we load the exported library back as said before. The
+``import_tree_logic`` is just write ``import_tree_row_ptr_`` and
+``import_tree_child_indices_`` into stream.
+
+After this step, we will pack it into a symbol
+``runtime::symbol::tvm_dev_mblob`` that can be recovered in the dynamic
+libary.
+
+Now, we complete the serialization part. As you have seen, we could
+support arbitrary modules to import ideally.
+
+****************
+Deserialization
+****************
+
+The entrance API is ``tvm.module.load``. This function
+is to call ``_LoadFromFile`` in fact. If we dig it a little deeper, this is
+``Module::LoadFromFile``. In our example, the file ``deploy.so``,
+according to the function logic, we will call ``module.loadfile_so`` in
+``dso_library.cc``. The key is here:
+
+.. code:: c++
+
+   // Load the imported modules
+   const char* dev_mblob = reinterpret_cast<const char*>(lib->GetSymbol(runtime::symbol::tvm_dev_mblob));
+   Module root_mod;
+   if (dev_mblob != nullptr) {
+   root_mod = ProcessModuleBlob(dev_mblob, lib);
+   } else {
+   // Only have one single DSO Module
+   root_mod = Module(n);
+   }
+
+As said before, we will pack the blob into the symbol
+``runtime::symbol::tvm_dev_mblob``. During deserialization part, we will
+inspect it. If we have it, we will call ``ProcessModuleBlob``, whose logic like this:
+
+.. code:: c++
+
+   READ(blob_size)
+   READ(blob_type_key)
+   for (size_t i = 0; i < blob_size; i++) {
+       if (blob_type_key == "_lib") {
+         // construct dso module using lib
+       } else if (blob_type_key == "_import_tree") {
+         // READ(_import_tree_row_ptr)
+         // READ(_import_tree_child_indices)
+       } else {
+         // call module.loadbinary_blob_type_key, such as module.loadbinary_cuda
+         // to restore.
+       }
+   }
+   // Using _import_tree_row_ptr and _import_tree_child_indices to
+   // restore module import relationship. The first module is the
+   // root module according to our invariance as said before.
+   return root_module;
+
+After this, we will set the ``ctx_address`` to be the ``root_module`` so
+that allow lookup of symbol from root (so all symbols are visible).
+
+Finally, we complete the deserialization part.