引言

有將近兩個(gè)月沒有學(xué)習(xí)一些新東西，更新一下博客了。一直在忙公司的一個(gè)項(xiàng)目，是做一款支持LSTM和RNN的通用架構(gòu)加速IP。自己恰好負(fù)責(zé)指令編譯工作，雖然開始的指令比較粗糙，沒有一套完整的編譯器架構(gòu)。但是這其中也去思考了一下基于FPGA加速器的編譯器架構(gòu)。在FPGA深度學(xué)習(xí)加速器中，編譯器除了需要自動(dòng)化生成指令外，還要優(yōu)化指令的結(jié)構(gòu)，來最大化加速器性能。TVM是一個(gè)支持GPU、CPU、FPGA指令生成的開源編譯器框架，雖然在我們自己的加速IP上無法直接拿過來用，但是其中的很多方法和思想還是很值得借鑒的。TVM最大的特點(diǎn)是基于圖和算符結(jié)構(gòu)來優(yōu)化指令生成，最大化硬件執(zhí)行效率。其中使用了很多方法來改善硬件執(zhí)行速度，包括算符融合、數(shù)據(jù)規(guī)劃、基于機(jī)器學(xué)習(xí)的優(yōu)化器等。它向上對(duì)接Tensorflow、Pytorch等深度學(xué)習(xí)框架，向下兼容GPU、CPU、ARM、TPU等硬件設(shè)備。

1. 整體結(jié)構(gòu)

TVM是一個(gè)端到端的指令生成器。它從深度學(xué)習(xí)框架中接收模型輸入，然后進(jìn)行圖的轉(zhuǎn)化和基本的優(yōu)化，最后生成指令完成到硬件的部署。整個(gè)架構(gòu)是基于圖描述結(jié)構(gòu)，不論是對(duì)指令的優(yōu)化還是指令生成，一個(gè)圖結(jié)構(gòu)清晰的描述了數(shù)據(jù)流方向，操作之間的依賴關(guān)系等。基于機(jī)器學(xué)習(xí)的優(yōu)化器是優(yōu)化過程中的重點(diǎn)，指令空間很大，通過優(yōu)化函數(shù)來尋找最優(yōu)值是一個(gè)很合理的想法。它的主要特點(diǎn)如下：

1） 基于GPU、TPU等硬件結(jié)構(gòu)，將張量運(yùn)算作為一個(gè)基本的算符，通過把一個(gè)深度學(xué)習(xí)網(wǎng)絡(luò)描述成圖結(jié)構(gòu)來抽象出數(shù)據(jù)計(jì)算流程。在這樣的圖結(jié)構(gòu)基礎(chǔ)上，更方便記性優(yōu)化。同時(shí)能夠有更好的向上向下兼容性，同時(shí)支持多種深度學(xué)習(xí)框架和硬件架構(gòu)。

2） 巨大的優(yōu)化搜索空間。在優(yōu)化圖結(jié)構(gòu)方面，其不再局限于通過某一種方式，而是通過機(jī)器學(xué)習(xí)方法來搜索可能的空間來最大化部署效率。這種方式雖然會(huì)導(dǎo)致編譯器較大的計(jì)算量，但是更加通用。

TVM提供了一個(gè)非常簡單的端到端用戶接口，通過調(diào)用TVM的API可以很方便的進(jìn)行硬件部署。比如：

以上就是將Keras的模型輸入到TVM，指定部署的硬件GPU，然后進(jìn)行優(yōu)化和代碼生成。

TVM也提供了Java、C++和Python界面供用戶調(diào)用。

2. 圖結(jié)構(gòu)基本優(yōu)化

圖結(jié)構(gòu)是大多數(shù)深度學(xué)習(xí)框架中普遍采用的描述方式。這種圖是一種高層次的描述，將一個(gè)張量運(yùn)算用一個(gè)算符描述，而不是拆分的更細(xì)。這樣更有利于優(yōu)化，而且也符合GPU、TPU的硬件架構(gòu)，在這些芯片中計(jì)算核算力很大，通?？梢酝瓿梢粋€(gè)較大的計(jì)算，比如卷積、矩陣運(yùn)算等。以下是一個(gè)卷積計(jì)算的例子，整個(gè)計(jì)算圖包含了2D卷積、ReLu，dense、softmas等。這樣的圖結(jié)構(gòu)也正好符合FPGA加速器的結(jié)構(gòu)，在FPGA中也是用一個(gè)計(jì)算核來專門計(jì)算某個(gè)大的計(jì)算。TVM圖中節(jié)點(diǎn)描述一個(gè)張量數(shù)據(jù)或者算符，而邊表示了不同計(jì)算的依賴關(guān)系。

