news 2026/4/24 11:39:53

【TVM教程】模块序列化指南

作者头像

张小明

前端开发工程师

1.2k 24
文章封面图
【TVM教程】模块序列化指南

TVM 现已更新到 0.21.0 版本,TVM 中文文档已经和新版本对齐。

Apache TVM 是一个深度的深度学习编译框架,适用于 CPU、GPU 和各种机器学习加速芯片。更多 TVM 中文文档可访问 →https://tvm.hyper.ai/

在部署 TVM 运行时模块时,无论目标是 CPU 还是 GPU,TVM最终只需要一个动态共享库(dynamic shared library)。实现这一点的关键就在于统一的模块序列化机制。本文将介绍 TVM 模块序列化的格式标准与实现细节。

序列化(Serialization)

入口 API 为tvm.module.Moduleexport_library。在此函数内部,我们会执行以下步骤:

  1. 收集所有 DSO 模块(例如 LLVM 模块和 C 模块)。
  2. 在获得 DSO 模块后,调用save函数将它们保存到文件。
  3. 随后检查是否存在已导入的模块(imported modules),例如 CUDA、OpenCL 等。这里对模块类型不做限制。
    如果存在导入模块,我们将创建一个名为devc.o/dev.cc的文件(用于将这些导入模块的二进制数据打包进最终的动态库中),然后调用_PackImportsToLLVM_PackImportsToC来执行模块序列化。
  4. 最后,调用fcompile,其内部会调用_cc.create_shared,生成动态共享库。

备注

  1. 对于 C 源码模块(CSourceModule),我们会将它们编译并与 DSO 模块一同进行链接。
  2. 是否使用_PackImportsToLLVM_PackImportsToC取决于 TVM 是否启用了 LLVM。它们本质上实现的是相同的目标。

序列化底层机制与格式标准

序列化主要发生在_PackImportsToLLVM_PackImportsToC中。它们都会调用SerializeModule来序列化 runtime module。在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

最后,除非我们的模块中仅有一个 DSO 模块并且它位于根位置,否则我们会写入一个键_import_tree。该键用于在重新加载导出的库时恢复模块导入关系,如前文所述。import_tree_logic的内容则是将import_tree_row_ptr_import_tree_child_indices_写入到流中。

在上述步骤完成后,我们会将最终结果打包进一个符号runtime::symbol::tvm_ffi_library_bin,该符号可在动态库中恢复。

现在,我们已经完成序列化部分。正如你所看到的,我们理论上可以支持导入任意模块。

反序列化

入口 API 是tvm.runtime.load。实际上,该函数会调用_LoadFromFile。 如果进一步展开,可以看到其对应的是Module::LoadFromFile

在我们的示例中,文件是deploy.so。根据其函数逻辑,我们会在dso_library.cc中调用module.loadfile_so,关键代码如下:

// Load the imported modules const char* library_bin = reinterpret_cast<const char*>( lib->GetSymbol(runtime::symbol::tvm_ffi_library_bin)); Module root_mod; if (library_bin != nullptr) { root_mod = ProcessLibraryBin(library_bin, lib); } else { // Only have one single DSO Module root_mod = Module(n); }``` 如前所述,我们会将 blob 打包进符号 `runtime::symbol::tvm_ffi_library_bin`· 中。 在反序列化阶段,我们会检查它。如果存在 `runtime::symbol::tvm_ffi_library_bin`,我们将调用 `ProcessLibraryBin`,其逻辑如下: ```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;

完成上述步骤后,我们会将ctx_address设置为root_module, 以便能够从根模块查找符号(使所有符号可见)。

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

版权声明: 本文来自互联网用户投稿,该文观点仅代表作者本人,不代表本站立场。本站仅提供信息存储空间服务,不拥有所有权,不承担相关法律责任。如若内容造成侵权/违法违规/事实不符,请联系邮箱:809451989@qq.com进行投诉反馈,一经查实,立即删除!
网站建设 2026/4/23 12:36:34

科研写作工具深度测评:7个平台功能与典型应用场景

工具核心特点速览 工具名称 核心优势 适用场景 数据支撑 aibiye 全流程覆盖降重优化 从开题到答辩的一站式需求 支持20万字长文逻辑连贯 aicheck 院校规范适配模板化输出 国内本硕博论文框架搭建 覆盖90%高校格式要求 秒篇 3分钟文献综述生成 紧急补文献章节 知…

作者头像 李华
网站建设 2026/4/23 11:15:18

Python基于Vue的网络安全产品推广平台 django flask pycharm

目录 这里写目录标题目录项目介绍项目展示详细视频演示技术栈文章下方名片联系我即可~解决的思路开发技术介绍性能/安全/负载方面python语言Django框架介绍技术路线关键代码详细视频演示收藏关注不迷路&#xff01;&#xff01;需要的小伙伴可以发链接或者截图给我 项目介绍 …

作者头像 李华
网站建设 2026/4/23 15:30:10

学术论文高效工具盘点:7大平台功能解析与适用场景

工具核心特点速览 工具名称 核心优势 适用场景 数据支撑 aibiye 全流程覆盖降重优化 从开题到答辩的一站式需求 支持20万字长文逻辑连贯 aicheck 院校规范适配模板化输出 国内本硕博论文框架搭建 覆盖90%高校格式要求 秒篇 3分钟文献综述生成 紧急补文献章节 知…

作者头像 李华
网站建设 2026/4/23 14:01:32

自动化智能体与测试用例生成

关注 霍格沃兹测试学院公众号&#xff0c;回复「资料」, 领取人工智能测试开发技术合集每天重复写着相似的测试用例&#xff0c;翻阅上百页的需求文档寻找测试点&#xff0c;为了一个边界值绞尽脑汁……这是不是你的日常&#xff1f;好消息是&#xff0c;AI自动化的时代已经到来…

作者头像 李华
网站建设 2026/4/23 12:25:13

收藏!刷招聘软件时的迟疑?AI大模型才是程序员的新底气

刷着招聘APP的你&#xff0c;是否也曾突然陷入迟疑&#xff1f; 屏幕上密密麻麻的“大模型工程师”“AIGC应用开发工程师”岗位&#xff0c;技能要求写得愈发细致具体&#xff0c;从模型微调、Prompt工程到落地部署&#xff0c;条条直指AI领域。反观自己简历上引以为傲的“微服…

作者头像 李华