TVMNN编译Compiler栈
TVMNN編譯Compiler棧
內(nèi)容綱要
- 前言
- 調(diào)研目標
- TVM介紹
- TVM源碼架構
i. FrontEnd
ii. Relay
iii. BackEnd - VTA實現(xiàn)原理及設計思想提煉
i. 整體結構
ii. VTA Hardware
a.
a. 指令集
b. 數(shù)據(jù)流
c. 控制流
b. VTA Config
c. Pyng HLS
d. 硬件設計思想提煉
e. Chisel Scalar
f. SIM C++
g. Xilinx Scripts
iii. VTA JIT
a. Driver
b. Runtime
c. TVM Runtime Library
iv. VTA Complier - 參考文獻
前言
深度學習/神經(jīng)網(wǎng)絡應用日益廣泛,多終端部署形成常態(tài)。從CPU、ARM、GPU到專用的神經(jīng)網(wǎng)絡加速器/深度學習處理器,不同的終端/不同的體系結構引起神經(jīng)網(wǎng)絡的碎片化;每一款設備特別是專用的加速芯片部署深度學習是一件費力不討好的事情;同時近年來,雖然CNN加速器在學術上和產(chǎn)業(yè)上火熱,但不同的研究人員/企業(yè)在CNN加速器上使用不同的指令集,不同的體系結構,沒有統(tǒng)一的標準使開發(fā)人員在終端上部署。所以一款支持前端,各種后端,統(tǒng)一結構/指令集,擴展方便的NN編譯棧在NN生態(tài)上的構建有著舉足輕重的影響力,這就是TVM的魅力。
??現(xiàn)在主流的深度學習訓練框架是Caffe/PyTorch/TensorFlow/MxNet等,對CPU/CUDA支持得很好。如果想把訓練好的神經(jīng)網(wǎng)絡部署到其它的終端設備,這就帶了幾個挑戰(zhàn): - 主流框架不支持ARM/FPGA/ASIC
- 嵌入式終端不需要訓練功能,對前向推理的速度有極大的要求
- 嵌入式終端性能/內(nèi)存/存儲有限,主流框架的臃腫不適合部署
- 終端指令集,架構沒有統(tǒng)一標準,開發(fā)部署難度很大
??這時需要一款軟件編譯棧,上接前端主流深度學習框架,后接各種終端;同時滿足輕量級,高性能,高度擴展與靈活性,開發(fā)容易等要求。TVM 是深度學習系統(tǒng)的編譯器堆棧。旨在縮小以算力為重點的深度學習框架與以性能和效率為重點的硬件后端之間的差距。TVM與深度學習框架協(xié)同工作,為不同的后端提供端到端編譯。TVM支持主流的深度學習前端框架,包括TensorFlow, MXNet, PyTorch, Keras, CNTK;同時能夠部署到寬泛的硬件后端,包括CPUs, server GPUs, mobile GPUs, and FPGA-based accelerators。
??本文深度分析TVM的源代碼(主要是FPGA-based accelerator/VTA方面),總結TVM的幾個實現(xiàn)特點: - 對接前端的神經(jīng)網(wǎng)絡模型配置文件,將前端網(wǎng)絡轉(zhuǎn)為自主設計的AST Graph/Relay IR;后端編譯都基于這個Graph/Relay IR。這類似于實現(xiàn)一個通用的編譯器,將C++/JAVA/Python等語言轉(zhuǎn)換為自主的編譯型語言IR;然后基于IR實現(xiàn)一個高性能的后端程序部署。
- 實現(xiàn)NN的Tensor Operator Libray:不同的后端采用不同的實現(xiàn)方式。
- 根據(jù)后端硬件將前端網(wǎng)絡編譯為Complied PackedFunc生成動態(tài)鏈接庫,用于后端執(zhí)行。
- RunTime和Driver加載Complied PackedFunc調(diào)用硬件執(zhí)行推理計算。
- 在HalideIR/LLVM框架/思想上實現(xiàn)編譯系統(tǒng)。
??前端的分離方式便于TVM兼容主流的深度學習框架,后端的編譯運行分離使后端只需要部署一個輕量級的TVM(只包括RunTime和Driver)。這種思想能夠兼顧通用性與性能。同時TVM在VTA(NN加速器)上的設計了一套通用的類RISC指令集,體系結構。TVM的整體思想與最終目標就是通用,形成一個深度學習全棧生態(tài);但目前TVM實現(xiàn)的效果如何,讓走近源代碼分析一番。
調(diào)研目標
本文著重分析TVM的Relay IR 層,編譯,VTA(NN加速器)源代碼;ARM/CUDA部分目前不在源代碼分析范圍內(nèi)。
? 明確TVM層次結構,掌握TVM的用法,工作原理,支持的軟硬件框架
? 分析TVM源碼設計與實現(xiàn),對TVM的工作框架有清晰地、層次化的認識
? 研究TVM論文、文檔與源碼,提煉出TVM的設計思想;包括整體設計思想,Relay、VTA,并細化到JIT、Complier、Hardware Design等設計理念。
? 評估TVM從前端到后端的性能,著重于VTA硬件的實現(xiàn)指標
? 整體評估基于TVM設計加速器的可行性
? 具體分析TVM的細節(jié)實現(xiàn),包括VTA Hardware,JIT Diver,Complier等。
? 評估設計或修改加速器所需的軟件改動,包括修改難度、工作量
? 綜合以上,得出基于TVM編譯棧生態(tài)設計自己的處理器,部署神經(jīng)網(wǎng)絡的可行性及成熟度分析。
TVM介紹
??TVM是一個端到端的深度學習工具鏈,能夠高效地把前端深度學習模型部署到CPUs、GPUs和專用的加速器上。TVM具有強大的兼容能力和擴展性,支持主流的深度學習前端框架,包括TensorFlow, MXNet, PyTorch, Keras, CNTK;同時能夠部署到寬泛的硬件后端,包括CPUs, server GPUs, mobile GPUs, and FPGA-based accelerators。
TVM設計架構
TVM架構如下:
- 導入前端深度學習模型 CNN model configuration and parameters,生成計算圖
- 重構原始計算圖,基于operator-level優(yōu)化圖的計算
- 基于Tensor張量化計算圖,并根據(jù)后端進行硬件原語級優(yōu)化
- 根據(jù)優(yōu)化目標探索搜索空間,找到最優(yōu)解
- 生成硬件所需指令、數(shù)據(jù)并部署到硬件上
Tvm注意點:
? TVM只做前向計算包括量化,不對網(wǎng)絡進行訓練
? 硬件是預先可配置的,編譯優(yōu)化部署的時候硬件需保持固定。
? 可以部署TVM Runtime System到嵌入式設備上(SOC),實現(xiàn)CPU與VTA之間的自動交互和實時運行環(huán)境。
TVM源碼架構
TVM軟件架構
??TVM 源代碼由三層構成,包括FrontEnd接口、Relay優(yōu)化和BackEnd部署。目前著重研究BackEnd VTA層級和與硬件相關的Relay部分。設計上三層應該互不影響,各自獨立;但源代碼分析中,Relay層糅合了與硬件無關的圖優(yōu)化和與VTA相關的調(diào)度生成,沒有分得很開。
FrontEnd
??支持主流的深度學習前端框架,包括TensorFlow, MXNet, PyTorch, Keras, CNTK。目前TVM可以繼承到PyTorch框架中優(yōu)化、訓練,而不是單純地調(diào)用CNN模型接口。
Relay
??根據(jù)具體硬件對原始計算圖進行重構、張量優(yōu)化、數(shù)據(jù)重排等圖優(yōu)化操作。源代碼分析中,Relay層比較雜,干的事情比較多,既對接上層的圖優(yōu)化又對接硬件的調(diào)度器。
Relay及Tensorlization示例
BackEnd
后端支持ARM、CUDA/Metal/OpenCL及加速器VTA。
- 生成硬件所需指令流與數(shù)據(jù)打包
- 一個CPU與VTA的交互式運行環(huán)境:包括driver、JIT。
VTA實現(xiàn)原理及設計思想提煉
整體結構
VTA層次結構
VTA源代碼架構有三個部分組成:
- 硬件實現(xiàn):包括硬件源代碼、硬件配置文件
- 編譯器:以operator或func級打包計算圖生成調(diào)度器,以及編譯出硬件執(zhí)行指令
- JIT:軟硬件交互運行環(huán)境
? Driver:CPU與VTA的通信與控制驅(qū)動
? Runtime:CPU對VTA的實時運行控制
? TVM runtime library:VTA runtime的基本庫,主要是句柄的定義
VTA Hardware
??為了實現(xiàn)硬件的通用化計算,VTA硬件參考RISC指令集,按照Fetch—>Load—>Compute—>Store模式,將所有操作劃分到這種粒度的計算,不設計專用復雜的計算模式。在犧牲一定性能的情況下,盡可能兼容更多的深度學習網(wǎng)絡。
VTA硬件體系結構
指令集
VTA指令分為四大類:
- Fetch:從DRAM取指到片上指令SRAM
? load 指令buffer
? GEMM指令buffer
? Store指令buffer - LOAD:從DRAM取數(shù)據(jù)到片上SRAM
? inp buffer和wgt buffer
? uop buffer (源碼實現(xiàn)在compute 中)
? Acc buffer放bias(源碼實現(xiàn)在compute 中) - Compute:加載SRAM數(shù)據(jù)到計算單元并執(zhí)行GEMM或ALU計算
? ALU:執(zhí)行vector的計算Min、Max、ADD、MUL、SH
? GEMM:執(zhí)行Tensor乘法 - Store:從SRAM存儲數(shù)據(jù)到DRAM
VTA 指令結構
數(shù)據(jù)流
??VTA relies on dependence FIFO queues between hardware modules to synchronize the execution of concurrent tasks. The figure below shows how a given hardware module can execute concurrently from its producer and consumer modules in a dataflow fashion through the use of dependence FIFO queues, and single-reader/single-writer SRAM buffers. Each module is connected to its consumer and producer via read-after-write (RAW) and write-after-read (WAR) dependence queues。
VTA依賴于硬件模塊之間的依賴FIFO隊列來同步并發(fā)任務的執(zhí)行。下圖顯示了給定硬件模塊如何通過使用依賴FIFO隊列和單讀寫器SRAM緩沖區(qū),以數(shù)據(jù)流方式從其生產(chǎn)者和消費者模塊并發(fā)執(zhí)行。每個模塊通過讀后寫(RAW)和讀后寫(WAR)依賴隊列連接到其使用者和生產(chǎn)者。
Dataflow and Dependency
??The pseudo-code above describes how a module executes a given instruction predicated on dependences with other instructions. First, the dependence flags within each instruction are decoded in hardware. If the instruction has an incoming RAW dependences, execution is predicated upon receiving a RAW dependence token from the producer module. Similarly, if the task has an incoming WAR dependence, execution is predicated upon receiving a WAR dependence token from the consumer module. Finally when the task is done, we check for outgoing RAW and WAR dependences, and notify the consumer and producer modules respectively.
上面的偽代碼描述了一個模塊如何根據(jù)與其它指令的依賴關系,執(zhí)行給定的指令。首先,在硬件中對每條指令中的依賴標志進行解碼。如果指令具有傳入的原始依賴項,則在從生產(chǎn)者模塊接收到原始依賴項令牌時,將斷言執(zhí)行。類似地,如果任務具有傳入的WAR依賴性,則在從使用者模塊接收到WAR依賴性令牌時,就斷言執(zhí)行。最后,當任務完成時,檢查輸出的RAW和WAR依賴,并分別通知使用者和生產(chǎn)者模塊。
控制流
按照Fetch—>Load—>Compute—>Store模式去計算:
- CPU把數(shù)據(jù)和指令放到DRAM
- Fetch和Load分別將指令和數(shù)據(jù)從DRAM搬運到SRAM
- CPU會提前計算好一部分AGU地址uop放到DRAM,然后被搬運到uop Sram
- Compute譯碼指令根據(jù)指令參數(shù),執(zhí)行對應的GEMM或ALU計算;計算結果放回到acc/out sram。
- Store將計算結果從SRAM放回到DRAM
- CPU取出計算結果并準備下次計算數(shù)據(jù),返回第一步
Inp buffer和Out buffer之間沒有數(shù)據(jù)交換,固定buffer in/out。控制邏輯會處理數(shù)據(jù)/指令依賴以及memory latency hiding。
VTA Config
??VTA主要對硬件的數(shù)據(jù)位寬,SRAM 大小,計算陣列大小進行配置,而不能更改大的計算架構,數(shù)據(jù)流,控制流等;同時TVM根據(jù)已配置好的VTA執(zhí)行編譯工作,而不會在編譯階段在線生成硬件代碼。
VTA配置
Pyng HLS
??VTA采用高層次綜合C++實現(xiàn)Xilinix Pynq FPGA的部署,包含配置文件,Vivado HLS 腳本,HLS C++硬件實現(xiàn)三部分組成。
Pynq FPGA部署
Hw_spec和VTA Implementation構成HLS 實現(xiàn)的主體:
- hw_spec.h: 數(shù)據(jù)及指令定義
? 根據(jù)VTA Config 定義數(shù)據(jù)寬度,GEMM規(guī)模,SRAM大小
? 定義指令layout(指令關鍵字段的排列順序,如下面GEMM layout所示)
?
// GEMM Layout
// ________________|type|
// arg 0: opcode | opcode_T |
// arg 1: pop_prev_dependence | bool |
// arg 2: pop_next_dependence | bool |
// arg 3: push_prev_dependence | bool |
// arg 4: push_next_dependence | bool |
// arg 5: reset_reg | bool |
// arg 6: uop_bgn | uop_idx_T |
// arg 7: uop_end | uop_idx_T |
// arg 8: iteration count ax0 | loop_T |
// arg 9: iteration count ax1 | loop_T |
// arg a: accum idx factor ax0 | acc_idx_T |
// arg b: accum idx factor ax1 | acc_idx_T |
// arg c: input idx factor ax0 | inp_idx_T |
// arg d: input idx factor ax1 | inp_idx_T |
// arg e: weight idx factor ax0 | wgt_idx_T |
// arg f: weight idx factor ax1 | wgt_idx_T - VTA Implementation:硬件實現(xiàn)包括Stream、fetch、load、compute、store五大模塊及VTA頂層控制。
VTA HLS C++ 實現(xiàn)
? Stream:定義CPU與VTA之間的數(shù)據(jù)接口/數(shù)據(jù)流。
? fetch:將指令從DRAM根據(jù)譯碼分類分別搬運到SRAM三個buffer中:load_queue,gemm_queue,store_queue。
? load:從load_queue中讀取load指令,譯碼并計算對應地址;從DRAM上搬運inp和weight到SRAM上inp_mem和wgt_mem中。同時數(shù)據(jù)的padding工作通過在數(shù)據(jù)傳輸時對inp_mem寫零完成。
? compute:從gemm_queue中讀取計算指令,譯碼計算對應地址;從DRAM上搬運uop(計算所需的部分數(shù)據(jù),權重地址)和bias到SRAM;從SRAM中加載數(shù)據(jù)和權重,根據(jù)譯碼計算GEMM或者ALU。其中計算的核心數(shù)據(jù)流為:
? Input SRAM—>Input Register
? Input Register—> Computing Unit
? Result—>Out Register/Out SRAM
? Out Register—> Out SRAM
? Store:將計算結果從SRAM—>DRAM
? VTA:以上模塊的控制邏輯實現(xiàn)。
? 封裝stream:包括instuction queue和dependency queue。
? 實例化SRAM
? 取值
? 根據(jù)依賴調(diào)度存儲和計算(依賴參考上文數(shù)據(jù)流中的dependence)
硬件設計思想提煉
VTA的核心思想,將計算劃分到一個通用的細膩度的計算結構Operation Wrapper:
? Load Inst:加載Operation對應指令
? Decode:譯碼
? AGU:計算Operation所需要的DRAM或SRAM地址
? Operation:利用Operation對應的硬件計算資源執(zhí)行操作
??無論是Fetch、load、compute、store等較大粒度的操作都可以利用OperationWrapper分解到更細粒度的四個步驟,這和傳統(tǒng)的設計方法不同。舉例,傳統(tǒng)設計是將芯片分為DMA—>Load—>Compute—>Store四個模式,每個模式單獨去設計一套盡量通用的操作單元,去匹配軟件算法;同時每個模式不會再分解到一個更通用的細膩度結構;OperationWrapper將Fetch—>Load—>Compute—>Store中每個模式都分解為Load—>Decode—>AGU—>Operation,即使是傳統(tǒng)認為Load/Store已經(jīng)是最小粒度,TVM依然再次分解到OperationWrapper。這有點像RISC中再RISC的思想。
VTA 硬件設計思想
??另外,傳統(tǒng)方法設計一個統(tǒng)一的AGU負責所有的地址計算;而TVM通過OperationWrapper將統(tǒng)一的AGU分解到每個模塊,讓每個模塊負責一個更細粒度的AGU。這種方式能夠加強AGU的通用性。
??VTA這種OperationWrapper的思想,將硬件進一步分解到更加通用的細粒度模式,比起傳統(tǒng)方法通用型更強。
??VTA的第二個核心思想,利用Stream傳輸并通過同步控制解決指令與數(shù)據(jù)的依賴性(有點類似操作系統(tǒng)中的設計理念),讓CPU與VTA的通訊、傳輸、控制變得更向操作系統(tǒng)級靠攏,可能也是為了TVM的通用性考慮。
參考文獻
TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
VTA: A Hardware-Software Blueprint for Flexible Deep Learning Specialization
TVM Tutorials
TVM GitHub
總結
以上是生活随笔為你收集整理的TVMNN编译Compiler栈的全部內(nèi)容,希望文章能夠幫你解決所遇到的問題。
- 上一篇: TVM适配NN编译Compiler缺陷
- 下一篇: PaddlePaddle推理部署