基于圖結(jié)構(gòu)，TVM采用了很多圖優(yōu)化策略。包括算符融合，將可以在硬件上用一個(gè)算符完成的多個(gè)連續(xù)運(yùn)算合并；常量折疊，將可以預(yù)先計(jì)算的數(shù)據(jù)放在編譯器中完成，減少硬件計(jì)算；存儲(chǔ)規(guī)劃，預(yù)先為中間數(shù)據(jù)分配存儲(chǔ)空間來儲(chǔ)存中間值，避免中間數(shù)據(jù)無法存儲(chǔ)在片上而增加片外存儲(chǔ)開銷；數(shù)據(jù)規(guī)劃，重新排列數(shù)據(jù)有利于硬件計(jì)算。

1） 算符融合

TVM中將運(yùn)算劃分為4種：1對(duì)1運(yùn)算，比如加法、點(diǎn)乘；降運(yùn)算，比如累加；復(fù)雜運(yùn)算，比如2D卷積，融合了乘法和累加；不透明的，比如分類、數(shù)據(jù)排列等，這些不能被融合。算符融合可以減少存儲(chǔ)開銷，實(shí)現(xiàn)pipeline，特別是在FPGA中更有利。比如我們目前的項(xiàng)目中是開發(fā)一款通用RNN架構(gòu)IP，其中涉及到矩陣乘法，加法。其中加法就可以融合到矩陣乘中，這樣就減少了單獨(dú)加法模塊計(jì)算開銷以及讀寫cache開銷。

2） 數(shù)據(jù)規(guī)劃

以我們XRNN來說，片上有一個(gè)矩陣運(yùn)算陣列，由于陣列大小固定，一次計(jì)算矩陣大小也是固定的。比如計(jì)算一個(gè)32x32對(duì)應(yīng)32x1的矩陣向量乘法，那么就要求權(quán)重和向量必須要按照32倍數(shù)進(jìn)行對(duì)齊，這就需要對(duì)權(quán)重?cái)?shù)據(jù)等進(jìn)行規(guī)劃。

3. 張量計(jì)算

TVM中使用的張量描述語言是透明的，可以根據(jù)硬件需要進(jìn)行修改。這樣更加靈活和有利于進(jìn)行優(yōu)化。但是這樣可能增加了編譯器優(yōu)化的復(fù)雜性。TVM描述例子如下：

描述算符中包含結(jié)果大小，計(jì)算方式。但是這其中沒有涉及到循環(huán)結(jié)構(gòu)和更多數(shù)據(jù)操作細(xì)節(jié)。TVM采用了Halide思想，通過使用schedule來對(duì)張量計(jì)算進(jìn)行等價(jià)變換，從中計(jì)算出執(zhí)行效率最高的schedule結(jié)構(gòu)。整個(gè)schedule流程如下圖：

從中可以看出，TVM除了采用了Halide的schedule方式外，還增加了三種針對(duì)GPU和TPU的schedule方式：specile memory scope，tensorization，latency hiding。這些schedule方式可以對(duì)一個(gè)張量運(yùn)算進(jìn)行等價(jià)變換，產(chǎn)生多種代碼結(jié)構(gòu)，從中選擇出最有利于硬件執(zhí)行的代碼結(jié)構(gòu)。

3.1 并行優(yōu)化

并行計(jì)算是提高硬件執(zhí)行效率的重要一步，因?yàn)橹T如卷積、矩陣計(jì)算等都是大量的可以并發(fā)進(jìn)行的計(jì)算，如何優(yōu)化并行結(jié)構(gòu)對(duì)改善硬件性能很關(guān)鍵。在這里需要考慮兩點(diǎn)問題：一個(gè)是并行度，另外一個(gè)是數(shù)據(jù)共享。如果數(shù)據(jù)不共享那么就會(huì)增加數(shù)據(jù)讀寫消耗。而盡量利用可共享數(shù)據(jù)，則需要盡心設(shè)計(jì)一個(gè)計(jì)算結(jié)構(gòu)。

