随机数的计时

前几节用指令数比较了不同的生成方法。本节用第三章的方法实测：vrng 自身的时序，以及在 kernel 中连续生成时，各种方法每个 TC VREG 要多少周期。

vrng 的时序

本小节实验源码、输出。

与第三章第 4 节相同，用 LccProbe 插入手写片段，setup 中先 setrngseed。读数 (R1 − R0, R2 − R0)：

vrng 的发射间隔是 8 个周期：一条 vrng 让每个生成器连续走 8 步（第 1 节），每个周期一步。两个槽共用同一个生成器，换到 va1 不能并行。间隔中的空闲周期可以发射其他向量运算。按第三章第 4 节的模型，它与 XLU 的转置相似：一个单元每 8 个周期接收一次，期间不挡住其他指令。

由此，硬件生成器的吞吐上限是每 8 个周期 4 KiB，即每周期 512 字节，按 1.05 GHz 约 540 GB/s。

在 kernel 中连续生成

本小节实验源码、输出。

kernel 在 pl.loop 中反复生成，每次 rows 行，结果异或进同一个 buffer，使编译器不能删去任何一次生成：

pltpu.prng_seed(1)
acc_vmem[...] = jnp.zeros((rows, 128), jnp.uint32)

@pl.loop(0, steps)
def _(step: jax.Array) -> None:
    acc_vmem[...] = acc_vmem[...] ^ generate(rows, step)

每种方法编译两个 kernel，分别生成 64 个和 576 个 TC VREG，各用设备上的 LCC 读出 kernel 的周期数（第一章第 3 节的 KernelClock），相减后除以多生成的 512 个，消去种子混合、DMA 等固定开销。实验只改 generate 和 rows，十种方法见源码中的 METHODS。

同时统计清单：pl.loop 不展开，循环体在清单中只出现一次；与循环体只生成全零的 kernel 相比，多出的计算指令除以每次的 TC VREG 数，就是每个 TC VREG 的计算指令。两个向量 ALU 槽每周期各发射一条，指令数的一半是只受 ALU 限制时的周期下限。

这里的计算指令包括异或进 buffer 等循环本身的指令，比第 4、5 节单独统计的变换略多。按瓶颈可以分成三类：

生成器受限。 vrng、均匀分布、伯努利都在 8.4–8.6 个周期，ALU 下限只有 1–3 个周期：瓶颈是 vrng 每 8 个周期一条的发射间隔，变换的指令完全藏在间隔里。指数分布与 Gumbel 分布的 10.0、11.5 个周期超过 8，也远高于 ALU 下限；vlog2 的发射间隔是 2 个周期（第三章第 4 节），每个 TC VREG 一到两条，也不是瓶颈。多出的部分来自循环体的结构，下一小节把它拆开。

ALU 受限。 两种正态分布与 Philox 的实测与 ALU 下限相差不到 4 个周期：两个向量 ALU 槽几乎每个周期都在发射，生成器反而闲着。要让这类变换更快，只能减少指令，例如选择不需要正弦、余弦的方法（第 5 节）。

依赖与槽位受限。 threefry2x32 的实测比 ALU 下限多 70% 以上。原因之一是槽位：tpuasm 的指令索引中，vshll、vshrl 只出现在 va1，不能在 va0 发射。统计每次 64 行的 kernel 段（含少量循环外的指令）：va1 上有 vshll 208 条、vshrl 201 条，va0 上有 vadd 310 条、vor 179 条、vxor 144 条。两个槽的指令数并不均衡，每个 TC VREG 的 va0 就有约 80 条，下限实际上是 80 而不是 70。其余的差距来自 hash 每一轮的依赖链（第 4 节）和溢出到 TC VMEM 的读写。每次只生成 16 行时更慢：每次只有 2 个 TC VREG 可以交错执行，依赖链无法被掩盖，而 fold_in 的固定开销只分摊到 2 个 TC VREG 上。

循环体：生成器之外的周期

本小节实验源码、输出，模型见第三章第 4 节。

为什么指数分布每个 TC VREG 多出约 1.5 个周期？脚本编译与上一小节相同的 kernel，从清单中取出 pl.loop 的循环体：从循环的标号到回跳分支之后的延迟槽，去掉计数与回跳的标量指令（只含标量指令的 bundle 变成空 bundle，仍占一个周期）。然后做两件事：

• 用 LccProbe 把循环体重复 4 次插入载体，直接测它的周期数。循环体会改写全部 32 个 TC VREG，其中 v20–v30 是载体保存标量寄存器的位置（第一章第 3 节），所以 setup 先把它们存进 TC VMEM，片段在最后一次读数之后再读回。
• 把同样的 bundle 序列交给第三章第 4 节的发射模型，分别在不考虑与考虑 vrng 的 8 周期间隔时预测。

