原文

GPU-Accelerated WHIR Proving on Apple Silicon — miha-stopar (2026-04-30)

Apple Silicon 上での GPU アクセラレーション WHIR 証明: クライアントサイド Metal コンピュートからのベンチマークと教訓

謝辞

セクション4のApple M3 MacBookおよびiPhoneのベンチマーク結果を提供してくれたMoven Tsai氏、そしてWHIRに関する議論をしてくれたAlex Kuzmin氏に感謝します。


TL;DR

Apple Silicon GPU上でMetalコンピュートシェーダーを用いて[WHIR]証明者 (Prover) を高速化し、M1チップ上で高度に最適化されたCPUコード(SIMD + リンク時最適化 (LTO) + target-cpu=native)と比較して最大2.03倍の高速化、補足のApple M3 MacBookでの実行では最大2.58倍、そしてWHIR Bench iOSアプリ(Metal Apple A19 GPU)からの疎なサンプルでは約1.4〜2.3倍の高速化を達成しました(セクション4を参照)。GPUパイプラインは、数論変換 (NTT)、ビット反転、Poseidon2 マークルツリー (Merkle tree) ハッシュ、およびプルーフ・オブ・ワーク (PoW) グラインディングを単一のコマンドバッファ (command buffer) 提出に融合し、Apple Siliconのユニファイドメモリ (Unified Memory) アーキテクチャを活用しています。この実装はオープンソースであり、Apple Siliconを搭載したあらゆるMacで動作し、同じアプリでiPhone上でのオンデバイステストも可能です。

主な発見:

  • 多項式のサイズが2^20(100万以上の係数)以上の場合、テストしたすべての構成でGPUが優位
  • 融合されたDFT+MerkleパイプラインはCPUの往復を回避し、最大の性能向上を提供
  • Apple SiliconのユニファイドメモリはPCIe転送コストを排除するが、より微妙なトレードオフを導入する:共有モードバッファ(CPU+GPUアクセス可能)は、キャッシュコヒーレンス (Cache Coherence) のオーバーヘッドにより、GPU管理バッファよりもGPUコンピュートで遅い。我々のパイプラインはハイブリッドアプローチを使用。
  • Poseidon2 Merkleカーネルが支配的なコスト(GPU時間の約58%)。Xcode GPUプロファイラは高いALU (算術論理演算ユニット) 使用率と明らかなストールがないことを示しており、Appleが正確な整数スループット仕様を公開していないため確認できないものの、このワークロードではハードウェアの限界に近いことを示唆。
  • コンパイラ最適化(LTO、target-cpu=native)によりCPUベースラインが約25%改善され、GPUが打ち破るのが難しくなったが、絶対的なエンドツーエンド性能は向上

リポジトリ: github.com/privacy-ethereum/whir-p3-metal

我々の作業は、Plonky3ライブラリを使用したWHIRプロトコルのRust実装であるtcoratger/whir-p3に基づいています。このコードベースにMetal GPUアクセラレーションを追加しました。(系譜: WizardOfMenlo/whirwhir-p3whir-p3-metal。)

開示. Metal GPU実装(RustおよびMSL)、サポートツール、およびこの記事の散文は、人間の指示、レビュー、ベンチマークの下で、AIコーディングアシスタントであるClaude(Anthropic)からの実質的な支援を受けて作成されました。暗号学的健全性は、公開されているWHIR / Plonky3の仕様およびアップストリームコードに従います。統合、性能に関する主張、または解釈におけるいかなる誤りも、プロジェクトメンテナーの責任です。


1. 動機: クライアントサイド証明が重要な理由

Ethereumの透明性にはプライバシーコストが伴います。すべてのトランザクションは永続的に可視であり、チェーン分析ツールは匿名アドレス (pseudonymous addresses) を実在の身元にリンクできます。ゼロ知識証明 (Zero-knowledge proofs) はプライバシーを回復できますが、証明生成をサーバーに委任すると、サーバーがプライベートな入力を参照するため、その目的が損なわれます。

