逐元素运算

逐元素运算是 TensorCore 向量单元最基本的工作：一条指令对一个 TC VREG 的 8×128 个位置同时做同一种运算。本节要建立的是一张“价目表”：哪些运算是一条硬件指令，哪些要用多条指令拼出来，拼出来的结果精度如何，各要多少周期。设计算法时，这张表决定了哪些写法便宜、哪些写法昂贵。

写法：同一个 kernel 框架，只换中间一行

本节的实验都用同一个两输入 kernel，放在 elementwise_common.py 中。相对于第 1 节的最小 kernel，它多一个输入 y，输出的 shape 和 dtype 由运算本身推出，计算换成任意函数 f：

out = jax.eval_shape(f, jax.ShapeDtypeStruct(x.shape, x.dtype), jax.ShapeDtypeStruct(y.shape, y.dtype))
...
def kernel(x_hbm: Ref, y_hbm: Ref, o_hbm: Ref, x_vmem: Ref, y_vmem: Ref, o_vmem: Ref, sem: Ref) -> None:
    pltpu.async_copy(x_hbm, x_vmem, sem).wait()
    pltpu.async_copy(y_hbm, y_vmem, sem).wait()
    o_vmem[...] = f(x_vmem[...], y_vmem[...])
    pltpu.async_copy(o_vmem, o_hbm, sem).wait()

jax.eval_shape 只推导函数输出的 shape 和 dtype，不实际计算。各实验脚本把一组 f 依次交给这个框架，例如 lambda x, y: jnp.exp(x)，每个 f 单独编译一个 kernel，统计清单中的计算指令（不计 load/store、DMA 和标量指令）。

kernel 内可以直接用 jnp 和 jax.lax 的逐元素函数。它们由 Mosaic 降低成 TensorCore 指令，降低不了的会在编译时报错。

f32：一条指令的运算

本小节实验源码、输出。

输入是 [-4, 4) 和 [0.25, 4) 中的均匀随机数，误差与 float64 参考比较，以 ULP（相差几个 f32 可表示数）计：

加减乘的结果都是正确舍入的。-x 和 abs(x) 没有专用指令，分别用减法和按位与实现。

x * y + 1 是一条 vmul 加一条 vadd。TPU v4 的向量单元没有乘加融合（FMA），乘积先舍入一次，加法再舍入一次。乘积接近 −1 时发生相消，结果相对误差可达 21 ULP，而一次舍入的 FMA 不会有这个问题。依赖 FMA 精度的算法（例如某些补偿求和）在这里要重新设计。

第 5 节已经提到，vmul.8x128.f32 只能在 va0 槽发射，vadd.8x128.f32 只能在 va1 槽发射。一个 bundle 可以同时做一次 f32 乘法和一次 f32 加法，但不能同时做两次乘法。

周期数：结果要几个周期才可用

本小节实验源码、输出。

第 3 节测到，相互依赖的 f32 加法和乘法每条要 2 个周期。这里对向量 ALU 的各类指令逐一测：把 16 条同样的指令串成一条链，每一条用上一条写出的 v11，读 R2 − R0。结果下一个周期可用时，16 个 bundle 互不等待，读数是 16 + 13 = 29；要 2 个周期时是 2 × 16 + 11 = 43。

'vadd.8x128.f32': ('va1: vadd.8x128.f32 v11, v11, v10',),
'vadd.8x128.s32': ('va0: vadd.8x128.s32 v11, v11, v10',),
'vgt.8x128.f32 → vsel（两条一轮）': ('va0: vgt.8x128.f32 vm0, v11, v10', 'va0: vsel.8x128 v11, vm0, v11, v10'),

规律很简单：要做浮点舍入或重新排布数据的指令是 2 个周期，其余都是 1 个周期。f32 的 vmax、vmin 只是比较和选择，不需要舍入，是 1 个周期。两条一轮的链也符合这张表：“比较、选择”和“读入、整数加”每轮 2 个周期（读数 45，即 32 个 bundle 互不等待），“f32 加法、再取最大值”每轮 3 个周期（读数 60）。

指令相同，先后不同，周期数也不同。 延迟只在下一条指令依赖它时才要付。同一实验把 8 条互不依赖的 vadd.f32 和写出各自结果的 8 条 vst 用两种顺序排列：

16 个 bundle 完全一样，只是顺序不同，相差 7 个周期。紧跟的排法里，每条 vst 都要等前一个 bundle 的加法结果多 1 个周期；向量一侧按顺序发射（第 3 节），它等的时候后面的指令也不能发射。分开的排法里，轮到第一条 vst 时它要的结果早已就绪。第 3 节的清单中，编译器把 vst 排在对应的 vmul 之后两个 bundle，就是在避开这个等待。手写或改写清单时，把依赖一个结果的指令往后挪、中间填上不依赖它的工作，是最基本的排法。

f32：超越函数与 EUP

exp(x) 的清单：

