TLDR
Rust で実装した Wasm インタープリタで PGO 試したら 0〜10% ぐらい速くなった
Profile-Guided Optimization (PGO) とは
PGO はコンパイラの最適化の手法のうちの1つですが，プログラムを実行した後に行うという点で少し他と異なります．
- まずは普段どおりの最適化オプションでプログラムをビルドする
- プログラムを実環境で動かし，profile data を取る
- 再度プログラムをビルドしなおす．ただし，ここの最適化では 2. で収集した profile data に基づいてインライン化やレジスタ割り当てなどを行う
- より最適化されたプログラムが出来上がる
このように，実世界でどう実行されるかをフィードバックした最適化をかけて再ビルドすることで，プログラムのビルド時で完結していた従来の最適化よりも進んだ最適化をかけられるようにするというものです．
最近だと，Facebook が BOLT という大規模アプリ向けの最適化フレームワークを研究開発していたり，Google が llvm-propeller というこれも大規模アプリ向けのフレームワークを研究開発していたりで，最適化の分野では熱いトピックだと思ってます．あと Android のネイティブコードでも PGO が使えるみたいです．
この PGO の仕組みは LLVM に既に実装されており，Clang などで利用可能で，Clang のユーザマニュアルもあります．上記の 2. において profile data を収集するために Linux の perf などを使う方法も考えられますが，LLVM を使っていれば，より正確なデータを効率的に収集し，それを元にした最適化をかけるところまで全部やってくれます．Rust の PGO サポートは LLVM のこの仕組みに完全に乗っかる形になっています．
Rust で PGO する
最近 Rust で PGO が既にサポートされているのを rustc のガイドで見かけて，僕の観測範囲ではまだ誰も試してなさそうだったので試してみました．
基本的には下記の短いドキュメントに非常にミニマルな main.rs 単一ファイルでのやり方が書いてあります．
Profile Guided Optimization - The rustc book
ただ，小さすぎるプログラムだと効果が出なさそうなのと，自分の開発してるプログラムでどれぐらい効果があるものなのか興味があったので，今回は wain プロジェクトを使ってやることにしました．
wain は僕が趣味で開発している WebAssembly インタープリタで，no dependency, Safe Rust で約1.5万行ほどの実装規模です．小さいですが，小規模というほどでもない感じですね． 外部依存は無いですが，複数 crate に分割して workspace で管理しています．
準備
PGO は Stable Rust でサポートされているようですが，一応最新の LLVM でやりたかったので nightly を使っています． rustup を使って LLVM 関連のツールをインストールします．
rustup component add llvm-tools-preview
llvm-tools-preview コンポーネントをインストールしてもパスは通されないので，~/.rustup/toolchains/{toolchain}-{triple}/lib/rustlib/{triple}/bin にパスを自前で通しておきます．
# nightly なら export PATH=$HOME/.rustup/toolchains/nightly-x86_64-apple-darwin/lib/rustlib/x86_64-apple-darwin/bin:$PATH # stable なら export PATH=$HOME/.rustup/toolchains/stable-x86_64-apple-darwin/lib/rustlib/x86_64-apple-darwin/bin:$PATH
まずは profile を有効にしてビルドする
profile data を収集するために，まずは profile data を自動生成してくれるオプションを有効にした実行ファイルをビルドします．
CARGO_BUILD_RUSTFLAGS="-Cprofile-generate=$(pwd)/pgo-data" cargo build --release
$CARGO_BUILD_RUSTFLAGS を使うことで，Cargo が rustc を使ってコンパイルする際に追加でコンパイルオプションを指定できます．
-Cprofile-generate=$(pwd)/pgo-data
というコンパイルオプションにより，カレントディレクトリの pgo-data というディレクトリに profile data が収集されていきます．
ここで注意なのですが，cargo rustc コマンドは使えません．今回は依存 crate を含むバイナリ全体を最適化したいので，依存 crate のコードでも profile data を収集する必要があります．そのため，依存 crate をビルドする際にも -Cprofile-generate オプションを指定するのですが，cargo rustc は依存 crate のビルド時には rustc に指定したオプションを渡しません．
上記 cargo build により，普段どおり ./target/release/ に最適化された実行ファイルが出来上がります（今回は ./target/release/wain）
profile data を収集する
今回はお試しなので，単に examples ディレクトリにある Wasm ファイルを順番に実行することにします
./target/release/wain ./examples/brainfxxk.wasm ./target/release/wain ./examples/mandelbrot.wasm ./target/release/wain ./examples/mt19937.wasm ./target/release/wain ./examples/n_queens.wasm ./target/release/wain ./examples/pi.wasm ...
これによって，./pgo-data 内に .profraw という拡張子のバイナリファイルが生成されます．.profraw ファイルは実行ファイルと動的リンクライブラリの分だけ生成されるので，例えば2つライブラリを動的リンクしている実行バイナリであれば，3つの .profraw ファイルが生成されます．
wain は動的リンクライブラリを使っていないので，1つだけ .profraw が生成されます．
> ls ./pgo-data default_13294579539908499303_0.profraw
./target/release/wain に異なる引数を与えて新しい Wasm コードを実行するたびに，収集した profile data が自動でこのファイルに追加されていきます．各 Wasm ファイルを ./target/release/wain で実行しながら default_13294579539908499303_0.profraw を見ると，徐々にファイルサイズが大きくなっていくのが分かります．
profile data を1つのファイルにマージする
前節で説明したように，profile data を収集した .profraw ファイルは複数になることがあります．また，今回は手元の PC でしか profile data を収集していませんが，例えば複数の異なる本番環境で実行した profile data を1箇所に持ってきたりすると必然的に profile data は複数ファイルになります．
これらを llvm-profdata コマンドで1つのファイルにマージします．llvm-profdata コマンドは準備の節で説明した llvm-tools-preview コンポーネントに含まれています．
llvm-profdata merge -o merged.profdata ./pgo-data
このコマンドを走らせるだけで，マージした profile data ファイル merged.profdata を生成してくれます．
profile data を元にした最適化を有効にしてビルド
収集した merged.profdata をコンパイラに食わせて，PGO を行います．
CARGO_BUILD_RUSTFLAGS="-Cprofile-use=$(pwd)/merged.profdata" cargo build --release
このように -Cprofile-use オプションで指定してやるだけです．後は rustc と LLVM がすべてやってくれます．-Cprofile-generate の場合と同様に，cargo rustc は使えないので注意してください．依存 crate を含むプログラム全体で PGO を有効にします．
これで ./target/release/wain に PGO を有効にしてビルドされた実行ファイルが生成されます．通常のリリースビルドでサイズが 716KB，PGO を有効にすると 783KB でわずかにサイズが増えていました．
どれぐらい速くなったのか測る
PGO を有効にした wain と通常のリリースモードでビルドした wain を用意します．
mv ./target/release/wain wain-pgo cargo build --release mv ./target/release/wain wain-orig
profile data の収集に使った examples ディレクトリにある Wasm ファイルを hyperfine を使って比較します．hyperfine は比較したいコマンドを与えると良い感じに実行・計測して比較してくれるすごいやつです．
hyperfine './wain-pgo ./examples/brainfxxk.wasm' './wain-orig ./examples/brainfxxk.wasm' hyperfine './wain-pgo ./examples/mandelbrot.wasm' './wain-orig ./examples/mandelbrot.wasm' hyperfine './wain-pgo ./examples/mt19937.wasm' './wain-orig ./examples/mt19937.wasm' hyperfine './wain-pgo ./examples/n_queens.wasm' './wain-orig ./examples/n_queens.wasm' hyperfine './wain-pgo ./examples/pi.wasm' './wain-orig ./examples/pi.wasm' ...
計測結果
examples/brainfxxk.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/brainfxxk.wasm |
25.0 ± 1.9 | 23.6 | 40.1 | 1.06 ± 0.10 |
./wain-pgo examples/brainfxxk.wasm |
23.5 ± 1.2 | 22.0 | 28.3 | 1.00 |
examples/mandelbrot.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/mandelbrot.wasm |
234.1 ± 5.5 | 227.3 | 248.3 | 1.09 ± 0.03 |
./wain-pgo examples/mandelbrot.wasm |
214.6 ± 3.0 | 211.1 | 220.4 | 1.00 |
examples/mt19937.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/mt19937.wasm |
74.4 ± 1.7 | 72.6 | 79.9 | 1.07 ± 0.03 |
./wain-pgo examples/mt19937.wasm |
69.7 ± 1.6 | 67.8 | 75.2 | 1.00 |
examples/n_queens.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/n_queens.wasm |
9.2 ± 0.6 | 8.5 | 12.9 | 1.06 ± 0.11 |
./wain-pgo examples/n_queens.wasm |
8.7 ± 0.6 | 8.0 | 12.1 | 1.00 |
examples/nbodies.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/nbodies.wasm |
131.1 ± 2.7 | 128.2 | 139.1 | 1.12 ± 0.03 |
./wain-pgo examples/nbodies.wasm |
117.3 ± 2.3 | 114.9 | 125.0 | 1.00 |
examples/pi.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/pi.wasm |
115.7 ± 2.8 | 113.2 | 124.6 | 1.04 ± 0.04 |
./wain-pgo examples/pi.wasm |
111.6 ± 2.7 | 109.2 | 120.3 | 1.00 |
examples/primes.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/primes.wasm |
2.6 ± 0.4 | 2.2 | 5.1 | 1.01 ± 0.21 |
./wain-pgo examples/primes.wasm |
2.6 ± 0.4 | 2.3 | 6.0 | 1.00 |
examples/quicksort.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/quicksort.wasm |
2.5 ± 0.4 | 2.0 | 5.4 | 1.00 |
./wain-pgo examples/quicksort.wasm |
2.5 ± 0.5 | 2.0 | 6.3 | 1.02 ± 0.25 |
examples/sqrt.wasm
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/sqrt.wasm |
2.6 ± 0.4 | 2.1 | 5.2 | 1.00 ± 0.20 |
./wain-pgo examples/sqrt.wasm |
2.6 ± 0.4 | 2.1 | 5.2 | 1.00 |
examples/brainfxxk.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/brainfxxk.wat |
25.3 ± 1.0 | 24.0 | 29.2 | 1.08 ± 0.07 |
./wain-pgo examples/brainfxxk.wat |
23.4 ± 1.1 | 22.3 | 30.9 | 1.00 |
examples/mandelbrot.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/mandelbrot.wat |
234.5 ± 5.0 | 229.0 | 244.1 | 1.08 ± 0.03 |
./wain-pgo examples/mandelbrot.wat |
217.2 ± 3.5 | 212.6 | 223.2 | 1.00 |
examples/mt19937.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/mt19937.wat |
75.2 ± 1.5 | 73.6 | 81.1 | 1.06 ± 0.04 |
./wain-pgo examples/mt19937.wat |
70.6 ± 2.2 | 68.4 | 77.9 | 1.00 |
examples/n_queens.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/n_queens.wat |
9.8 ± 0.7 | 9.0 | 12.3 | 1.05 ± 0.10 |
./wain-pgo examples/n_queens.wat |
9.3 ± 0.6 | 8.6 | 11.9 | 1.00 |
examples/nbodies.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/nbodies.wat |
132.2 ± 3.2 | 128.0 | 139.8 | 1.12 ± 0.03 |
./wain-pgo examples/nbodies.wat |
117.9 ± 2.0 | 115.7 | 122.6 | 1.00 |
examples/pi.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/pi.wat |
115.8 ± 2.8 | 112.6 | 122.0 | 1.04 ± 0.03 |
./wain-pgo examples/pi.wat |
111.3 ± 2.3 | 109.2 | 118.4 | 1.00 |
examples/primes.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/primes.wat |
2.9 ± 0.4 | 2.5 | 6.1 | 1.01 ± 0.22 |
./wain-pgo examples/primes.wat |
2.9 ± 0.4 | 2.4 | 5.6 | 1.00 |
examples/quicksort.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/quicksort.wat |
2.8 ± 0.4 | 2.4 | 6.2 | 1.00 |
./wain-pgo examples/quicksort.wat |
2.8 ± 0.4 | 2.3 | 6.6 | 1.00 ± 0.19 |
examples/sqrt.wat
| Command | Mean [ms] | Min [ms] | Max [ms] | Relative |
|---|---|---|---|---|
./wain-orig examples/sqrt.wat |
2.9 ± 0.4 | 2.4 | 5.3 | 1.01 ± 0.19 |
./wain-pgo examples/sqrt.wat |
2.8 ± 0.4 | 2.4 | 5.3 | 1.00 |
ほぼ変わらないものもありますが，0〜10% ほど速くなっていることが分かります．実行対象がほぼ決まっているプログラムであれば，プログラムに直接手を入れること無く最大 10% の高速化が得られると考えるとかなりお得に見えます．
ただし，もちろん profile data 収集時に使わなかった入力では逆に遅くなる可能性もあるので，実運用時には profile data をどう上手く収集するかがカギの1つになりそうです．
ちなみに毎回手でコマンドを打つのが面倒だったので数十行のシェルスクリプトを書いてやってました．
https://gist.github.com/rhysd/1324f6726f8905f4e71c0f549cb35d97
実際に実行したコマンドを知りたい場合は参照してみてください
