模块序列化介绍

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

模块导出示例

让我们先为构建一个ResNet-18 GPU 工作负载作为例子。

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模块,我们将调用保存函数将它们保存到文件中。

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。它将需要模块做一些初始化工作,如标记模块索引。然后我们可以使用它的SerializeModule来序列化模块。

为了更好地理解,让我们更深入地研究这个类的实现。 

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,可以通过模块->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,以便允许从根目录查找符号(这样所有符号都是可见的)。

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

  • 0
    点赞
  • 1
    收藏
    觉得还不错? 一键收藏
  • 0
    评论

“相关推荐”对你有帮助么?

  • 非常没帮助
  • 没帮助
  • 一般
  • 有帮助
  • 非常有帮助
提交
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值