{ va0: vmul.8x128.f32 v1, 1.4426950216293335, v0 }   # x × log2(e)
{ va0: vpow2.8x128.f32 erf, v1 }                     # 2 的幂，结果送入 erf
{ vr0: vpop.8x128 v2, erf }                          # 从 erf 取回结果

超越函数由独立的超越函数单元 EUP 计算。vpow2 在 va0 槽发射，但它的目的操作数不是 TC VREG，而是结果队列 erf；结果要另用一条 vpop 从 vr0 或 vr1 槽取回。

周期数：EUP 的延迟与发射间隔

本小节实验源码、输出。

EUP 是本章遇到的第一个“提交—取回”的单元，它有两个时间参数，测法以后各节还会反复使用（tpuasm_tools.show_latency）。

延迟。 在提交与取回之间插入 d − 1 个 vnop：

show_latency(probe, f'EUP {function} → vpop erf', f'va0: {function}.8x128.f32 erf, v10', 'vr0: vpop.8x128 v11, erf')

d ≤ 6 时读数都是 20：vpop 在等结果，向量一侧停在这里，中间放几个 vnop 都一样。d = 7 起每多一个 vnop 多一个周期，说明结果已经就绪。20 减去固定的 13 个周期，从提交到可以取回是 7 个周期。五种 EUP 指令的读数完全相同。

发射间隔。 连续提交 k 次，看每多一次多几个周期：

第一行每多一次提交多 2 个周期：EUP 每 2 个周期接受一次提交。 第二行，16 次提交占 32 个周期，16 次取回紧随其后每周期一条。

这两个数给出了使用 EUP 的两种情形。只算一个 TC VREG 时，7 个周期是空等。连续计算许多个 TC VREG 时，提交每 2 个周期一次，延迟被后面的提交填满，EUP 每 2 个周期完成一个 TC VREG：它的吞吐是向量 ALU 上一条普通指令的一半。

本节实验中出现的 EUP 指令有 vpow2、vlog2、vrcp、vrsqrt、vtanh，每条都配一条 vpop。其他函数由它们加上普通向量运算拼成：

这张表中有两种不同的取舍：

• 倒数、平方根倒数和除法：EUP 只给出近似值，编译器再用几次乘加做牛顿迭代，把误差修正到 1–2 ULP。多出来的 vweird（判断是否为无穷或 NaN，结果写进掩码）、veq、vsel 等指令处理无穷、零和 NaN：这些输入不能进入迭代，直接选用特殊值。
• exp、log、tanh：直接使用 EUP 的近似结果，不做修正，误差为几十到几百 ULP。log 的 ULP 数格外大，是因为 y 接近 1 时结果接近 0，约 1e-4 的绝对误差折合成很多个 ULP。

这两种取舍的依据是 EUP 指令本身的精度。实验源码、输出用 tpuasm 直接执行五条 EUP 指令，不加任何修正：先由 1024 个随机的 u32 拼出 [1, 2) 中的 y 和 [−4, 4) 中的 x，再对每条指令各提交、取回一次，与 float64 的参考值比较。

f32 的尾数是 24 位，而 EUP 只给出 12–18 位有效的结果。vrcp 的 16 位经过一次牛顿迭代（r + r × (1 − y × r)，清单中的 2 条 vmul、1 条 vsub、1 条 vadd）翻倍，超过 24 位，所以 1 / y 只需要一步修正就达到 1 ULP；vrsqrt 同理。vpow2、vlog2、vtanh 没有这样便宜的修正公式，编译器直接使用原始结果，上表的误差就是 exp2、log、tanh 的误差。

sin 不使用 EUP，而是范围约化加多项式，共 198 条指令，换来 2 ULP 的精度。用它之前要先想清楚是否值得。

对照实验源码、输出：原生 XLA 对 x / y、exp、log、tanh、sin、x * y + 1 生成的计算指令和误差与 Pallas 完全相同，两者用的是同一套降低规则。需要更高精度时，例如要求 exp 或 log 达到 f32 的完整精度，就必须自己在 EUP 结果上追加修正步骤，XLA 和 Pallas 都不会自动做。

整数：没有整数乘法器

本小节实验源码、输出。

输入是随机的 uint32（脚本第二部分把同样的 bit 解释为 int32），结果与 CPU 上的 JAX 逐元素比较，全部一致：

加减、位运算、移位、位计数各是一条指令。但 32 bit 整数乘法没有对应的硬件指令。编译器把两个操作数拆成若干段，每段小到可以精确地放进 f32 的 24 bit 尾数，用 6 次 f32 乘法算出部分积，再转回整数、移位相加，共 33 条指令。整数除法和取余更贵，是逐位的长除法，约 330 条指令。

这个事实直接影响算法设计。例如第四章介绍的计数器式随机数生成器 Philox 以 32 bit 整数乘法为核心，在 TPU v4 上它的每次乘法都要付出这 33 条指令的开销；而只用加法、异或和移位的生成器则便宜得多。

另外两点：

