本篇文章譯自英文文檔 Making your Hardware Accelerator TVM-ready with UMA

作者是 Michael J. Klaiber,Christoph Gerum,Paul Palomero Bernardo。

更多 TVM 中文文檔可訪問 →TVM 中文站

本節(jié)介紹通用模塊化加速器接口（UMA）。UMA 提供了一個易用的 API 來將新的硬件加速器集成到 TVM 中。

本教程詳細(xì)介紹了如何利用 UMA 使得你的硬件加速器可直接用于 TVM。雖然這個問題沒有萬能的解決方案，但 UMA 旨在提供一個穩(wěn)定的純 Python API，從而將許多種類的硬件加速器集成到 TVM 中。

本教程將通過三個逐漸復(fù)雜的用例來介紹 UMA API。這些用例引入了三個模擬加速器 Vanilla、Strawberry 和 Chocolate，并用 UMA 將它們集成到 TVM 中。

Vanillau200b

Vanilla 是一個由 MAC 數(shù)組組成的簡單加速器，沒有內(nèi)部存儲器。它只能處理 Conv2D 層，所有其他層都在 CPU 上執(zhí)行，同時也協(xié)調(diào) Vanilla。 CPU 和 Vanilla 共享內(nèi)存。

Vanilla 的 C 接口 vanilla_conv2dnchw(...) 用于執(zhí)行 Conv2D 操作（包括 same-padding），它接收指向輸入特征圖、權(quán)重和結(jié)果的指針，以及 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 為新的加速器創(chuàng)建帶有 API（UMA-API）調(diào)用的代碼骨架。

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 數(shù)據(jù)流模式（DFPattern）。

def conv2d_pattern():
    pattern = is_op("nn.conv2d")(wildcard(), wildcard())
    pattern = pattern.has_attr({"strides": [1, 1]})
 return pattern

為了將輸入計算圖的 Conv2D 算子映射到 Vanilla 的底層函數(shù)調(diào)用 vanilla_conv2dnchw(...)，在 VanillaAcceleratorBackend 中注冊了 TIR pass VanillaAcceleratorConv2DPass（稍后討論）。

Codegen

文件 vanilla_accelerator/codegen.py 定義了靜態(tài) C 代碼，它被添加到生成的結(jié)果 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)

構(gòu)建神經(jīng)網(wǎng)絡(luò)并在 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）生成輸出文件。

輸出結(jié)果：

Generated files are in /tmp/tvm-debug-mode-tempdirs/2022-07-13T13-26-22___x5u76h0p/00000

查看生成的文件：

輸出結(jié)果：

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 中，可以看到生成的代碼調(diào)用了 Vanilla 的 C-API，然后執(zhí)行了一個 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;
}

Strawberryu200b

即將上線

Chocolateu200b

即將上線

征求社區(qū)意見u200b

若本教程不適合你的加速器，請將你的需求添加到 TVM 論壇中的 UMA 帖子 中。我們很樂意通過擴(kuò)展本教程來提供更多指導(dǎo)，例如如何利用 UMA 接口使得更多種類的 AI 硬件加速器可直接用于 TVM。

參考u200b [UMA-RFC]UMA：通用模塊化加速器接口，TVM RFC，2022 年 6 月。

[DFPattern]Relay 中的模式匹配

下載 Python 源代碼：uma.py

下載 Jupyter Notebook：uma.ipynb

以上就是該文檔的全部內(nèi)容，點擊查看更多 TVM 中文文檔。