在 Apple Silicon 上利用 GPU 加速 WHIR 證明生成
我們使用 Metal 計算著色器在 Apple Silicon GPU 上加速了 WHIR 證明器，透過利用統一記憶體架構和融合內核流水線，實現了比高度優化的 CPU 代碼快達 2.58 倍的效能提升。
致謝
感謝 Moven Tsai 提供第 4 節中的 Apple M3 MacBook 與 iPhone 基準測試結果，以及 Alex Kuzmin 對 WHIR 的討論。
內容提要 (TL;DR)
我們利用 Metal 計算著色器（compute shaders）在 Apple Silicon GPU 上加速了 WHIR 證明器。在 M1 晶片上，相較於高度優化的 CPU 代碼（SIMD + LTO + target-cpu=native），實現了 高達 2.03 倍的加速；在補充的 Apple M3 MacBook 測試中，加速達 2.58 倍；而在 WHIR Bench iOS 應用程式的稀疏樣本測試中（Metal Apple A19 GPU；見第 4 節），加速約為 1.4-2.3 倍。GPU 流水線將 NTT（數論變換）、位元反轉（bit-reversal）、Poseidon2 Merkle 樹雜湊以及工作量證明（PoW）研磨（grinding）整合到單次命令緩衝區提交中，充分利用了 Apple Silicon 的統一記憶體架構。該實作已開源，可在任何搭載 Apple Silicon 的 Mac 上運行，並提供同款 App 用於 iPhone 裝置端測試。
核心發現：
- 在多項式規模 >= 2^20（100 萬個以上係數）的所有測試配置中，GPU 均勝過 CPU。
- 融合的 DFT+Merkle 流水線避免了 CPU 與 GPU 之間的往返通訊，提供了最大的效能提升。
- Apple Silicon 的統一記憶體消除了 PCIe 傳輸成本，但引入了更微妙的權衡：由於快取一致性（cache coherence）開銷，共享模式緩衝區（CPU 與 GPU 均可存取）在 GPU 計算上的速度慢於 GPU 管理的緩衝區。我們的流水線採用了混合方法。
- Poseidon2 Merkle 核心是主要的成本來源（約佔 GPU 時間的 58%）。Xcode GPU 分析器顯示 ALU 利用率很高且無明顯停頓，這表明我們已接近該工作負載的硬體極限，儘管 Apple 未發佈精確的整數吞吐量規格以供證實。
- 編譯器優化（LTO、target-cpu=native）將 CPU 基準效能提升了約 25%，這使得 GPU 更難以超越，但也提升了絕對的端到端效能。
程式碼庫：github.com/privacy-ethereum/whir-p3-metal
我們的工作基於 tcoratger/whir-p3，這是一個使用 Plonky3 函式庫實作 WHIR 協議的 Rust 專案。我們在此基礎上增加了 Metal GPU 加速。（演進路徑：WizardOfMenlo/whir → whir-p3 → whir-p3-metal。）
聲明：Metal GPU 實作（Rust 和 MSL）、支援工具以及本文內容是在 AI 程式碼助手 Claude (Anthropic) 的實質協助下，在人工指導、審查和基準測試下完成的。密碼學安全性遵循已發佈的 WHIR / Plonky3 規範及上游代碼；任何整合錯誤、效能主張或解釋上的失誤均由專案維護者負責。
1. 動機：為何客戶端證明至關重要
以太坊的透明度是以隱私為代價的：每筆交易都永久可見，且鏈上分析工具可以將匿名地址與真實身份關聯。零知識證明可以恢復隱私，但將證明生成委託給伺服器會違背初衷——伺服器會看到你的私有輸入。
真正的隱私需要 客戶端證明：用戶在自己的裝置上生成證明。這對於以下場景至關重要：
- 私密支付 – 隱藏金額、交易對手和交易模式。
- 身份識別 – 在不洩露憑證本身的情況下證明憑證事實（年齡、國籍、成員身份）（如 ZK Email、Anon Aadhaar）。
- 投票 – 匿名參與 DAO 和治理（如 Semaphore）。
障礙在於效能。在消費級硬體上進行證明必須足夠快，才能滿足互動式使用的需求。伺服器端 GPU 證明器（資料中心 GPU 上的 CUDA）可實現巨大的加速，但客戶端裝置面臨不同的限制：散熱限制、共享記憶體頻寬、較少的 GPU 核心數，且沒有獨立的 VRAM。
客戶端 GPU 的機遇
現代手機和筆記型電腦配備了效能日益強大的 GPU。Apple Silicon 的 M 系列和 A 系列晶片在 CPU 和 GPU 之間共享統一記憶體，消除了主導資料中心 GPU 證明的 PCIe 傳輸瓶頸。這種架構差異為精細的 CPU-GPU 協作創造了獨特的機會。
多個專案正在探索這一領域：
- Mopro – Metal MSM 加速（v2 報告），WebGPU 欄位運算基準測試顯示在小欄位上相較於 BN254 有 100 倍以上的吞吐量。
- ICICLE Metal – 適用於 Apple Metal 的 MSM 和 NTT 原語，加速高達 5 倍（v3.6）。
- zkSecurity / Stwo WebGPU – 透過 WebGPU 計算著色器在瀏覽器中將 Circle STARK 的整體證明速度提升了 2 倍。
- Ligetron – 用於跨平台證明的 WebGPU SHA-256 和 NTT。
- FibRace – 大規模行動端基準測試（6,000 多名參與者，210 萬個證明）顯示大多數現代智慧型手機可在 5 秒內完成證明。
我們的工作專注於一個具體且具備實際意義的目標：使用 Apple 的 Metal API 進行原生 GPU 計算，加速 WHIR 多項式承諾方案，該方案基於雜湊且具備後量子安全性。
2. WHIR 背景
WHIR（Gal Arnon, Weizmann Institute; Alessandro Chiesa, EPFL; Giacomo Fenzi, EPFL; Eylon Yogev, Bar-Ilan University；發表於 EUROCRYPT 2025）是一種針對 Reed-Solomon 碼的鄰近性交互式預言機證明（Interactive Oracle Proof of proximity）。它是 FRI、STIR 和 BaseFold 的高效替代方案，其驗證速度顯著更快（數百微秒，而其他方案為毫秒級）。
對於證明過程，主要的成本在於：
- NTT（數論變換） – 在擴展域上進行多項式求值。
- Merkle 樹構建 – 用於多項式承諾的 Poseidon2 雜湊。
- 工作量證明（PoW）研磨 – 尋找滿足雜湊難度目標的隨機數（Fiat-Shamir 安全性）。
這三者都是高度並行的，非常適合 GPU 計算。證明器執行多輪 STIR，每輪涉及一次 NTT、一次 Merkle 承諾和一次 PoW 研磨。在各輪之間，CPU 執行跨輪次的順序 sumcheck 操作。
我們的實作基於 tcoratger/whir-p3，它使用 Plonky3 函式庫在 BabyBear 欄位（質數 p = 2^31 - 2^27 + 1，蒙哥馬利形式）上進行欄位算術和 WHIR 協議實作。
3. GPU 架構與實作
3.1 為何選擇 Metal（而非 WebGPU 或 CUDA）
- 統一記憶體：Apple Silicon 在 CPU 和 GPU 之間共享物理記憶體。這消除了 WebGPU 和 CUDA 證明專案中提到的 PCIe 傳輸瓶頸。然而，如第 6 節所述，「統一」並不意味著「免費」——存在重要的快取一致性權衡。
- 低調度開銷：Metal 命令緩衝區可以在 CPU 上構建，同時前一個任務在 GPU 上執行，從而實現緊密的流水線化。
- 原生效能：Metal 著色語言 (MSL) 在構建時編譯為 AIR (Apple Intermediate Representation)，並在加載時編譯為裝置特定的機器碼。沒有運行時著色器編譯開銷。
- iOS 相容性：同樣的 Metal 代碼可在 iPhone 和 iPad 上運行，支援行動端基準測試。
權衡之處在於平台鎖定在 Apple 生態。對於跨平台部署，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 個葉節點雜湊 + 首次 Merkle 壓縮
│
├── Poseidon2 Merkle 壓縮（所有剩餘的樹層級）
│ 用於小層級的 SIMD shuffle 核心（避免使用執行緒組記憶體）
│
└── [單次 GPU 等待] → 結果存於 CPU 可存取的記憶體中
與單純的實作相比，這種融合在每輪 STIR 中減少了 3-4 個 CPU-GPU 同步點。
3.3 核心優化（30 次迭代）
我們經歷了 30 次優化迭代。影響最顯著的包括：
| 優化項目 | 影響 |
|---|---|
| 1-4: 基礎 Metal NTT + Merkle 核心 | 建立基準 GPU 路徑 |
| 5-8: Radix-16 DIF，共享記憶體蝴蝶運算 | NTT 加速 2-3 倍 |
| 9-12: 融合 DFT+Merkle 流水線，零拷貝 I/O | 端到端加速 1.5-2 倍 |
| 13-16: Poseidon2 4 葉融合核心，SIMD Merkle | Merkle 加速 1.3 倍 |
| 17-20: GPU PoW 研磨，零拷貝 EF 轉換 | 助力 PoW 主導的配置 |
| 21-24: 擴展域零拷貝，逐輪融合 | 大規模 n 加速 1.2 倍 |
| 25-28: LTO, target-cpu=native, 分析引導優化 | 絕對效能提升 10% |
| 29-30: R32 DIF 核心，跳過打包轉置 | 每輪節省 5-15ms |
最重要的經驗是：融合操作以避免 CPU-GPU 往返通訊，比優化單個核心更重要。融合流水線（DFT+Merkle 單次命令緩衝區）通常比非融合路徑快 15-30%。
3.4 Metal 中的蒙哥馬利算術
BabyBear 欄位運算在整個 GPU 核心中均使用蒙哥馬利形式。蒙哥馬利乘法將昂貴的基於除法的模還原替換為較便宜的乘法和移位操作。關鍵在於，對於像 BabyBear 這樣的 31 位元質數 (p < 2^31)，兩個元素相乘產生的乘積可容納於 62 位元中。標準方法需要 64 位元算術來保存此中間結果。然而，如果你能獲取 32x32→64 乘法結果的高 32 位元，蒙哥馬利還原演算法僅需使用 32 位元乘法即可實作。
這正是 Metal 的 mulhi(a, b) 內置函數發揮關鍵作用的地方：它在單條指令中返回 64 位元乘積 a * b 的高 32 位元。如果沒有 mulhi，你需要使用多個 32 位元操作（4 條以上指令）來模擬 64 位元算術，這是在不支援原生 mulhi 的 GPU 上會發生的情況。
inline uint bb_mont_mul(uint a, uint b) {
uint lo = a * b; // a*b 的低 32 位元
uint q = lo * BB_MONT_NINV; // 蒙哥馬利商數: lo * (-p^{-1}) mod 2^32
uint hi = mulhi(a, b); // a*b 的高 32 位元 (關鍵內置函數)
uint qn_hi = mulhi(q, BB_P); // q*p 的高 32 位元
uint t = hi - qn_hi; // 結果在 [0, 2p) 區間
return (t >= BB_P) ? t - BB_P : t;
}
整個蒙哥馬利乘法僅使用 2 次 mul + 2 次 mulhi + 1 次減法 + 1 次條件減法——共六個 32 位元整數運算。Apple GPU 核心具有原生 32 位元整數 ALU，因此每個運算都對應單條硬體指令。
為何這對效能很重要：BabyBear 是一個 31 位元質數，因此所有欄位元素都能放入單個 32 位元 GPU 暫存器中。相比之下，像 BN254（254 位元）這樣的橢圓曲線欄位，每個元素需要 8 個 32 位元整數組成的肢體（limbs），單次欄位乘法需要約 64 次帶進位的乘加運算。這就是為什麼小欄位證明系統（BabyBear, Mersenne-31）在 GPU 上的速度遠快於大欄位系統（BN254, BLS12-381）的根本原因。
Mopro 的欄位運算基準測試量化了這一點：在 Apple M3 晶片上，BabyBear/M31 欄位乘法透過 Metal 可達到約 112 GOP/s（每秒十億次運算），而 BN254 欄位乘法僅能達到 <1 GOP/s——差距達 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, release 版本，開啟 LTO (thin) + codegen-units=1 + target-cpu=native
- 方法論：每個配置運行 3 次，報告 中位數
- 基準 (Baseline)：高度優化的 CPU 路徑，使用 Plonky3 的 Radix2DFT（含 NEON SIMD）、rayon 並行化，以及相同的 LTO/native 設置
參數說明
- n = 變數數量（多項式有 2^n 個係數）
- fold = 每輪 STIR 的折疊因子
- rate = 初始對數反向速率（RS 碼速率 = 1/2^rate）
測試結果
共測試了 29 種配置。「最佳 GPU」取 GPU、FUSED 和 GRIND 模式中的最小值。
n=20 (100 萬個係數)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 1 | 267 | 171 | 1.56x |
| 1 | 2 | 453 | 290 | 1.56x |
| 1 | 3 | 897 | 483 | 1.86x |
| 2 | 1 | 127 | 75 | 1.70x |
| 2 | 2 | 217 | 122 | 1.78x |
| 2 | 3 | 414 | 210 | 1.98x |
| 4 | 1 | 49 | 35 | 1.41x |
| 4 | 2 | 89 | 53 | 1.70x |
| 4 | 3 | 200 | 120 | 1.67x |
n=22 (400 萬個係數)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 1 | 1174 | 579 | 2.03x |
| 1 | 2 | 1938 | 1092 | 1.77x |
| 1 | 3 | 3459 | 1934 | 1.79x |
| 2 | 1 | 441 | 287 | 1.54x |
| 2 | 2 | 805 | 484 | 1.66x |
| 2 | 3 | 1676 | 897 | 1.87x |
| 3 | 1 | 206 | 138 | 1.49x |
| 3 | 2 | 381 | 248 | 1.54x |
| 3 | 3 | 897 | 611 | 1.47x |
| 4 | 1 | 166 | 128 | 1.30x |
| 4 | 2 | 322 | 215 | 1.50x |
| 4 | 3 | 661 | 530 | 1.25x |
| 6 | 1 | 141 | 98 | 1.44x |
| 6 | 2 | 410 | 327 | 1.25x |
| 6 | 3 | 1763 | 1320 | 1.34x |
n=22, rate > 3 (擴展測試)
當 rate > 3 時，求值域和中間緩衝區增長迅速。Metal 調度路徑應用了保守的預設值（見 gpu_dft.rs）：大約為 log_n <= 24 且 GPU 可用存儲為 1 GiB (WHIR_GPU_MAX_LOG_N / WHIR_GPU_MAX_TOTAL_BYTES)。若不調高這些值，n=22 且速率較高時通常會保留在 CPU 執行，或觸及與下方 n=24 註釋中相同的上限。
下表使用與主表相同的硬體和方法論（每種模式 3 次運行，報告中位數；最佳 GPU = gpu、gpu_fused 和 gpu_grind 中位數的最小值）。運行前請設置：
export WHIR_GPU_MAX_LOG_N=28
export WHIR_GPU_MAX_TOTAL_BYTES=4294967296 # 4 GiB
./bench.sh 22 1 5 # 範例：單個 (n, fold, rate)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 5 | 15226 | 9300 | 1.64x |
| 2 | 5 | 6327 | 3644 | 1.74x |
| 3 | 5 | 8694 | 3779 | 2.30x |
| 4 | 5 | 7546 | 4568 | 1.65x |
| 2 | 6 | 17142 | 8408 | 2.04x |
為何部分 (n, fold, rate) 組合失敗或被省略：這與上述封閉的 29 種配置網格不同；rate > 3 的 GPU 測試是實驗性的。
- 定義域與記憶體上限：如果緩衝區超過預設的 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 變異極大）。高 fold 值會增加輪次和 Fiat–Shamir 研磨次數；結合大定義域，加速路徑的預測性遠低於 fold 1–4。
- 驅動程式與穩定性：在激進的設置下，GPU 進程在某些機器上可能會異常終止（例如信號退出）——這也是預設值保持保守的另一個原因。
在這些參數下，當 PoW 佔主導地位時，GRIND 通常提供最佳的 GPU 中位數；但在某些行（例如此處的 fold=1, rate=5），FUSED 仍然勝出。
n=24 (1600 萬個係數)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 1 | 4153 | 2463 | 1.69x |
| 2 | 1 | 1814 | 1049 | 1.73x |
| 3 | 1 | 890 | 531 | 1.68x |
| 4 | 1 | 978 | 694 | 1.41x |
| 6 | 1 | 588 | 405 | 1.45x |
n=24 且 rate>=2 超過了 GPU 定義域上限（2^25 個元素），未進行測試。
補充：Apple M3 MacBook
下表在不同機器上重複相同的參數含義，供使用較新 Apple Silicon 的讀者參考。最佳 GPU 是每行三種 GPU 模式（GPU, FUSED, GRIND）中的最小值，與 M1 表格一致。
- 硬體：Apple M3 (MacBook, 統一記憶體)
- 軟體：Rust 1.95.0 (2026-04-14), macOS 15.7.5 (arm64), release 版本，採用與 M1 測試相同的優化等級（LTO, 基準測試採用原生 CPU 代碼生成）
- 方法論：每個配置運行 3 次，報告每列的中位數；加速比 = CPU 中位數 / 最佳 GPU 中位數
- 日期：2026-04-27
n=20 (100 萬個係數)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 1 | 202.4 | 109.8 | 1.84x |
| 1 | 2 | 349.3 | 174.7 | 2.00x |
| 1 | 3 | 616.3 | 290.3 | 2.12x |
| 2 | 1 | 83.4 | 51.9 | 1.61x |
| 2 | 2 | 138.8 | 78.8 | 1.76x |
| 2 | 3 | 262.2 | 136.9 | 1.92x |
| 4 | 1 | 30.8 | 28.1 | 1.10x |
| 4 | 2 | 55.2 | 39.2 | 1.41x |
| 4 | 3 | 113.5 | 85.5 | 1.33x |
n=22 (400 萬個係數)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 1 | 744.0 | 394.3 | 1.89x |
| 1 | 2 | 1557.9 | 617.9 | 2.52x |
| 1 | 3 | 3006.7 | 1164.2 | 2.58x |
| 2 | 1 | 704.9 | 173.4 | 4.07x |
| 2 | 2 | 618.0 | 297.3 | 2.08x |
| 2 | 3 | 1210.2 | 552.3 | 2.19x |
| 3 | 1 | 172.7 | 97.3 | 1.77x |
| 3 | 2 | 316.5 | 164.4 | 1.93x |
| 3 | 3 | 617.3 | 385.6 | 1.60x |
| 4 | 1 | 121.3 | 79.9 | 1.52x |
| 4 | 2 | 226.1 | 141.0 | 1.60x |
| 4 | 3 | 510.6 | 382.9 | 1.33x |
| 6 | 1 | 108.8 | 72.4 | 1.50x |
| 6 | 2 | 335.3 | 242.5 | 1.38x |
| 6 | 3 | 1573.3 | 1064.1 | 1.48x |
在 (n=22, fold=2, rate=1) 處，M3 數據顯示 4.07 倍 的最佳 GPU 加速；該行相對於網格其餘部分以及 M1 的對應單元格（1.54x）是一個離群值。在得出強烈結論前，應將其視為值得重新測試的對象（可能受 PoW 變異、熱狀態或不尋常的 CPU 中位數影響）。
n=24 (1600 萬個係數)
| fold | rate | CPU (ms) | 最佳 GPU (ms) | 加速比 |
|---|---|---|---|---|
| 1 | 1 | 3390.6 | 1432.3 | 2.37x |
| 2 | 1 | 1538.3 | 688.7 | 2.23x |
| 3 | 1 | 708.5 | 370.8 | 1.91x |
| 4 | 1 | 809.7 | 549.3 | 1.47x |
| 6 | 1 | 438.5 | 333.7 | 1.31x |
補充：iPhone (WHIR Bench, Apple A19 GPU)
iOS App 報告的是單次點擊下單個 GPU 模式的實際時間（而非 Mac bench.sh 的 FUSED/GRIND 變體網格）。下表是配置的稀疏集合。加速比為每行的 CPU 時間除以 GPU 時間。
- App：WHIR Bench (專案 ios/ 目標)
- GPU：Metal 報告為 Apple A19 GPU
- 單位：毫秒 (裝置端 WHIR Bench 運行結果)
| n | fold | rate | CPU (ms) | GPU (ms) | 加速比 |
|---|---|---|---|---|---|
| 20 | 1 | 1 | 240 | 133 | 1.8x |
| 20 | 2 | 2 | 196 | 105 | 1.9x |
| 20 | 4 | 3 | 207 | 121 | 1.7x |
| 22 | 1 | 1 | 1068 | 486 | 2.2x |
| 22 | 1 | 2 | 1942 | 886 | 2.2x |
| 22 | 2 | 1 | 482 | 238 | 2.0x |
| 22 | 3 | 2 | 437 | 246 | 1.8x |
| 22 | 4 | 3 | 739 | 517 | 1.4x |
| 24 | 1 | 1 | 4618 | 2019 | 2.3x |
| 24 | 2 | 1 | 2017 | 976 | 2.1x |
| 24 | 3 | 1 | 1055 | 540 | 2.0x |
| 24 | 4 | 1 | 1197 | 845 | 1.4x |
核心觀察
在主 M1 網格中，對於 n >= 20 的所有 29 種測試配置，GPU 均快於 CPU。加速比從 1.25 倍到 2.03 倍不等，在低 fold 值和高速率下表現最佳。在 M3 補充網格（第 4 節）中，所有列出的配置在最佳 GPU 模式下也都勝過 CPU，加速比多在 1.10x 到 2.58x 之間，另有一個 (n=22, fold=2, rate=1) 的離群值達 4.07x（見該處註釋）。第 4 節中的 iPhone (A19) WHIR Bench 樣本同樣顯示 GPU 在所有列出的行中均快於 CPU（~1.4-2.3x），且呈現相同的定性模式：fold=4 的行差距最小。n=22, rate > 3 是一個較小的單獨記錄集合（見表）：在我們能測得中位數的配置中 GPU 仍然勝出，但其他高倍率組合如前所述會出現上限失敗、超時或崩潰。
低 fold 值可獲得最佳加速比（在 fold=1-2 時為 1.5-2.0x），因為 NTT 和 Merkle 樹主導了運行時間——這正是我們加速的操作。較高的 fold 值減少了每輪 NTT 的元素數量，但增加了輪次和 PoW 研磨次數，將工作量轉移到 sumcheck（GPU 並行性較難發揮作用之處——見第 7 節討論）。
速率（rate）的增加會增加工作量，且通常會提高加速比，因為定義域擴展（2^rate 倍以上的點）產生了更多的 NTT 和 Merkle 工作，放大了 GPU 的優勢。
CPU 基準效能非常強勁。 Plonky3 的 BabyBear 實作使用了具備 4 路打包運算的 NEON SIMD 內置函數。結合 LTO 和 target-cpu=native，CPU 路徑在我們的優化過程中提升了約 25%（源於 Rust 編譯器更新和構建設置）。這使得 GPU 加速更難達成，但也使其更有意義——我們擊敗的是一個真正經過優化的基準線。
5. 時間去哪兒了
代表性配置的分析細目（n=24, fold=3, rate=1, GPU 總時間約 537ms）：
| 組件 | 時間 | 佔比 | 備註 |
|---|---|---|---|
| GPU Poseidon2 Merkle | ~315 ms | 58% | 受限於計算 (蒙哥馬利乘法) |
| GPU DFT (NTT) | ~75 ms | 14% | 受限於記憶體頻寬 |
| CPU sumcheck | ~80 ms | 15% | 外部 crate，經 SIMD 優化 |
| CPU 約束組合 | ~45 ms | 8% | 部分由 GPU 卸載 |
| GPU 讀回 + 調度 | ~24 ms | 4% | 已使用零拷貝 |
Poseidon2 Merkle 佔主導地位。 每個 Poseidon2 置換（寬度 16）需要 25 輪 S-box + MDS 矩陣運算，其中每個 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 crate 中，並使用 SIMD 優化的多項式算術。sumcheck 協議在各輪之間是順序執行的（每輪都需要驗證者的隨機挑戰才能開始下一輪）。在每一輪內部，計算是可以並行的。我們沒有嘗試對輪內計算進行 GPU 加速；這是否有效取決於數據傳輸開銷（每輪將多項式移至 GPU 並移回）是否會抵消計算加速。對於我們的問題規模，輪內計算耗時約 5-15ms，而 GPU 調度 + 讀回開銷約 1-2ms，因此可能存在適度的改進空間。這仍是未來的工作。
6. 統一記憶體：全貌
GPU 證明專案最常提到的瓶頸是 CPU-GPU 數據傳輸。在獨立 GPU 上，透過 PCIe 傳輸 64MB 的多項式需要約 4ms (16 GB/s)——這與 NTT 計算本身的時間相當。在 Apple Silicon 上，CPU 和 GPU 存取相同的物理 DRAM，因此完全消除了 PCIe 成本。
然而，統一記憶體並非沒有權衡。Metal 提供兩種緩衝區存儲模式，選擇哪種至關重要：
共享模式 (MTLResourceStorageModeShared)
CPU 和 GPU 均可讀寫緩衝區。硬體必須維持快取一致性——當 GPU 讀取 CPU 最近寫入的數據時，必須刷新 CPU 的快取，以便 GPU 看到最新值，反之亦然。這種一致性是有效能成本的：
- 從共享緩衝區讀取的速度比從 GPU 管理的緩衝區慢。
- 記憶體控制器必須在 CPU 和 GPU 快取之間進行協調。
- 對於大型緩衝區，快取刷新會增加約 0.1-0.5ms。
管理模式 (MTLResourceStorageModePrivate)
僅 GPU 可存取緩衝區。不需要快取一致性，因此 GPU 可以獲得完整的記憶體頻寬而無競爭。這對於計算密集型核心來說速度更快。
我們的混合方法
我們透過分析發現，GPU 管理的緩衝區在計算密集型核心（NTT、Merkle）中速度明顯更快。我們的流水線採用：
- 輸入：CPU 數據存於共享緩衝區（零拷貝——CPU 寫入多項式係數，GPU 直接讀取）。
- 中間計算：NTT 將數據從共享輸入緩衝區複製到 GPU 管理的緩衝區，用於所有蝴蝶運算階段。這種複製是單次 GPU 側的 blit 操作，以全記憶體頻寬運行。
- 位元反轉輸出：最終排列將結果寫回共享緩衝區，以便 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 演算法。輸入為位元反轉順序，輸出為自然順序。每個蝴蝶運算階段讀取兩個元素，將其中一個乘以旋轉因子（twiddle factor），然後寫回兩個結果。記憶體存取步長從最小（相鄰元素）開始，並在最終階段增長到 N/2。後期階段的大步長導致 GPU 上的快取利用率不佳。
頻域抽取 (DIF)：DIT 的「反向」。輸入為自然順序，輸出為位元反轉順序。記憶體步長從最大 (N/2) 開始並縮小到 1。對於 GPU 工作負載，DIF 通常更可取，因為最初的階段（處理最多全域記憶體數據的階段）具有最大的步長——而當多個蝴蝶運算並行調度時，帶有 radix-16 蝴蝶模式的大步長恰好能很好地對應 GPU 記憶體合併（memory coalescing）。
Stockham 自動排序 (Stockham, 1966)：一種異地（out-of-place）變體，透過在兩個緩衝區之間交替來避免單獨的位元反轉排列。每個階段從一個緩衝區讀取並以不同順序寫入另一個緩衝區，因此最終階段的輸出已經是正確順序。缺點是 2 倍的記憶體使用量。
我們選擇 DIF 作為主要演算法，原因如下：
- DIF 的輸出是位元反轉的，但我們本來就需要位元反轉順序來構建 Merkle 樹（葉節點必須處於求值域順序）。我們將位元反轉與最後一個 NTT 階段融合，消除了單獨的排列步驟。
- DIF 自然地與 radix-16/32 分解結合（見下文）。
- DIF 的步長模式與 Apple GPU 的記憶體子系統配合良好。
我們也實作了 Stockham 進行對比。對於額外緩衝區能放入 GPU 快取的小型 NTT (< 2^16)，它具有競爭力，但對於我們的負載規模（2^20 到 2^25），由於記憶體壓力較低，DIF 始終更快。
基數選擇：為何選擇 radix-16（和 radix-32）
基數為 R 的蝴蝶運算一次處理 R 個元素，每階段執行 R log_R(N)/log_2(N) 次運算，但僅需 log_R(N) 次全域記憶體存取，而非 log_2(N) 次。權衡如下：
- Radix-2: 每個蝴蝶運算：1 乘 + 1 加。2^24 次方需 24 次存取。
- Radix-4: 每個蝴蝶運算：3 乘 + 8 加。2^24 次方需 12 次存取。
- Radix-8: 每個蝴蝶運算：~17 乘 + 加。2^24 次方需 8 次存取。
- Radix-16: 每個蝴蝶運算：~43 乘 + 加。2^24 次方需 6 次存取。
- Radix-32: 每個蝴蝶運算：~100 乘 + 加。2^24 次方需 ~5 次存取。
由於我們的 NTT 受限於記憶體頻寬（每次存取都要讀寫整個陣列），將存取次數減半大致能使運行時間減半。代價是每次存取需要更多的 ALU 運算，但在頻寬受限的 NTT 存取期間 GPU ALU 未被充分利用，因此這種交換是有利的。
我們使用 radix-16 作為主力，因為：
- 每個蝴蝶運算 16 個元素 = 16 個暫存器，這能舒適地放入 Apple GPU 的暫存器檔案中。
- 每個執行緒處理 16 個元素，內部包含 4 「層」 radix-2 蝴蝶運算（因為 16 = 2^4）。
- Radix-32（5 個內部層，32 個暫存器）也有效且能為 2^24 節省一次存取，但對於較小的 NTT 收益遞減。
直觀地說，單個 radix-16 蝴蝶運算在內部透過 4 個階段處理一組 16 個元素：
- 階段 0 (步長 8): [0,8] [1,9] [2,10] [3,11] [4,12] [5,13] [6,14] [7,15]
- 階段 1 (步長 4): [0,4] [1,5] [2,6] [3,7] [8,12] [9,13] [10,14] [11,15]
- 階段 2 (步長 2): [0,2] [1,3] [4,6] [5,7] [8,10] [9,11] [12,14] [13,15]
- 階段 3 (步長 1): [0,1] [2,3] [4,5] [6,7] [8,9] [10,11] [12,13] [14,15]
每對 [i,j] 都是一個 radix-2 蝴蝶運算：a' = a + twb, b' = a - twb。所有 4 個階段都在暫存器中完成（無全域記憶體存取），然後將 16 個結果寫回。這意味著每次全域記憶體存取都完成了 4 個蝴蝶運算階段的工作。
四步 FFT：嘗試後放棄
四步 FFT 演算法 (Bailey, 1990) 將大型 1D NTT 分解為較小的 2D 子變換：
- 將輸入視為 N1 x N2 矩陣（行優先）。
- 執行 N2 個獨立的大小為 N1 的 NTT（行變換）。
- 乘以旋轉因子。
- 轉置矩陣。
- 執行 N1 個獨立的大小為 N2 的 NTT（列變換）。
其吸引力在於行/列 NTT 足夠小，可以完全放入執行緒組共享記憶體（Apple GPU 上為 32KB），從而在蝴蝶運算階段避免全域記憶體存取。只有轉置需要全域記憶體存取。
在實踐中，這比我們的扁平 DIF 方法慢了 20%，原因如下：
- 矩陣轉置並非免費——它需要一次全域記憶體存取，且合併模式不佳。
- 行/列 NTT 期間的共享記憶體銀行衝突（每行 16 個元素 = 步長 16 存取 = Apple GPU 32 銀行共享記憶體的最差情況銀行衝突）。
- 額外的旋轉因子乘法步驟增加了另一次全域記憶體往返。
- Apple GPU 共享記憶體每個執行緒組僅 32KB，將子 NTT 大小限制在 2^13 = 8192 個元素。
失敗的嘗試
- 將 DFT 讀回與 Merkle GPU 工作重疊：我們嘗試使用 Metal 事件讓 CPU 從共享緩衝區讀取 NTT 結果，同時 GPU 在單獨的管理緩衝區上開始 Merkle 雜湊。這反而更慢，因為 CPU 讀取與 GPU 競爭統一記憶體頻寬。在 Apple Silicon 上，記憶體匯流排是共享的，因此在重度 GPU 計算期間進行 CPU 讀取會有效地竊取頻寬週期。
- GPU sumcheck：sumcheck 協議 (Lund, Fortnow, Karloff, Nisan, 1992) 具有順序輪次結構：每一輪，證明器在超立方體上計算單變數多項式，驗證者發送隨機挑戰，證明器為下一輪「折疊」超立方體。這種輪次間的依賴關係本質上是串行的。
我們沒有實作或測試 GPU sumcheck。順序結構使其原則上不適合 GPU 並行化，但我們尚未針對具體的工作負載規模進行經驗驗證。輪內計算（在第 i 輪對 2^(n-i) 個元素求和）是高度並行的，對於具有大型超立方體的早期輪次，可能會從 GPU 加速中受益。這仍然是一個開放性問題。Thaler 2022 調查報告對 sumcheck 協議及其優化提供了詳盡的論述。
實際挑戰
- PoW 導致的基準測試變異：工作量證明研磨具有指數分佈的完成時間（搜尋隨機數）。對於 PoW 密集型配置，單次運行的基準測試結果可能波動 50% 以上。使用 3 次運行的中位數可顯著減少這種噪聲。
- 編譯器優化對 CPU 的幫助大於 GPU：LTO 和 target-cpu=native 提升了 CPU 效能約 25%，但對 GPU 核心效能影響微乎其微（Metal 著色器與 Rust 代碼分開編譯）。儘管提升了絕對效能，這卻縮小了 GPU/CPU 的比率。
- 行動端熱調節：在筆記型電腦（尤其是手機）上長時間運行基準測試會觸發熱調節（thermal throttling），降低 GPU 時脈頻率。我們的基準測試使用 3 次運行的中位數來緩解此問題，但實際的持續效能可能會較低。
8. 與相關工作對比
| 專案 | 目標 | 協議 | 加速比 | 欄位 | API | 來源 |
|---|---|---|---|---|---|---|
| 本工作 | Apple M1 GPU | WHIR/Poseidon2 | 1.3-2.0x vs CPU | BabyBear (31-bit) | Metal | repo |
| 本工作 (M3) | Apple M3 GPU | WHIR/Poseidon2 | 1.1-2.6x vs CPU | BabyBear (31-bit) | Metal | 同上 |
| 本工作 (iPhone) | Apple A19 GPU | WHIR/Poseidon2 | ~1.4-2.3x vs CPU | BabyBear (31-bit) | Metal | 同上 |
| Mopro Metal MSM v2 | Apple GPU | MSM (BN254) | 40-100x vs v1 | BN254 (254-bit) | Metal | 報告 |
| ICICLE Metal v3.6 | Apple GPU | MSM, NTT | 高達 5x | 多種 | Metal | 部落格 |
| ICICLE-Stwo (CUDA) | 資料中心 GPU | Circle STARK | 3.25-7x vs CPU SIMD | M31 (31-bit) | CUDA | 部落格 |
| zkSecurity Stwo WebGPU | 瀏覽器 GPU | Circle STARK | 整體 2x | M31 | WebGPU | 報告 |
| Ligetron | 瀏覽器/原生 | SHA-256, NTT | N/A (開發中) | 多種 | WebGPU | 代碼 |
我們的加速數據 (1.3-2.0x) 比專注於 MSM 的專案更保守，原因如下：
- 我們測試的是端到端證明，包括受限於 CPU 的 sumcheck 輪次。僅 GPU 部分（NTT + Merkle）相較於非 SIMD CPU 代碼顯示出 3-4 倍的加速。
- 我們的 CPU 基準非常強勁——Plonky3 的 BabyBear 使用 NEON SIMD 4 路打包算術，加上 LTO 和原生 CPU 代碼生成。
- WHIR 的工作負載以雜湊為主（Poseidon2 Merkle 約佔運行時間的 58%），其算術強度高，但與 MSM 相比並行化縮減有限。
請注意，Mopro MSM v2 的加速（相較於 v1 為 40-100 倍）是與其自身的 v1 相比，而非優化後的 CPU 基準。若與 Arkworks CPU MSM 相比，加速比會更適中。同樣地，ICICLE-Stwo 的 3.25-7 倍是與 Stwo 的 CPU SIMD 後端相比。
9. 未來方向
高元 Merkle 樹 (Higher-arity Merkle trees)
Poseidon2 Merkle 核心主導了運行時間 (58%)。採用 4 元或 8 元樹將使雜湊調用次數減少 2-3 倍，從而按比例加速 GPU 流水線。這需要在 WHIR 承諾方案中進行協議級別的更改。
較新的 Apple Silicon
我們主要發佈的數據基於 M1 (2020)。M3 MacBook 的測試（第 4 節）顯示了相同的定性趨勢，且在大規模 n 下通常具有更高的加速比（例如 n=24, fold=1, rate=1 時為 2.37x，而 M1 為 1.69x）。iPhone (A19) 上的 WHIR Bench 在採樣單元格中也處於類似區間（n=24, fold=1, rate=1 時為 2.3x）。M4 Max 擁有 40 個 GPU 核心（M1 為 8 個）和 273 GB/s 記憶體頻寬（M1 為 68 GB/s）。我們預期在新晶片上會有進一步的提升；程式碼庫包含 bench.sh 腳本和 iOS App，方便進行跨裝置基準測試——歡迎貢獻數據。
WebGPU 後端
為了實現跨平台覆蓋，WebGPU (wgpu) 後端可以在 WGSL 中重用相同的核心演算法。主要的代價是失去共享/管理記憶體的區分（WebGPU 記憶體模型較簡單）以及 WGSL 缺乏 mulhi 內置函數（需要模擬）。Mopro 的欄位運算基準測試表明，WebGPU 在欄位運算上可達到 Metal 效能的 50-80%。
GPU sumcheck
sumcheck 協議是目前主要的 CPU 瓶頸（約佔運行時間的 15%）。雖然輪次間的依賴本質上是串行的，但對於極大規模的實例，輪內計算（在大型超立方體上求和）可能會從 GPU 並行化中受益。我們尚未對此進行基準測試，這仍是一個開放性問題。Thaler 2022 調查報告對此提供了深入探討。
相關文章
其他收藏 · 0