• 无符号的 maximum 无法编译（failed to legalize operation 'arith.maxui'）。有符号的 maximum 是一次比较加一次选择。需要无符号比较时，可以先把两数都异或 0x80000000 再做有符号比较。
• x << y 中，脚本写的是 x << (y & 31)：移位量来自数据时，要自己保证它在 0–31 之间。

比较、选择与向量掩码

本小节实验源码、输出。

jnp.where(x > y, x, y) 是两条指令：

vgt.8x128.f32 vm0, v0, v1      # 比较结果写进掩码寄存器 vm0
vsel.8x128    v2, vm0, v0, v1  # 按掩码逐元素选择

比较指令的目的操作数是向量掩码寄存器 vm0–vm7，每个元素一个 bit；vsel 按掩码在两个 TC VREG 之间逐元素选择。比较的一侧可以是立即数，where(x > 0, x, 0.1 * x) 就是一条 vgt、一条 vmul、一条 vsel。

掩码之间的逻辑运算有专用指令 vmand、vmor、vmxor、vmneg，在 misc 槽执行。但编译器不一定用它们：where((x > 0) & (y > 0), x, y) 被改写成两次嵌套的 vsel，where((x > 0) | ~(y > 0), x, y) 则用了 vmor 和 vmxor。

TPU v4 只有 8 个向量掩码寄存器。实验在函数中先算出 count 个掩码，再依次使用它们，迫使所有掩码同时活跃：

掩码超过 8 个时，编译器把一部分掩码转存到 TC VREG 中：

vimm.8x128.s32 v31, 0                  # 先把一个 TC VREG 清零
vsel.8x128 v31, vm6, 0xffffffff, v31   # 掩码为 1 的位置写全 1，掩码转存为 32 bit 数据
...
vsmask vm6, v31                        # 需要时再从 TC VREG 恢复成掩码

每转存一个掩码要多一条 vimm、一条 vsel、一条 vsmask，还要占用一个 TC VREG；另有一个掩码被重新比较了一次（13 条 vgt）。这是第 5 节所说的掩码开销。需要大量条件的算法，应尽量让每个掩码产生后立刻用掉。研究报告 58 记录了这种转存的更多细节。

周期数：价目表的最后一栏

本小节实验源码、输出。

把前面各种运算放进第 3 节的 TC VMEM 边界，每种运算测两种大小：1 个 TC VREG（[8,128]）和 128 个 TC VREG（[1024,128]）。前者看单独算一次要多少周期，后者减去前者再除以 127，是连续处理时每个 TC VREG 的周期数。原生 XLA 与 Pallas kernel 各测一次，结果逐元素相同；两边周期数不同的只有最后注明的两行。

for name, f, used in OPS:
    timings = tpuasm_tools.compare_in_vmem(clock, f'{name}，[{rows},128]', f, [tuple(inputs[key] for key in used)] * 2)

每个数包含外壳固定的 39 个周期。sin(x) 和 int32 a // b 两行是 Pallas 的读数，XLA 分别是 150、12342 和 242、25058，其余各行两边相同（x / y 在 128 个 TC VREG 时 XLA 是 941）。

连续处理时每个 TC VREG 的周期数，由最紧的那种资源决定。表中出现了三种：

• 读写内存的槽。 一个 bundle 只有一个 vld 槽。x * 2 每个 TC VREG 读一次，每个周期完成一个；两个输入的运算每个 TC VREG 要读两次，就是 2 个周期，与中间做的是加法、乘法还是比较加选择无关。这一类运算的瓶颈不在计算，而在把数据读进寄存器。
• EUP 的发射间隔。 exp、log、tanh 各只有一条 EUP 指令，每个 TC VREG 2 个周期，正是上文测得的间隔。
• 向量 ALU 的两个槽。 指令多的运算，每个 TC VREG 的周期数接近计算指令数的一半：sqrt 13 条是 6.0，1 / y 15 条是 6.5，int32 乘法 33 条是 16.5，sin 198 条是 98.5。va0、va1 两个槽每个周期最多发射两条指令，编译器把它们几乎填满了。

单独算一个 TC VREG 时，周期数没有这样的比例：1 / y 只比 x + y 多 8 个周期，sin 的 198 条指令是 110 个周期。只有一个 TC VREG 时，决定周期数的是最长的那条依赖链，不是指令的总数。

这张表把本节开头的“价目表”补完整了。几条直接的用法：

• 整数乘法每个 TC VREG 16.5 个周期，是加法的 8 倍；整数除法是加法的 100 倍。算法里能用移位和加法代替的，就不要用乘除。
• 1 / y 与 x / y 的精度是 1 至 2 ULP，每个 TC VREG 6.5 至 7 个周期；直接用 vrcp 的近似结果只要 2 个周期。精度要求不高时（上文 vrcp 约 16 位有效），手写可以省掉迭代。
• 逐元素运算上 Pallas 与 XLA 没有差别，两边用的是同一套降低规则。逐元素的 kernel 要比 XLA 快，只能靠别处：与前后的计算融合，少搬几次数据。