這樣的兩層排程法儘管為系統設計帶來了一些簡潔，但在實際的部署中，兩層排程器互相不感知會導致幾個問題：執行時的排程開銷很大，運算元間並行性沒有被有效利用，以及忽視了運算元內和運算元間兩種並行性的相互影響。
圖1：(a) 傳統低效率的排程方案，(b) 最佳化後的排程方案
為了能夠打破這種僵局，微軟亞洲研究院和北京大學、上海科技大學合作提出了一種可以成倍甚至幾十倍地提升深度學習計算速度的編譯框架 RAMMER。研究員們將原資料流圖中的運算元解析為 rOperator 並將其分解為更小的排程單元 rTask，將底層的硬體抽象為由多個虛擬執行單元（virtualized execution units, vEU）組成的 vDevice。在這套新的抽象下，使用者可以透過更細的 rTask 粒度將資料流圖排程到多個 vEU 之上，兼顧了計算任務中的兩種並行性與底層計算資源的協調。整個排程方案在編譯期生成並“靜態”對映到硬體計算單元上，因此可以天然地消除掉許多原本存在的排程開銷。
圖2：(a) 傳統深度學習框架，(b)RAMMER 深度學習編譯框架
儘管上述介紹中用了不少 CUDA 的概念，但不難發現 RAMMER 整個設計是硬體通用的。它可以很好地適配到諸如 GPU、IPU 等主流深度學習加速器上。研究員們在 NVIDIA GPU、AMD GPU 和 GraphCore 上都評估了這套編譯技術所能取得的效能收益。與 TensorRT 相比，RAMMER 在部分模型上實現了最高3.1倍的效能加速。
圖3：在 NVIDIA V100 GPU 上批次大小設定為1的端到端的模型推理時間對比
RAMMER 背後是微軟亞洲研究院在過去一年多時間裡打造的一套名為 NNFusion（https://github.com/microsoft/nnfusion）的深度神經網路編譯器（DNN compiler）。NNFusion 能夠將現有的模型編譯為對應裝置的可高效執行的原始碼，同時支援使用者自行替換核心實現或自動從外部匯入高效能的核心實現。為了方便地與現有的程式碼庫和 GPU 程式設計模型相容，RAMMER 採用的是原始碼轉換的方式，而不是像 TVM、Tensor Comprehension [1] 一樣定義新的計算抽象需要使用者提供運算元的計算邏輯。
我們先來了解一下什麼是 rTask。rTask 是組成 rOperator 的互相獨立的更小的任務單元，也是 RAMMER 抽象中最小的排程單元。在 NVIDIA GPU 上，對於使用者提供的核心 CUDA 實現，透過一個小的解析器，可以將每個執行緒塊（thread block）轉化為一個 rTask。所以 rTask 可以利用原本核心實現中的語義，雖然這樣 rTask 在實現上是與原本的程式設計模型耦合的，但是也大幅度降低了所需的工作負擔。
那麼要如何建立 vDevice 和 vEU 呢？目前，根據硬體的特性再配合簡單的啟發式搜尋就可以建立 vDevice 和 vEU 了。如果 V100 中有80個 SM，每個 SM 最多能夠執行32個執行緒塊，那麼就可以建立一個包含有2560個 vEU 的 vDevice，而後根據 rTask 所構成的資料流圖與 vDevice 的情況，透過一些簡單的策略（譬如直接將 rTask 平鋪上 vEU）就能夠生成足夠高效的排程計劃。
由於硬體和程式設計模型的限制，GPU 執行時的分發器（dispatcher）和排程器（scheduler）並不對使用者開放可程式設計介面，所以研究員們採用了持續執行緒（persistent thread）[2]的方式，巧妙地以一個相對小的開銷將 vEU 與 SM 繫結起來。這樣就可以將排程計劃“靜態”對映給硬體裝置了。
一篇好的系統論文，不僅可以最佳化效能，更要闡明一個問題。起初，研究員們只是想改善一個具體的神經網路推理時 GPU 利用率偏低的問題（出於對延遲的保障，很多場景下設定小的批次大小其實是標準做法），而除了最佳化運算元實現以外，樸素的想法就是將多個運算元一同交給 GPU 裝置同時執行，但這並不是一個新的問題。CUDA 很早就引入了流的概念對其提供支援，GPU 社群之前也有一些效果不錯的工作（concurrent kernel execution [3]、elastic kernel [4]等），那麼在深度神經網路的場景下，為什麼大家對這個問題的認知不足？
像上文提到的 MXNet 中有依賴引擎一樣，TensorFlow 開發早期也有支援多個流的嘗試。但是到後來都接近棄置了，主要原因可能有以下幾個方面：
不同的 CUDA 流在執行時採用空間分片（spatial multiplexing）的方法來排程不同流佇列（stream queue）上的運算元，粒度更粗而彼此之間又極易產生相互干擾影響最終效能 [5]。
GPU 的 SM 在不斷增加。現在 Ampere GA100 中有128個 SM，但幾年前 Kepler GK180 中僅有15個 SM，所以在早期，無論是 GPU 社群還是 DNN 的框架開發，在現有的 GPU 程式設計模型下都已經形成了硬體對於運算元間並行性並沒有太多加速潛力的印象。
早期的神經網路結構比較簡單，如 AlexNet 等本身在運算元間並行性上也沒有更多發揮的空間。但隨著 AutoML 的出現，網路結構趨於複雜，此外也有ResNext [6]、ResNeSt [7] 等工作引入了新的神經網路設計模式，這個問題正變得更重要。
只是將運算元間並行性挖掘起來會是一個好的效能最佳化，但不足以成為一個好的系統工作。在之前的工作中，微軟亞洲研究院的研究員們已經完成了初步的實現並且在一些模型上獲得了較好的加速效果，但是當時還沒有完全對問題進行清楚地定義，而且因為沒有 NNFusion 程式碼庫的支援，實驗相對簡陋，沒有取得很好的反饋。
重新定義一個問題和定位一個工作並不是在用不同的寫法來寫“茴”字。之前的工作只是在做一個廣義上的核心融合，並沒有設立起 rTask 和 vEU 的抽象。而此次確定本質的問題在於原本系統中兩層排程的差距以後，新的抽象很快探明瞭更大的最佳化空間：首先是將原本的透過成本模型來選擇子圖進行融合的問題，轉變為了以更細粒度下的排程和資源分配問題。而得益於絕大部分情況下，神經網路計算的特徵（DFG, 運算元和張量）在編譯時間是已知的，因此可以將排程的開銷移交給編譯器，這既提升了搜尋的效率也簡化了系統設計。
更重要的是，讓運算元間並行性與運算元內並行性相互影響這個問題走進研究員們的視野。舉個例子，如果對於同一個運算元有兩種核心實現，其中一個比另一個多消耗三倍的資源（如 CUDA Cores、Shared Memory 等），但是隻取得兩倍的加速，這在平行計算中是很常見的一個現象。而在此前單個運算元獨佔整個硬體的情況下，毫無疑問會選擇更快的實現。而研究員們的實驗表明，在運算元間和運算元內兩種並行性協同排程的情況下，選擇資源“價效比”最高的實現而非“最快”往往是更優的選擇。這其實挑戰了之前許多生成高效能運算元的工作如 AutoTVM [8] 等的一個基本假設，單個運算元獨佔整個硬體表現出的計算效能是否真的是效能調優的金標準？顯然，子圖替換TASO [9] 加上高效能核心實現 TVM 兩個“最優”相結合，並沒有帶來真的最優。
研究員們基於新的抽象，僅嘗試了簡單的策略就在一些場景下獲得了超過現有 SOTA 的效能。在此，歡迎大家基於這個抽象進一步嘗試更多排程策略，來探索對於一個資料流圖（或者其子圖）搜尋運算元間和運算元內並行性相互影響下的更高效能的整體實現。
NNFusion 現已在 GitHub 開源：https://github.com/microsoft/nnfusion。目前已經發布了0.1 版本, 該版本支援 TensorFlow 和 ONNX 在內的主流模型格式以及 CUDA GPU 等裝置，並提供了豐富的效能最佳化策略，支援端到端的模型到原始碼的 AOT 編譯來消除執行時的開銷，同時還消除了對第三方庫或框架的依賴。如果你有更深入的研發需求，可以直接修改 NNFusion 生成的程式碼來進行模型的定製化最佳化。
歡迎前往體驗 NNFusion，也期待你可以在 NNFusion 中貢獻真知灼見，一起“壓榨”加速器的效能！