真のプライバシーにはクライアントサイド証明 (client-side proving) が必要です。ユーザーは自分のデバイスで証明を生成します。これは以下の点で重要です。

  • プライベートな支払い – 金額、取引相手、取引パターンを隠す
  • ID – 資格情報自体(年齢、市民権、メンバーシップ)を明かすことなく、資格情報に関する事実を証明する(ZK EmailAnon Aadhaar
  • 投票 – DAO (分散型自律組織) およびガバナンスへの匿名参加(Semaphore

障壁となるのは性能です。消費者向けハードウェアでの証明は、インタラクティブな使用に十分な速さである必要があります。サーバーサイドGPU証明者(データセンターGPU上のCUDA)は劇的な高速化を達成しますが、クライアントデバイスには異なる制約があります。熱制限 (thermal limits)、共有メモリ帯域幅 (shared memory bandwidth)、より少ないGPUコア数、およびディスクリートVRAM (discrete VRAM) がありません。

クライアントサイドGPUの機会

現代のスマートフォンやラップトップには、ますます高性能なGPUが搭載されています。Apple SiliconのMシリーズおよびAシリーズチップは、CPUとGPU間でユニファイドメモリを共有しており、データセンターGPU証明を支配するPCIe転送ボトルネック (PCIe transfer bottleneck) を排除します。このアーキテクチャの違いは、CPUとGPUのきめ細かな連携に独自の機会を生み出します。

いくつかのプロジェクトがこの分野を探求しています。

  • Mopro – Metal 多点スカラー乗算 (MSM) アクセラレーション(v2の解説)、WebGPU 有限体演算 (Field Ops) ベンチマークでBN254と比較して小規模な有限体で100倍以上のスループットを示す
  • ICICLE Metal – Apple Metal用のMSMおよびNTTプリミティブ、最大5倍の高速化(v3.6)
  • zkSecurity / Stwo WebGPU – WebGPUコンピュートシェーダーを介したブラウザでのCircle STARKsの全体的な証明速度2倍向上
  • Ligetron – クロスプラットフォーム証明用のWebGPU SHA-256およびNTT
  • FibRace – 大規模なモバイルベンチマーク(6,000人以上の参加者、210万件の証明)で、ほとんどの最新スマートフォンが5秒未満で証明できることを示す

我々の作業は、特定の、実用的に関連性の高いターゲットに焦点を当てています。ハッシュベースでポスト量子安全なWHIR多項式コミットメントスキームを、AppleのMetal APIを使用してネイティブGPUコンピュートで高速化することです。


2. WHIR: 背景

WHIR(Gal Arnon, Weizmann Institute; Alessandro Chiesa, EPFL; Giacomo Fenzi, EPFL; Eylon Yogev, Bar-Ilan University; EUROCRYPT 2025で発表)は、リード・ソロモン符号 (Reed-Solomon codes) の近接性に関するインタラクティブオラクル証明 (Interactive Oracle Proof) です。これは、FRISTIR、およびBaseFoldの効率的な代替として機能し、特に高速な検証 (Verification)(他の代替案がミリ秒単位であるのに対し、数百マイクロ秒)が特徴です。

証明 (Proving) における支配的なコストは以下の通りです。

  1. 数論変換 (NTT) – 拡張領域での多項式評価
  2. マークルツリー (Merkle tree) 構築 – 多項式コミットメントのためのPoseidon2ハッシュ
  3. プルーフ・オブ・ワーク (PoW) グラインディング – ハッシュ難易度ターゲットを満たすノンスを見つける(Fiat-Shamirセキュリティ)

これら3つはすべて大規模に並列化可能であり、GPUコンピュートによく適合します。証明者 (Prover) は複数のSTIRラウンドを実行し、各ラウンドにはNTT、Merkleコミットメント、およびPoWグラインドが含まれます。ラウンド間では、CPUがラウンド間でシーケンシャルなサムチェック (Sumcheck) 演算を実行します。

我々の実装は、Plonky3ライブラリの有限体演算 (field arithmetic) とBabyBear有限体 (BabyBear field)(素数 p = 2^31 - 2^27 + 1、モンゴメリ形式 (Montgomery Form))上のWHIRプロトコル実装を使用するtcoratger/whir-p3に基づいています。


3. GPUアーキテクチャと実装

3.1 なぜMetalなのか(WebGPUやCUDAではないのか)

  • ユニファイドメモリ (Unified Memory): Apple SiliconはCPUとGPU間で物理メモリを共有します。これにより、WebGPUやCUDA証明プロジェクトで報告されているPCIe転送ボトルネックが解消されます。しかし、セクション6で議論するように、「ユニファイド」は「無料」を意味しません。重要なキャッシュコヒーレンス (Cache Coherence) のトレードオフがあります。
  • 低いディスパッチオーバーヘッド: Metalコマンドバッファ (command buffer) は、以前の作業がGPUで実行されている間にCPUで構築できるため、密なパイプライン処理が可能です。
  • ネイティブ性能: Metal Shading Language (MSL) はビルド時にAIR (Apple Intermediate Representation) にコンパイルされ、ロード時にデバイス固有のマシンコードにコンパイルされます。ランタイムシェーダーコンパイルのオーバーヘッドはありません。
  • iOS互換性: 同じMetalコードがiPhoneおよびiPadで実行され、モバイルベンチマークが可能です。

トレードオフはAppleへのプラットフォームロックイン (platform lock-in) です。クロスプラットフォーム展開の場合、WebGPU(wgpu経由)が代替手段となりますが、ある程度の性能コストがかかります。

3.2 パイプラインアーキテクチャ

GPUパイプラインは、複数のステージを単一のMetalコマンドバッファに融合します。

CPU入力バッファ (ゼロコピー共有メモリ)
    │
    ├── Radix-16/32 DIF-NTTステージ (GPU管理バッファ上でインプレース)
    │       R32, R16, R8, R4, R2バタフライカーネルを使用
    │       最終ステージ用共有メモリカーネル (スレッドグループあたり最大4096要素)
    │
    ├── ビット反転順列 (最終NTTステージと融合)
    │       ゼロコピーバッファに直接書き戻し
    │
    ├── Poseidon2リーフハッシュ (幅16, 8+13+4ラウンド, x^7 S-box)
    │       4リーフ融合カーネル: 4リーフをハッシュ + 最初のマークル圧縮を1ディスパッチで実行
    │
    ├── Poseidon2マークル圧縮 (残りの全ツリーレベル)
    │       小レベル用SIMDシャッフルカーネル (スレッドグループメモリを回避)
    │
    └── [単一GPU待機] → CPUアクセス可能メモリに結果

この融合により、ナイーブな実装と比較して、STIRラウンドごとに3〜4回のCPU-GPU同期ポイントが解消されます。

3.3 主要な最適化(30回のイテレーション)

我々は30回の最適化イテレーションを実施しました。最も影響の大きかったものは以下の通りです。

#最適化影響
1-4基本的なMetal NTT + MerkleカーネルベースラインGPUパス
5-8Radix-16 DIF、共有メモリバタフライNTT速度2-3倍向上
9-12融合DFT+Merkleパイプライン、ゼロコピーI/O (zero-copy I/O)エンドツーエンドで1.5-2倍向上
13-16Poseidon2 4リーフ融合カーネル、SIMD MerkleMerkle速度1.3倍向上
17-20GPU PoWグラインディング、ゼロコピー拡張体 (Extension Field) 変換 (EF conversions)PoWが支配的な構成で有効
21-24拡張体ゼロコピー、ラウンドごとの融合大規模なnで1.2倍向上
25-28LTO、target-cpu=native、プロファイリングに基づく絶対性能10%向上
29-30R32 DIFカーネル、パックトランスポーズバイパスラウンドあたり5-15ms節約

最も重要な教訓は、個々のカーネルを最適化するよりも、CPU-GPU間の往復を避けるために操作を融合することの方が重要であるということです。融合されたパイプライン(DFT+Merkle用の単一コマンドバッファ)は、非融合パスを通常15〜30%上回ります。

3.4 Metalにおけるモンゴメリ演算

BabyBear有限体 (BabyBear field) の演算は、GPUカーネル全体でモンゴメリ形式 (Montgomery Form) を使用します。モンゴメリモジュラ乗算 (Montgomery Modular Multiplication) は、高価な除算ベースのモジュラ削減 (modular reduction) を、より安価な乗算とシフト操作に置き換えます。重要な洞察は、BabyBearのような31ビット素数(p < 2^31)の場合、2つの要素を乗算すると62ビットに収まる積が生成されるという点です。標準的なアプローチでは、この中間結果を保持するために64ビット演算が必要になります。しかし、モンゴメリ削減アルゴリズムは、32x32→64乗算の上位32ビットにアクセスできる場合、32ビット乗算のみを使用して実装できます。

ここでMetalのmulhi(a, b)組み込み関数 (intrinsic) が重要になります。これは、64ビット積a * bの上位32ビットを単一命令で返します。mulhiがない場合、複数の32ビット操作(4つ以上の命令)を使用して64ビット演算をエミュレートする必要があります。これは、ネイティブのmulhiサポートがないGPUで発生することです。

inline uint bb_mont_mul(uint a, uint b) {
    uint lo = a * b;               // low 32 bits of a*b
    uint q  = lo * BB_MONT_NINV;   // Montgomery quotient: lo * (-p^{-1}) mod 2^32
    uint hi = mulhi(a, b);         // high 32 bits of a*b (the critical intrinsic)
    uint qn_hi = mulhi(q, BB_P);   // high 32 bits of q*p
    uint t = hi - qn_hi;           // result in [0, 2p)
    return (t >= BB_P) ? t - BB_P : t;
}

モンゴメリ乗算全体は、2つのmul + 2つのmulhi + 1つの減算 + 1つの条件付き減算、つまり6つの32ビット整数演算です。Apple GPUコアはネイティブの32ビット整数ALUを備えているため、これらのそれぞれが単一のハードウェア命令にマッピングされます。

これが性能に重要な理由: BabyBearは31ビット素数であるため、すべての有限体要素が単一の32ビットGPUレジスタに収まります。対照的に、BN254(254ビット)のような楕円曲線有限体は、要素ごとに32ビット整数の8リム (limb) を必要とし、単一の有限体乗算には、キャリー伝播 (carry propagation) を伴う約64回の乗算加算演算が必要です。これが、小規模有限体証明システム(BabyBear、Mersenne-31)が、大規模有限体システム(BN254、BLS12-381)よりもGPUで劇的に高速である根本的な理由です。

Moproの有限体演算ベンチマークはこれを定量化しています。Apple M3チップでは、BabyBear/M31の有限体乗算はMetal経由で約112ギガ演算/秒 (GOP/s) を達成しますが、BN254の有限体乗算は1ギガ演算/秒未満であり、100倍以上の差があります。このギャップは、BN254が有限体乗算あたり約100倍多くの32ビットALU演算を必要とし、GPUコアが基本的に32ビットマシンであるために存在します。


4. ベンチマーク結果

セットアップ

  • ハードウェア: Apple M1(8 GPUコア、16GBユニファイドメモリ、68.25 GB/s帯域幅)
  • ソフトウェア: Rust nightly 1.97.0、macOS 15.5、LTO (thin) + codegen-units=1 + target-cpu=native を使用したリリースビルド
  • 方法論: 各構成で3回実行し、中央値を報告
  • ベースライン: Plonky3のRadix2DFTとNEON SIMD、Rayon並列処理 (Rayon parallelism)、および同じLTO/ネイティブ設定を使用した高度に最適化されたCPUパス

パラメータ

  • n = 変数の数(多項式は2^n個の係数を持つ)
  • fold = STIRラウンドごとの折りたたみ係数
  • rate = 開始対数逆レート (log inverse rate)(RS符号レート (RS code rate) = 1/2^rate)

結果

29の構成をテストしました。「Best GPU」は、GPU、FUSED、GRINDモードの最小値です。

n=20(100万係数)

foldrateCPU (ms)Best GPU (ms)Speedup
112671711.56x
124532901.56x
138974831.86x
21127751.70x
222171221.78x
234142101.98x
4149351.41x
4289531.70x
432001201.67x

n=22(400万係数)

foldrateCPU (ms)Best GPU (ms)Speedup
1111745792.03x
12193810921.77x
13345919341.79x
214412871.54x
228054841.66x
2316768971.87x
312061381.49x
323812481.54x
338976111.47x
411661281.30x
423222151.50x
436615301.25x
61141981.44x
624103271.25x
63176313201.34x

n=22、rate > 3(拡張)

rate > 3の場合、評価ドメインと中間バッファは急速に増加します。Metalディスパッチパスは保守的なデフォルトを適用します(gpu_dft.rsを参照):おおよそ**log_n <= 24および1 GiBのGPU適格ストレージ(WHIR_GPU_MAX_LOG_N / WHIR_GPU_MAX_TOTAL_BYTES)。これらを上げない限り、高レートのn=22は通常CPUに留まる**か、以下のn=24の注記と同じクラスの上限に達します。

以下の行は、メインテーブルと同じハードウェアと方法論を使用しています(モードごとに3回実行し、中央値を報告。「Best GPU」はgpugpu_fusedgpu_grindの中央値の最小値)。実行前に以下を設定してください。

export WHIR_GPU_MAX_LOG_N=28
export WHIR_GPU_MAX_TOTAL_BYTES=4294967296   # 4 GiB
./bench.sh 22 1 5   # example: single (n, fold, rate)
foldrateCPU (ms)Best GPU (ms)Speedup
151522693001.64x
25632736441.74x
35869437792.30x
45754645681.65x
261714284082.04x

一部の(n, fold, rate)ペアが失敗または省略される理由。 これは上記の29構成の閉じたグリッドと同じではありません。rate > 3でのGPUは実験的です。

  • ドメインとメモリの上限 (Domain and memory caps): バッファがデフォルトのlog_nまたはバイト予算を超えると、コードは(設計上)GPUを拒否し、CPUを使用します。上記の環境変数による上書きはベンチマークのためだけに緩和されます。十分なユニファイドメモリの余裕が必要です。
  • (n=22, fold=1, rate=6): 我々の実行では、すべてのGPUモードが使用可能なタイミングなしで終了しました(GPU証明の成功タイミングなし)。一方、CPUは完了しました。これは、非常に大きなアロケーションとMetalパイプラインの相互作用、またはエッジケースのレイアウトにおけるバグの可能性があります。ここでは根本原因の調査は行っていません。
  • (n=22, fold=6, rate=5): CPUはプローブでO(10²)秒で完了しましたが、GPU実行は実用的な実行時間内に終了しませんでした(ハングまたはGPUパスでの極端なPoW分散 (PoW variance))。高いfoldはラウンド数とFiat–Shamirグラインディングを増加させます。大きなドメインと組み合わせると、高速化されたパスはfold 1〜4の場合よりもはるかに予測不能になります。
  • ドライバーと安定性: 積極的な設定では、GPUプロセスが一部のマシンで異常終了する(例:シグナル終了)ことがあります。これもデフォルトが保守的である理由の一つです。

これらのパラメータでは、PoWが支配的な場合、GRINDが最良のGPU中央値を提供することがよくあります。FUSEDは一部の行(例えば、ここではfold=1、rate=5)で依然として優位です。

n=24(1600万係数)

foldrateCPU (ms)Best GPU (ms)Speedup
11415324631.69x
21181410491.73x
318905311.68x
419786941.41x
615884051.45x

n=24 rate>=2はGPUドメイン上限 (domain cap)(2^25要素)を超過するため、テストされていません。

補足: Apple M3 MacBook

以下の表は、新しいApple Siliconを使用している読者向けに、上記のパラメータと同じ意味を別のマシンで繰り返したものです。Best GPUは、M1の表と同様に、各行の3つのGPUモード(GPUFUSEDGRIND)の最小値です。

  • ハードウェア: Apple M3(MacBook、ユニファイドメモリ)
  • ソフトウェア: Rust 1.95.0 (2026-04-14)、macOS 15.7.5 (arm64)、M1実行時と同じクラスの最適化(LTO、ベースライン用のネイティブCPUコード生成)を施したリリースビルド
  • 方法論: 各構成で3回実行し、各列の中央値を報告。高速化 = CPU中央値 / Best GPU中央値
  • 日付: 2026-04-27

n=20(100万係数)

foldrateCPU (ms)Best GPU (ms)Speedup
11202.4109.81.84x
12349.3174.72.00x
13616.3290.32.12x
2183.451.91.61x
22138.878.81.76x
23262.2136.91.92x
4130.828.11.10x
4255.239.21.41x
43113.585.51.33x

n=22(400万係数)

foldrateCPU (ms)Best GPU (ms)Speedup
11744.0394.31.89x
121557.9617.92.52x
133006.71164.22.58x
21704.9173.44.07x
22618.0297.32.08x
231210.2552.32.19x
31172.797.31.77x
32316.5164.41.93x
33617.3385.61.60x
41121.379.91.52x
42226.1141.01.60x
43510.6382.91.33x
61108.872.41.50x
62335.3242.51.38x
631573.31064.11.48x

(n=22, fold=2, rate=1) のM3シートでは、Best GPUの高速化が4.07倍と示されています。この行は、グリッドの他の部分やM1のn=22 fold=2 rate=1セル(1.54倍)と比較して外れ値です。強い結論を出す前に、再実行する価値があると考えてください(PoW分散、熱状態 (thermal state)、または特異なCPU中央値のいずれか)。

n=24(1600万係数)

foldrateCPU (ms)Best GPU (ms)Speedup
113390.61432.32.37x
211538.3688.72.23x
31708.5370.81.91x
41809.7549.31.47x
61438.5333.71.31x

補足: iPhone(WHIR Bench、Apple A19 GPU)

iOSアプリは、タップごとに単一のGPUモードの実行時間(ミリ秒単位)を報告します(Macのbench.shグリッドのFUSED/GRINDバリアントではありません)。以下の表は、疎な構成セットです(M1の29セルグリッド全体ではありません)。高速化は、各行のCPU時間とGPU時間の比率です。

  • アプリ: WHIR Bench(プロジェクトios/ターゲット)
  • GPU: MetalはApple A19 GPUを報告
  • 単位: ミリ秒(オンデバイスWHIR Bench実行)
nfoldrateCPU (ms)GPU (ms)Speedup
20112401331.8x
20221961051.9x
20432071211.7x
221110684862.2x
221219428862.2x
22214822382.0x
22324372461.8x
22437395171.4x
2411461820192.3x
242120179762.1x
243110555402.0x
244111978451.4x

主な観察

上記の主要なM1グリッドでは、n >= 20のテストされた29の構成すべてでGPUがCPUよりも高速です。高速化は1.25倍から2.03倍の範囲で、低いfold値と高いrateで最良の結果が得られます。セクション4のM3補足グリッドでは、リストされたすべての構成でBest GPUがCPUを上回り、ほとんどの高速化は1.10倍から2.58倍の間であり、さらに1つの外れ値**(n=22, fold=2, rate=1)4.07倍です(そこの注記を参照)。セクション4のiPhone(A19)**WHIR Benchサンプルも、リストされたすべての行でGPUがCPUよりも高速であることを示しており(約1.4〜2.3倍)、fold=4の行が最も差が小さいという同じ定性的なパターンが見られます。n=22, rate > 3はより小さい、個別に文書化されたセットです(表を参照)。GPUは、中央値を測定できた構成では依然として優位ですが、他の高レートのペアは、前述のように上限に達したり、タイムアウトしたり、クラッシュしたりします。

低いfold値で最高の高速化が得られます(fold=1-2で1.5-2.0倍)。これは、NTTとMerkleツリーが実行時間を支配し、これらが我々が高速化した操作そのものであるためです。高いfold値(fold=4-6)は、ラウンドあたりのNTT要素数を減らしますが、ラウンド数とPoWグラインディングを増やし、サムチェック (Sumcheck) への作業シフトを引き起こします(GPU並列処理の活用がより困難になります。セクション7の議論を参照)。

CPUベースラインは非常に強力です。 Plonky3のBabyBear実装は、4ワイドパック演算を伴うNEON SIMD組み込み関数 (intrinsics) を使用しています。LTOとtarget-cpu=nativeと組み合わせることで、CPUパスは我々の最適化作業中に約25%改善されました(Rustコンパイラの更新とビルド設定による)。これにより、絶対性能は向上したものの、GPUの高速化を達成するのがより困難になりました。


5. 時間がどこに使われているか

プロファイリング breakdown for a representative configuration (n=24, fold=3, rate=1, total GPU time ~537ms):

コンポーネント時間%注記
GPU Poseidon2 Merkle約315 ms58%コンピュートバウンド(モンゴメリ乗算)
GPU DFT (NTT)約75 ms14%メモリ帯域幅制限
CPU サムチェック (Sumcheck)約80 ms15%外部クレート、SIMD最適化済み
CPU 制約結合約45 ms8%部分的にGPUオフロード済み
GPU 読み戻し + ディスパッチ約24 ms4%ゼロコピーは既に使用中

Poseidon2 Merkleが支配的です。 各Poseidon2順列(幅16)は、S-box + MDS行列演算の25ラウンドを必要とし、各S-boxはx^7 = x * x * x * x * x * x * x(6回のモンゴメリ乗算)です。数百万のリーフをハッシュし、完全な二分木を圧縮する必要があるため、これはGPU実行時間の約58%を占めます。

GPUはどの程度飽和しているか? XcodeのMetal GPUプロファイラは、Poseidon2カーネルで高いALU使用率(85%以上の占有率)と顕著なメモリスタールがないことを示しています。このカーネルはコンピュートバウンドであり、メモリバウンドではありません。しかし、AppleはGPUコアの正確な整数ALUスループット仕様を公開していないため、理論上のピークに対する厳密な割合を計算することはできません。言えることは、プロファイラは明らかな最適化の機会(無駄なサイクル、メモリボトルネック、完全な占有率)を示しておらず、カーネルをさらに最適化しようとする試み(ループアンローリング、レジスタプレッシャ削減)は測定可能な改善をもたらさなかったということです。これは、このワークロードではハードウェアの限界に近いことを強く示唆していますが、公開されている仕様を持つNVIDIAハードウェアのように、正確なFLOP/s計算でそれを証明することはできません。

NTTはメモリ帯域幅制限(コンピュート制限ではない)です。これは、バタフライ演算が単純(ペアあたり1回の乗算+1回の加算)であるにもかかわらず、大きなストライドで非シーケンシャルなメモリアドレスにアクセスするためです。我々のradix-16/32カーネルは、グローバルメモリパスの数を減らします(セクション7を参照)。

CPUサムチェック (Sumcheck) が残りのボトルネックです。これはPlonky3のp3-whirクレートにあり、SIMD最適化された多項式演算を使用しています。サムチェックプロトコルはラウンド間でシーケンシャルです(各ラウンドは、次のラウンドが始まる前に検証者 (Verifier) のランダムなチャレンジを必要とします)。各ラウンド内では、計算は並列化可能です。我々はラウンド内の計算をGPUアクセラレートしようとはしませんでした。これが役立つかどうかは、データ転送オーバーヘッド(多項式をGPUに移動して戻す各ラウンド)が計算速度の向上を上回るかどうかにかかっています。我々の問題サイズでは、ラウンド内の計算は5〜15ミリ秒かかり、GPUディスパッチ+読み戻しオーバーヘッドは1〜2ミリ秒なので、わずかな改善の余地があるかもしれません。これは今後の課題です。


6. ユニファイドメモリ: 全体像

GPU証明プロジェクトで最も一般的に報告されるボトルネックは、CPU-GPUデータ転送です。ディスクリートGPUでは、PCIe経由で64MBの多項式を転送するのに約4ms(16 GB/s)かかります。これはNTT計算自体に匹敵します。Apple Siliconでは、CPUとGPUが同じ物理DRAMにアクセスするため、このPCIeコストは完全に排除されます。

しかし、ユニファイドメモリ (Unified Memory) にもトレードオフがないわけではありません。Metalは2つのバッファストレージモードを提供しており、その選択は大きく影響します。

共有モード (MTLResourceStorageModeShared)

CPUとGPUの両方がバッファを読み書きできます。ハードウェアはキャッシュコヒーレンス (Cache Coherence) を維持する必要があります。GPUがCPUが最近書き込んだデータを読み取る場合、GPUが最新の値を見るためにCPUのキャッシュをフラッシュする必要があり、その逆も同様です。このコヒーレンスにはパフォーマンスコストがかかります。

  • 共有バッファからのGPU読み取りは、GPU管理バッファからの読み取りよりも遅い
  • メモリコントローラはCPUとGPUのキャッシュ間の仲介をする必要がある
  • 大きなバッファの場合、キャッシュフラッシュに約0.1〜0.5msが追加される

管理モード (MTLResourceStorageModePrivate)

GPUのみがバッファにアクセスできます。キャッシュコヒーレンスは不要なため、GPUは競合なしでフルメモリ帯域幅を得られます。これはGPU集約的な計算でより高速です。

我々のハイブリッドアプローチ

プロファイリングを通じて、GPU管理バッファが計算負荷の高いカーネル(NTT、Merkle)で著しく高速であることを発見しました。我々のパイプラインは以下を使用します。

  1. 入力: CPUデータは共有バッファに存在します(ゼロコピー – CPUは多項式係数を書き込むだけで、GPUはそれを直接読み取ります)。
  2. 中間計算: NTTは、共有入力バッファからGPU管理バッファにデータをコピーし、すべてのバタフライステージを実行します。このコピーは、フルメモリ帯域幅で実行される単一のGPU側blitです。
  3. ビット反転出力: 最終的な順列は結果を共有バッファに書き戻すため、CPUはコピーなしでMerkle結果を読み取ることができます。

このハイブリッドアプローチは、両方の利点を提供します。CPUにはゼロコピー入出力、計算負荷の高い中間ステージにはフルGPU帯域幅です。GPU管理メモリへの初期コピーは64MBバッファで約0.3msかかりますが、より高速なNTT実行で約2〜3ms節約できます。

さらに、GPU実行中にCPUとGPUが同じ共有バッファ領域に同時に安全にアクセスすることはできません。コマンドバッファの実行中にはハードウェアコヒーレンスが存在しないためです。CPUは、結果を読み取る前にGPUが完了を通知するのを待つ必要があります。これが、我々の融合パイプラインがすべてを単一のコマンドバッファで送信し、最後に一度待機する理由です。CPUの読み取りと進行中のGPU作業をオーバーラップさせようとしましたが、それは遅くなりました(セクション7を参照)。


7. 学んだ教訓

NTTアルゴリズムの選択: DIT、DIF、Stockham

数論変換 (NTT) は、FFTの有限体アナログです。いくつかのアルゴリズムバリアントがあり、それぞれ異なるメモリアクセスパターンを持ちます。これはGPU性能にとって重要な考慮事項です。

時間間引き (DIT) (Cooley-Tukey, 1965): 古典的なFFTアルゴリズム。入力はビット反転順序、出力は自然順序です。各バタフライステージは2つの要素を読み取り、1つをトゥイドルファクター (twiddle factor) で乗算し、2つの結果を書き込みます。メモリアクセスストライドは小さく(隣接要素)始まり、最終ステージではN/2まで増加します。後のステージでの大きなストライドは、GPUでのキャッシュ利用率を低下させます。

周波数間引き (DIF): DITの「逆」。入力は自然順序、出力はビット反転順序です。メモリストライドは大きく(N/2)始まり、1まで縮小します。GPUワークロードの場合、DIFはしばしば好まれます。なぜなら、(グローバルメモリを介して最も多くのデータを処理する)最初のステージが最大のストライドを持ち、複数のバタフライが並行してディスパッチされるときに、radix-16バタフライパターンを持つ大きなストライドがGPUのメモリコヒーシング (memory coalescing) にうまくマッピングされるためです。

ストックハム自動ソート (Stockham auto-sort) (Stockham, 1966): 個別のビット反転順列を回避するために2つのバッファを交互に使用するアウトオブプレースバリアント。各ステージは一方のバッファから読み取り、もう一方のバッファに異なる順序で書き込むため、最終ステージの出力はすでに正しい順序になっています。欠点はメモリ使用量が2倍になることです。

我々はいくつかの理由でDIFを主要なアルゴリズムとしました。

  1. DIFの出力はビット反転順序ですが、Merkleツリー(リーフは評価ドメイン順序である必要があります)のためにいずれにせよビット反転順序が必要です。ビット反転を最後のNTTステージと融合させ、個別の順列パスを排除します。
  2. DIFはradix-16/32分解と自然にペアになります(下記参照)。
  3. DIFのストライドパターンはApple GPUのメモリサブシステムとうまく機能します。

比較のためにStockhamも実装しました。追加のバッファがGPUキャッシュに収まる小さなNTT(< 2^16)では競争力がありましたが、我々のワークロードサイズ(2^20から2^25)では、メモリプレッシャが低いためDIFが一貫して高速でした。

Radix選択: なぜradix-16(およびradix-32)なのか

radix-Rバタフライは一度にR個の要素を処理し、ステージあたりR log_R(N)/log_2(N)回の演算を実行しますが、log_2(N)回のグローバルメモリパスではなくlog_R(N)回しか必要としません。トレードオフは以下の通りです。

Radix-2:  各バタフライ: 1 mul + 1 add.  2^24の場合のパス数: 24
Radix-4:  各バタフライ: 3 mul + 8 add.  2^24の場合のパス数: 12
Radix-8:  各バタフライ: 約17 mul + add. 2^24の場合のパス数:  8
Radix-16: 各バタフライ: 約43 mul + add. 2^24の場合のパス数:  6
Radix-32: 各バタフライ: 約100 mul + add. 2^24の場合のパス数:  約5

我々のNTTはメモリ帯域幅制限であるため(各パスが配列全体を読み書きするため)、パス数を半分にすると実行時間もおおよそ半分になります。コストはパスあたりのALU作業が増えることですが、帯域幅制限のあるNTTパスではGPU ALUは十分に活用されていないため、このトレードは有利です。

我々はradix-16を主力として使用しています。

  • バタフライあたり16要素 = 16レジスタ。Apple GPUのレジスタファイルに快適に収まる
  • 各スレッドは、内部的に4つの「レイヤー」のradix-2バタフライ(16 = 2^4のため)で16要素を処理
  • Radix-32(内部5レイヤー、32レジスタ)も機能し、2^24の場合に1パス節約できますが、小さなNTTでは収穫逓減が見られます。

視覚的には、単一のradix-16バタフライは、16要素のブロックを内部的に4つのステージで処理します。

Stage 0 (stride 8):  [0,8] [1,9] [2,10] [3,11] [4,12] [5,13] [6,14] [7,15]
Stage 1 (stride 4):  [0,4] [1,5] [2,6] [3,7]   [8,12] [9,13] [10,14] [11,15]
Stage 2 (stride 2):  [0,2] [1,3] [4,6] [5,7]   [8,10] [9,11] [12,14] [13,15]
Stage 3 (stride 1):  [0,1] [2,3] [4,5] [6,7]   [8,9] [10,11] [12,13] [14,15]

各ペア[i,j]はradix-2バタフライです。a' = a + tw*bb' = a - tw*b。これら4つのステージはすべてレジスタ内で発生し(グローバルメモリへのアクセスなし)、その後16個の結果が書き戻されます。これは、各グローバルメモリパスが4つのバタフライステージ分の作業を行うことを意味します。

4ステップFFT (four-step FFT): 試行と断念

4ステップFFTアルゴリズム(Bailey, 1990)は、大きな1D NTTを小さな2Dサブ変換に分解します。

  1. 入力をN1 x N2行列(行優先)として解釈
  2. サイズN1のN2個の独立したNTTを実行(行変換)
  3. トゥイドルファクターで乗算
  4. 行列を転置
  5. サイズN2のN1個の独立したNTTを実行(列変換)

魅力的な点は、行/列NTTがスレッドグループ共有メモリ (threadgroup shared memory)(Apple GPUでは32KB)に完全に収まるほど小さく、バタフライステージ中のグローバルメモリアクセスを回避できることです。転置のみがグローバルメモリパスを必要とします。

実際には、これは我々のフラットなDIFアプローチよりも20%遅かったです。その理由は以下の通りです。

  • 行列転置は無料ではありません。コヒーシングパターンが悪いグローバルメモリパスが必要です。
  • 行/列NTT中の共有メモリバンク競合 (bank conflicts)(行あたり16要素 = ストライド16アクセス = Apple GPUの32バンク共有メモリで最悪のバンク競合)
  • 追加のトゥイドル乗算パスが別のグローバルメモリ往復を追加
  • Apple GPUの共有メモリはスレッドグループあたり32KBしかなく、サブNTTサイズを2^13 = 8192要素に制限

うまくいかなかったこと

  1. DFT読み戻しとMerkle GPU作業のオーバーラップ。 Metalイベントを使用して、GPUが別の管理バッファでMerkleハッシュを開始している間に、CPUが共有バッファからNTT結果を読み取れるように試みました。これは、CPUの読み取りがユニファイドメモリ帯域幅をGPUと競合したため、遅くなりました。Apple Siliconでは、メモリバスが共有されているため、GPUの重い計算中のCPU読み取りは、実質的に帯域幅サイクルを奪います。
  2. GPUサムチェック (Sumcheck)。 サムチェックプロトコル(Lund, Fortnow, Karloff, Nisan, 1992)はシーケンシャルなラウンド構造を持っています。各ラウンドで、証明者 (Prover) はハイパーキューブ (hypercube) 上で単変量多項式を計算し、検証者 (Verifier) はランダムなチャレンジを送信し、証明者は次のラウンドのためにハイパーキューブを「折りたたみ」ます。このラウンド間の依存性は本質的に直列です。 我々はGPUサムチェックを実装もベンチマークもしていません。シーケンシャルな構造は、原則としてGPU並列処理には不向きですが、我々の特定のワークロードサイズでこれが経験的に検証されたわけではありません。ラウンド内の計算(ラウンドiで2^(n-i)要素を合計する)は完全に並列化可能 (Embarrassingly Parallel) であり、非常に大規模なインスタンスの初期ラウンドではGPUアクセラレーションの恩恵を受ける可能性があります。これは未解決の疑問です。Thaler 2022の調査は、サムチェックプロトコルとその最適化について徹底的な考察を提供しています。

実用的な課題

  • PoWによるベンチマークのばらつき。 プルーフ・オブ・ワーク (PoW) グラインディングは、指数分布の完了時間(ランダムなノンスの探索)を持ちます。PoWが重い構成では、単一実行のベンチマークは50%以上変動する可能性があります。3回の実行の中央値を使用することで、このノイズを大幅に低減できます。
  • コンパイラ最適化はGPUよりもCPUに効果的。 LTOとtarget-cpu=nativeはCPU性能を約25%向上させましたが、GPUカーネル性能にはほとんど影響がありませんでした(MetalシェーダーはRustコードとは別にコンパイルされるため)。これにより、絶対性能は向上したものの、GPU/CPU比率は縮小しました。
  • モバイルでのサーマルスロットリング (Thermal throttling)。 ラップトップ(特にスマートフォン)での長時間のベンチマーク実行は、サーマルスロットリングを引き起こし、GPUクロック速度を低下させる可能性があります。我々のベンチマークでは、これを軽減するために3回の実行の中央値を使用していますが、実際の持続性能は低くなる可能性があります。

8. 関連研究との比較

プロジェクトターゲットプロトコル高速化有限体APIソース
本研究Apple M1 GPUWHIR/Poseidon2CPU比1.3-2.0倍BabyBear (31ビット)Metalリポジトリ
本研究 (M3実行)Apple M3 GPUWHIR/Poseidon2CPU比1.1-2.6倍(§4参照; 1セル約4.1倍)BabyBear (31ビット)Metal同上
本研究 (iPhone)Apple A19 GPUWHIR/Poseidon2CPU比約1.4-2.3倍(疎なWHIR Benchサンプル、§4)BabyBear (31ビット)Metal同上
Mopro Metal MSM v2Apple GPUMSM (BN254)v1比40-100倍BN254 (254ビット)Metal解説、コード
ICICLE Metal v3.6Apple GPUMSM, NTT最大5倍複数Metalブログ、ドキュメント
ICICLE-Stwo (CUDA)データセンターGPUCircle STARKCPU SIMD比3.25-7倍M31 (31ビット)CUDAブログ
zkSecurity Stwo WebGPUブラウザGPUCircle STARK全体で2倍M31WebGPU解説、PR
Ligetronブラウザ/ネイティブSHA-256, NTTN/A (WIP)複数WebGPUコード

我々の高速化の数値(1.3-2.0倍)は、MSMに焦点を当てたプロジェクトよりも控えめです。その理由は以下の通りです。

  1. 我々はエンドツーエンドの証明をベンチマークしています。これにはCPUバウンドのサムチェック (Sumcheck) ラウンドが含まれます。GPUのみの部分(NTT + Merkle)は、非SIMD CPUコードと比較して3-4倍の高速化を示します。
  2. 我々のCPUベースラインは非常に強力です。Plonky3のBabyBearは、NEON SIMD 4ワイドパック演算に加え、LTOとネイティブCPUコード生成を使用しています。
  3. WHIRのワークロードはハッシュが支配的です(Poseidon2 Merkleが実行時間の約58%を占める)。これは高い算術強度 (arithmetic intensity) を持ちますが、MSMと比較して並列処理の削減が限定的です。

Mopro MSM v2の高速化(v1比40-100倍)は、最適化されたCPUベースラインではなく、彼ら自身のv1と比較していることに注意してください。Arkworks CPU MSMと比較すると、高速化はより控えめです。同様に、ICICLE-Stwoの3.25-7倍は、StwoのCPU SIMDバックエンドと比較しています。


9. 今後の方向性

高アリティ (arity) Merkleツリー

Poseidon2 Merkleカーネルが実行時間(58%)を支配しています。4アリティまたは8アリティのツリーは、ハッシュ呼び出しの数を2〜3倍削減し、GPUパイプラインを比例的に高速化するでしょう。これにはWHIRコミットメントスキームのプロトコルレベルの変更が必要です。

新しいApple Silicon

我々の主要な公開数値はM1(2020年)のものです。M3 MacBookでの実行(セクション4)は、同じ定性的な結果を示し、大規模なnではしばしばより高い高速化を示しています(例えば、n=24, fold=1, rate=1で2.37倍に対し、M1の同じセルでは1.69倍)。iPhone(A19、セクション4)のWHIR Benchは、サンプリングされたセルで同様の範囲(n=24, fold=1, rate=1で2.3倍)に収まります。M4 Maxは40 GPUコア(M1の8コアに対し)と273 GB/sのメモリ帯域幅(68 GB/sに対し)を持っています。新しいチップからのさらなる性能向上を期待しており、リポジトリにはbench.shスクリプトとiOSアプリが含まれており、クロスデバイスベンチマークを容易にしています。貢献を歓迎します。

WebGPUバックエンド

クロスプラットフォーム対応のために、WebGPU(wgpu)バックエンドは、WGSLで同じカーネルアルゴリズムを再利用できます。主なコストは、共有/管理メモリの区別を失うこと(WebGPUはより単純なメモリモデルを持つため)と、WGSLにmulhi組み込み関数がないこと(エミュレーションが必要)です。Moproの有限体演算ベンチマークは、WebGPUがMetalの有限体演算性能の約50〜80%を達成することを示唆しています。

GPUサムチェック (Sumcheck)

サムチェックプロトコルは、残りの主要なCPUボトルネック(実行時間の約15%)です。ラウンド間の依存性は本質的に直列ですが、ラウンド内の計算(大きなハイパーキューブ (hypercube) 上での合計)は、非常に大規模なインスタンスの場合、GPU並列処理の恩恵を受ける可能性があります。我々はこれをベンチマークしておらず、未解決の疑問として残っています。

1投稿 - 1参加者

トピック全文を読む