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

现代计算机诞生，如何编译更快、更小的代码问题随之出现。

编译优化是成本收益比最高的优化手段，更好的代码优化可以显著降低大型数据中心应用程序的操作成本。编译代码的大小对于部署在安全引导分区上的移动和嵌入式系统或软件来说是至关重要的，因为编译后的二进制文件必须符合严格的代码大小预算。随着这一领域的进步，越来越复杂的启发式方法严重挤压有限的系统空间，阻碍了维护和进一步的改进。

最近的研究表明，机器学习可以通过用机器学习策略取代复杂的启发式方法，在编译器优化中释放更多的机会。然而，在通用的、行业级编译器中采用机器学习策略仍然是一个挑战。

为了解决这个问题，谷歌两位高级工程师钱云迪、Mircea Trofin 提出了“MLGO，一个机器学习指导的编译器优化框架”，这是第一个工业级的通用框架，用于将机器学习技术系统地集成到 LLVM（一个开源的工业编译器基础设施，在构建关键任务、高性能软件时无处不在）中。

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

MLGO 使用强化学习训练神经网络来做出决策，以取代 LLVM 中的启发式算法。根据作者描述，LLVM 上有两处 MLGO 优化：

1）通过内联减少代码量；

2）通过寄存器分配提高代码性能。

这两种优化都可以在 LLVM 资源库中获得，并已在生产中部署。

1 MLGO是如何工作的？

内联（Inlining）有助于通过做出能够删除冗余代码的决策来减少代码大小。在下面的示例中，调用者函数 ​​<span style="font-size: 15px;">foo()</span>​​调用被调用者函数 ​​<span style="font-size: 15px;">bar()</span>​​，而 ​​<span style="font-size: 15px;">bar()</span>​​本身又调用了 ​​<span style="font-size: 15px;">baz()</span>​​。内联这两个调用站点将返回一个简单的 ​​<span style="font-size: 15px;">foo()</span>​​函数，该函数将减小代码大小。

图注：内联通过删除冗余代码来减少代码大小

在實際程式碼中，有成千上萬的函數相互調用，因此構成了一個調用圖（Call graph）。在內聯階段，編譯器遍歷（traverses）所有調用者-被調用者對的調用圖，並決定是否內聯一個調用者-被調用者對。這是一個連續的決策過程，因為先前的內聯決策會改變呼叫圖，影響後面的決策和最終的結果。在上面的範例中，呼叫圖<span style="font-size: 15px;">#foo()</span># → <span style="font-size: 15px;">bar()</span> → <span style="font-size: 15px;">baz()</span>#需要在兩邊做出「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 的改進在部署後持續了幾個月，顯示模型的可推廣性。

3 總結

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

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

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

以上是記憶體減少3%-7%！谷歌提出用於編譯器最佳化的機器學習框架 MLGO的詳細內容。更多資訊請關注PHP中文網其他相關文章！