メモリを3%～7%削減！ Google がコンパイラ最適化のための機械学習フレームワーク MLGO を提案

現代のコンピュータの出現により、より高速でより小さなコードをコンパイルする方法が問題になりました。

コンパイル最適化は、最もコスト効率の高い最適化方法です。コードの最適化を向上させることで、大規模なデータセンターアプリケーションの運用コストを大幅に削減できます。コンパイルされたバイナリは厳格なコードサイズバジェットを満たす必要があるため、コンパイルされたコードのサイズは、モバイルシステムや組み込みシステム、またはセキュアブートパーティションに展開されるソフトウェアにとって重要です。分野が進歩するにつれて、ますます複雑化するヒューリスティックによって限られたシステムスペースが著しく圧迫され、メンテナンスやさらなる改善が妨げられます。

最近の研究では、機械学習によって複雑なヒューリスティックを機械学習戦略に置き換えることで、コンパイラの最適化における新たな機会が生まれることが示されています。しかし、汎用の業界グレードのコンパイラーに機械学習戦略を採用することは、依然として課題となっています。

この問題を解決するために、Google の 2 人の上級エンジニア、Yundi Qian 氏と Mircea Trofin 氏が「機械学習ガイド付きコンパイラ最適化フレームワーク MLGO」を提案しました。これは、機械学習テクノロジーを LLVM (ミッションクリティカルな高性能ソフトウェアの構築に広く使用されているオープンソースの産業用コンパイラインフラストラクチャ) に体系的に統合する、初の産業グレードの汎用フレームワークです。

論文アドレス: https://arxiv.org/pdf/2101.04808.pdf

MLGO は、LLVM のヒューリスティックアルゴリズムに代わる、強化学習を使用してニューラルネットワークをトレーニングし、意思決定を行います。著者によると、LLVM には 2 つの MLGO 最適化があります。

1) インライン化によってコードの量を削減します。

2) レジスタ割り当てを通じてコードのパフォーマンスを向上します。

両方の最適化は LLVM リポジトリで利用可能であり、本番環境に導入されています。

1 MLGO はどのように機能しますか?

インライン化は、冗長なコードを削除する決定を行うことで、コードサイズを削減するのに役立ちます。次の例では、呼び出し元関数foo()先関数bar() bar() baz()これら 2 つの呼び出しサイトをインライン化すると、単純なfoo()が返され、コードサイズが削減されます。

図1: インライン化により冗長なコードを削除してコードサイズを削減

実際のコードでは、何千もの関数が互いに呼び出し合って呼び出しグラフを形成します。インライン化フェーズでは、コンパイラはすべての呼び出し元と呼び出し先のペアの呼び出しグラフを走査し、呼び出し元と呼び出し先のペアをインライン化するかどうかを決定します。これは継続的な意思決定プロセスであり、以前のインライン化の決定によって呼び出しグラフが変更され、後の決定と最終結果に影響します。上記の例では、呼び出しグラフfoo() bar() baz() 。

MLGO 以前は、インライン/非インラインの決定はヒューリスティックによって行われていましたが、時間の経過とともに改善がますます困難になっていました。 MLGO はヒューリスティックなアプローチを機械学習モデルに置き換えます。呼び出しグラフのトラバース中に、コンパイラはグラフ内の関連する機能 (つまり、入力) を入力して、特定の呼び出し元と呼び出し先のペアをインライン化するかどうかについてニューラルネットワークのアドバイスを求め、呼び出しグラフ全体がトラバースされるまで、決定を順番に実行します。

図 1: インライン化中の MLGO の図。「#bbs」、「#users」、および「callsite height」は、呼び出し元と呼び出し先のペア機能のインスタンスです。

MLGO は、ポリシー勾配と進化戦略アルゴリズムを使用して、決定ネットワークで RL トレーニングを実行します。最適な決定に関する真実は存在しませんが、オンライン RL はトレーニング済みのポリシーを使用して、トレーニングと実行コンパイルを反復し、データを収集してポリシーを改善します。特に、現在トレーニング中のモデルが与えられた場合、コンパイラーはインライン化フェーズ中にモデルを参照して、インライン化/非インライン化の決定を行います。コンパイルされると、順次的な意思決定プロセス (状態、アクション、報酬) のログが生成されます。このログはトレーナーに渡され、モデルが更新されます。満足のいくモデルが得られるまでこのプロセスが繰り返されます。

図 1: トレーニング中のコンパイラの動作- コンパイラはソースコード foo.cpp をオブジェクトファイル foo.o にコンパイルし、一連の最適化を実行します。その 1 つがインラインパスです。

