硬件决定上限,软件实现上限。即使是顶尖的硬件也需要匹配优秀的软件栈才能发挥最大功效。如今 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
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 ↩︎
Making your Hardware Accelerator TVM-ready with UMA — tvm 0.17.dev0 documentation. (不详). 取读于 2024年4月20日, 从 https://tvm.apache.org/docs/tutorial/uma.html ↩︎
Install from Source—Tvm 0.17.dev0 documentation. (不详). 取读于 2024年4月20日, 从 https://tvm.apache.org/docs/install/from_source.html#install-from-source ↩︎