【TVM 教程】TVM 代码库实例讲解

💡 原文中文,约7300字,阅读约需18分钟。
📝

内容提要

Apache TVM 是一个支持 CPU、GPU 和加速芯片的深度学习编译框架。本文介绍了 TVM 的代码库结构、编译管道及 C++ 与 Python 的互操作性,并通过向量加法的实现示例,帮助开发者快速上手。

🎯

关键要点

  • Apache TVM 是一个支持 CPU、GPU 和加速芯片的深度学习编译框架。
  • TVM 代码库结构包括 src、src/relay、python 和 src/topi 等子目录。
  • src 目录包含用于算子编译和部署的 C++ 代码。
  • python 目录是 Python 前端,用于包装 C++ 函数和对象。
  • 用户通过 relay.build(...) 调用图结构编译,涉及算子查找、生成计算表达式和调度。
  • TVM 中 C++ 和 Python 之间的互操作性是双向的,C++ 代码可以调用 Python 函数。
  • 向量加法示例展示了如何直接调用底层 TVM API。
  • Tensor 对象与 Operation 对象相关联,能够跟踪操作之间的依赖关系。
  • 调度通过 tvm.te.create_schedule() 函数创建,包含多个阶段的信息。
  • tvm.build() 函数用于生成目标机器代码,分为降级和代码生成两个步骤。
  • 降级过程通过边界推断和初始循环嵌套结构的创建实现。
  • 代码生成使用目标特定的代码生成器生成机器代码。
  • 返回的模块可以在 TVM 的 NDArray 对象上被调用,自动管理设备内存和内存传输。
  • 首次调用已编译模块时,会获取可用于内核调用的 PackedFunc。

延伸问答

Apache TVM 是什么?

Apache TVM 是一个支持 CPU、GPU 和加速芯片的深度学习编译框架。

TVM 的代码库结构包含哪些主要目录?

TVM 的代码库结构包括 src、src/relay、python 和 src/topi 等子目录。

如何在 TVM 中实现向量加法?

通过定义张量 A 和 B,然后使用 tvm.te.compute() 函数计算 C,最后调用 tvm.build() 生成目标代码。

TVM 中 C++ 和 Python 的互操作性是如何实现的?

TVM 中 C++ 和 Python 的互操作性是双向的,C++ 代码可以调用 Python 函数,反之亦然。

在 TVM 中,如何创建调度?

使用 tvm.te.create_schedule() 函数创建调度,该函数包含多个阶段的信息。

tvm.build() 函数的作用是什么?

tvm.build() 函数用于生成目标机器代码,分为降级和代码生成两个步骤。

➡️

继续阅读