現(xiàn)代計算機誕生，如何編譯更快、更小的代碼問題隨之出現(xiàn)。

編譯優(yōu)化是成本收益比最高的優(yōu)化手段，更好的代碼優(yōu)化可以顯著降低大型數(shù)據(jù)中心應用程序的操作成本。編譯代碼的大小對于部署在安全引導分區(qū)上的移動和嵌入式系統(tǒng)或軟件來說是至關重要的，因為編譯后的二進制文件必須符合嚴格的代碼大小預算。隨著這一領域的進步，越來越復雜的啟發(fā)式方法嚴重擠壓有限的系統(tǒng)空間，阻礙了維護和進一步的改進。

最近的研究表明，機器學習可以通過用機器學習策略取代復雜的啟發(fā)式方法，在編譯器優(yōu)化中釋放更多的機會。然而，在通用的、行業(yè)級編譯器中采用機器學習策略仍然是一個挑戰(zhàn)。

為了解決這個問題，谷歌兩位高級工程師錢云迪、Mircea Trofin 提出了“MLGO，一個機器學習指導的編譯器優(yōu)化框架”，這是第一個工業(yè)級的通用框架，用于將機器學習技術系統(tǒng)地集成到 LLVM（一個開源的工業(yè)編譯器基礎設施，在構建關鍵任務、高性能軟件時無處不在）中。

論文地址：https://arxiv.org/pdf/2101.04808.pdf

MLGO 使用強化學習訓練神經(jīng)網(wǎng)絡來做出決策，以取代 LLVM 中的啟發(fā)式算法。根據(jù)作者描述，LLVM 上有兩處 MLGO 優(yōu)化：

1）通過內(nèi)聯(lián)減少代碼量；

2）通過寄存器分配提高代碼性能。

這兩種優(yōu)化都可以在 LLVM 資源庫中獲得，并已在生產(chǎn)中部署。

1 MLGO是如何工作的？

內(nèi)聯(lián)（Inlining）有助于通過做出能夠刪除冗余代碼的決策來減少代碼大小。在下面的示例中，調(diào)用者函數(shù) ??foo()??調(diào)用被調(diào)用者函數(shù) ??bar()??，而 ??bar()??本身又調(diào)用了 ??baz()??。內(nèi)聯(lián)這兩個調(diào)用站點將返回一個簡單的 ??foo()??函數(shù)，該函數(shù)將減小代碼大小。

圖注：內(nèi)聯(lián)通過刪除冗余代碼來減少代碼大小

在實際代碼中，有成千上萬的函數(shù)相互調(diào)用，因此構成了一個調(diào)用圖（Call graph）。在內(nèi)聯(lián)階段，編譯器遍歷（traverses）所有調(diào)用者-被調(diào)用者對的調(diào)用圖，并決定是否內(nèi)聯(lián)一個調(diào)用者-被調(diào)用者對。這是一個連續(xù)的決策過程，因為以前的內(nèi)聯(lián)決策會改變調(diào)用圖，影響后面的決策和最終的結果。在上面的例子中，調(diào)用圖??foo()?? → ??bar()?? → ??baz()??需要在兩條邊上做出“yes”的決定，以使代碼大小減少。

在MLGO之前，內(nèi)聯(lián)/非內(nèi)聯(lián)的決定是由啟發(fā)式方法做出的，隨著時間的推移，這種方法越來越難以改進。MLGO用一個機器學習模型代替了啟發(fā)式方法。在調(diào)用圖的遍歷過程中，編譯器通過輸入圖中的相關特征（即輸入）來尋求神經(jīng)網(wǎng)絡對是否內(nèi)聯(lián)特定的調(diào)用者-被調(diào)用者對的建議，并按順序執(zhí)行決策，直到遍歷整個調(diào)用圖為止。

圖注：內(nèi)聯(lián)過程中MLGO的圖示，“ # bbs”、“ # users”和“ callsite height”是調(diào)用者-被調(diào)用者對特性的實例

MLGO 使用策略梯度和進化策略算法對決策網(wǎng)絡進行 RL 訓練。雖然沒有關于最佳決策的基本事實，但在線 RL 使用經(jīng)過培訓的策略在培訓和運行匯編之間進行迭代，以收集數(shù)據(jù)并改進策略。特別是，考慮到當前訓練中的模型，編譯器在內(nèi)聯(lián)階段咨詢模型，以做出內(nèi)聯(lián)/不內(nèi)聯(lián)的決策。編譯完成后，它產(chǎn)生一個順序決策過程的日志（狀態(tài)、行動、獎勵）。然后，該日志被傳遞給訓練器以更新模型。這個過程不斷重復，直到得到一個滿意的模型為止。

圖注：訓練期間的編譯器行為——編譯器將源代碼foo.cpp編譯成對象文件foo.o，并進行了一系列的優(yōu)化，其中一個是內(nèi)聯(lián)通道。

