當前位置：首頁 > 人文社科 > 生活经验 >内容正文

生活经验

TVM适配NN编译Compiler缺陷

發布時間：2023/11/28 生活经验 41 豆豆

生活随笔收集整理的這篇文章主要介紹了 TVM适配NN编译Compiler缺陷小編覺得挺不錯的,現在分享給大家,幫大家做個參考.

TVM適配NN編譯Compiler缺陷
內容綱要

前言
TVM針對VTA的編譯流程
i. 自定義VTA架構：TVM的缺陷與性能瓶頸
TVM缺陷與瓶頸
i. 缺陷一：SRAM配置靈活性差
ii. 缺陷二：計算陣列配置僵硬
iii. 缺陷三：網絡支持少
TVM源碼修改之靜態調度搜索算法
前言
前文NN編譯棧之TVM研究報告深度分析TVM的源碼結構，編譯器特點。本文介紹TVM的當前缺陷以及如何修改源代碼彌補缺陷并適配自己開發的神經網絡加速器。不久會在GitHub上開源自己的適配修改工作并向TVM倉庫提交新的版本。
??現在主流的深度學習訓練框架是Caffe/PyTorch/TensorFlow/MxNet等，對CPU/CUDA支持得很好。如果想把訓練好的神經網絡部署到其它的終端設備，這就帶了幾個挑戰：
主流框架不支持ARM/FPGA/ASIC
嵌入式終端不需要訓練功能，對前向推理的速度有極大的要求
嵌入式終端性能/內存/存儲有限，主流框架的臃腫不適合部署
終端指令集，架構沒有統一標準，開發部署難度很大
??TVM 是深度學習系統的編譯器堆棧。旨在縮小以算力為重點的深度學習框架，以性能和效率為重點的硬件后端之間的差距。 TVM與深度學習框架協同工作，為不同的后端提供端到端編譯。TVM支持主流的深度學習前端框架，包括TensorFlow, MXNet, PyTorch, Keras, CNTK；同時能夠部署到寬泛的硬件后端，包括CPUs, server GPUs, mobile GPUs, and FPGA-based accelerators。
??如果想自己設計一款深度學習處理器（VTA）并兼容TVM編譯棧，那么該如何開發？TVM編譯器的擴展性如何？如何修改TVM編譯器，適配自定義的深度學習處理器？
TVM針對VTA的編譯流程
?? TVM定義了VTA的指令集，體系結構。為了實現硬件的通用化計算，VTA硬件參考RISC指令集，按照Fetch—>Load—>Compute—>Store模式，將所有操作劃分到這種粒度的計算，不設計專用復雜的計算模式。

VTA硬件體系結構
??同時，TVM定義了VTA硬件的可重構約束。VTA主要對硬件的數據位寬，SRAM 大小，計算陣列大小進行配置，而不能更改大的計算架構，數據流，控制流等。

VTA配置
??假設VTA配置已經固定下來，TVM會按照如下流程編譯VTA：

定義前端網絡：對接MxNet。

#“tvm/vta/tutorials/frontend/deploy_resnet_on_vta.py”

Name of Gluon model to compile

The `<code>start_pack</code><code> and </code><code>stop_pack</code>` labels indicate where

to start and end the graph packing relay pass: in other words

where to start and finish offloading to VTA.

model = “resnet18_v1”
start_pack = “nn.max_pool2d”
stop_pack = “nn.global_avg_pool2d”
…
dtype_dict = {“data”: ‘float32’}
shape_dict = {“data”: (env.BATCH, 3, 224, 224)}

Get off the shelf gluon model, and convert to relay

gluon_model = vision.get_model(model, pretrained=True)
2. 將前端網絡轉為自己設計的AST Graph/Relay IR：“relay.frontend.from_mxnet“函數

#"tvm/vta/tutorials/frontend/deploy_resnet_on_vta.py"

Start front end compilation

mod, params = relay.frontend.from_mxnet(gluon_model, shape_dict)

Update shape and type dictionary

shape_dict.update({k: v.shape for k, v in params.items()})
dtype_dict.update({k: str(v.dtype) for k, v in params.items()})
3. 量化網絡：這里應該采用的post-quantization方法，不需要重新訓練或autotune。

#"tvm/vta/tutorials/frontend/deploy_resnet_on_vta.py"

Perform quantization in Relay

