模型序列化简介

模型序列化简介#

在部署 TVM 运行时模块时，无论是 CPU 还是 GPU，TVM 只需要一个单一的动态共享库。关键在于统一的模块序列化机制。本文档将介绍 TVM 模块序列化格式标准及其实现细节。

模块导出示例#

首先，以 GPU 为例构建 ResNet-18 的工作负载。

from tvm import relay
from tvm.relay import testing
from tvm.contrib import utils
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 = utils.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.runtime.load_module(path_lib)
assert loaded_lib.type_key == "library"
assert loaded_lib.imported_modules[0].type_key == "cuda"

序列化#

入口 API 是 tvm.module.Module 的 export_library。在这个函数内部，将执行以下步骤：

  1. 收集所有 DSO 模块（LLVM 模块和 C 模块）

  2. 一旦拥有 DSO 模块，便可调用 save 函数将它们保存到文件中。

  3. 接着，将检查是否导入了模块，例如 CUDA、OpenCL 或任何其他模块。不限制模块类型。一旦导入了模块，将创建名为 devc.o / dev.cc 的文件，（以便能够将二进制 blob 数据嵌入到动态共享库中），然后调用函数 _PackImportsToLLVM 或 _PackImportsToC 来完成模块序列化。

  4. 最后，将调用 fcompile，它将调用 _cc.create_shared 来获取动态共享库。

备注

  1. 对于 C 源模块，将编译它们并将它们链接在一起。

  2. 最后，根据在 TVM 中是否启用 LLVM，使用 _PackImportsToLLVM 或 _PackImportsToC。实际上，它们实现了相同的目标。

序列化与格式标准的背后原理"#

如前所述，将会执行序列化工作于 _PackImportsToLLVM 或 _PackImportsToC 函数中。这两个函数都会调用 SerializeModule 来对运行时模块进行序列化。在 SerializeModule 函数中，首先构建辅助类 ModuleSerializer。这个类会接收 module 并做一些初始化工作，比如标记模块索引。之后可以使用它的 SerializeModule 方法来对模块进行序列化。

为了更深入地理解，对这个类的实现进行更细致的探讨。

以下代码用于构建 ModuleSerializer：

explicit ModuleSerializer(runtime::Module mod) : mod_(mod) {
  Init();
}
private:
void Init() {
  CreateModuleIndex();
  CreateImportTree();
}

在 CreateModuleIndex() 函数中，使用深度优先搜索（DFS）来检查模块间的导入关系，并为它们创建索引。请注意，根模块的位置固定为 0。在示例中，模块之间的关系如下：

llvm_mod:imported_modules
  - cuda_mod

故而 LLVM 模块索引为 0，CUDA 模块索引为 1

构建模块索引之后，将尝试构建导入树（CreateImportTree()），该树在重新加载导出的库时用于恢复模块导入关系。在我们的设计中，使用 CSR 格式存储导入树，每一行代表父索引，子索引对应其子节点的索引。在代码中，用 import_tree_row_ptr_ 和 import_tree_child_indices_ 来表示它们。

初始化完成后，可以使用 SerializeModule 函数对模块进行序列化。在它的功能逻辑中，假设序列化的格式如下：

binary_blob_size
binary_blob_type_key
binary_blob_logic
binary_blob_type_key
binary_blob_logic
...
_import_tree
_import_tree_logic

binary_blob_size 是在这一步序列化过程中将要处理的二进制数据块的数量。在示例中，将有三个二进制数据块被创建，分别对应于 LLVM 模块、CUDA 模块和 _import_tree。

binary_blob_type_key 是模块的二进制类型键。对于 LLVM/C++ 模块，其二进制类型键是 _lib。而对于 CUDA 模块，则是 cuda，可以通过调用 module->type_key() 获取该键。

binary_blob_logic 是指处理 blob 的逻辑。对于大多数 blob（如 CUDA、OpenCL），将调用 SaveToBinary 函数将 blob 序列化为二进制格式。然而，对于 LLVM/C 模块，我们只会写入 _lib 来表示这是动态共享对象模块。

备注

是否需要实现 SaveToBinary 虚函数，这取决于模块的使用方式。例如，如果在重新加载动态共享库时需要模块中的信息，那么应该实现它。以 CUDA 模块为例，当加载动态共享库时，需要将其二进制数据传递给 GPU 驱动，因此应实现 SaveToBinary 来序列化其二进制数据。但对于主机模块（如 DSO），在加载动态共享库时不需要其他信息，因此无需实现 SaveToBinary。然而，如果将来希望记录 DSO 模块的一些元信息，也可以为 DSO 模块实现 SaveToBinary。

最终，将编写关键的 _import_tree 函数，除非模块只包含一个 DSO 模块并且它位于根目录中。这个函数用于在重新加载导出的库时重建模块导入关系，正如之前所述。import_tree_logic 的功能仅是将 import_tree_row_ptr_ 和 import_tree_child_indices_ 写入流中。

完成此步骤后，将它打包进名为 runtime::symbol::tvm_dev_mblob 的符号中，该符号可以在动态库中恢复。

现在，完成了序列化部分。正如你所看到的，理想情况下可以支持任意模块的导入。

反序列化#

入口API为 tvm.runtime.load。此函数实际上是调用 _LoadFromFile。如果更深入地研究，这其实是 Module::LoadFromFile。在示例中，文件名为 deploy.so，根据函数逻辑，将在 dso_library.cc 中调用 module.loadfile_so。关键点在此：

// 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);
}

如前所述，将 blob 打包到符号 runtime::symbol::tvm_dev_mblob 中。在反序列化部分，将检查它。如果有 runtime::symbol::tvm_dev_mblob，将调用 ProcessModuleBlob，其逻辑如下：

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;

之后，将 ctx_address 设置为 root_module，这样可以从根目录查找符号（因此所有符号都是可见的）。

最后，完成反序列化部分。