Firefox、WebAssemblyのコンパイル処理が75倍速に。SpiderMonkeyのJITコンパイラ改善で

FirefoxのJavaScriptエンジンでありWebAssemblyの実行エンジンでもある「SpiderMonkey」の開発チームは、WebAssemblyのコンパイル処理を従来よりも最大で75倍高速にする改善を行ったことを、ブログ「75x faster: optimizing the Ion compiler backend | SpiderMonkey JavaScript/WebAssembly Engine」で明らかにしました。

二段目のコンパイラがメモリを使いすぎている

この性能改善のきっかけは、マイクロソフトがオープンソースで公開している機械学習ライブラリ「ONNX Runtime」のWebAssembly版をSpiderMonekyで実行した際に、最適化のためのコンパイルで大量のメモリと長い時間がかかっていることが判明したことでした。

SpiderMonekyによるWebAssemblyの実行はJITコンパイラによって行われますが、そのJITコンパイラは2段階に分かれています。

一段階目は「Baselineコンパイラ」です。これはWebAssemblyモジュールを読み込むとまずは迅速にコンパイルし、すぐにWebAssemblyモジュールを実行できるようにします。

二段目は、より最適化された高速なネイティブコードを生成する「Ionコンパイラ」です。Ionコンパイラが働くことで、WebAssemblyの実行速度がさらに向上する仕組みになっています。

前述のブログの著者であるJan de Mooij氏の環境においてONNX Runtimeを実行すると、このIonコンパイラによるコンパイルに4GBのメモリが消費され、コンパイルが完了するまでに5分かかることが分かりました。

これは最適化のためのコンパイルとしてもオーバーヘッドが大きすぎるということで、改善が始まったわけです。

5分かかっていたコンパイルが4秒に高速化

結果としてSpiderMonkeyの開発チームはIonコンパイラの改善に成功しました。

その成果は、これまで5分かかっていたIonコンパイラによるONNX Runtimeのコンパイル時間がわずか3.9秒に短縮され、75倍もの速度向上が実現できたことが報告されました。

さらにアドビが公開しているPhotoshopオンライン版のデモも、4分かかっていたコンパイルが14秒になったとのことです。

SpiderMonkey開発チームは今後、WebAssemblyモジュール全体をコンパイルする代わりに関数ごとのコンパイルやインライン化などのテクニックを投入することなどで、Ionコンパイラをさらに高速化するとしています。

Ionコンパイラ、4つの改善点

今回のIonコンパイラの改善は、主に4つの処理に関して行われたことが説明されています。コンパイラの改善に関する解説は本記事の筆者（新野）の正確な理解が追いつかないところが多いため、ここでは参考程度に4つのポイントとなる部分の引用と日本語訳を紹介するに留めさせていただきます。

1つ目の改善は、レジスタへの割り当てに関する部分。

In Ion’s register allocator, each VirtualRegister has a list of LiveRange objects. We were using a linked list for this, sorted by start position. This caused quadratic behavior when allocating registers: the allocator often splits live ranges into smaller ranges and we’d have to iterate over the list for each new range to insert it at the correct position to keep the list sorted. This was very slow for virtual registers with thousands of live ranges.

Ionのレジスタアロケータでは、各バーチャルレジスタがライブレンジオブジェクトのリストを持っている。私たちはこの処理のために、開始位置でソートされたリンクされたリストを使用していたが、そのせいでレジスタを割り当てる際の処理量が二次曲線的になっていた。というのも、アロケータはしばしばライブレンジを小さな範囲に分割するので、あるリストをソートされた状態のままで正しい位置に挿入するには、新しい範囲ごとにリストの処理を繰り返さなければならなかったのだ。何千ものライブレンジを持つ仮想レジスタでは、この処理は非常に遅いものだった。

これは複数のライブレンジのデータ構造をリンクリストからベクターに変更するなどによって改善されました。

次の課題は、スケールしないアルゴリズムによる性能低下でした。

The next problem that stood out in performance profiles was the Dominator Tree Building compiler pass, in particular a function called ComputeImmediateDominators. This function determines the immediate dominator block for each basic block in the MIR graph.

The algorithm we used for this (based on A Simple, Fast Dominance Algorithm by Cooper et al) is relatively simple but didn’t scale well to very large graphs.

パフォーマンスプロファイル分析で目立った次の問題は、Dominator Tree Buildingコンパイラのパス、特にComputeImmediateDominatorsと呼ばれる関数だ。この関数はMIRグラフの各基本ブロックの即時ドミネーターブロックを決定する。これに使用したアルゴリズム（Cooperらの「A Simple, Fast Dominance Algorithm」）は比較的単純だが、非常に大きなグラフにはうまくスケールしなかったのだ。

これはアルゴリズムをLLVMやJuliaのコンパイラに採用されているSemi-NCAに変更することで改善されました。

3つ目の問題は、レジスタアロケータが用いるビットセットのメモリ消費が大きくなること。

For each basic block, the register allocator allocated a (dense) bit set with a bit for each virtual register. These bit sets are used to check which virtual registers are live at the start of a block.

各基本ブロックに対して、レジスタアロケータは各バーチャルレジスタのビットを持つ（密な）ビットセットを割り当てた。 これらのビットセットは、ブロックの開始時にどのバーチャルレジスタが生きているかをチェックするために使用される。

For the largest function in the ONNX Wasm module, this used a lot of memory: 199477 virtual registers x 132856 basic blocks is at least 3.1 GB just for these bit sets! Because most virtual registers have short live ranges, these bit sets had relatively few bits set to 1.

ONNX Wasmモジュールの最大の関数では、これは多くのメモリ、すなわち199477仮想レジスタ×132856基本ブロックのビットセットによる少なくとも3.1GB！ を使用した。ほとんどのバーチャルレジスタはライブレンジが短いため、これらのビットセットは1にセットされるビットが比較的少なかったのだ。

この問題はビットセットを小さな領域に格納できる新しいデータ構造に最適化したことで解決されました。

4つ目の問題は、特定の関数の動作が遅かった点でした。

The last issue that showed up in profiles was a function in the register allocator called createMoveGroupsFromLiveRangeTransitions. After the register allocator assigns a register or stack slot to each live range, this function is responsible for connecting pairs of live ranges by inserting moves.

最後にプロファイルに現れた問題は、createMoveGroupsFromLiveRangeTransitionsと呼ばれるレジスタアロケータ内の関数だった。レジスタアロケータが各ライブレンジにレジスタまたはスタックスロットを割り当てた後、この関数はムーブを挿入してライブレンジのペアを接続する役割を果たす。

この関数の動作が遅い原因は、二次関数的に処理量が増えるループがいくつもあったためでした。このうち主要な2つのループを処理量が線形に増えるように最適化して解決されました。

（追記 10/21 10:22。タイトルを変更しました。当初、WebAssemblyの実行速度が75倍速と表記していましたが、正しくはコンパイル処理が75倍速になったので表現を改めました。お詫びして訂正します）