陳天奇的tvm更新了：VTA，開源AI晶元棧

科技 07-14

陳天奇發表於 tvm社區

量子位搬運 | 公眾號 QbitAI

華盛頓大學陳天奇博士近日在tvm社區介紹了VTA(Versatile Tensor Accelerator)項目，量子位全文搬運如下。

問題：不只是硬體設計

2016年的秋天，我收到我們系硬體的教授Luis的郵件邀請暢談合作計劃，當時正在醞釀TVM計劃的我這樣認識了同系的法國同學Thierry，開始了我們將近兩年的合作。AI晶元和硬體加速是深度學習發展寫下一個大的方向，而如何設計專用加速晶元無疑是大家都感興趣的話題，Thierry也不例外。加速晶元專用性強，反而可以使得邏輯設計更加趨於簡單，在設計思路上面也可以更加開放，不用受到傳統指令集和功能完整性因素的制約，正好適合在研究前沿的高校同學一展身手。

然而，當我們的合作真正開始的時候，才發現問題比原來想的要困難的多：深度學習晶元的主要問題遠遠超過了硬體本身。深度學習硬體真正的目標，是希望大家都可以把自己的深度學習的模型跑在AI晶元上。而加速器帶來的特殊定製導致對於每個加速器，我們都需要重新設計一套驅動，以及上面的軟體棧—-包含編譯，優化和部署模塊。雖然以高校研究人員的能力可以設計出不亞於工業級別的晶元設計，但是他還需要重頭打造所有的這些模塊，才可以讓大家都真正的把自己的模型運行在他的加速器之上。

軟體棧設計本身的巨大壓力成為了讓我們的項目完成的一道大大的鴻溝。而反觀現在的行業現狀，不論是高校還是公司，大家都面臨的同樣的一個問題 — 深度學習晶元從來不只是硬體的問題，而是涉及到一個從硬體，到驅動，到編譯優化部署，再到模型本身的全棧問題。

一直以來我們不斷在思考這個問題：如何降低AI晶元設計門檻，讓每一個人都可以玩深度學習硬體加速和系統開發。

今天，由Thierry領銜設計，我合作參與的項目VTA(Versatile Tensor Accelerator)作為TVM stack的一部分正式亮相，也給出了對於這個問題的一個答案。

VTA是什麼

VTA的全名叫做Versatile Tensor Accelerator，直譯過來為靈活的張量加速器。首先，VTA是一個完全開源的深度學習加速器。但是VTA不光包含了加速器設計本身，完整的驅動，tvm編譯的整合和直接從tvm前端python編譯部署深度學習模型的完整開源工具鏈。本次發布的VTA包含了模擬器和FPGA部署模塊。當然因為設計本身完全開源，有興趣的同學也可以直接基於目前的設計進行流片設計。

作為tvm社區的一部分，VTA對於不同的人有著不同的意義：

對於硬體設計者而言，VTA就像一個硬體和軟體棧的藍圖。為不論是基於VTA擴展定製，或者是給自己的加速器增加TVM的支持提供詳細的參照。硬體設計者也可以通過tvm快速的直接在自己的硬體上運行深度學習程序，減少開放開銷。VTA本身也會不斷演化發布各種可定製設計來支持各種應用的需求。對於硬體設計研究者而言，VTA提供了一套從軟體到硬體完全開源的測試平台。

對於深度學習的從業人員來說，新的深度學習優化往往需要考慮到加速硬體的特性。如何讓系統支持未來的各種加速器優化成為了系統設計人員頭疼的問題。VTA的設計總結了現在的深度學習加速器的特性，並且會不斷更新反應最新的研究成果。系統設計者可以通過VTA為目標來完成優化深度學習加速器的一個基石。這也是TVM社區加速器支持發展的路線之一。同時完全開放的設計帶來了更多軟硬體結合設計的機會。對於深度學習演算法同學來說，有一個開放的硬體平台可以讓大家更好地研究如量化，低精度，壓縮和設計針對加速硬體優化的模型。

設計細節: 精簡硬體和軟硬體協同設計

VTA的設計在同時發布的論文中有詳細解說，麻雀雖小，五臟俱全，基本反應了TPU數據流風格的AI加速器設計經驗。我就講一些個人的體會。VTA的設計總結了目前已經有的深度學習晶元的設計經驗。最後設計的結果是我和Thierry在兩年之內不斷討論凝練的結果。

