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

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

90后的陳天奇是機(jī)器學(xué)習(xí)領(lǐng)域著名的青年華人，在2020年秋季加入CMU任助理教授，成為加入CMU的年輕華人學(xué)者之一。

近日，陳天奇團(tuán)隊在arxiv上發(fā)表了針對深度學(xué)習(xí)開發(fā)的編譯器「CORTEX」，側(cè)重解決深度學(xué)習(xí)模型的優(yōu)化問題，能夠為遞歸神經(jīng)網(wǎng)絡(luò)生成高效的編碼。

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

深度學(xué)習(xí)的應(yīng)用瓶頸

深度學(xué)習(xí)模型在生產(chǎn)中越來越多地被用作應(yīng)用程序的一部分，例如個人助理，自動駕駛汽車和聊天機(jī)器人。

這些應(yīng)用對模型的推理延遲提出了嚴(yán)格的要求。

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

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

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

優(yōu)化深度學(xué)習(xí)模型一般分兩步進(jìn)行：

• 高層次的圖優(yōu)化，如內(nèi)核融合；

• 低層次的內(nèi)核優(yōu)化，如依賴庫中的內(nèi)核優(yōu)化。

這種方法對于遞歸深度學(xué)習(xí)模型，通常會表現(xiàn)出顯著的性能。

技術(shù)突破

遞歸深度學(xué)習(xí)模型的編譯器「CORTEX」旨在構(gòu)建一個編譯器和運(yùn)行基礎(chǔ)結(jié)構(gòu)，以通過不規(guī)則和動態(tài)的控制流，甚至不規(guī)則的數(shù)據(jù)結(jié)構(gòu)訪問也能有效地編譯ML應(yīng)用程序。

這種編譯器方法和不依賴其他庫，能夠?qū)崿F(xiàn)端到端優(yōu)化，在不同backend上的推理延遲最多降低了14倍。

CORTEX包括一個運(yùn)行框架，這個框架可以處理具有遞歸數(shù)據(jù)結(jié)構(gòu)遍歷的模型。

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

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

最后，論文總結(jié)了貢獻(xiàn)：

• 設(shè)計了CORTEX，一個基于編譯器的框架，可實現(xiàn)端到端優(yōu)化和高效的代碼生成，以遞歸深度學(xué)習(xí)模型的低延遲推理。

• 作為設(shè)計的一部分，擴(kuò)展了張量編譯器，并提出了遞歸模型的新調(diào)度原語和優(yōu)化。

• 對提出的框架進(jìn)行原型設(shè)計，并根據(jù)最新的遞歸深度學(xué)習(xí)框架對其進(jìn)行評估，并獲得了顯著的性能提升（ 在Nvidia GPU，Intel和ARM CPU上最高可達(dá)14倍）。

編譯器API之遞歸模型

CORTEX需要具有模型計算的端到端視圖，以便執(zhí)行諸如內(nèi)核融合之類的優(yōu)化。

輸入程序需要包含有關(guān)在模型中執(zhí)行的張量操作的足夠信息，以在將其降低到ILIR時進(jìn)行調(diào)度。

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

除了進(jìn)行RA計算外，用戶還需要提供有關(guān)輸入數(shù)據(jù)結(jié)構(gòu)的基本信息，例如每個節(jié)點(diǎn)的最大子代數(shù)以及數(shù)據(jù)結(jié)構(gòu)的類型（序列，樹或DAG）。

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

降低循環(huán)的次數(shù)

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

他們還具體化了張量rnn，它存儲了計算結(jié)果。三個張量中的每個張量都存儲每個遞歸調(diào)用的數(shù)據(jù)，在這種情況下，這些數(shù)據(jù)相當(dāng)于每個樹節(jié)點(diǎn)。

不規(guī)則循環(huán)IR（ILIR）

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

該計算表示為運(yùn)算符的DAG，每個運(yùn)算符都通過消耗先前生成的或輸入的張量來生成張量?？梢栽谡{(diào)度原語的幫助下執(zhí)行優(yōu)化，例如循環(huán)切片，循環(huán)展開，矢量化等。

當(dāng)中間張量存儲在暫存器中時，用非仿射表達(dá)式索引它會導(dǎo)致張量稀疏。這種稀疏填充的張量會占用過多的內(nèi)存， 其中A1的一半未使用。

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

展望未來

過去有關(guān)機(jī)器學(xué)習(xí)編譯器的以及深度學(xué)習(xí)的工作建議在機(jī)器學(xué)習(xí)模型中支持有效地動態(tài)執(zhí)行是非常必要的。

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

陳天奇團(tuán)隊認(rèn)為，像在ILIR中一樣，擴(kuò)大當(dāng)今使用的高度專業(yè)化的機(jī)器學(xué)習(xí)框架和技術(shù)的范圍也很重要（在不損害其優(yōu)化靜態(tài)前饋模型的能力的前提下）。

同時，他們也介紹了團(tuán)隊的未來開發(fā)方向：

• 應(yīng)用這些研究來開發(fā)類似的技術(shù)來訓(xùn)練和服務(wù)具有潛在非遞歸動態(tài)控制流的模型。

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