with relay.quantize.qconfig(global_scale=8.0,
skip_conv_layers=[0]):
relay_prog = relay.quantize.quantize(mod[“main”], params=params)
4. 圖打包和常量展開：預先計算常量節點，減少參數量和計算量。

#"tvm/vta/tutorials/frontend/deploy_resnet_on_vta.py"

Perform graph packing and constant folding for VTA target

if target.device_name == “vta”:
assert env.BLOCK_IN == env.BLOCK_OUT
relay_prog = graph_pack(
relay_prog,
env.BATCH,
env.BLOCK_OUT,
env.WGT_WIDTH,
start_name=start_pack,
stop_name=stop_pack)
5. VTA編譯，生成Complied PackedFunc動態鏈接庫。

#"tvm/vta/tutorials/frontend/deploy_resnet_on_vta.py"

with vta.build_config():
graph, lib, params = relay.build(
relay_prog, target=target,
params=params, target_host=env.target_host)
VTA編譯這步是研究的重點，分析TVM如何針對特定硬件進行編譯，步驟如下：
? 將網絡分解為針對VTA硬件的卷機計算圖:
? 卷機topi.nn.conv2d
? 移位topi.right_shift
? bias相加topi.add
? 溢出處理my_clip

  # tvm/vta/tests/python/integration/test_benchmark_topi_conv2d.py

Define base computation schedule

with target:
res = topi.nn.conv2d(
data, kernel, (wl.hstride, wl.wstride), (wl.hpad, wl.wpad), (1, 1),
layout, env.acc_dtype)
res = topi.right_shift(res, 8)
res = topi.add(res, bias)
res = my_clip(res, 0, (1 << env.OUT_WIDTH - 1) - 1)
? 來查看VTA的卷積是如何定義的：
? 傳入數據（data，kernel）及卷機形狀定義（strides，padding）
? 定義占位符（類似TensorFlow）reduce_axis
? 卷積張量定義tvm.compute
# tvm/vta/python/vta/top/vta_conv2d.py
@autotvm.register_topi_compute(topi.nn.conv2d, ‘vta’, ‘direct’)
def _declaration_conv2d(cfg,
data,
kernel,
strides,
padding,
dilation,
layout,
out_dtype):
“”" Packed conv2d function."""
if not is_packed_layout(layout):
raise topi.InvalidShapeError()
assert dilation == (1, 1)
d_i = tvm.reduce_axis((0, kshape[2]), name=‘d_i’)
d_j = tvm.reduce_axis((0, kshape[3]), name=‘d_j’)
k_o = tvm.reduce_axis((0, ishape[1]), name=‘k_o’)
k_i = tvm.reduce_axis((0, ishape[-1]), name=‘k_i’)
hstride, wstride = strides
res = tvm.compute(
oshape,
lambda b_o, c_o, i, j, b_i, c_i: tvm.sum(
pad_data[b_o, k_o, ihstride+d_i, jwstride+d_j, b_i, k_i].astype(out_dtype) *
kernel[c_o, k_o, d_i, d_j, c_i, k_i].astype(out_dtype),
axis=[k_o, d_i, d_j, k_i]),
name=“res”, tag=“conv2d_dense”)

cfg.add_flop(2 * np.prod(topi.util.get_const_tuple(oshape)) *
kshape[2] * kshape[3] * ishape[1] * ishape[-1])

return res
6. 生成卷積層的調度方案
# tvm/vta/tests/python/integration/test_benchmark_topi_conv2d.py

Derive base schedule

s = topi.generic.schedule_conv2d_nchw([res])
? 看一下生成的調度器是什么樣子：
? 綁定數據與對應buffer
? DMA搬運
? 循環分片：包括Channels、Height、Width
? 地址生成：數據加載存儲地址
? 虛擬線程：掩蓋訪存延遲

# tvm/vta/tests/python/integration/test_benchmark_topi_conv2d.py