訓練后的策略被嵌入到編譯器中，在編譯過程中提供內(nèi)聯(lián)/非內(nèi)聯(lián)的決策。與訓練場景不同的是，該策略不生成日志。TensorFlow 模型被嵌入 XLA AOT ，它將模型轉(zhuǎn)換為可執(zhí)行代碼。這避免了TensorFlow運行時的依賴性和開銷，最大限度地減少了在編譯時由ML模型推理引入的額外時間和內(nèi)存成本。

圖注：生產(chǎn)環(huán)境中的編譯器行為

我們在一個包含30k 模塊的大型內(nèi)部軟件包上培訓了大小內(nèi)聯(lián)策略。訓練后的策略在編譯其他軟件時可以推廣，并減少了3% ~ 7% 的時間和內(nèi)存開銷。除了跨軟件的通用性之外，跨時間的通用性也很重要，軟件和編譯器都在積極開發(fā)之中，因此訓練有素的策略需要在合理的時間內(nèi)保持良好的性能。我們在三個月后評估了該模型在同一組軟件上的性能，發(fā)現(xiàn)只有輕微的退化。

圖注：內(nèi)聯(lián)大小策略大小減少百分比，x 軸表示不同的軟件，y 軸表示減小的百分比。“Training”是訓練模型的軟件，“InfraX”是不同的內(nèi)部軟件包。

MLGO 的內(nèi)聯(lián)換大小訓練已經(jīng)在 Fuchsia 上部署，F(xiàn)uchsia 是一個通用的開源操作系統(tǒng)，旨在為不同的硬件和軟件生態(tài)系統(tǒng)提供動力，其中二進制大小是關鍵。在這里，MLGO 顯示 C++ 翻譯單元的大小減少了6.3%。

2 寄存器分配

作為一個通用框架，我們使用 MLGO 來改進寄存器分配（Register allocation）通道，從而提高 LLVM 中的代碼性能。寄存器分配解決了將物理寄存器分配給活動范圍（即變量）的問題。

隨著代碼的執(zhí)行，不同的活范圍在不同的時間完成，釋放出的寄存器供后續(xù)處理階段使用。在下面的例子中，每個 "加法 "和 "乘法 "指令要求所有操作數(shù)和結果都在物理寄存器中。實時范圍x被分配到綠色寄存器，并在藍色或黃色寄存器的實時范圍之前完成。x 完成后，綠色寄存器變得可用，并被分配給活范圍t。

在代碼執(zhí)行過程中，不同的活范圍在不同的時間完成，釋放出的寄存器供后續(xù)處理階段使用。在下面的例子中，每個“加法”和“乘法”指令要求所有操作數(shù)和結果都在物理寄存器中?；顒臃秶?x 被分配到綠色寄存器，并在藍色或黃色寄存器的實時范圍之前完成。x 完成后，綠色寄存器變得可用，并被分配給活范圍 t 。

圖注：寄存器分配示例

當分配活動范圍 q 時，沒有可用的寄存器，因此寄存器分配通道必須決定哪個活動范圍可以從其寄存器中“驅(qū)逐”，以便為 q 騰出空間。這被稱為“現(xiàn)場驅(qū)逐”問題，是我們訓練模型來取代原始啟發(fā)式算法的決策。在這個例子中，它將 z 從黃色寄存器中驅(qū)逐出去，并將其賦給 q 和 z 的前半部分。

我們現(xiàn)在考慮實際范圍 z 的未分配的下半部分。我們又有一個沖突，這次活動范圍 t 被驅(qū)逐和分割，t 的前半部分和 z 的最后一部分最終使用綠色寄存器。Z 的中間部分對應于指令 q = t * y，其中沒有使用 z，因此它沒有被分配給任何寄存器，它的值存儲在來自黃色寄存器的堆棧中，之后被重新加載到綠色寄存器中。同樣的情況也發(fā)生在 t 上。這給代碼增加了額外的加載/存儲指令，降低了性能。寄存器分配算法的目標是盡可能地減少這種低效率。這被用作指導 RL 策略訓練的獎勵。

與內(nèi)聯(lián)大小策略類似，寄存器分配(regalloc-for-Performance)策略在 Google 內(nèi)部一個大型軟件包上進行了培訓，并且可以在不同的軟件上通用，在一組內(nèi)部大型數(shù)據(jù)中心應用程序上每秒查詢次數(shù)（QPS）提高了0.3% ~ 1.5% 。QPS 的改進在部署后持續(xù)了幾個月，顯示該模型的可推廣性。

3 總結

MLGO使用強化學習訓練神經(jīng)網(wǎng)絡來作決策，是一種機器學習策略取代復雜的啟發(fā)式方法。作為一個通用的工業(yè)級框架它將更深入、更廣泛應用于更多環(huán)境，不僅僅在內(nèi)聯(lián)和寄存器分配。

MLGO可以發(fā)展為：1）更深入，例如增加更多的功能，并應用更好的 RL 算法；2）更廣泛，可應用于內(nèi)聯(lián)和重新分配之外的更多優(yōu)化啟發(fā)式方法。

作者對 MLGO 能夠為編譯器優(yōu)化領域帶來的可能性充滿熱情，并期待著它的進一步采用和研究界未來的貢獻。