Go 的 simd/archsimd 实验：用纯 Go SIMD 加速 ChaCha20

Go 标准库里的 ChaCha20 在 amd64 上一直缺专门的优化，因为 TLS 中实际使用的 ChaCha20-Poly1305 已有汇编实现。这导致我在 Go 里做 ChaCha20-BLAKE3 时，即使在支持 AVX-512 的 AMD EPYC Zen5 上也跑不快。

试了 Go 的实验性 simd 和 simd/archsimd 包来回折腾一天，最终给 amd64（AVX2/AVX-512）和 arm64（NEON）写出了纯 Go 的 ChaCha20 实现。在 AMD EPYC 9555P 上，AVX-512 路径比标准库快约 7.5 倍；在 Linux/arm64 上则能到标准库的约 97%。代码在 stdx-go/chacha20，仓库明确标注请勿用于生产环境。

两个平台关键吞吐数据（MB/s，均 0 分配）：

64 B 的小块在 amd64 上反而变慢，属于 SIMD 启动/转置开销；数据块到 512 B 以上收益就比较明显。

SIMD 加速两条路

ChaCha20 是教科书级的 SIMD 流密码，设计时就是为了让性能随 SIMD 寄存器宽度增长。常规做法有两种：

1. 为算法寻找可并行的依赖子结构，通常很依赖算法本身，复杂。
2. 把输入切成若干片，每片包含 X 个独立加密块，X 等于可用 lane 数，然后并行处理这些块。

ChaCha20 一个块是 16 个 32 位字，共 512 位 / 64 字节。若运行在 256 位向量上，lane 数就是 256 / 32 = 8，即一次并行处理 8 个块。输入只有达到 8 × 64 = 512 字节以上，SIMD 才能开始回本。

simd/archsimd：固定宽度、直接映射指令

archsimd 给的是固定大小的向量类型，方法直接映射到具体 SIMD 指令。某个方法在当前平台没有对应指令时，会退化成模拟执行——很慢。所以用之前必须确认你要调的方法在该 CPU 上是真实指令。

一个 AVX-512 版本核心结构：

// 旋转计数的向量只加载一次，跨轮次留在寄存器
rotl16 := archsimd.LoadUint32x16Array(&rotlCounts[0])
rotl12 := archsimd.LoadUint32x16Array(&rotlCounts[1])
rotl8  := archsimd.LoadUint32x16Array(&rotlCounts[2])
rotl7  := archsimd.LoadUint32x16Array(&rotlCounts[3])

for r := 0; r < 10; r++ {
    w0, w4, w8,  w12 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w0, w4, w8, w12)
    w1, w5, w9,  w13 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w1, w5, w9, w13)
    w2, w6, w10, w14 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w2, w6, w10, w14)
    w3, w7, w11, w15 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w3, w7, w11, w15)

    w0, w5, w10, w15 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w0, w5, w10, w15)
    w1, w6, w11, w12 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w1, w6, w11, w12)
    w2, w7, w8,  w13 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w2, w7, w8, w13)
    w3, w4, w9,  w14 = quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, w3, w4, w9, w14)
}

// add back 后做两级 4×4 转置，把 word-major 状态
// 转成按 64 字节块连续输出的布局

quarterRound 本身很短，但注意它接收的是 4 个独立向量，不是寄存器数组：

func quarterRoundAVX512(rotl16, rotl12, rotl8, rotl7, a, b, c, d archsimd.Uint32x16) (archsimd.Uint32x16, archsimd.Uint32x16, archsimd.Uint32x16, archsimd.Uint32x16) {
    // a += b; d ^= a; d <<<= 16
    a = a.Add(b)
    d = d.Xor(a)
    d = d.RotateLeft(rotl16)

    // c += d; b ^= c; b <<<= 12
    c = c.Add(d)
    b = b.Xor(c)
    b = b.RotateLeft(rotl12)

    // a += b; d ^= a; d <<<= 8
    a = a.Add(b)
    d = d.Xor(a)
    d = d.RotateLeft(rotl8)

    // c += d; b ^= c; b <<<= 7
    c = c.Add(d)
    b = b.Xor(c)
    b = b.RotateLeft(rotl7)

    return a, b, c, d
}

相比之下，Rust 版是直接用 intrinsics：

_mm512_add_epi32($a, $b);
_mm512_xor_si512($d, $a);
_mm512_rol_epi32($d, 16);

完全可控、与指令一一对应，但可读性差很多。Go 的 archsimd 抽象在行为上接近 intrinsics，写法上更接近普通代码。

这套 API 是实验性的，必须加环境变量才能用，例如：

GOEXPERIMENT=simd go test -bench=. -benchmem ./chacha20

不设 GOEXPERIMENT=simd 无法导入。

simd 包：宽度无关，但 ChaCha20 上并不“可移植”

simd 是对向量宽度做抽象的包，编译时会选择平台对应的 SIMD 指令，不支持则模拟执行。例如给初始常量做广播：

i0 := simd.BroadcastUint32s(0x61707865) // Uint32s
i1 := simd.BroadcastUint32s(0x3320646e)

我也试过用它做 ChaCha20，结果不理想：至少在这个算法里，很难写出真正跨指令集的代码。不同架构能直接干的事不一样，位旋转就是最明显的例子——有的有单指令，有的没有。最终代码里仍免不了平台判断，或者说平台相关分支。

更希望的是能在 simd 与 archsimd 类型之间互换：上层逻辑用宽窄无关 API 写，需要精确控制时下沉到固定宽度。目前做不到。

寄存器压力与溢出

不同架构的 SIMD 寄存器数量不一样：AVX2 只有 16 个向量寄存器，AVX-512 和 NEON 都是 32 个。写 SIMD 时得把状态装进显式变量，并保证不要因为函数调用或中间操作导致寄存器被 spilling 到栈上。

AVX-512 路径中 16 个状态字加 4 个旋转计数向量，共 20 个 ZMM 寄存器，还能塞下。但如果移植到 AVX2，寄存器容量和数量都会紧很多。通用经验：写完看编译器生成的汇编，确认没有栈溢出；再看优化报告确认热函数被内联。Go 的优化器在这种极端使用场景下不会总按你预想的方式生成代码。

经验总结

• 每次都要看生成的汇编，追踪寄存器 spill。这一条往往决定性能拐点。
• 确认用的每个 SIMD 方法是真指令还是模拟执行，后者慢到没有意义。
• 在编辑器里开优化详情，确认核心函数被内联，否则函数边界上的拷贝会毁掉寄存器分配。
• simd/archsimd 仍是实验性 API，生产环境不要用。
• 对多 SIMD 架构做抽象很难。Go 团队这两个包已经给出了一个比较可用的方向。

ChaCha20 在 Go 里现在已经能跑到 7.5 GB/s 级别。下一步的瓶颈显然在 BLAKE3：同一台机器上 Go 实现大概只有 5452 MB/s，Rust 版在 EPYC 4245P 上能到约 13000 MB/s。用 archsimd 重写 BLAKE3 应该是把整体吞吐推到 4–5 GB/s 的关键。

另一个未完成的想法：如果 Go 能把 SIMD 寄存器当数组处理，比如：

var state [16]archsimd.Uint32x16

就不必手写 i0 到 i15 共 16 个独立变量了。但 Go 编译器目前很难让函数参数按寄存器传这种大数组，实现复杂度不低。类似的，还希望补上 _mm512_rol_epi32 这类相对常见却缺失的方法。

测机：AMD EPYC 9555P 64 核虚拟机，CPU flags 含 avx512f avx512dq avx512bw avx512vl avx512vbmi vaes 等。