本篇文章譯自英文文件 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 中文文件。