ARM64 での gearhash の高速化

ARM64版gearhashの高速化（2倍高速に）

tl;dr: バージョン 0.1.4 以降、gearhash クレートはNEONバックエンドを獲得し、典型的なチャンクサイズにおいてARM64で約2倍高速になりました。aarch64では自動的に選択され、後方互換性もあるため、クレートの利用者は更新するだけで済みます。どのようにしてこれが達成されたかの詳細に興味がある場合は、続きを読むか、最終結果へ直接スキップしてください。

すべてはどのように始まったか

2019年末、私は個人的なバックアップシステムを構築しており、その一部としてコンテンツ定義チャンク分割（content-defined chunking）という技術に興味を持ちました。その中心的なアイデアは、固定のチャンク境界でファイルを分割するのではなく、スライディングウィンドウハッシュ関数をファイル全体に適用し、ハッシュが特定の値になったときにチャンク境界をトリガーするというものです。欠点は、これにより可変長チャンクが分布することですが、利点は、ファイルの中間にバイトシーケンスが挿入または削除されても、チャンク分割がはるかに堅牢になることです。

いずれにせよ、その過程で私はFastCDC論文に出会いました。その構成要素はGEARローリングハッシュです。私は速いものが好きなので、論文で公開されているシリアルアルゴリズムをSIMDアルゴリズムに変換する方法をかなり長い時間かけて検討しました。その結果を、SSE4.2とAVX2の最適化を備えた小さなRustクレートであるgearhashとして公開しました。

クレートを作成した当時、ARM64は最適化する価値のあるターゲットではありませんでした。AWSは1年間ARM64インスタンスを提供していましたが、Cortex-A72コアをベースにした最初の世代のGraviton A1ファミリーだけで、汎用コンピューティングではなくスケールアウトワークロード向けにマーケティングされていました。汎用コンピューティングで競争力のあるコアを備えた最初の世代であるGraviton2は、私の最初のコミットと同じ月にre:Inventで発表され、2020年5月に一般提供が開始されました。Appleは2020年11月にM1を発表しました。

今日まで、多くのことが変わりました。AppleはARM64を一般消費者向けハードウェアの主流に押し上げました。AWSはさらに数世代のGravitonをリリースし、3年連続で追加したCPU容量の半分以上がGravitonであると述べています。GitHub Actionsは2025年に公開リポジトリ向けに無料のARM64ランナーを追加しました。これらのマシンのすべてで、gearhash クレートはスカラーループにフォールバックしていました。

さらに、gearhash は当初私自身以外にほとんど本番ユーザーがいませんでしたが、その後、Hugging Faceの大型ファイル用ストレージプロトコルであるXetクライアントのコア部分となり、Git LFSに代わってデフォルトになりました。gearhash にとっては、これは1日あたり1万から2万回のダウンロードが行われていることを意味します。クレートへのこの新たな関心が、どこまでプッシュできるかを探求するモチベーションを見つけるのに役立ちました。

並列化を可能にした洞察

ギアハッシュカーネルは、64ビット符号なし整数に対するシリアル関数として定義されています。

hash = (hash << 1).wrapping_add(table[byte as usize]);

この2つのプロパティが、これをベクトル化することを困難にしています。

1. シリアルな依存関係チェーンです。 各バイトのハッシュは前のバイトに依存します。単一ストリームから抽出できるデータ並列性はありません。
2. テーブルルックアップはギャザーです。 256 × 8バイトは2 KBで、レジスタ内のパーミュートには大きすぎます。各バイトは実際のロードを必要とします。

これについてしばらく頭を悩ませた後、最初のプロパティに関する観察を行いました。ハッシュは64ビット幅で、バイトごとに1ビット左シフトされるため、64バイト後には開始値は完全にシフトアウトされます。これは、バッファの任意のオフセットから hash = 0 でハッシュを開始し、64バイトでウォームアップすれば、それ以降のハッシュは先頭からのパスとビット単位で同一であることを意味します。

これにより、チャンクをストリップに分割できます。レーン0には実際の入力ハッシュをシードし、他のすべてのレーンにはそのストリップの前の64バイトをハッシュしてシードし、すべてのストリップを同期して実行します。レーンが一致を報告した場合、最も早い一致がどれかを特定するだけで済みます。実装の複雑さの大部分はここにありました。

