密码学 - Stwo WebGPU - Xm9Z2P8Ssis1W Rv1O0Uda

概要

• 我们以 Stwo Poseidon AIR 作为基准，分析了 proving 的执行时间，并在 WebGPU 中实现了最耗时的操作：trace 生成、trace 插值、trace 扩展和计算组合多项式。

• 即使在考虑了各种开销（例如 CPU 到 GPU 的内存复制）之后，这些使用 WebGPU 的操作也比 wasm 构建上的 Simd 实现快大约 5 倍。

• 由于这些操作大约占总 proving 执行时间的 65%，即使在当前的部分实现下，整体执行速度也快了 2 倍。

• 与其他 Stwo GPU 实现一样，将整个证明过程卸载到 GPU 似乎是最有效的方法。

  • 然而，由于 WebGPU 的内存限制（受到浏览器和移动设备的约束），可能还需要实施其他措施，例如动态 proving（基于设备限制在 CPU 和 WebGPU 之间切换）和递归 proving。

• 在主要任务中，尚未移植到 WebGPU 的剩余组件是 commitment 和 FRI 代码。

  • 在此之后，下一步将是将尽可能多的 proving 逻辑迁移到 WebGPU。

Stwo 性能分析

flowchart LR
    L1["主要影响<br/>>15%"] --- L2["显著影响<br/>5-15%"] --- L3["中等影响<br/>1-5%"] --- L4["轻微影响<br/><1%"]

    classDef red fill:#ff0000,color:white
    classDef orange fill:#ffa500,color:black
    classDef gray fill:#808080, color:white;
    classDef white fill:#ffffff, color:black;

    class L1 red
    class L2 orange
    class L3 gray
    class L4 white

Poseidon n=15 (并行)

flowchart TD
    A["总 Poseidon 证明<br/>13.78s | 100%"] --> B["预计算 twiddle 因子<br/>9.01ms | 0.07%"]
    A --> C["常量<br/>1.04ms | 0.01%"]
    A --> D["Trace<br/>5.19s | 37.66%"]
    A --> E["交互<br/>331ms | 2.40%"]
    A --> F["证明<br/>8.24s | 59.80%"]

    F --> G["组合<br/>6.80s | 49.35%"]
    F --> H["计算域外列<br/>997ms | 7.24%"]
    F --> I["计算 FRI 商<br/>261ms | 1.89%"]
    F --> J["承诺<br/>105ms | 0.76%"]
    F --> K["研磨<br/>13.3ms | 0.10%"]

    G --> L["生成<br/>6.73s | 48.84%"]
    G --> M["承诺<br/>68.7ms | 0.50%"]

    L --> N["扩展<br/>5.66s | 41.07%"]
    L --> O["约束逐点求值<br/>1.04s | 7.55%"]
    L --> P["约束插值<br/>22.8ms | 0.17%"]

    D --> Q["生成<br/>865ms | 6.28%"]
    D --> R["承诺插值<br/>1.51s | 10.96%"]
    D --> S["承诺<br/>2.82s | 20.46%"]

    S --> T["承诺<br/>2.82s | 20.46%"]
    T --> U["扩展<br/>2.58s | 18.72%"]
    T --> V["Merkle<br/>243ms | 1.76%"]

    %% Color definitions based on impact:
    %% Major Impact (>15%): red
    %% Significant Impact (5-15%): orange
    %% Moderate Impact (1-5%): gray
    %% Minor Impact (<1%): white
    classDef red fill:#ff0000, color:white;
    classDef orange fill:#ffa500, color:black;
    classDef gray fill:#808080, color:white;
    classDef white fill:#ffffff, color:black;

    class A,D,F,G,L,N,S,T,U red;
    class H,O,Q,R orange;
    class E,I,V gray;
    class B,C,J,K,M,P white;

细节

基准测试结果（扩展 trace + 计算约束多项式）

Repo: https://github.com/zksecurity/stwo/commits/gpu-extend-trace/

• 在 Macbook M3 Pro 上运行

WASM (release build) 执行时间 (ms)

Native build (parallel feature on + optimized) 执行时间 (ms)

基准测试步骤

基准测试步骤 ### 运行 WASM

1. 更改 test_poseidon_prove_wasm_gpu_cpu 中的 LOG_N_INSTANCES
2. 更改 compute_composition_polynomial.rs 中的 N_ROWS
   1. LOG_N_INSTANCES=12 => N_ROWS=32
   2. LOG_N_INSTANCES=13 => N_ROWS=64
   3. …
3. 更改 compute_composition_polynomial.wgsl 中的 N_ROWS
   1. 与 compute_composition_polynomial.rs 相同的值
4. 运行 wasm-pack test --chrome --release

运行 Native CPU & GPU (parallel features on & release)

1. 更改 test_simd_poseidon_prove 中的 LOG_N_INSTANCES
2. 更改 compute_composition_polynomial.rs 中的 N_ROWS
   1. LOG_N_INSTANCES=12 => N_ROWS=32
   2. LOG_N_INSTANCES=13 => N_ROWS=64
   3. …
3. 更改 compute_composition_polynomial.wgsl 中的 N_ROWS
   1. 与 compute_composition_polynomial.rs 相同的值
4. 运行 cargo test --release --features parallel test_simd_poseidon_prove -- --nocapture

• 原文链接： hackmd.io/XM9z2p8sSIS1w_...
• 登链社区 AI 助手，为大家转译优秀英文文章，如有翻译不通的地方，还请包涵～