一開始Thierry給出的設計，比較偏向於複雜指令集(CISC)和固定的流水線，但我們發現RISC的微指令變成加上CISC訪存和單元張量指令混合的風格可以做出更好的設計。原因非常簡單，如果把碰到的問題都通過硬體的手段來解決，勢必會給硬體增加很多負擔。以深度學習應用為例，一般TPU類的加速器會有一個矩陣乘法核心，而卷積等操作都會轉化到這個矩陣乘法核心來做。CISC的設計方式會使得我們的指令集裡面直接有一個conv2d的卷積操作，然後由硬體來把這個指令解析成微指令。但是這麼做並不是最好的，因為光是卷積，就需要支持stride, dilation，如果對於醫學圖像應用，我們需要支持三維卷積。因為細微的應用差別就需要重新改變硬體顯然並不是我們想要的。這一點其實也體現在了深度學習加速器的學術研究上 — 大家都會說我們的硬體設計了XX功能解決了XX問題。比如我之前在會議上有看到專門用於處理deconvolution的GAN加速器。

但是事實上，並非所有的問題都適合在硬體解決。針對剛才的conv2d的例子，雖然深度學習晶元的核心是矩陣乘法或者卷積單元，但是如何利用這些簡單的基本單元組合成我們想要的操作的微指令本身，其實可以直接暴露給軟體來可編程。這樣的設計讓硬體「越來越笨」，而把更多的任務交給軟體也來接。這正是RISC的思想。當然天下沒有免費的午餐，更加靈活的設計需要更多的軟體的支持，這也是VTA在tvm stack中協同設計的原因。

當然，深度學習加速器肯定和傳統的RISC不同 —傳統的CPU的指令為單個寄存器指令，一定程度上VTA的loop源語類似於比較靈活的表達loop的CISC。但是把大的任務打碎成基本的張量計算單元的這樣的思想是一致的。這樣的設計涉及到更多的軟硬體協同，但是一旦處理的當，也可以使得加速器變得更加靈活。比如VTA目前的設計可以直接支持是誰三維卷積，對於之前的GAN加速器的例子也都可以直接通過編程實現。

VTA裡面最讓我喜歡的設計是顯式的控制流依賴操作。利用這些控制流依賴我們可以直接在軟體層面實現動態的流水線，達到掩蓋內存開銷的目的。有趣的是，當我們回過頭來看這個過程時候，在VTA設計的這個過程我們不斷重新發現了計算機體系結構教科書裡面要解決的問題。比如內存讀寫掩蓋的這個技術一定程度上可以看成是SMT(simultaneous multithreading)，我們現在的系主任是主要的發明人之一。為了用好這個設計，我們必須在軟體層面做出創新，但是反過來使得VTA可以擁有更加靈活的流水線依賴。

除了控制流之外，我們還直接為VTA設計了一個完整的JIT驅動，並且可以通過RPC直接在自己筆記本上遠程運行FPGA加速器實現，真正做到用python去編寫加速器程序。這些設計無一包含了軟體和硬體協同設計的痕迹，這也需要軟體和硬體專家的同時努力。我相信未來的深度加速器也會往「更加聰明的軟體，更加笨的硬體」，和更多軟硬體協同方向發展。

未來展望

作為TVM社區加速器支持的重要一環，我相信VTA可以使得大家都可以玩加速器的目標更近一步。當然VTA的設計並非完美，我們也會不斷的完善和擴展硬體和軟體層，同時結合TVM社區的其他模塊一起來做聯合的優化。

圍繞這個問題，這個方向還留下包括硬體抽象，協同設計，模型表示等許多非常有趣的問題等著大家去探索,也歡迎更多的小夥伴加入TVM社區。

資源鏈接

VTA英文博客：

https://tvm.ai/2018/07/12/vta-release-announcement.html

技術文檔：

https://docs.tvm.ai/vta/

代碼倉庫：

https://github.com/dmlc/tvm/tree/master/vta

有興趣的小夥伴可以直接下載代碼玩，目前支持FPGA運行。即使沒有FPGA，也可以直接通模擬器運行編譯出來的加速器代碼，具體可以參考官方文檔教程。

—完—