NEONへの移植開始

SSE4.2実装からの直接移植から始めました。aarch64::uint64x2_t は2つの64ビットレーンで、x86_64::__m128i と同じなので、SSE4.2の構造はほぼ機械的にマッピングされます。

マッピングされなかった唯一のものはマスク抽出です。NEONには pmovmskb のような等価物がないため、レーン比較結果をスカラーレジスタに取得するには、ナローイングシフトとムーブが必要で、これを小さな movemask ヘルパーでラップしました。

結果は残念なものでした。0.92倍、スカラーコードよりも遅かったです。

その理由を理解するには、ARM64でのループ内レイテンシを確認する必要があります。イテレーションごとに、NEONバージョンは次のことを行います。

add.2d   v1, v1, v1      ; h << 1, LLVMは自身への加算として発行します
add.2d   v1, v1, v_g     ; h + g

Appleコアでは、これらはそれぞれ約2サイクル（Dougall JohnsonのM1テーブルによる）なので、イテレーションあたり約4サイクルです。イテレーションは2バイト（レーンあたり1バイト）をカバーするため、バイトあたり約2サイクルになります。

スカラーバージョン hash = (hash << 1) + table[b] は、単一のシフトレジスタ加算 add x0, x1, x0, lsl #1 にコンパイルされ、レイテンシは約2サイクルです。これもバイトあたり約2サイクルです。

これは、ベクトルバージョンがクリティカルパスあたりの作業量がスカラーバージョンと同じであることを意味しますが、それに加えてロードとマスク抽出のコストを支払う必要があります。これでは勝てません。

NEONで勝つには、依存関係チェーン自体を短くする必要があります。

チェーンの短縮

チェーンが2バイトあたり2オペレーションであるなら、2バイトあたりの更新を2ステップ実行し、掛け合わせることで、なぜ4バイトあたりのオペレーションにしないのでしょうか。

h₁ = (h << 1) + g₀
h₂ = (h << 2) + (g₀ << 1) + g₁

これにより、h₂ は単一のシフトと単一の加算を通じて h に依存します。ただし、G = (g₀ << 1) + g₁ を事前に計算しておく必要があります。G は h には依存せず、テーブルルックアップのみに依存するため、クリティカルパスの外にあります。

結果：0.92倍 → 1.13倍。改善しましたが、まだ期待される2倍には程遠いです。

ディスアセンブリに隠された理由

add.2d  v2, v1, v1      ; h << 1
add.2d  v2, v3, v2      ; h₁ = (h<<1) + g₀
shl.2d  v1, v1, #2      ; h << 2
add.2d  v3, v3, v3      ; g₀ << 1
add.2d  v1, v1, v4      ; (h<<2) + g₁      <-- hチェーン上
add.2d  v1, v3, v1      ; ... + (g₀<<1)    <-- hチェーン上でも

LLVMが単に再結合したことが判明しました。私は (h << 2) + (G₀ + G₁) と書きましたが、((h << 2) + G₁) + G₀ を発行しました。これはもちろん合法的な変換ですが、依存関係チェーンに2番目の加算を戻してしまいます。

やや不正な修正

コンパイラが合計を再結合するのを止めることはできませんが、1つを（見えなくするよう）試みることはできます。結合された項は2つのテーブルルックアップから構築され、それらは anyway 一般用途レジスタに到着するため、結合はそこで行うことができます。

let (t00, t01) = (table[b00 as usize], table[b01 as usize]);
let (t10, t11) = (table[b10 as usize], table[b11 as usize]);

// 2つのテーブルエントリをスカラーレジスタで結合すると、ベクトルオペランドが
// 不透明になり、コンパイラが以下の加算をループで持ち越される`h`チェーン上の
// 2つの依存ベクトル加算に再結合するのを防ぎます。
let g1 = vcombine_u64(
    vcreate_u64((t00 << 1).wrapping_add(t01)),
    vcreate_u64((t10 << 1).wrapping_add(t11)),
);

let h2 = vaddq_u64(vshlq_n_u64::<2>(h), g1);

ディスアセンブリを確認すると、チェーン上では shl.2d → add.2d のみが見られました。

結果：1.13倍 → 1.46倍。ようやく意味のある高速化が始まりましたが、まだ十分ではありません！