print(vta.lower(s, [data, kernel, bias, res], simple_mode=True))
##調度器部分代碼
// attr [res_conv] storage_scope = “local.acc_buffer”
// attr [data_buf] storage_scope = “local.inp_buffer”
// attr [kernel_buf] storage_scope = “local.wgt_buffer”
produce res {
vta.coproc_dep_push(3, 2)
vta.coproc_dep_push(3, 2)
for (ic.outer, 0, 8) {//input channel
// attr [iter_var(vta, , vta)] coproc_scope = 1
vta.coproc_dep_pop(2, 1)
produce data_buf {
VTALoadBuffer2D(tvm_thread_context(VTATLSCommandHandle()), data, ((((ic.outer196) + (i2.outer98)) + (max((1 - (i2.outer7)), 0)14)) - 14), 14, ((9 - max((1 - (i2.outer7)), 0)) - max(((i2.outer7) - 6), 0)), 14, 1, max((1 - (i2.outer7)), 0), 1, max(((i2.outer7) - 6), 0), 0, 2)
}
produce kernel_buf {
VTALoadBuffer2D(tvm_thread_context(VTATLSCommandHandle()), kernel, ((i1.outer.outer1152) + (ic.outer9)), 9, 8, 72, 0, 0, 0, 0, 0, 1)
}
vta.coproc_dep_push(1, 2)
// attr [iter_var(vta, , vta)] coproc_scope = 1
vta.coproc_dep_pop(2, 1)
produce data_buf {
VTALoadBuffer2D(tvm_thread_context(VTATLSCommandHandle()), data, ((((ic.outer196) + (i2.outer98)) + (max((1 - (i2.outer7)), 0)14)) - 14), 14, ((9 - max((1 - (i2.outer7)), 0)) - max(((i2.outer7) - 6), 0)), 14, 1, max((1 - (i2.outer7)), 0), 1, max(((i2.outer7) - 6), 0), 144, 2)
}
produce kernel_buf {
VTALoadBuffer2D(tvm_thread_context(VTATLSCommandHandle()), kernel, (((i1.outer.outer1152) + (ic.outer9)) + 576), 9, 8, 72, 0, 0, 0, 0, 72, 1)
}
? 然后看一下調度器是什么計算的
? 加載默認的分片配置文件cfg
? 綁定數據與對應buffer:set_scope
? DMA：compute_at和cache_read
? 循環分片：apply
? 虛擬線程：split
? 指令生成:pragma

# tvm/vta/python/vta/top/vta_conv2d.py

@autotvm.register_topi_schedule(topi.generic.schedule_conv2d_nchw, ‘vta’, ‘direct’)
def _schedule_conv2d(cfg, outs):
# 綁定數據與對應buffer
cdata = s.cache_read(data, env.inp_scope, [conv2d_stage])
ckernel = s.cache_read(kernel, env.wgt_scope, [conv2d_stage])

DMA cache數據

s[conv2d_stage].set_scope(env.acc_scope)

tile 循環分片

x_bo, x_co, x_i, x_j, x_bi, x_ci = s[output].op.axis
x_co0, x_co1 = cfg[‘tile_co’].apply(s, output, x_co)
x_i0, x_i1 = cfg[‘tile_h’].apply(s, output, x_i)
x_j0, x_j1 = cfg[‘tile_w’].apply(s, output, x_j)
s[output].reorder(x_bo, x_i0, x_co0, x_j0, x_co1, x_i1, x_j1, x_bi, x_ci)
store_pt = x_j0

set all compute scopes DMA

for tensor in cache_read_ewise:
s[tensor].compute_at(s[output], store_pt)
s[tensor].pragma(s[tensor].op.axis[0], env.dma_copy)

virtual threading along output channel axes 虛擬線程

if cfg[‘oc_nthread’].val > 1:
_, v_t = s[output].split(x_co0, factor=cfg[‘oc_nthread’].val)
s[output].reorder(v_t, x_bo)
s[output].bind(v_t, tvm.thread_axis(“cthread”))

Use VTA instructions 指令生成

s[cdata].pragma(s[cdata].op.axis[0], env.dma_copy)
s[ckernel].pragma(s[ckernel].op.axis[0], env.dma_copy)
s[conv2d_stage].tensorize(x_bi, env.gemm)
s[output].pragma(x_co1, env.dma_copy)

return s
7. 到這一步Complied PackedFunc已經生成好了，可以給VTA硬件編程使用

#"tvm/vta/tutorials/frontend/deploy_resnet_on_vta.py"

vta.reconfig_runtime(remote)
vta.program_fpga(remote, bitstream=None)
8. 前向測試:resnet18_v1分類結果并統計訪存計算數據。

resnet18_v1

Execution statistics:
inp_load_nbytes : 48864512
wgt_load_nbytes : 1695547392
acc_load_nbytes : 6723584
uop_load_nbytes : 36
out_store_nbytes: 1680896
gemm_counter : 1655808
alu_counter : 286160

resnet18_v1 prediction
#1: tiger cat
#2: Egyptian cat
#3: tabby, tabby cat
#4: lynx, catamount
#5: weasel
??總的來說，TVM針對前端網絡和VTA的編譯，調用流程如下。其中重點是調度方案的生成，對不同架構的VTA適配最需要的就是定義對應架構的調度方案。

定義前端網絡
將前端網絡轉為自己設計的AST Graph/Relay IR
量化網絡
圖打包和常量展開
VTA編譯，生成Complied PackedFunc動態鏈接庫
生成卷積層的調度方案
VTA硬件編程
前向測試
自定義VTA架構：TVM的缺陷與性能瓶頸
??自定義深度學習處理器VTA架構，可以分為兩種設計方案：
魔改VTA架構，指令集：需要對TVM編譯器作大規模源碼更改，目前TVM的針對VTA的編譯是固定指令集與架構，擴展性較差。
保持VTA架構不變，修改配置：修改硬件的數據位寬，SRAM 大小，計算陣列大小進行配置，而不能更改大的計算架構，數據流，控制流等。
本文的目的是分析TVM的缺陷與瓶頸，所以先從最簡單的第二種方案研究。VTA架構配置更改在"tvm/vta/config/vta_config.json",參考前文進行對應的修改。
{
“TARGET” : “sim”,
“HW_VER” : “0.0.1”,
“LOG_INP_WIDTH” : 3,
“LOG_WGT_WIDTH” : 3,
“LOG_ACC_WIDTH” : 5,
“LOG_BATCH” : 0,
“LOG_BLOCK” : 4,
“LOG_UOP_BUFF_SIZE” : 15,
“LOG_INP_BUFF_SIZE” : 15,
“LOG_WGT_BUFF_SIZE” : 18,
“LOG_ACC_BUFF_SIZE” : 17
}
TVM缺陷與瓶頸
缺陷一：SRAM配置靈活性差
??VTA配置SRAM大小的時候不是任意可配的，假如固定計算位寬和計算規模，理論上講BUFF_SIZE只要大于一個最小值就可以。但其實不是這樣，BUFF_SIZE必須固定，計算方式本文推定如下。總的來說，SRAM配置不夠靈活，如果芯片/FPGA資源過小或過大，都不好兼容TVM。
VTA_UOP_BUFF_SIZE=VTA_LOG_UOP_BUFF_DEPTH+VTA_LOG_UOP_WIDTH-3=15
VTA_INP_BUFF_SIZE=VTA_LOG_INP_BUFF_DEPTH+VTA_LOG_BLOCK_IN+VTA_LOG_INP_WIDTH-3=
VTA_WGT_BUFF_DEPTH=VTA_LOG_WGT_BUFF_DEPTH+VTA_LOG_BLOCK_OUT + VTA_LOG_BLOCK_IN + VTA_LOG_WGT_WIDTH - 3
VTA_ACC_BUFF_DEPTH=VTA_LOG_ACC_BUFF_DEPTH+VTA_LOG_BATCH + VTA_LOG_BLOCK_OUT + VTA_LOG_ACC_WIDTH - 3
VTA_LOG_UOP_BUFF_DEPTH必須固定為15，沒有什么道理可講。TVM固定UOP_BUFF_DEPTH必須為8192
LOG_INP_BUFF_SIZE/LOG_ACC_BUFF_SIZE必須保持上面的計算方式，不可大于計算值。
LOG_WGT_BUFF_SIZE可小于計算值，靈活性還算可以。
缺陷二：計算陣列配置僵硬
??如果修改計算陣列，比如16×16到32×32或者64×64,那么VTA編譯依然按照16×16的調度方案編譯新的架構，修改的配置毫無意義。同時新的架構性能比16×16還差。如下圖所示，32×32訪存比16×16高出95倍，帶來極大的性能損失。
resnet18_v1:16x16 架構
Execution statistics:
inp_load_nbytes : 5549568
wgt_load_nbytes : 12763136
acc_load_nbytes : 30720
uop_load_nbytes : 1540
out_store_nbytes: 1680896
gemm_counter : 6623232
alu_counter : 572320
resnet18_v1 prediction for sample 0
#1: tiger cat
#2: Egyptian cat
#3: tabby, tabby cat
#4: lynx, catamount
#5: weasel
resnet18_v1:32x32架構
Execution statistics:
inp_load_nbytes : 48864512
wgt_load_nbytes : 1695547392
acc_load_nbytes : 6723584
uop_load_nbytes : 36
out_store_nbytes: 1680896
gemm_counter : 1655808
alu_counter : 286160
resnet18_v1 prediction for sample 0
#1: tiger cat
#2: Egyptian cat
#3: tabby, tabby cat
#4: lynx, catamount
#5: weasel
缺陷三：網絡支持少
??雖然TVM宣稱支持主流的深度學習框架及大量神經網絡模型，但是當前TVM（*本文指VTA，ARM/CUDA沒有深入研究）支持/適配的網絡較少，目前本文在不改變任何代碼的情況下只跑通了默認的resnet18_v1。支持其它網絡模型較少的原因如下：
? Relay量化目前只支持有限的操作類型
? VTA目前只支持有限的操作類型
? Relay量化Bug：如果部署resnet50_v1，會報如下錯誤：

1 Check failed: lhs->dtype == dtype (int8 vs. int32)

Separate quantization code base into different files: partition.cc, annotate.cc, realize.cc
Change rewrite_for_vta to extra partition pass and enable it by default
Change annotation.force_cast(x) to annotation.cast_hint(x, dtype)
Remove qconfig.store_lowbit_output and enable it by default
Fixed accuracy of models like mobilenet：
resnet18_v1(8-16bit): 69.29%
resnet18_v1(8-32bit): 69.29%
resnet34_v1: 73.33%
resnet50_v1: 74.78%
resnet101_v1: 75.66%
mobilenetv2_1.0: 66.64%
TVM源碼修改之靜態調度搜索算法
??VTA的架構配置僵化與性能瓶頸原因在于TVM編譯器僵化的調度方案。僵化原因如上文VTA編譯流程所說，VTA采用寫死的分片配置文件cfg，無論怎么修改架構，分片方案都不變。所以性能瓶頸出現。

tvm/vta/python/vta/top/vta_conv2d.py

@autotvm.register_topi_schedule(topi.generic.schedule_conv2d_nchw, 'vta', 'direct')
def _schedule_conv2d(cfg, outs):
??有解決方案嗎？TVM的設計方案是AutoTVM，一種在線的調度空間搜索算法：

修改VTA配置
聯機FPGA：不支持仿真模式
對調度方案進行迭代，大概1000次（可配置）
找尋在所有迭代方案中性能最優的調度器。
??這種方案不支持仿真模式，因為仿真迭代速度太慢（仿真 VS FPGA類似CPU VS GPU）。本文對這種方案表示不解，原因如下：
如果是做芯片ASIC，設計階段的在線迭代會很麻煩。
如果硬件VTA的架構固定，調度方案是確定解。完全可以實現一套靜態的調度方案，根據VTA配置，在編譯階段就可以確定最優的調度。
??本文實現了這種靜態調度搜索算法，針對前端網絡與VTA硬件在編譯階段搜索最優的調度器。實現性能的大幅提升，不需要FPGA這種在線AutoTVM迭代方案。編譯時間的增加可忽略不計。

resnet18_v1:32x32架構靜態調度搜索算法
Execution statistics:
inp_load_nbytes : 5743104
wgt_load_nbytes : 12730368
acc_load_nbytes : 29696
uop_load_nbytes : 829
out_store_nbytes: 1680896
gemm_counter : 1655808
alu_counter : 286160
resnet18_v1 prediction for sample 0
#1: tiger cat
#2: Egyptian cat
#3: tabby, tabby cat
#4: lynx, catamount
#5: weasel
??如上圖搜索，本文實現的搜索算法實現32×32架構比16×16結構性能提高四倍，訪存保持不變。徹底解決TVM編譯器的固有缺陷與性能瓶頸。不久會更新解釋算法原理與實現細節，并在GitHub上開源自己的靜態調度搜索算法并向TVM倉庫提交新的版本。

總結

以上是生活随笔為你收集整理的TVM适配NN编译Compiler缺陷的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。

上一篇：自动微分基本理论
下一篇： TVMNN编译Compiler栈