为zkVM中的RISC-V提供1024个寄存器

本文探讨了在zkVM中,RISC-V的32个寄存器限制导致大量spill操作,增加证明成本的问题。作者通过让Claude修改LLVM,生成了支持1024个寄存器的RISC-V变体(RISCV-X),实验表明可消除函数内spill,减少约14%的trace单元,但调用约定导致的跨函数spill仍然存在。与powdr的crush方案对比,后者通过无限寄存器和帧分离更好地解决了此问题。结论:扩展寄存器对内部寄存器压力有效,但无法解决调用约定问题;AI辅助实验效率高。

OpenVMSP1ZisK 这类通用 zkVM 都选择 RISC-V 作为其主要的源 ISA,因为它简单。它是一个现代、精简、易于理解的 ISA,拥有成熟的工具链。但 RISC-V 是为硬件设计的:当寄存器溢出到缓存仅消耗几个周期时,32 个通用寄存器是一个合理的数量。

在 zkVM 中,内存更昂贵,而通常的编译器工具链意识不到这种截然不同的成本模型。每一次加载和存储都需要被证明,并且对于计算密集型程序来说,内存一致性约束的证明成本很高。那个在硬件上无害的 32 寄存器限制,在这里变成了一个问题。

powdr,我们构建了 crush,它从 WebAssembly 编译到一个具有无限寄存器和零溢出的自定义 ISA,从而比 RISC-V 实现更快的证明时间。我最近在 EthProofs Beast Mode Day 展示了 crushpowdr-wasm。我的幻灯片可以在这里找到。我们还发布了一篇关于 crush 和 powdr-wasm 的专门文章,所以我不会在这里重点讲这个。

除了从 WebAssembly 编译到自定义 ISA 之外,至少还有另外两种类似的、我们之前已经排除的尝试消除寄存器溢出的方法:

  • 为自定义 ISA 构建一个完整的 LLVM 后端。从 WebAssembly 开始让我们可以复用覆盖大部分编译流程的大量工具,所以我们更倾向于 WebAssembly。
  • 为 zkVM 定制 RISC-V。这也需要修改 LLVM 和工具链,结果更难预测,感觉像一种 hack。

在 Beast Mode Day 期间,我讨论了一些这些未被选中的想法,尤其对第二点感兴趣。如果我们保留 RISC-V,但尝试给它无限(或更多)的寄存器,会发生什么? LLVM IR 内部已经使用了无限的虚拟寄存器。寄存器分配器的工作是将它们映射到有限的物理寄存器集。如果我们给它 1024 个寄存器而不是 32 个,对大多数实际函数来说,它应该能够在没有溢出的情况下成功。我们选择了 1024 个寄存器,因为它可以放在 10 位中,这比 RISC-V 在指令编码中为寄存器保留的 5 位多了一倍,我们认为将格式加倍是一种允许更多寄存器的简单 hack。

去年我们第一次想到这个的时候,作为周末实验还不可行,但幸运的是,现在我们有了 AI 可以依赖。我让 Claude 在 LLVM 中实现这一点,并将其连接到 OpenVM。实验成功了,我们得到了自定义 RISC-V(Claude 将其命名为 RISCV-X)的证明。

RISC-V 上的寄存器溢出是什么样的

基准测试使用了 tiny_sha3(Saarinen 的 FIPS 202 参考 C 实现,提交 dcbb319,与上游源代码字节一致)以及一个独立驱动程序,该程序在一个起始为 32 个零字节的缓冲区上运行 1000 次 SHA3-256 迭代。我们比较了标准 RISC-V 版本与我们的 RISCV-X 扩展的指令数、zkVM 追踪单元和证明时间。

这是 sha3_keccakf(Keccak 置换函数)编译到标准 RISC-V 的开头部分:

sha3_keccakf:
    addi sp, sp, -272
    sw   ra,  268(sp)          # 被调用者保存的寄存器溢出
    sw   s0,  264(sp)
    sw   s1,  260(sp)
    sw   s2,  256(sp)
    sw   s3,  252(sp)
    sw   s4,  248(sp)
    sw   s5,  244(sp)
    sw   s6,  240(sp)
    sw   s7,  236(sp)
    sw   s8,  232(sp)
    sw   s9,  228(sp)
    sw   s10, 224(sp)
    sw   s11, 220(sp)          # 13 个被调用者保存的寄存器溢出
    li   t3, 0
    lw   t2, 0(a0)
    lw   t4, 4(a0)
    lui  a1, %hi(.L__const.sha3_keccakf.keccakf_rotc)
    addi a1, a1, %lo(.L__const.sha3_keccakf.keccakf_rotc)
    addi a2, a1, 4
    sw   a2, 12(sp)            # 立即溢出常量指针
    addi a1, a1, 96
    sw   a1, 8(sp)             # 溢出另一个
    ...

