【按：在深度學習「紅透」半邊天的同時，當前很多深度學習框架卻面臨著共同的性能問題：被頻繁調用的代數運算符嚴重影響模型的執行效率。本文中，微軟亞洲研究院研究員薛繼龍將為大家介紹能夠輕鬆玩轉計算性能的「加速神器」——內核融合，探討內核融合在加速深度學習上的主要方法以及當前面臨的主要挑戰。】
如今，較為常見的深度學習框架（如CNTK、TensorFlow和Caffe2等）都會將一個深度學習的模型抽象成為一個由一些基本運算符（Operator）組成的有向無環的數據流圖（DAG），然後再由下層計算引擎按照某一種拓撲序來依次調度並執行這些節點對應的內核函數，從而完成一個模型的執行。為了能夠支持在不同的硬體上進行計算，一個Operator往往會對應多個內核函數的實現，例如，GPU上的內核函數是由CUDA或者一些GPU的函數庫（如cuDNN、cuBLAS等）提供的操作組合而成。
為了提供較好的靈活性，大多深度學習框架中的Operator都是定義在了代數運算符這個粒度上，例如向量的加、減、乘、除和矩陣乘法等等，一般的計算框架都會有幾百甚至上千個Operator。由於這些運算符的抽象粒度較低，所以一個真實的訓練模型的數據流圖往往會包括數千個節點，這些節點在GPU上的執行就會變成數千次GPU上的內核執行。這些粒度較小的內核函數在提供了靈活性的同時，其頻繁的調用也成為當前影響許多深度學習框架性能的一個重要因素，其帶來的性能開銷主要體現在：數據流圖的調度開銷，GPU內核函數的啟動開銷，以及內核函數之間的數據傳輸開銷。
解決這些性能問題的一個直接方法就是內核融合(Kernel Fusion)。所謂內核融合，就是將一個計算圖中的節點所對應的內核函數融合成一個函數，使得整個數據流圖只需要通過一次函數調用即可完成，從而減小平臺調度和內核啟動帶來的開銷。並且，通過合理地設計不同內核函數的輸入輸出數據的放置（例如使用GPU上的共享內存或寄存器），可以極大地提高數據傳輸效率，從而提升整體計算性能。
為了展示內核融合能夠帶來的好處，我們對比了一個80步長的單樣本LSTM網絡在TensorFlow上的模型推理（inference）時間和我們手工將所有計算融合併優化在同一個內核函數中的計算時間（圖1）。可以看出，在相同的GPU上，融合的內核函數比TensorFlow上基於圖的計算可以快40倍左右。當然，這裡的TensorFlow與手動融合的內核的性能差距除了來源於上述性能開銷外，還包括TensorFlow本身的框架開銷。
LSTM網絡在TensorFlow上的執行時間和手工融合內核的執行時間對比
然而，為每一個計算圖的內核函數進行手工融合併不是一種可以擴展的方法。因此，研究自動化的內核融合成為最近的一個熱點，例如TensorFlow中的XLA項目就是要將給定的計算圖生成硬體設備相關的機器碼，再如NNVM-Fusion是DMLC社區為加速MXNet所提出的相關項目，還有最近比較流行的動態圖計算框架PyTorch也開始逐漸引入內核融合的技術來提升性能。
自動化的內核生成一般包括以下幾個步驟：1. 圖優化。即在進行內核融合之前，首先對計算圖進行分析並應用一系列與硬體無關的優化策略，從而在邏輯上降低運行時的開銷，常見的類似優化策略包括常數摺疊（constant folding）、公共子表達式消除（common subexpression elimination）等；2. 檢測融合子圖。即在給定數據流圖中，找出一些可以被融合的圖節點，這些節點往往是一段連續的子圖。3. 代碼生成。在給定一個融合子圖，為其生成一份內核函數代碼。這裡可以直接生成與硬體相關的代碼，也可以先生成到一個統一的中間表示層（intermediate representation），如LLVM，然而再由相應的編譯器將其編譯到與針對特定硬體的執行代碼，TensorFlow的XLA就採用了後者的方法。4. 圖的修改。即將融合後的內核所對應的Operator替換之前的子圖，並插入原來的數據流圖中。整個流程如圖2所示。
內核融合在數據流圖計算框架中的應用流程
然而，自動化的內核融合併生成高效的內核代碼還存在著許多挑戰，如何解決內核間跨線程的數據同步和如何實現高效的線程模型及任務劃分都是非常重要的問題。
內核間的數據同步
當前，在GPU上的內核融合技術大部分都只支持element-wise的Operator，如PyTorch和NNVM-Fusion。其主要原因是由於CUDA採用的是SIMT（單指令多線程）的編程模型，這使得融合element-wise的操作更加容易。例如，在圖3的示意圖中，如果我們想將y1=x1+x2 和h=sigmoid(y1) 兩個計算表達式進行融合，那只需要讓每個線程都處理輸入向量中的一個元素並且執行相同的表達式h=sigmoid(x1+x2) 即可，在這種情況中，由於融合後的計算邏輯都在相同的一個線程內完成，所以前一個計算輸出的結果可以通過寄存器或共享內存直接傳到下一個計算的輸入中。
GPU上的element-wise內核融合示意圖
然而，如果我們想將兩個矩陣乘法、或更加複雜的計算（如卷積操作）融合在一起，就需要引入數據之間的同步機制，即前一個內核完成的結果需要傳播到下一個內核的部分或全部線程中。這時，若想融合這樣的內核，我們必須有較為靈活的同步機制。然而，在CUDA 8.0之前，CUDA只支持同一個線程塊內的計算同步，其無法滿足融合的需求。Shucai Xiao等人早期提出一種能支持全局跨線程塊的同步機制，但其需要對計算的資源有一定的假設，即要求線程塊個數要小於SM的個數。最近，在Nvidia發布的最新版CUDA 9.0中首次提出了Cooperative Groups的概念，其可以靈活地支持不同粒度上的線程同步，這將會使得在GPU上的更加複雜的內核融合變得容易，也同時為實現更加高效的融合提供了更多空間。
線程模型與任務劃分
內核融合中另一個挑戰是如何優化任務的劃分，從而充分發揮GPU的計算和片上存儲性能。我們知道，執行一個GPU的內核函數，不僅需要指定內核函數的計算算法，還需要指定其調度邏輯，即如何分配線程塊的大小和數量等等。通常，這需要有經驗的程式設計師根據計算算法的特性仔細地設計每一個內核的調度邏輯。然而，在本文介紹的內核融合的場景中，我們需要系統能夠根據當前使用的GPU架構快速、自動化地生成調度邏輯。因此，目前一種研究趨勢是採用來自於MIT的Halide項目的思想，即通過將計算算法和調度邏輯進行抽象並分離，然後採用一些搜索算法來找到較優的調度方案，從而自動生成最終的執行代碼。Halide項目是針對圖像處理所設計的編譯系統。目前，像來自DMLC的TVM項目、以及MIT的Taco項目都採用該思想並針對深度學習庫進行自動化的優化，目前大部分這些項目還都還處在較早期階段。
未經允許不得轉載：DOIT » 內核融合：GPU深度學習的「加速神器」