考虑 vrng 间隔的模型与 LCC 实测完全相同；不考虑时少 0.7–1 个周期。LCC 测的是去掉了循环计数和回跳的循环体，比上一小节的整个 kernel 少约 0.4 个周期。

指数分布的循环体（输出中有完整的 65 个 bundle）是这样排的：

[ 3] { va0: vrng.8x128.u32 v23 }
[ 4] { va1: vshrl.8x128.s32 v24, v23, 0x9 }
[ 5] { va0: vor.8x128.u32 v25, 0x3f800000, v24 }
[ 6] { va1: vadd.8x128.f32 v26, -1.0, v25 }
[ 7] { va1: vsub.8x128.f32 v27, 1.0, v26 ; vld: vld.8x128 v26, [vmem:0x10] }
[ 8] { va0: vlog2.8x128.f32 erf, v27 ; va1: vrng.8x128.u32 v28 }
...
[48] { va0: vrng.8x128.u32 v20 ; vr0: vpop.8x128 v19, erf }
...
[55] { va0: vlog2.8x128.f32 erf, v28 ; vr0: vpop.8x128 v29, erf }
[56] { va0: vmul.8x128.f32 v30, 0.6931471824645996, v29 }
[57] { va1: vsub.8x128.f32 v0, 0.0, v30 }
[58] { va0: vxor.8x128.u32 v1, v0, v31 }
[59] { vst: vst.8x128 [vmem:0x30], v1 }
[60] { vr0: vpop.8x128 v2, erf }
[61] { va0: vmul.8x128.f32 v3, 0.6931471824645996, v2 }
[62] { va1: vsub.8x128.f32 v5, 0.0, v3 }
[63] { va0: vxor.8x128.u32 v6, v5, v4 }
[64] { vst: vst.8x128 [vmem:0x38], v6 }

每个 TC VREG 的变换是一条链：vrng → 右移 9 位 → 或上 0x3F800000 得到 [1, 2) 中的浮点数 → 减 1 → 1 − u → vlog2 提交 → 7 个周期后取回 → 乘 ln 2 → 取负 → 异或进 buffer → 写回。编译器把相邻两个 TC VREG 的链错开 5–7 个 bundle，交错排列。从清单上看，8 条 vrng 之间相隔 5–7 个 bundle，比 8 个周期的间隔短：编译器排指令时没有考虑这个间隔，每条 vrng 都要在向量一侧多等 1–3 个周期，后面所有的 bundle 一起推迟（第三章第 4 节的按序发射）。

于是一次迭代的时间可以拆成三段：

• 第一条 vrng 之前：读入 buffer 的 vld 等 3 个 bundle。
• 8 条 vrng 之间：7 个间隔，按 8 个周期计共 56 个周期。均匀分布与指数分布都一样，与编译器排得多紧无关。
• 最后一条 vrng 之后：最后一个 TC VREG 的整条链，加上倒数第二个的 vlog2 取回与写回。指数分布有 16 个 bundle，其中 [55] 的 vlog2 到 [60] 的取回还要再等 2 个周期；均匀分布的链只有 5 条运算加一次写回，这一段只有 6 个 bundle。

pl.loop 的迭代之间不重叠：下一次迭代的 vrng 要等这一次的最后一个 bundle 发射之后才发射。所以首尾两段的周期每次迭代都要付一次，摊到 8 个 TC VREG 上，就是指数分布比 8 多出的部分；Gumbel 分布的链多一次 vlog2 和几条运算，尾段更长，多出的也更多。

验证这个解释的办法是只改每次生成的行数：64 行改为 256 行，循环体中的 vrng 从 8 条变为 32 条，首尾两段的长度不变，摊到 32 个 TC VREG 上。实测每个 TC VREG 从 9.59 降到 8.40 个周期，接近 vrng 的间隔 8。要在 kernel 中以接近生成器上限的速度产生需要较长变换的分布，每次迭代应当生成足够多的行；代价是寄存器与 TC VMEM 的占用随之增加。

结论

对比第二章第 2 节：HBM 读入一个 TC VREG 约需 4 个周期。所以：

• 用 vrng 在 kernel 内生成均匀分布或伯努利掩码，每个 TC VREG 约 8.5 个周期（指数分布这类短变换，每次迭代生成足够多的行时也能接近这个数），比从 HBM 读入预先生成的随机数慢一倍；但它不占 HBM 带宽，可以与其他读写、计算并行，而生成器在 8 个周期的间隔中只占用极少的 ALU。
• 正态分布这类变换受 ALU 限制，会与 kernel 的其他向量计算争用两个 ALU 槽；在计算密集的 kernel 中，它的真实开销就是这些指令本身。
• threefry2x32 每个 TC VREG 约 120 个周期，比 HBM 读入慢约 30 倍，只在需要与 jax.random 逐位一致、或需要按位置寻址时才值得；使用时每次至少生成几十行，让多个 TC VREG 交错掩盖依赖链。

这些计时补充了研究报告 57 未测的时序部分。