JITコンパイルを5μsで実行する

歴史的に、高速なJITコンパイルは秘術でした。高速なJITコンパイラを書くには、アセンブリ言語の書き方を知る必要がありました。例を挙げると、現在、独自のJITコンパイラを持つ本番環境対応のデータベースは存在しません。それらはすべてLLVMを使用するか、C/C++コードを生成しています。これらのどちらの選択肢もコンパイル時間が長くなるという欠点があり、適用範囲を制限しています。 現在では、AIの助けを借りることで、アセンブリを直接ターゲットにすることで、高速なコンパイル時間を持つJITコンパイラを簡単に書けるようになりました。これは、新しいデータベースが古いデータベースを改善する機会でもあります。 pgrustを構築する際、当初はJITコンパイラの実装は非常に難しいだろうと考えていました。しかし、最終的にはAIの支援のおかげで予想よりもはるかに簡単であることがわかり、pgrustが非常に高速である理由の一部となっています。 pgrustのJITコンパイラは、約5μsでコードをコンパイルします。これにより、サブセットだけでなく、すべてのSQLクエリをJITコンパイルすることが可能になります。この記事では、独自の高速JITコンパイラを構築する方法を順を追って説明します。例として、JITコンパイルを使用する簡単な正規表現エンジンを構築します。 なぜJITコンパイルなのか JITコンパイルとは、実行時にコンパイル済みコードを生成する、つまり「ジャストインタイム」にコンパイルする手法です。適切に行われれば、パフォーマンスが劇的に向上し、多くの場合2〜5倍、時にはそれ以上になります。 JITコンパイルの主なユースケースは、実行時に得られる情報がプログラムの動作を劇的に変更する場合です。これは、プログラミング言語インタプリタで特に一般的です。インタプリタは実行時に実行するコードを受け取ります。 JITコンパイラは、プログラミング言語以外の分野、例えばデータの解析などでも役立ちます。実行時に解析するデータのスキーマがわからない場合があり、JITはその助けになります。 まず、おもちゃの正規表現エンジンの実装から始めましょう。シンプルにするために、リテラル文字列と繰り返し（つまり正規表現の*）の2つの機能のみをサポートします。パーサーはスキップし、正規表現をすでに解析されたRustの構造体として表現します。これは、apples b(an)*のような文字列をサポートできることを意味しますが、代替やルックビハインドなどはサポートしません。 コードでは、これは非常にシンプルです。3種類のノードがあります：リテラル文字列ノード、繰り返しノード、そして2つのノードの組み合わせである連結ノードです。これは最終的に次のようになります。 enum Node { Literal(&'static str), Concatenation(Box<Node>, Box<Node>), Repetition(Box<Node>), } fn literal(text: &'static str) -> Node { Node::Literal(text) } fn concatenation(left: Node, right: Node) -> Node { Node::Concatenation(Box::new(left), Box::new(right)) } fn repetition(body: Node) -> Node { Node::Repetition(Box::new(body)) } 私たちの正規表現エンジンのインタプリタを書くのも同様に簡単です。 fn match_node(node: &Node, input: &[u8], pos: usize, next: &dyn Fn(usize) -> bool) -> bool { match node { Node::Literal(text) => { let literal = text.as_bytes(); input[pos..].starts_with(literal) && next(pos + literal.len()) } Node::Concatenation(left, right) => { match_node(left, input, pos, &|left_end| { match_node(right, input, left_end, next) }) } Node::Repetition(body) => { match_node(body, input, pos, &|body_end| { match_node(node, input, body_end, next) }) || next(pos) } } } fn interp_match(regex: &Node, input: &str) -> bool { let bytes = input.as_bytes(); match_node(regex, bytes, 0, &|pos| pos == bytes.len()) } この正規表現エンジンは非常にシンプルです。コードは20行未満ですが、パフォーマンスはどうでしょうか。 比較のために、正規表現専用に実装された手書きコードと比較します。例として、正規表現b(an)*を使用します。手書きコードは次のようになります。 fn handwritten_b_an_star(input: &str) -> bool { let bytes = input.as_bytes(); let mut pos = 0; if pos == bytes.len() || bytes[pos] != b'b' { return false; } pos += 1; while pos < bytes.len() { if bytes[pos] != b'a' { return false; } pos += 1; if pos == bytes.len() || bytes[pos] != b'n' { return false; } pos += 1; } true } （このコードを最適化してさらに高速化する方法はありますが、ここでは比較対象として十分です） 私がこれらの2つをいくつかの例でベンチマークすると、手書きバージョンがインタプリタの10〜20倍高速であることがわかりました。明らかに改善の余地がたくさんあります。 次に、JITコンパイルを使用して、手書きバージョンと同等のパフォーマンスを持つ汎用正規表現エンジンをどのように使用できるかを見てみましょう。 JITコンパイルの方法 コードをJITコンパイルするには、2つのステップがあります。まず、実行したいコードのアセンブリを生成します。コードができたら、そのアセンブリコードをプログラム内の他のコードと同じように呼び出せる関数にパッケージ化します。 アセンブリを生成するために、コピー＆パッチと呼ばれるアプローチのバリアントを使用します。このアイデアは、JITコンパイルしたいさまざまな操作のために、アセンブリに一連のテンプレートがあるということです。これらのテンプレートは「ステンシル」と呼ばれます。操作をJITコンパイルしたい場合は、関連するステンシルを取得し、操作の具体性に基づいて小さな調整を加えます。実際のステンシルに記入するのと非常によく似ています。 これらの記入済みステンシルをいくつか連結することで、手書きバージョンと同等のパフォーマンスを持つプログラムを実行時に構築できます。ここで進む道は次のとおりです。まず、b(an)*に対して生成したいARM64コードを見ていきます。次に、繰り返される命令シーケンスを再利用可能なステンシルに変換し、正規表現ASTからそれらのステンシルを埋め込んで結合するエミッターを書き、最後に生成された命令を実行可能メモリにコピーしてRustが通常の関数のように呼び出せるようにします。 これがどのように機能するかを説明するために、生成されたコードから始めてJITコンパイラ自体に向かって逆行するのが最も簡単です。再び、正規表現「b(an)*」を扱います。 設計上の決定をいくつか示します。 バックトラッキングのためにスタックを使用します。スタックは、正規表現で行き詰まった場合に遷移すべき状態を追跡します。 マッチング対象の文字列はヌルバイトで終わります。これは、文字列の終わりに到達した場合、文字比較はすべて自動的に失敗することを意味します。これにより、どの時点でも長さ比較を行う必要がなくなります。 プログラムの状態には、次のレジスタを使用します。 x0 – 文字列内の現在の位置と戻り値 x1 – バックトラッキングに使用するスタックのトップ x2 – バックトラッキングに使用するスタックのボトム（スタックが空かどうかを判断するために必要です） x9 – 一時変数として使用 プログラムへの入力として、次のものが渡されます。 x0 – 文字列の先頭へのポインタ x1 – スタックに使用する場所へのポインタ 生成されたARM64 それらの準備ができたので、生成されたアセンブリを部分ごとに見ていきましょう。これはmacOSのARM64に特化しています。まず、プログラムを初期化するプロローグがあります。これが行うことはすべて、スタックのトップとスタックのボトムを渡された値に設定してスタックを初期化することです。 0: aa0103e2 mov x2, x1 次に、文字bをチェックするコードがあります。文字がbでない場合、フォールバックロジックを処理するコードブロックにジャンプします。そうでない場合は、文字列内の位置を進めます。 ; CHAR 'b' 4: 39400009 ldrb w9, [x0] ; 現在の入力バイトをロード 8: 7101893f cmp w9, #0x62 ; 'b'か？ c: 54000281 b.ne 0x5c ; いいえ -> フォールバックブロック 10: 91000400 add x0, x0, #1 ; はい -> 入力を進める 次に、繰り返し（an）*のコードがあります。繰り返しの場合、バックトラッキングを行う必要があります。ここでバックトラックする場合、ループの終了に直接ジャンプすることを意味します。これは、ループの後の命令のアドレスと文字列内の位置の両方をスタックに保存する必要があることを意味します。 14: d2800989 movz x9, #0x004c ; 再開アドレスを構築 18: f2a00009 movk x9, #0x0000, lsl #16 ; = 0x1_0000_004c 1c: f2c00029 movk x9, #0x0001, lsl #32 ; （ループ終了） 20: f2e00009 movk x9, #0x0000, lsl #48 ; 24: a8810029 stp x9, x0, [x1], #16 ; (終了、位置)をスタックにプッシュ