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

陳天奇在AI圈的名氣，好比李小龍?jiān)谖湫g(shù)界的地位，且都是少年有成。

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

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

重點(diǎn)來(lái)了，這個(gè)編譯器不依賴其他庫(kù)，能夠端到端地測(cè)量?jī)?yōu)化性能，最高降低了14倍推理延遲。

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

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

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

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

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

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

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

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

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

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

技術(shù)突破

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

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

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

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

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

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

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

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

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

編譯器API之遞歸模型

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

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

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

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

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

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

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

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

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

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

該計(jì)算表示為運(yùn)算符的DAG，每個(gè)運(yùn)算符都通過(guò)消耗先前生成的或輸入的張量來(lái)生成張量。可以在調(diào)度原語(yǔ)的幫助下執(zhí)行優(yōu)化，例如循環(huán)切片，循環(huán)展開，矢量化等。

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

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

展望未來(lái)

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

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

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

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

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

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