TVM提出了memory scope的概念，其將數(shù)據(jù)計(jì)算進(jìn)行可并行和不可并行分類，對(duì)于可以并行計(jì)算的，就可以使用多線程來并行計(jì)算，而不可并行，則需要等待被依賴數(shù)據(jù)計(jì)算完成。比如一個(gè)矩陣乘法例子：

在XRNN中也會(huì)遇到類似的問題，對(duì)于不同的計(jì)算比如矩陣乘法和激活函數(shù)，如果沒有依賴的話，就是可以并行執(zhí)行的。

3.2 存儲(chǔ)讀寫優(yōu)化

在FPGA中讀寫cache或者external ddr也是一筆開銷。如何將存儲(chǔ)讀寫開銷降低也有利于提高硬件執(zhí)行效率。比如在我們的XRNN中，save數(shù)據(jù)到ddr會(huì)消耗很多時(shí)間，而這個(gè)數(shù)據(jù)下一次又會(huì)被利用，同時(shí)又增加了load的時(shí)間。如果將數(shù)據(jù)緩存到片上，那么就減少了load和save開銷。還有比如片上cache也會(huì)來回讀寫數(shù)據(jù)，如果可以將一個(gè)計(jì)算完成的數(shù)據(jù)直接送給下一個(gè)計(jì)算核，實(shí)現(xiàn)流水，那么讀寫cache開銷也節(jié)省了。

對(duì)于隱藏存儲(chǔ)讀寫開銷還有一種方法，在下一次計(jì)算開始前而這一次計(jì)算正在進(jìn)行的時(shí)候，就進(jìn)行片外數(shù)據(jù)加載，就將加載和計(jì)算重合了，減少額外加載消耗的時(shí)間。

4. 自動(dòng)優(yōu)化器

TVM在豐富的schedule方式基礎(chǔ)上，提出了一個(gè)機(jī)器學(xué)習(xí)模型來尋找最優(yōu)化的schedule結(jié)構(gòu)。其包含兩部分：一部分是基于schedule方式產(chǎn)生所有可能的計(jì)算結(jié)構(gòu)；另外一個(gè)是機(jī)器學(xué)習(xí)代駕模型來預(yù)測(cè)性能。

Schedule空間是巨大的，它可能產(chǎn)生很多種計(jì)算流結(jié)構(gòu)，而對(duì)其中進(jìn)行探索，找到最合適的結(jié)構(gòu)。這會(huì)產(chǎn)生大量的計(jì)算。不過在我們的XRNN結(jié)構(gòu)中，因?yàn)槭艿接布?nèi)核的限制，還是可以對(duì)空間進(jìn)行縮減的。只采用最可能影響性能schedule方式。比如矩陣乘、加法、點(diǎn)乘、激活這些不同的計(jì)算可以進(jìn)行并發(fā)和非并發(fā)安排，這樣的空間相對(duì)小，有利于加速編譯器生成工作。

機(jī)器學(xué)習(xí)代價(jià)模型主要考慮了不同操作的延時(shí)來預(yù)測(cè)性能：存儲(chǔ)訪問方式，數(shù)據(jù)重利用，pipeline等。TVM中對(duì)代價(jià)函數(shù)的求解不是通過隨機(jī)統(tǒng)計(jì)方式，而是使用實(shí)時(shí)的配置數(shù)據(jù)進(jìn)行訓(xùn)練，周期性更替代碼結(jié)構(gòu)。人為的探尋更合理的優(yōu)化結(jié)構(gòu)，然后提供給模型，讓其不斷更替。這種方式避免了探索大量schedule空間的時(shí)間消耗，同時(shí)也能更接近于實(shí)際情況。

總結(jié)

以上簡要介紹了TVM的整體架構(gòu)和基本方法，其實(shí)還挺符合FPGA加速硬件結(jié)構(gòu)的。很多方法還是可以借用的。而且TVM是一個(gè)兼容更廣的編譯器架構(gòu)，針對(duì)我們自身的FPGA特性，也會(huì)有很多不一樣的設(shè)計(jì)。

審核編輯：湯梓紅