然而，為每一個(gè)計(jì)算圖的內(nèi)核函數(shù)進(jìn)行手工融合并不是一種可以擴(kuò)展的方法。因此，研究自動(dòng)化的內(nèi)核融合成為最近的一個(gè)熱點(diǎn)，例如TensorFlow中的XLA項(xiàng)目就是要將給定的計(jì)算圖生成硬件設(shè)備相關(guān)的機(jī)器碼，再如NNVM-Fusion是DMLC社區(qū)為加速M(fèi)XNet所提出的相關(guān)項(xiàng)目，還有最近比較流行的動(dòng)態(tài)圖計(jì)算框架PyTorch也開始逐漸引入內(nèi)核融合的技術(shù)來提升性能。

自動(dòng)化的內(nèi)核生成一般包括以下幾個(gè)步驟：1. 圖優(yōu)化。即在進(jìn)行內(nèi)核融合之前，首先對(duì)計(jì)算圖進(jìn)行分析并應(yīng)用一系列與硬件無關(guān)的優(yōu)化策略，從而在邏輯上降低運(yùn)行時(shí)的開銷，常見的類似優(yōu)化策略包括常數(shù)折疊（constant folding）、公共子表達(dá)式消除（common subexpression elimination）等；2. 檢測融合子圖。即在給定數(shù)據(jù)流圖中，找出一些可以被融合的圖節(jié)點(diǎn)，這些節(jié)點(diǎn)往往是一段連續(xù)的子圖。3. 代碼生成。在給定一個(gè)融合子圖，為其生成一份內(nèi)核函數(shù)代碼。這里可以直接生成與硬件相關(guān)的代碼，也可以先生成到一個(gè)統(tǒng)一的中間表示層（intermediate representation），如LLVM，然而再由相應(yīng)的編譯器將其編譯到與針對(duì)特定硬件的執(zhí)行代碼，TensorFlow的XLA就采用了后者的方法。4. 圖的修改。即將融合后的內(nèi)核所對(duì)應(yīng)的Operator替換之前的子圖，并插入原來的數(shù)據(jù)流圖中。整個(gè)流程如圖2所示。

內(nèi)核融合在數(shù)據(jù)流圖計(jì)算框架中的應(yīng)用流程

然而，自動(dòng)化的內(nèi)核融合并生成高效的內(nèi)核代碼還存在著許多挑戰(zhàn)，如何解決內(nèi)核間跨線程的數(shù)據(jù)同步和如何實(shí)現(xiàn)高效的線程模型及任務(wù)劃分都是非常重要的問題。

內(nèi)核間的數(shù)據(jù)同步

當(dāng)前，在GPU上的內(nèi)核融合技術(shù)大部分都只支持element-wise的Operator，如PyTorch和NNVM-Fusion。其主要原因是由于CUDA采用的是SIMT（單指令多線程）的編程模型，這使得融合element-wise的操作更加容易。例如，在圖3的示意圖中，如果我們想將y1=x1+x2 和h=sigmoid(y1) 兩個(gè)計(jì)算表達(dá)式進(jìn)行融合，那只需要讓每個(gè)線程都處理輸入向量中的一個(gè)元素并且執(zhí)行相同的表達(dá)式h=sigmoid(x1+x2) 即可，在這種情況中，由于融合后的計(jì)算邏輯都在相同的一個(gè)線程內(nèi)完成，所以前一個(gè)計(jì)算輸出的結(jié)果可以通過寄存器或共享內(nèi)存直接傳到下一個(gè)計(jì)算的輸入中。

GPU上的element-wise內(nèi)核融合示意圖

然而，如果我們想將兩個(gè)矩陣乘法、或更加復(fù)雜的計(jì)算（如卷積操作）融合在一起，就需要引入數(shù)據(jù)之間的同步機(jī)制，即前一個(gè)內(nèi)核完成的結(jié)果需要傳播到下一個(gè)內(nèi)核的部分或全部線程中。這時(shí)，若想融合這樣的內(nèi)核，我們必須有較為靈活的同步機(jī)制。然而，在CUDA 8.0之前，CUDA只支持同一個(gè)線程塊內(nèi)的計(jì)算同步，其無法滿足融合的需求。Shucai Xiao等人早期提出一種能支持全局跨線程塊的同步機(jī)制，但其需要對(duì)計(jì)算的資源有一定的假設(shè)，即要求線程塊個(gè)數(shù)要小于SM的個(gè)數(shù)。最近，在Nvidia發(fā)布的最新版CUDA 9.0中首次提出了Cooperative Groups的概念，其可以靈活地支持不同粒度上的線程同步，這將會(huì)使得在GPU上的更加復(fù)雜的內(nèi)核融合變得容易，也同時(shí)為實(shí)現(xiàn)更加高效的融合提供了更多空間。

線程模型與任務(wù)劃分

內(nèi)核融合中另一個(gè)挑戰(zhàn)是如何優(yōu)化任務(wù)的劃分，從而充分發(fā)揮GPU的計(jì)算和片上存儲(chǔ)性能。我們知道，執(zhí)行一個(gè)GPU的內(nèi)核函數(shù)，不僅需要指定內(nèi)核函數(shù)的計(jì)算算法，還需要指定其調(diào)度邏輯，即如何分配線程塊的大小和數(shù)量等等。通常，這需要有經(jīng)驗(yàn)的程序員根據(jù)計(jì)算算法的特性仔細(xì)地設(shè)計(jì)每一個(gè)內(nèi)核的調(diào)度邏輯。然而，在本文介紹的內(nèi)核融合的場景中，我們需要系統(tǒng)能夠根據(jù)當(dāng)前使用的GPU架構(gòu)快速、自動(dòng)化地生成調(diào)度邏輯。因此，目前一種研究趨勢是采用來自于MIT的Halide項(xiàng)目的思想，即通過將計(jì)算算法和調(diào)度邏輯進(jìn)行抽象并分離，然后采用一些搜索算法來找到較優(yōu)的調(diào)度方案，從而自動(dòng)生成最終的執(zhí)行代碼。Halide項(xiàng)目是針對(duì)圖像處理所設(shè)計(jì)的編譯系統(tǒng)。目前，像來自DMLC的TVM項(xiàng)目、以及MIT的Taco項(xiàng)目都采用該思想并針對(duì)深度學(xué)習(xí)庫進(jìn)行自動(dòng)化的優(yōu)化，目前大部分這些項(xiàng)目還都還處在較早期階段。

xiesc