该函数以 13 个被调用者保存的寄存器溢出开始,然后开始加载 Keccak 状态和常量指针并立即溢出它们,因为它已经没有寄存器了。Keccak 状态是 25 × 64 位通道(在 RV32 上是 50 × 32 位)。该函数必须将所有通道加上临时变量和常量指针一起在 32 个架构寄存器中穿梭。

在整个 sha3_keccakf 中进行统计:

  • 共 784 条指令
  • 111 次栈存储 + 111 次栈加载 = 222 次内存操作
  • 272 字节的栈帧

在 zkVM 中,这 222 次内存操作中的每一次都是纯粹的开销:它们之所以存在,仅仅是因为我们用完了寄存器。

在具有 1024 个寄存器的 RISC-V 上运行相同的代码

Claude 在 LLVM 的 RISC-V 后端中添加了一个 +xregs1024 子目标特性:1024 个 GPR、64 位指令编码(以适应 10 位寄存器字段),以及一个除 ra 外所有寄存器都作为调用者保存的调用约定。我们使用了相同的 tiny_sha3 C 源码,并在 llc 调用中添加了 -mattr=+xregs1024

以下是相同的 sha3_keccakf 函数序言:

sha3_keccakf:                  # 完全没有栈帧
    li   a1, 0
    lw   t1, 0(a0)
    lw   t2, 4(a0)
    lui  a6, %hi(.L__const.sha3_keccakf.keccakf_rotc)
    addi a6, a6, %lo(.L__const.sha3_keccakf.keccakf_rotc)
    lui  a2, %hi(.L__const.sha3_keccakf.keccakf_piln+4)
    addi a2, a2, %lo(.L__const.sha3_keccakf.keccakf_piln+4)
    li   a3, 64
    li   a4, 32
    addi a5, a6, 4
    addi a6, a6, 96
    lui  a7, %hi(.L__const.sha3_keccakf.keccakf_rndc)
    addi a7, a7, %lo(.L__const.sha3_keccakf.keccakf_rndc)
    li   t0, 24
    j    .LBB0_2
    ...

没有序言,也没有溢出。所有常量都位于寄存器中,循环计数器位于寄存器中,在循环体内部,50 个状态字位于 x32x80 中。

该函数的完整数据:

  • 561 条指令(从 784 条减少,减少了 28%)
  • 0 次栈存储,0 次栈加载(从 222 减少,100% 消除)
  • 0 字节的栈帧
  • 使用了 80 个不同的寄存器——25 个标准寄存器加上约 55 个扩展寄存器(x32x91),这与预期大致相符:50 个 Keccak 状态字加上临时变量。

因此,sha3_keccakf 需要大约 80 个寄存器才能在没有溢出的情况下运行,当你给 LLVM 寄存器分配器这么多寄存器时,它确实会使用它们。这在一定程度上是预期的,但看到寄存器分配如此通用仍然令人欣慰。标准的 32 寄存器 RISC-V 迫使它在整个函数中通过内存来回穿梭值。

在 OpenVM 中运行它

有趣的部分是当 zkVM 尝试证明这一点时会发生什么。由于更宽的寄存器无法放入标准的 32 位 RISC-V 指令格式中,因此二进制文件对每条指令使用了 64 位编码。

Claude 必须修改 OpenVM 以接受新的二进制格式:一个新的转译器扩展,用于解码 64 位指令,内部扩展寄存器字节偏移量(电路代码中进行了 u8 → u16 清理),以及扩展的寄存器文件。指令集本身没有改变。我们使用相同的操作码、相同的语义,只是更宽的寄存器标识符。

然后,我们将两个二进制文件输入 OpenVM(v1),并为一个基准测试生成了实际证明,该基准测试在一个起始为全零的 32 字节缓冲区上连续运行 SHA3-256 1000 次。最终的哈希值是 52cf48e88ce4dea40f272b6aaf083675ade26504a0129f51ec30204a2fdb1c5b。基础版本和扩展版本的 ELF 都产生这个结果,与 Python 的 hashlib.sha3_256 匹配。

具体到 OpenVM 访客,我们使用的是独立 C 语言,并直接与 ld.lld 链接;通过我们修改的后端连接完整的 Rust + OpenVM 工具链需要一个自定义的 rustc sysroot,这需要更多的工作,我认为当前的结果已经足以证明这个实验。

证明指标

