LSTM網絡在TensorFlow上的執(zhí)行時間和手工融合內核的執(zhí)行時間對比

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

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

內核融合在數(shù)據(jù)流圖計算框架中的應用流程

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

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

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

GPU上的element-wise內核融合示意圖

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

線程模型與任務劃分

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

分享到

xiesc

相關推薦