內核融合：GPU深度學習的「加速神器」

​編者按：在深度學習「紅透」半邊天的同時，當前不少深度學習框架卻面臨着共同的性能問題：被頻繁調用的代數運算符嚴重影響模型的執行效率。

本文中，微軟亞洲研究院研究員薛繼龍將爲你們介紹可以輕鬆玩轉計算性能的「加速神器」——內核融合，探討內核融合在加速深度學習上的主要方法以及當前面臨的主要挑戰。

現在，較爲常見的深度學習框架（如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自己的框架開銷。

然而，爲每個計算圖的內核函數進行手工融合並非一種能夠擴展的方法。所以，研究自動化的內核融合成爲最近的一個熱點，例如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) 便可，在這種狀況中，因爲融合後的計算邏輯都在相同的一個線程內完成，因此前一個計算輸出的結果能夠經過寄存器或共享內存直接傳到下一個計算的輸入中。

然而，若是咱們想將兩個矩陣乘法、或更加複雜的計算（如卷積操做）融合在一塊兒，就須要引入數據之間的同步機制，即前一個內核完成的結果須要傳播到下一個內核的部分或所有線程中。這時，若想融合這樣的內核，咱們必須有較爲靈活的同步機制。然而，在CUDA 8.0以前，CUDA只支持同一個線程塊內的計算同步，其沒法知足融合的需求。Shucai Xiao等人早期提出一種能支持全局跨線程塊的同步機制，但其須要對計算的資源有必定的假設，即要求線程塊個數要小於SM的個數。最近，在Nvidia發佈的最新版CUDA 9.0中首次提出了Cooperative Groups的概念，其能夠靈活地支持不一樣粒度上的線程同步，這將會使得在GPU上的更加複雜的內核融合變得容易，也同時爲實現更加高效的融合提供了更多空間。

線程模型與任務劃分

內核融合中另外一個挑戰是如何優化任務的劃分，從而充分發揮GPU的計算和片上存儲性能。咱們知道，執行一個GPU的內核函數，不只須要指定內核函數的計算算法，還須要指定其調度邏輯，即如何分配線程塊的大小和數量等等。一般，這須要有經驗的程序員根據計算算法的特性仔細地設計每個內核的調度邏輯。然而，在本文介紹的內核融合的場景中，咱們須要系統可以根據當前使用的GPU架構快速、自動化地生成調度邏輯。所以，目前一種研究趨勢是採用來自於MIT的Halide項目的思想，即經過將計算算法和調度邏輯進行抽象並分離，而後採用一些搜索算法來找到較優的調度方案，從而自動生成最終的執行代碼。Halide項目是針對圖像處理所設計的編譯系統。目前，像來自DMLC的TVM項目、以及MIT的Taco項目都採用該思想並針對深度學習庫進行自動化的優化，目前大部分這些項目還都還處在較早期階段。

歡迎對內核融合感興趣的小夥伴們積極留言，和研究員互動起來吧！

做者介紹

薛繼龍，2016年博士畢業於北京大學，後加入微軟亞洲研究院，現爲系統組研究員，主要研究方向爲大規模分佈式計算系統，包括機器學習、圖計算系統等以及在GPU、RDMA等硬件上加速的優化。