聚合所有 AIR,来自两次运行的原始指标 JSON(1000 次 SHA3-256 迭代,4 个证明段,无递归的 STARK 应用证明):

基线(32 个寄存器) 扩展版(1024 个寄存器) 变化
已执行指令数 35,896,008 31,067,004 −13.5%
total_cells(分配的追踪单元) 4,049,774,760 3,853,796,520 −4.8%
total_cells_used(已填充的追踪单元) 3,275,912,933 2,829,414,701 −13.6%
总证明时间 130.8s 127.7s −2.4%

我们实际用于证明的追踪单元减少了约 14%。total_cells 的分配数字减少较少,因为要对齐到 2 的幂进行填充,而证明时间通常与最终追踪单元的节省量成正比。

节省来自于已证明的内存操作更少:基线中的每次溢出/重新加载都是一次加载和一次存储,这些必须在内存参数中得到处理,而现在它们消失了。这是一个真正的改进,但幅度不是很大。

详细分析可以在我们的指标查看器中看到。

溢出在哪里(以及不在哪里)

汇编级数据显示,仅在 sha3_keccakf 内部就消除了 222 次内存操作。端到端的减少只有约 14%,而 crush 相对于 RV32 可以实现高达 50% 的节省。为什么没有更多?

因为这次实验解决了问题中较容易的一半:基本块内的溢出。一个函数有大量活跃值却没有地方存放。给寄存器分配器一个更大的寄存器文件,它就不再需要穿梭了。

没有解决的是另一半:跨函数调用的溢出。即使有 1024 个调用者保存的寄存器,当一个函数调用另一个函数时,调用者必须保存它在调用返回后想要看到的任何活跃值。被调用者可以随意破坏所有内容。因此,如果存在跨越调用的活跃值,它们就会进入栈。特别是,ra 总是这样,因为每次 call 都会覆盖它。

这就是 crush 在结构上不同的地方。在 crush 中,每个函数通过帧指针获得无限寄存器空间的一个不重叠切片。调用者和被调用者从不共享寄存器(除了输出优化的预期情况),所以没有任何东西需要保存和恢复。这是一种真正的帧分离,而不是在共享寄存器文件之上的约定。

RISC-V 像每个低级 ISA 一样,有一个单一的扁平寄存器文件,在每个栈帧之间共享。调用约定是关于谁保存什么的合作约定,但它始终是在共享资源之上的约定。你可以让所有寄存器都是调用者保存的(我们就是这么做的),或者所有都是被调用者保存的,或者拆分它们,但你无法逃脱这样一个事实:如果两个函数想要使用相同的寄存器,调用边界就是活跃值必须移动的地方。并且,没有全局程序分析,寄存器分配器无法跨函数边界进行协调。

基于帧的寄存器划分,就像 crush 所做的那样,需要从一开始就设计到 ISA 中。你需要某种类似于帧指针的东西来移动寄存器窗口,或者真正无限多的具有 SSA 风格命名的寄存器。通过现有的调用约定机制将其改造到 RISC-V 上是行不通的,因为该机制本身就是导致跨调用点溢出的原因。

结论

结果

这个实验成功地完成了它设定的目标。函数内部的寄存器压力降为零。这对实际工作负载来说是一个有意义的改进(在 Keccak 上追踪单元减少了约 14%,在计算更密集的工作负载上可能更多)。然而,调用约定问题仍然存在(crush 很好地解决了这个问题),对此你需要一种不同类型的 ISA。

LLVM 🤝 Claude

当前 LLVM 的差异 包含不到 500 行代码。我之前认为这个实验用这么少的代码是不可能实现的,所以祝贺 LLVM!

Claude 在这项任务中令人印象深刻。它似乎对 LLVM 理解得很好,并在不到 10 分钟(!!!)内让一个小型 C 程序的原型运行起来。

感觉不可思议的是,我可以在一天之内发起这样一个实验,而一行代码都不用写(我仍然阅读了代码,并检查了计算出的哈希值是否正确)。我真的很兴奋,我们将能够仅凭原始思路和 Markdown 文件做出这些改进。

链接

所有代码和报告均由 Claude 编写。

  • 原文链接: hackmd.io/et5ECnWLQkOnU1...
  • 登链社区 AI 助手,为大家转译优秀英文文章,如有翻译不通的地方,还请包涵~
点赞 0
收藏 0
分享
本文参与登链社区写作激励计划 ,好文好收益,欢迎正在阅读的你也加入。

0 条评论

请先 登录 后评论
et5ECnWLQkOnU10mvPpUeQ
et5ECnWLQkOnU10mvPpUeQ
江湖只有他的大名,没有他的介绍。