ML Sys | Apache TVM 添加新 backend

硬件决定上限，软件实现上限。即使是顶尖的硬件也需要匹配优秀的软件栈才能发挥最大功效。如今 Machine Learning Compiler 主要有 Apache TVM 和 MLIR(Multi Level IR) 两种范式。让我们来研究一下 TVM 如何支持新的 backend。

Apache TVM 后端主要维护 GPU 和 CPU 平台，对于新后端文档不是很完善，想要搞懂还得深入理解 TVM 源码。

tvm 有俩种添加后端的接口^[1]

• Porcelain Layer: UMA (Universal Modular Accelerator Interface)
• Plumbing Layer: BYOC (Bring Your Own Codegen)

前者轻量纯 Python API 接口，后者则是 Python + C++。 UMA codegen 只支持生成 C 代码，想要高自由度不得不淌 BYOC 这趟浑水了。

Install

官网提供了一个 vanilla backend 的例子 ^[2]，但默认 pip 仓库的 pre-build 没有开启 UMA 功能，需要我们重新 build 安装。

主要流程按照官网教程 ^[3]，除了以下俩点：

• 编译 tvm 时在 config.cmake 里把 set(UMA ON) 和 set(MICRO ON) 选项打开。
• 安装 tvm 时用 develop 模式安装。 UMABackend 类中有个地方是根据相对路径查找内容，这里必须按照 develop 模式安装 Python package，防止将内容迁移虚拟环境目录下。切换到 python 目录下调用 python setup.py develop 安装

其余部分参考官网教程，应该就能跑通例程。

UMA 结构解析

TODO

BYOC

TODO

---

1. Tvm-rfcs/rfcs/0060_UMA_Unified_Modular_Accelerator_Interface.md at main · apache/tvm-rfcs. (不详). GitHub. 取读于 2024年4月17日, 从 https://github.com/apache/tvm-rfcs/blob/main/rfcs/0060_UMA_Unified_Modular_Accelerator_Interface.md ↩︎

2. Making your Hardware Accelerator TVM-ready with UMA — tvm 0.17.dev0 documentation. (不详). 取读于 2024年4月20日, 从 https://tvm.apache.org/docs/tutorial/uma.html ↩︎

3. Install from Source—Tvm 0.17.dev0 documentation. (不详). 取读于 2024年4月20日, 从 https://tvm.apache.org/docs/install/from_source.html#install-from-source ↩︎