本篇文章译自英文文档 Making your Hardware Accelerator TVM-ready with UMA
作者是 Michael J. Klaiber,Christoph Gerum,Paul Palomero Bernardo。
更多 TVM 中文文档可拜访 →TVM 中文站
本节介绍通用模块化加速器接口(UMA)。UMA 提供了一个易用的 API 来将新的硬件加速器集成到 TVM 中。
本教程具体介绍了如何利用 UMA 使得你的硬件加速器可间接用于 TVM。尽管这个问题没有万能的解决方案,但 UMA 旨在提供一个稳固的纯 Python API,从而将许多品种的硬件加速器集成到 TVM 中。
本教程将通过三个逐步简单的用例来介绍 UMA API。这些用例引入了三个模仿加速器 Vanilla、Strawberry 和 Chocolate,并用 UMA 将它们集成到 TVM 中。
Vanilla
Vanilla 是一个由 MAC 数组组成的简略加速器,没有外部存储器。它只能解决 Conv2D 层,所有其余层都在 CPU 上执行,同时也协调 Vanilla。CPU 和 Vanilla 共享内存。
Vanilla 的 C 接口 vanilla_conv2dnchw(…) 用于执行 Conv2D 操作(包含 same-padding),它接管指向输出特色图、权重和后果的指针,以及 Conv2D 的维度:oc、iw、ih、ic、kh 和 kw。
int vanilla_conv2dnchw(float* ifmap, float* weights, float* result, int oc, int iw, int ih, int ic, int kh, int kw);
脚本 uma_cli 为新的加速器创立带有 API(UMA-API)调用的代码骨架。
Vanilla 的应用形式如下:(–tutorial vanilla 增加了本局部教程所需的所有附加文件)
pip install inflection
cd $TVM_HOME/apps/uma
python uma_cli.py --add_hardware vanilla_accelerator --tutorial vanilla
uma_cli.py 在 vanilla_accelerator 目录中生成这些文件。
backend.py
codegen.py
conv2dnchw.cc
passes.py
patterns.py
run.py
strategies.py
Vanilla 后端
vanilla 生成的后端位于 vanilla_accelerator/backend.py 中:
class VanillaAcceleratorBackend(UMABackend):
"""VanillaAccelerator 的 UMA 后端。"""
def __init__(self):
super().__init__()
self._register_pattern("conv2d", conv2d_pattern())
self._register_tir_pass(PassPhase.TIR_PHASE_0, VanillaAcceleratorConv2DPass())
self._register_codegen(fmt="c", includes=gen_includes)
@property
def target_name(self):
return "vanilla_accelerator"
定义迁徙模式
为了指定 Conv2D 迁徙到 Vanilla,vanilla_accelerator/patterns.py 中将其形容为 Relay 数据流模式(DFPattern)。
def conv2d_pattern():
pattern = is_op("nn.conv2d")(wildcard(), wildcard())
pattern = pattern.has_attr({"strides": [1, 1]})
return pattern
为了将输出计算图的 Conv2D 算子映射到 Vanilla 的底层函数调用 vanilla_conv2dnchw(…),在 VanillaAcceleratorBackend 中注册了 TIR pass VanillaAcceleratorConv2DPass(稍后探讨)。
Codegen
文件 vanilla_accelerator/codegen.py 定义了动态 C 代码,它被增加到生成的后果 C 代码(由 gen_includes 中的 TVM 的 C-Codegen 生成)中,其目标是蕴含 Vanilla 的底层库 vanilla_conv2dnchw()。
def gen_includes() -> str:
topdir = pathlib.Path(__file__).parent.absolute()
includes = ""includes += f'#include "{topdir}/conv2dnchw.cc"'
return includes
如下面的 VanillaAcceleratorBackend 所示,用 self._register_codegen 可将其注册到 UMA。
self._register_codegen(fmt="c", includes=gen_includes)
构建神经网络并在 Vanilla 上运行
为了演示 UMA 的性能,将为单个 Conv2D 层生成 C 代码,并在 Vanilla 加速器上运行。文件 vanilla_accelerator/run.py 提供了一个应用 Vanilla 的 C-API 运行 Conv2D 层的 demo。
def main():
mod, inputs, output_list, runner = create_conv2d()
uma_backend = VanillaAcceleratorBackend()
uma_backend.register()
mod = uma_backend.partition(mod)
target = tvm.target.Target("vanilla_accelerator", host=tvm.target.Target("c"))
export_directory = tvm.contrib.utils.tempdir(keep_for_debug=True).path
print(f"Generated files are in {export_directory}")
compile_and_run(AOTModel(module=mod, inputs=inputs, outputs=output_list),
runner,
interface_api="c",
use_unpacked_api=True,
target=target,
test_dir=str(export_directory),
)
main()
运行 vanilla_accelerator/run.py,将以模型库格局(MLF)生成输入文件。
输入后果:
Generated files are in /tmp/tvm-debug-mode-tempdirs/2022-07-13T13-26-22___x5u76h0p/00000
查看生成的文件:
输入后果:
cd /tmp/tvm-debug-mode-tempdirs/2022-07-13T13-26-22___x5u76h0p/00000
cd build/
ls -1
codegen
lib.tar
metadata.json
parameters
runtime
src
若要评估生成的 C 代码,请查看 codegen/host/src/default_lib2.c。
cd codegen/host/src/
ls -1
default_lib0.c
default_lib1.c
default_lib2.c
在 default_lib2.c 中,能够看到生成的代码调用了 Vanilla 的 C-API,而后执行了一个 Conv2D 层:
TVM_DLL int32_t tvmgen_default_vanilla_accelerator_main_0(float* placeholder, float* placeholder1, float* conv2d_nchw, uint8_t* global_workspace_1_var) {vanilla_accelerator_conv2dnchw(placeholder, placeholder1, conv2d_nchw, 32, 14, 14, 32, 3, 3);
return 0;
}
Strawberry
行将上线
Chocolate
行将上线
征求社区意见
若本教程不适宜你的加速器,请将你的需要增加到 TVM 论坛中的 UMA 帖子 中。咱们很乐意通过扩大本教程来提供更多领导,例如如何利用 UMA 接口使得更多品种的 AI 硬件加速器可间接用于 TVM。
参考
[UMA-RFC]UMA:通用模块化加速器接口,TVM RFC,2022 年 6 月。
[DFPattern]Relay 中的模式匹配
下载 Python 源代码:uma.py
下载 Jupyter Notebook:uma.ipynb
以上就是该文档的全部内容,点击查看更多 TVM 中文文档。