レイテンシバウンドからスループットバウンドへ

一度に2ステップまでアンロールした結果に励まされ、4つの独立した状態（それぞれが依然として h から直接計算される）で同じことを試しました。

let h1 = vaddq_u64(vshlq_n_u64::<1>(h), g[0]);   // hk == (h << k) + g[k-1]
// ...
let h4 = vaddq_u64(vshlq_n_u64::<4>(h), g[3]);

これにより、バイトあたりの依存関係チェーンの長さが再び半分になるため、さらなる大きなステップを期待していました。しかし、測定してみると、まったく違いはありませんでした。この時点で、限界はもはやイテレーション間のレイテンシではなく、単にCPUのスループットである可能性が高いと疑いました。

その仮説に従い、命令数を減らすことに焦点を移しました。ループは現在、バイトあたり2回のロードを行っていました。1つはバイト自体のため、もう1つはそのテーブルエントリのためです。後者は避けられませんが、ここでは4つの連続するバイトロードを1つのアラインメントされていない32ビットロードに置き換えることができ、その後、各ワードから1バイトをシフトで剥がします。

結果：1.46倍 → 1.63倍。2倍の目標に向けて、さらに大きな一歩です。

4つの状態を1つのブランチでテスト

ロードが最適化された後、イテレーションあたりの次の大きな命令ブロックは境界テストでした。各ステップで、コードはハッシュをマスクしてゼロと比較することでチャンク境界をチェックします。

4つの独立した状態までアンロールしていましたが、それらを1つずつプローブしていました。これはベクトルユニットから4回の別々のムーブであり、それぞれにブランチが待機していました。

一般的なケースは一致しないため、実際に行えることは、4つのテストをベクトルユニット内で結合し、1回の移動で済ませることです。どちらのストリップの4つの位置にも境界がない場合、イテレーションは汎用レジスタへの1回のムーブと1回のブランチの後に続行できます。

let t = vandq_u64(
    vandq_u64(vtstq_u64(h1, maskv), vtstq_u64(h2, maskv)),
    vandq_u64(vtstq_u64(h3, maskv), vtstq_u64(h4, maskv)),
);

if movemask(t) == u64::MAX {
    i += UNROLL;
    continue;
}

このチェックが失敗するまれなケースでのみ、コードは4つの状態を1つずつ調べます。Xetクライアントが64 KiBチャンクに使用する16ビットマスクでは、これは約8192イテレーションごとに1回発生します。

結果：1.63倍 → 1.81倍。これには非常に満足しており、現時点ではここで止めました。

最終結果

すべてを組み合わせると、クレートの11ビットベンチマークマスクで、直接移植の0.92倍から、公開したバージョン0.1.4の1.81倍になります。

しかし、作業中に気づいたことですが、ベンチマークがいかにマスク密度に依存しているかは、後から考えると非常に明白です。これは、最適化されたパスにはスカラーパスにはない固定のコールコストがあり、マスクが疎であるほど境界とコールが多くなるためです。それがどれほど重要かを確認するために、4ビットから20ビットを設定したさまざまなマスクでベンチマークを実行しました。

［画像: Apple M4での平均チャンクサイズに対するスカラーパスとNEONパスのスループット。NEONは350バイト未満のチャンクではスカラーより遅く、64 KiB以上では4,350 MB/s近くで平坦化しますが、スカラーは全体で2,000 MB/s近くを維持します。］

クレート自体のベンチマークは1.8倍で、カーブの急峻な部分にあり、64 KiBから1 MiBの間で約2.17倍で平坦化するまで上昇し続けます。16ビットの行はXetクライアントが使用するマスクで、2.14倍です。

サイドノート：平均チャンクが約350バイトを下回ると、コールあたりのコストが利点を上回り、NEONパスはスカラーよりも徐々に遅くなります。このタイプの構成を使用することは期待しませんが、ビット数が少ないマスクに対してスカラーパスにフォールバックすることは、そのギャップを埋めるための安価な方法のように思えます。

次は何？

NEONがスカラーの約2倍高速になったので、x86バックエンドを再検討する時期が来たように感じます。誰が知っているでしょうか、NEON実装で学んだトリックのいくつかが引き継がれるかもしれません。続報をお待ちください！