內存減少3%-7%！谷歌提出用于編譯器優化的機器學習框架 MLGO

本文作者：我在思考中

2022-07-11 10:00

導語：谷歌整活！將更快、更小的編譯代碼置入通用型工業級框架。

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

作者 | 錢云迪、Mircea Trofin

編譯 | 劉冰一

編輯 | 陳彩嫻

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

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

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

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

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

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 使用經過培訓的策略在培訓和運行匯編之間進行迭代，以收集數據并改進策略。特別是，考慮到當前訓練中的模型，編譯器在內聯階段咨詢模型，以做出內聯/不內聯的決策。編譯完成后，它產生一個順序決策過程的日志（狀態、行動、獎勵）。然后，該日志被傳遞給訓練器以更新模型。這個過程不斷重復，直到得到一個滿意的模型為止。

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

訓練后的策略被嵌入到編譯器中，在編譯過程中提供內聯/非內聯的決策。與訓練場景不同的是，該策略不生成日志。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 的改進在部署后持續了幾個月，顯示該模型的可推廣性。