トレーニングされたポリシーはコンパイラーに埋め込まれ、コンパイルプロセス中にインライン/非インラインの決定を提供します。トレーニングシナリオとは異なり、この戦略ではログは生成されません。 TensorFlow モデルは XLA AOT に埋め込まれ、モデルを実行可能コードに変換します。これにより、TensorFlow ランタイムへの依存性とオーバーヘッドが回避され、コンパイル時の ML モデル推論によって発生する追加の時間とメモリコストが最小限に抑えられます。

図1: 実稼働環境でのコンパイラの動作

30,000 個のモジュールを含む大規模な内部パッケージで、大規模なインラインポリシーと小規模なインラインポリシーの両方をトレーニングしました。トレーニングされたポリシーは、他のソフトウェアをコンパイルするときに一般化でき、時間とメモリのオーバーヘッドが 3% ～ 7% 削減されます。ソフトウェア全体にわたる一般化可能性に加えて、時間全体にわたる一般化可能性も重要です。ソフトウェアとコンパイラーは両方とも活発に開発されているため、トレーニングされたポリシーは、妥当な時間にわたって良好なパフォーマンスを維持する必要があります。 3 か月後、同じソフトウェアセットでモデルのパフォーマンスを評価したところ、わずかな低下しか見られませんでした。

図のキャプション: インラインサイズポリシーのサイズ削減率。x 軸はさまざまなソフトウェアを表し、y 軸は削減率を表します。「Training」はモデルをトレーニングするソフトウェアであり、「InfraX」は別の内部ソフトウェアパッケージです。

MLGO のインラインサイズ変更トレーニングは、バイナリサイズが重要な多様なハードウェアおよびソフトウェアエコシステムを強化するために設計された汎用オープンソースオペレーティングシステムである Fuchsia に導入されています。ここで、MLGO は C++ 翻訳単位のサイズが 6.3% 削減されたことを示しています。

2 レジスタ割り当て

一般的なフレームワークとして、MLGO を使用してレジスタ割り当てパイプラインを改善し、LLVM のコードパフォーマンスを向上させます。レジスタ割り当ては、物理レジスタをアクティブスコープ (つまり変数) に割り当てる問題を解決します。

コードが実行されると、異なるライブ範囲が異なる時間に完了し、後続の処理段階で使用できるようにレジスタが解放されます。次の例では、「加算」命令と「乗算」命令のそれぞれで、すべてのオペランドと結果が物理レジスタに存在する必要があります。ライブスコープ x は緑のレジスタに割り当てられ、青または黄色のレジスタのライブスコープの前に完了します。 x が完了すると、緑色のレジスタが使用可能になり、ライブスコープ t に割り当てられます。

コード実行中、異なるライブ範囲が異なる時間に完了し、後続の処理段階で使用するためにレジスタが解放されます。次の例では、「加算」命令と「乗算」命令のそれぞれで、すべてのオペランドと結果が物理レジスタに存在する必要があります。ライブスコープ x は緑のレジスタに割り当てられ、青または黄色のレジスタのライブスコープの前に完了します。 x が完了すると、緑色のレジスタが使用可能になり、ライブスコープ t に割り当てられます。

図1: レジスタ割り当ての例

アクティブ範囲 q が割り当てられると、使用可能なレジスタがなくなるため、レジスタ割り当てパスでは、q のためのスペースを確保するために、どのアクティブ範囲をレジスタから「削除」できるかを決定する必要があります。これは「フィールドエビクション」問題と呼ばれ、元のヒューリスティックの意思決定を置き換えるようにモデルをトレーニングする場所です。この場合、z を黄色のレジスタから削除し、q と z の前半に割り当てます。

ここで、実際の範囲 z の割り当てられていない下半分を検討します。別の競合が発生し、今回はアクティブ範囲 t が削除されて分割され、t の前半と z の後半が緑のレジスタを使用することになります。 Z の中央部分は、命令 q = t * y に対応します。ここで、z は使用されないためにどのレジスタにも割り当てられず、その値は黄色のレジスタからスタックに格納され、後で緑のレジスタに再ロードされます。 t についても同じことが起こります。これにより、コードに余分なロード/ストア命令が追加され、パフォーマンスが低下します。レジスタ割り当てアルゴリズムの目的は、この非効率性を可能な限り削減することです。これは、RL ポリシーのトレーニングをガイドするための報酬として使用されます。

インラインサイズポリシーと同様に、レジスタ割り当て (regalloc-for-performance) ポリシーは、Google 内の大規模なソフトウェアパッケージでトレーニングされており、さまざまなソフトウェアに一般化できるため、一連の内部大規模データセンターアプリケーションで 1 秒あたりのクエリ数 (QPS) が 0.3% ～ 1.5% 向上します。 QPS の改善は導入後数か月間持続し、モデルの一般化可能性を示しました。