記憶體減少3%-7%！Google提出用於編譯器優化的機器學習框架 MLGO

• 2022 年 7 月 11 日
• AI

作者 | 錢雲迪、Mircea Trofin

編譯 | 劉冰一

現代電腦誕生，如何編譯更快、更小的程式碼問題隨之出現。

編譯優化是成本收益比最高的優化手段，更好的程式碼優化可以顯著降低大型數據中心應用程式的操作成本。編譯程式碼的大小對於部署在安全引導分區上的移動和嵌入式系統或軟體來說是至關重要的，因為編譯後的二進位文件必須符合嚴格的程式碼大小預算。隨著這一領域的進步，越來越複雜的啟發式方法嚴重擠壓有限的系統空間，阻礙了維護和進一步的改進。

最近的研究表明，機器學習可以通過用機器學習策略取代複雜的啟發式方法，在編譯器優化中釋放更多的機會。然而，在通用的、行業級編譯器中採用機器學習策略仍然是一個挑戰。

為了解決這個問題，Google兩位高級工程師錢雲迪、Mircea Trofin 提出了「MLGO，一個機器學習指導的編譯器優化框架」，這是第一個工業級的通用框架，用於將機器學習技術系統地集成到 LLVM（一個開源的工業編譯器基礎設施，在構建關鍵任務、高性能軟體時無處不在）中。

MLGO 使用強化學習訓練神經網路來做出決策，以取代 LLVM 中的啟發式演算法。根據作者描述，LLVM 上有兩處 MLGO 優化：

1）通過內聯減少程式碼量；

2）通過暫存器分配提高程式碼性能。

這兩種優化都可以在 LLVM 資源庫中獲得，並已在生產中部署。

1

MLGO是如何工作的？

內聯（Inlining）有助於通過做出能夠刪除冗餘程式碼的決策來減少程式碼大小。在下面的示例中，調用者函數 foo()調用被調用者函數 bar()，而 bar()本身又調用了 baz()。內聯這兩個調用站點將返回一個簡單的 foo()函數，該函數將減小程式碼大小。

圖註：內聯通過刪除冗餘程式碼來減少程式碼大小

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

在MLGO之前，內聯/非內聯的決定是由啟發式方法做出的，隨著時間的推移，這種方法越來越難以改進。MLGO用一個機器學習模型代替了啟發式方法。在調用圖的遍歷過程中，編譯器通過輸入圖中的相關特徵（即輸入）來尋求神經網路對是否內聯特定的調用者-被調用者對的建議，並按順序執行決策，直到遍歷整個調用圖為止。

圖註：內聯過程中MLGO的圖示，「 # bbs」、「 # users」和「 callsite height」是調用者-被調用者對特性的實例

MLGO 使用策略梯度和進化策略演算法對決策網路進行 RL 訓練。雖然沒有關於最佳決策的基本事實，但在線 RL 使用經過培訓的策略在培訓和運行彙編之間進行迭代，以收集數據並改進策略。特別是，考慮到當前訓練中的模型，編譯器在內聯階段諮詢模型，以做出內聯/不內聯的決策。編譯完成後，它產生一個順序決策過程的日誌（狀態、行動、獎勵）。然後，該日誌被傳遞給訓練器以更新模型。這個過程不斷重複，直到得到一個滿意的模型為止。

訓練後的策略被嵌入到編譯器中，在編譯過程中提供內聯/非內聯的決策。與訓練場景不同的是，該策略不生成日誌。TensorFlow 模型被嵌入 XLA AOT ，它將模型轉換為可執行程式碼。這避免了TensorFlow運行時的依賴性和開銷，最大限度地減少了在編譯時由ML模型推理引入的額外時間和記憶體成本。

圖註：生產環境中的編譯器行為

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

圖註：內聯大小策略大小減少百分比，x 軸表示不同的軟體，y 軸表示減小的百分比。「Training」是訓練模型的軟體，「InfraX」是不同的內部軟體包。

MLGO 的內聯換大小訓練已經在 Fuchsia 上部署，Fuchsia 是一個通用的開源作業系統，旨在為不同的硬體和軟體生態系統提供動力，其中二進位大小是關鍵。在這裡，MLGO 顯示 C++ 翻譯單元的大小減少了6.3%。

2

暫存器分配

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

隨著程式碼的執行，不同的活範圍在不同的時間完成，釋放出的暫存器供後續處理階段使用。在下面的例子中，每個 “加法 “和 “乘法 “指令要求所有操作數和結果都在物理暫存器中。實時範圍x被分配到綠色暫存器，並在藍色或黃色暫存器的實時範圍之前完成。x 完成後，綠色暫存器變得可用，並被分配給活範圍t。

在程式碼執行過程中，不同的活範圍在不同的時間完成，釋放出的暫存器供後續處理階段使用。在下面的例子中，每個「加法」和「乘法」指令要求所有操作數和結果都在物理暫存器中。活動範圍 x 被分配到綠色暫存器，並在藍色或黃色暫存器的實時範圍之前完成。x 完成後，綠色暫存器變得可用，並被分配給活範圍 t 。

圖註：暫存器分配示例

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

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

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

3

總結

MLGO使用強化學習訓練神經網路來作決策，是一種機器學習策略取代複雜的啟發式方法。作為一個通用的工業級框架它將更深入、更廣泛應用於更多環境，不僅僅在內聯和暫存器分配。

MLGO可以發展為：1）更深入，例如增加更多的功能，並應用更好的 RL 演算法；2）更廣泛，可應用於內聯和重新分配之外的更多優化啟發式方法。

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

參考鏈接：

//ai.googleblog.com/

github：//github.com/google/ml-compiler-opt

demo：//github.com/google/ml-compiler-opt/blob/main/docs/demo/demo.md雷峰網

雷峰網