西雅圖是美國的超一線城市，不僅孕育了科技巨頭微軟，華盛頓大學也是城市的招牌，華人武術宗師李小龍也畢業于此。

陳天奇在AI圈的名氣，好比李小龍在武術界的地位，且都是少年有成。

90后的陳天奇是機器學習領域著名的青年華人，在2020年秋季加入CMU任助理教授，成為加入CMU的年輕華人學者之一。

近日，陳天奇團隊在arxiv上發表了針對深度學習開發的編譯器「CORTEX」，側重解決深度學習模型的優化問題，能夠為遞歸神經網絡生成高效的編碼。

重點來了，這個編譯器不依賴其他庫，能夠端到端地測量優化性能，最高降低了14倍推理延遲。

深度學習的應用瓶頸

深度學習模型在生產中越來越多地被用作應用程序的一部分，例如個人助理，自動駕駛汽車和聊天機器人。

這些應用對模型的推理延遲提出了嚴格的要求。

因此，生產中使用了各種各樣的專用硬件，包括CPU，GPU和專用加速器，以實現低延遲推理。

對于具有遞歸和其他動態控制流的模型，減少推理延遲尤其困難。

陳天奇提出了一個解決方法，一種遞歸深度學習模型的編譯器「CORTEX」，在編譯階段就生成高效代碼。

優化深度學習模型一般分兩步進行：

• 高層次的圖優化，如內核融合；

• 低層次的內核優化，如依賴庫中的內核優化。

這種方法對于遞歸深度學習模型，通常會表現出顯著的性能。

技術突破

遞歸深度學習模型的編譯器「CORTEX」旨在構建一個編譯器和運行基礎結構，以通過不規則和動態的控制流，甚至不規則的數據結構訪問也能有效地編譯ML應用程序。

這種編譯器方法和不依賴其他庫，能夠實現端到端優化，在不同backend上的推理延遲最多降低了14倍。

CORTEX包括一個運行框架，這個框架可以處理具有遞歸數據結構遍歷的模型。

這個框架能夠從張量計算中完全解開遞歸遍歷，并進行部分推理，從而能夠快速地將整個張量計算放入到硬件加速器（目前支持CPU和GPU）。

這意味著編譯器可以生成一個高度優化的代碼，能夠避免加速器和主機CPU之間頻繁通信的開銷。

最后，論文總結了貢獻：

• 設計了CORTEX，一個基于編譯器的框架，可實現端到端優化和高效的代碼生成，以遞歸深度學習模型的低延遲推理。

• 作為設計的一部分，擴展了張量編譯器，并提出了遞歸模型的新調度原語和優化。

• 對提出的框架進行原型設計，并根據最新的遞歸深度學習框架對其進行評估，并獲得了顯著的性能提升（ 在Nvidia GPU，Intel和ARM CPU上最高可達14倍）。

編譯器API之遞歸模型

CORTEX需要具有模型計算的端到端視圖，以便執行諸如內核融合之類的優化。

輸入程序需要包含有關在模型中執行的張量操作的足夠信息，以在將其降低到ILIR時進行調度。

從而RA將輸入計算建模為運算符的DAG，其中每個運算符都指定為循環嵌套。

除了進行RA計算外，用戶還需要提供有關輸入數據結構的基本信息，例如每個節點的最大子代數以及數據結構的類型（序列，樹或DAG）。

這些信息在編譯期間使用，并且可以在運行時很容易驗證。

降低循環的次數

從RA到ILIR的降低本質上是從遞歸到迭代的降低。因此，陳團隊在降低過程中使所有臨時顯性張量。

他們還具體化了張量rnn，它存儲了計算結果。三個張量中的每個張量都存儲每個遞歸調用的數據，在這種情況下，這些數據相當于每個樹節點。

不規則循環IR（ILIR）

ILIR是張量編譯器使用的程序表示的擴展。因此，ILIR中分別指定了計算和優化。

該計算表示為運算符的DAG，每個運算符都通過消耗先前生成的或輸入的張量來生成張量。可以在調度原語的幫助下執行優化，例如循環切片，循環展開，矢量化等。

當中間張量存儲在暫存器中時，用非仿射表達式索引它會導致張量稀疏。這種稀疏填充的張量會占用過多的內存， 其中A1的一半未使用。

在這種情況下，可以通過循環迭代空間來索引張量。這種轉換還通過將間接內存訪問轉換為仿射訪問來降低索引成本。

展望未來

過去有關機器學習編譯器的以及深度學習的工作建議在機器學習模型中支持有效地動態執行是非常必要的。

CORTEX展示了一種行之有效的方法是利用過去在通用編譯中的工作，例如檢查執行器技術或稀疏的多面體框架。

陳天奇團隊認為，像在ILIR中一樣，擴大當今使用的高度專業化的機器學習框架和技術的范圍也很重要（在不損害其優化靜態前饋模型的能力的前提下）。

同時，他們也介紹了團隊的未來開發方向：

• 應用這些研究來開發類似的技術來訓練和服務具有潛在非遞歸動態控制流的模型。

• 將CORTEX集成到更高級別的編程抽象中。