Module Serialization 简介

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

序列化

入口 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 源代码模块，我们将编译它们并将它们与 DSO 模块链接在一起。

  2. 使用 _PackImportsToLLVM 还是 _PackImportsToC 取决于我们是否在 TVM 中启用了 LLVM。事实上，它们实现了相同的目标。

序列化和格式标准幕后原理

如前所述，我们将在 _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 是我们在此序列化步骤中将拥有的 blob 数量。在我们的示例中，将有三个 blob，它们分别是为 LLVM 模块、CUDA 模块和 _import_tree 创建的。

binary_blob_type_key 是模块的 blob 类型键。对于 LLVM / C 模块，其 blob 类型键为 _lib。对于 CUDA 模块，它是 cuda，可以通过 module->type_key() 获取。

binary_blob_logic 是 blob 的逻辑处理。对于大多数 blob（如 CUDA、OpenCL），我们将调用 SaveToBinary 函数将 blob 序列化为二进制。但是，对于像 LLVM / C 模块这样的模块，我们只会写入 _lib 以指示这是一个 DSO 模块。

注意

是否需要实现 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，以便允许从根目录查找符号（因此所有符号都是可见的）。

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