0
| 本文作者： 我在思考中 | 2022-07-11 10:00 |
作者 | 錢云迪、Mircea Trofin
編譯 | 劉冰一
現(xiàn)代計(jì)算機(jī)誕生，如何編譯更快、更小的代碼問題隨之出現(xiàn)。
編譯優(yōu)化是成本收益比最高的優(yōu)化手段，更好的代碼優(yōu)化可以顯著降低大型數(shù)據(jù)中心應(yīng)用程序的操作成本。編譯代碼的大小對于部署在安全引導(dǎo)分區(qū)上的移動(dòng)和嵌入式系統(tǒng)或軟件來說是至關(guān)重要的，因?yàn)榫幾g后的二進(jìn)制文件必須符合嚴(yán)格的代碼大小預(yù)算。隨著這一領(lǐng)域的進(jìn)步，越來越復(fù)雜的啟發(fā)式方法嚴(yán)重?cái)D壓有限的系統(tǒng)空間，阻礙了維護(hù)和進(jìn)一步的改進(jìn)。
最近的研究表明，機(jī)器學(xué)習(xí)可以通過用機(jī)器學(xué)習(xí)策略取代復(fù)雜的啟發(fā)式方法，在編譯器優(yōu)化中釋放更多的機(jī)會(huì)。然而，在通用的、行業(yè)級編譯器中采用機(jī)器學(xué)習(xí)策略仍然是一個(gè)挑戰(zhàn)。
為了解決這個(gè)問題，谷歌兩位高級工程師錢云迪、Mircea Trofin 提出了“MLGO，一個(gè)機(jī)器學(xué)習(xí)指導(dǎo)的編譯器優(yōu)化框架”，這是第一個(gè)工業(yè)級的通用框架，用于將機(jī)器學(xué)習(xí)技術(shù)系統(tǒng)地集成到 LLVM（一個(gè)開源的工業(yè)編譯器基礎(chǔ)設(shè)施，在構(gòu)建關(guān)鍵任務(wù)、高性能軟件時(shí)無處不在）中。
MLGO 使用強(qiáng)化學(xué)習(xí)訓(xùn)練神經(jīng)網(wǎng)絡(luò)來做出決策，以取代 LLVM 中的啟發(fā)式算法。根據(jù)作者描述，LLVM 上有兩處 MLGO 優(yōu)化：
1）通過內(nèi)聯(lián)減少代碼量；
2）通過寄存器分配提高代碼性能。
這兩種優(yōu)化都可以在 LLVM 資源庫中獲得，并已在生產(chǎn)中部署。
內(nèi)聯(lián)（Inlining）有助于通過做出能夠刪除冗余代碼的決策來減少代碼大小。在下面的示例中，調(diào)用者函數(shù) foo()調(diào)用被調(diào)用者函數(shù) bar()，而 bar()本身又調(diào)用了 baz()。內(nèi)聯(lián)這兩個(gè)調(diào)用站點(diǎn)將返回一個(gè)簡單的 foo()函數(shù)，該函數(shù)將減小代碼大小。
圖注：內(nèi)聯(lián)通過刪除冗余代碼來減少代碼大小
在實(shí)際代碼中，有成千上萬的函數(shù)相互調(diào)用，因此構(gòu)成了一個(gè)調(diào)用圖（Call graph）。在內(nèi)聯(lián)階段，編譯器遍歷（traverses）所有調(diào)用者-被調(diào)用者對的調(diào)用圖，并決定是否內(nèi)聯(lián)一個(gè)調(diào)用者-被調(diào)用者對。這是一個(gè)連續(xù)的決策過程，因?yàn)橐郧暗膬?nèi)聯(lián)決策會(huì)改變調(diào)用圖，影響后面的決策和最終的結(jié)果。在上面的例子中，調(diào)用圖foo() → bar() → baz()需要在兩條邊上做出“yes”的決定，以使代碼大小減少。
在MLGO之前，內(nèi)聯(lián)/非內(nèi)聯(lián)的決定是由啟發(fā)式方法做出的，隨著時(shí)間的推移，這種方法越來越難以改進(jìn)。MLGO用一個(gè)機(jī)器學(xué)習(xí)模型代替了啟發(fā)式方法。在調(diào)用圖的遍歷過程中，編譯器通過輸入圖中的相關(guān)特征（即輸入）來尋求神經(jīng)網(wǎng)絡(luò)對是否內(nèi)聯(lián)特定的調(diào)用者-被調(diào)用者對的建議，并按順序執(zhí)行決策，直到遍歷整個(gè)調(diào)用圖為止。
圖注：內(nèi)聯(lián)過程中MLGO的圖示，“ # bbs”、“ # users”和“ callsite height”是調(diào)用者-被調(diào)用者對特性的實(shí)例
MLGO 使用策略梯度和進(jìn)化策略算法對決策網(wǎng)絡(luò)進(jìn)行 RL 訓(xùn)練。雖然沒有關(guān)于最佳決策的基本事實(shí)，但在線 RL 使用經(jīng)過培訓(xùn)的策略在培訓(xùn)和運(yùn)行匯編之間進(jìn)行迭代，以收集數(shù)據(jù)并改進(jìn)策略。特別是，考慮到當(dāng)前訓(xùn)練中的模型，編譯器在內(nèi)聯(lián)階段咨詢模型，以做出內(nèi)聯(lián)/不內(nèi)聯(lián)的決策。編譯完成后，它產(chǎn)生一個(gè)順序決策過程的日志（狀態(tài)、行動(dòng)、獎(jiǎng)勵(lì)）。然后，該日志被傳遞給訓(xùn)練器以更新模型。這個(gè)過程不斷重復(fù)，直到得到一個(gè)滿意的模型為止。
foo.cpp編譯成對象文件foo.o，并進(jìn)行了一系列的優(yōu)化，其中一個(gè)是內(nèi)聯(lián)通道。訓(xùn)練后的策略被嵌入到編譯器中，在編譯過程中提供內(nèi)聯(lián)/非內(nèi)聯(lián)的決策。與訓(xùn)練場景不同的是，該策略不生成日志。TensorFlow 模型被嵌入 XLA AOT ，它將模型轉(zhuǎn)換為可執(zhí)行代碼。這避免了TensorFlow運(yùn)行時(shí)的依賴性和開銷，最大限度地減少了在編譯時(shí)由ML模型推理引入的額外時(shí)間和內(nèi)存成本。
圖注：生產(chǎn)環(huán)境中的編譯器行為
我們在一個(gè)包含30k 模塊的大型內(nèi)部軟件包上培訓(xùn)了大小內(nèi)聯(lián)策略。訓(xùn)練后的策略在編譯其他軟件時(shí)可以推廣，并減少了3% ~ 7% 的時(shí)間和內(nèi)存開銷。除了跨軟件的通用性之外，跨時(shí)間的通用性也很重要，軟件和編譯器都在積極開發(fā)之中，因此訓(xùn)練有素的策略需要在合理的時(shí)間內(nèi)保持良好的性能。我們在三個(gè)月后評估了該模型在同一組軟件上的性能，發(fā)現(xiàn)只有輕微的退化。
圖注：內(nèi)聯(lián)大小策略大小減少百分比，x 軸表示不同的軟件，y 軸表示減小的百分比?！癟raining”是訓(xùn)練模型的軟件，“InfraX”是不同的內(nèi)部軟件包。
MLGO 的內(nèi)聯(lián)換大小訓(xùn)練已經(jīng)在 Fuchsia 上部署，F(xiàn)uchsia 是一個(gè)通用的開源操作系統(tǒng)，旨在為不同的硬件和軟件生態(tài)系統(tǒng)提供動(dòng)力，其中二進(jìn)制大小是關(guān)鍵。在這里，MLGO 顯示 C++ 翻譯單元的大小減少了6.3%。
作為一個(gè)通用框架，我們使用 MLGO 來改進(jìn)寄存器分配（Register allocation）通道，從而提高 LLVM 中的代碼性能。寄存器分配解決了將物理寄存器分配給活動(dòng)范圍（即變量）的問題。
隨著代碼的執(zhí)行，不同的活范圍在不同的時(shí)間完成，釋放出的寄存器供后續(xù)處理階段使用。在下面的例子中，每個(gè) "加法 "和 "乘法 "指令要求所有操作數(shù)和結(jié)果都在物理寄存器中。實(shí)時(shí)范圍x被分配到綠色寄存器，并在藍(lán)色或黃色寄存器的實(shí)時(shí)范圍之前完成。x 完成后，綠色寄存器變得可用，并被分配給活范圍t。
在代碼執(zhí)行過程中，不同的活范圍在不同的時(shí)間完成，釋放出的寄存器供后續(xù)處理階段使用。在下面的例子中，每個(gè)“加法”和“乘法”指令要求所有操作數(shù)和結(jié)果都在物理寄存器中。活動(dòng)范圍 x 被分配到綠色寄存器，并在藍(lán)色或黃色寄存器的實(shí)時(shí)范圍之前完成。x 完成后，綠色寄存器變得可用，并被分配給活范圍 t 。
圖注：寄存器分配示例
當(dāng)分配活動(dòng)范圍 q 時(shí)，沒有可用的寄存器，因此寄存器分配通道必須決定哪個(gè)活動(dòng)范圍可以從其寄存器中“驅(qū)逐”，以便為 q 騰出空間。這被稱為“現(xiàn)場驅(qū)逐”問題，是我們訓(xùn)練模型來取代原始啟發(fā)式算法的決策。在這個(gè)例子中，它將 z 從黃色寄存器中驅(qū)逐出去，并將其賦給 q 和 z 的前半部分。
我們現(xiàn)在考慮實(shí)際范圍 z 的未分配的下半部分。我們又有一個(gè)沖突，這次活動(dòng)范圍 t 被驅(qū)逐和分割，t 的前半部分和 z 的最后一部分最終使用綠色寄存器。Z 的中間部分對應(yīng)于指令 q = t * y，其中沒有使用 z，因此它沒有被分配給任何寄存器，它的值存儲在來自黃色寄存器的堆棧中，之后被重新加載到綠色寄存器中。同樣的情況也發(fā)生在 t 上。這給代碼增加了額外的加載/存儲指令，降低了性能。寄存器分配算法的目標(biāo)是盡可能地減少這種低效率。這被用作指導(dǎo) RL 策略訓(xùn)練的獎(jiǎng)勵(lì)。
與內(nèi)聯(lián)大小策略類似，寄存器分配(regalloc-for-Performance)策略在 Google 內(nèi)部一個(gè)大型軟件包上進(jìn)行了培訓(xùn)，并且可以在不同的軟件上通用，在一組內(nèi)部大型數(shù)據(jù)中心應(yīng)用程序上每秒查詢次數(shù)（QPS）提高了0.3% ~ 1.5% 。QPS 的改進(jìn)在部署后持續(xù)了幾個(gè)月，顯示該模型的可推廣性。
MLGO使用強(qiáng)化學(xué)習(xí)訓(xùn)練神經(jīng)網(wǎng)絡(luò)來作決策，是一種機(jī)器學(xué)習(xí)策略取代復(fù)雜的啟發(fā)式方法。作為一個(gè)通用的工業(yè)級框架它將更深入、更廣泛應(yīng)用于更多環(huán)境，不僅僅在內(nèi)聯(lián)和寄存器分配。
MLGO可以發(fā)展為：1）更深入，例如增加更多的功能，并應(yīng)用更好的 RL 算法；2）更廣泛，可應(yīng)用于內(nèi)聯(lián)和重新分配之外的更多優(yōu)化啟發(fā)式方法。
作者對 MLGO 能夠?yàn)榫幾g器優(yōu)化領(lǐng)域帶來的可能性充滿熱情，并期待著它的進(jìn)一步采用和研究界未來的貢獻(xiàn)。
參考鏈接：
https://ai.googleblog.com/
github：https://github.com/google/ml-compiler-opt
demo：https://github.com/google/ml-compiler-opt/blob/main/docs/demo/demo.md雷峰網(wǎng)(公眾號：雷峰網(wǎng))
雷峰網(wǎng)
雷峰網(wǎng)版權(quán)文章，未經(jīng)授權(quán)禁止轉(zhuǎn)載。詳情見轉(zhuǎn)載須知。
