模块序列化简介
部署TVM运行时模块时,无论是CPU还是GPU,TVM都只需要一个动态共享库。关键是我们统一的模块序列化机制。本文档将介绍TVM模块序列化格式标准和实现细节。
模块导出例子
让我们首先为GPU构建一个ResNet-18工作负载作为示例。
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.runtime.load(path_lib)
assert loaded_lib.type_key == "library"
assert loaded_lib.imported_modules[0].type_key == "cuda"
序列化
入口API是【tvm.module.Module
】的【export_library
】。在此函数中,我们将执行以下步骤:
-
收集所有DSO模块(LLVM模块和C模块)
-
一旦有了DSO模块,我们将调用【
save
】函数将其保存到文件中。 -
接下来,我们将检查是否已导入模块,例如CUDA,OpenCL或其他任何模块。我们在这里不限制模块类型。导入模块后,我们将创建一个名为【
devc.o
】或者【dev.cc
】的文件 (以便可以将导入模块的二进制Blob数据嵌入到一个动态共享库中),然后调用函数【_PackImportsToLLVM
】或【_PackImportsToC
】进行模块序列化。 -
最后,我们调用【
fcompile
】,它
调用【_cc.create_shared
】以获取动态共享库。
注意
-
对于C源模块,我们将对其进行编译并将其与DSO模块链接在一起。
-
使用【
_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数。在我们的示例中,将分别为LLVM模块,CUDA模块和【_import_tree
】创建三个Blob 。
【binary_blob_type_key】
是模块的Blob类型键。对于LLVM / C模块,其Blob类型键为【_lib
】。对于CUDA模块,其Blob类型键为【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
】。
最后,除非模块只有一个DSO模块并且位于根目录中,否则我们将编写一个关键【_import_tree
】。如前所述,当我们重新加载导出的库时,它用于重建模块导入关系。该【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
】,我们将在【module.loadfile_so
】中调用【dso_library.cc
】。关键在这里:
// 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打包到符号【untime::symbol::tvm_dev_mblob
】中 r
。在反序列化部分,我们将对其进行检查。如果有【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
】,以便允许从根目录查找符号(这样所有符号都是可见的)。
最后,我们完成反序列化部分。