jangafx/FMAG 源码: https://github.com/jangafx/FMAG

FMAG

单指令 GPU 虚拟机及工具链

一个听起来很蠢但居然能跑起来的想法：一台只有一条指令的 GPU 虚拟机，这样我就再也不用担心分支发散问题了。有时你需要运行一个并非自己编写、直到运行时才见到的程序 —— 美术在编辑器中连线而成的节点图，一段用户生成的小型表达式内容。你不想只运行一次，而是希望每个元素都运行一次：同一程序同时应用于一百万个不同的输入，并在 GPU 上执行。理想情况下，你会把它交给着色器编译器，但在大多数场景下，你无法在运行时即时编译着色器，因此程序必须以数据形式存在，由 GPU 上的某个解释器来执行。

编写该解释器最直观的方式是使用操作码表：取一条指令，查看它的标签，然后跳转到实现 add、mul 或 sin 的对应代码。在 CPU 上这几乎零开销，没人会去考虑它，但在 GPU 上却是毒药。一个 warp 内的通道以锁步方式执行，只要不同通道之间的值到达控制流，就会发生发散，而操作码解释器恰好让这种发散无处不在。某些处理函数内部会根据每条通道的值进行分支 —— 这些值并不来自硬件指令，而是编译器用软件模拟的内建值，数量远超你的想象，再加上你自己编写的复杂操作。被解释程序自身的分支则更糟糕，因为它们使程序计数器变得数据相关，而一旦每条通道的计数器不一致，之后的每一次取指都会不一致。

通常，着色器甚至在循环开始之前就已经做出了数据相关的决策，你往往希望不同通道运行不同的程序 —— 通过线程 ID 对流进行索引，此时它们甚至取到的操作码都不同。单条指令消除了所有这些问题，因为程序计数器对所有通道都按相同方式递增。决策通过条件写入（guard）而非分支来实现，这样计算出来的值永远不会进入控制流；即使不同通道运行不同程序，每一步它们执行的指令仍然相同，区别仅在于操作数。因此，即使是每通道独立的程序流也不会发生发散。

这里的示例在所有通道上运行同一个程序，但整个设计正是为了确保上述所有情况都无发散。通道计算出的任何值都不会进入控制流。如果你愿意，你也可以让每条通道运行不同的程序。

那条指令需要是什么？

这样的指令长什么样？听起来像是一个模糊的概念，而且似乎不可能实现 —— 一条单一操作居然能替代任何程序。事实证明它可以，而且所需的条件出奇地少。

它必须能进行算术运算，因为归根结底，这些程序只是对输入做数学计算。它还必须能做出决策，因为一个不能分支的程序价值不大。关键的约束是：决策不能变成真正的分支，否则就会重蹈整个想法旨在避免的发散问题，因此它必须在值的之间做选择，而无需跳转。

融合乘加（FMA）覆盖了算术部分，而在其之上放置一个 guard（防护）则覆盖了决策部分。我将这种组合称为 FMAG，代表 “fused multiply-add if greater (than zero) / guard”。

单指令计算机已经有了名字：单指令集计算机（OISC，https://en.wikipedia.org/wiki/One-instruction_set_computer），不过通常被视为一种有趣但不实用的东西。据我所知，这是它少数几个真正有用的场景之一。

仓库内容

本仓库包含 VM 和为其提供输入的整套工具链，是一个小型 C 库。你可以用着色语言编写程序，也可以直接使用 SSA CFG IR 构建器 —— 因为核心思想是在运行时将它们即时编译为数据，而库会对你所用的任何一种方式进行优化并降级为平坦的 FMAG 流。一个参考 VM 可以运行该流，以便在 CPU 上测试。

此外，还有一个命令行驱动程序，以及 —— 对于一条指令的机器而言有点荒谬 —— 一个文本汇编器和对等的反汇编器，供那些能在 FMAG 指令中思考的受虐狂使用（目前只有我一个人这么干 xD）。

大致上，程序经过的路径如下：

• 着色语言（tools/sl.c）。这是一个示例前端，旨在展示如何使用构建器，而非你必须使用的东西。因此它放在 tools 中，完全基于公开库构建，而不是库的一部分。这是一个小型类型语言，支持 f32 到 f32x4，具备函数、局部变量、if/else、函数调用以及常用运算符，编译成普通的 IR 供后续处理。

• IR 构建器（fmag_ir.h）是一个 SSA（https://en.wikipedia.org/wiki/Static_single-assignment_form）CFG（https://en.wikipedia.org/wiki/Control-flow_graph）。你可以构建函数、块、常量以及 FMAG 操作，并通过调用、分支和返回将它们连接起来。常用操作（如 add、mul、select、abs 等）作为纯 FMAG 的语法糖提供。

• 优化与降级。这是大部分工作发生的地方，老实说也是整个仓库的绝大部分 —— 因为事实证明，在单条指令周围可以做的优化数量惊人。完整的优化遍（按流水线顺序）如下，每个遍的详细说明见 OPTIMIZATIONS.md：

  • 内联：将所有调用克隆到叶子位置，直到整个模块成为一个函数；拒绝递归。
  • 三元组识别：if/else 钻石结构，如果两个分支要么返回要么汇合到 merge 点，则直接转换为 select，同时将 merge 处的 PHI 节点重写为 select，并拼接掉转发块和空块。
  • if 转换：在此之后剩下的结构化控制流都被谓词化成一个块。
  • 常量折叠：常数算术运算加上代数恒等式，如 x*1、x*0、两个分支值相同的 select，以及已知为真或假的 guard。所有这些都共享一个常量池，相同的字面量被内部化。
  • guard 剥离：剥离 (x > 0) ? 1 : 0 这种符号布尔值及其 1 - ... 补码，使得比较可以直接作为 guard 使用。
  • 乘加融合：将产生 a*1 + 0 的 producer 折叠到消费它的操作数的乘法或加法槽中。
  • 交换律规范化：将孤立的常量因子移到 b 中，并将两个常量因子合并。
  • 公共子表达式消除。
  • 死代码消除。
  • 最小/最大值定向：翻转操作数，使得保留的值位于编号较低的寄存器中，从而省去一次移动。
  • 收缩与重关联：仅在构建器以快速数学模式创建时启用，合并常量因子和加数，丢弃 x * 0。
  • 降级期间的寄存器分配：偏向结果寄存器，在每个寄存器最后使用时回收它，重新利用即将消亡的目标寄存器作为后备，最后用一个临时寄存器将结果混洗到输出寄存器中，以打破可能的循环。

• 汇编器与反汇编器（fmag_as.h）。纯文本形式，每行一个 %rd = fmag %ra, b, c, guard，加上 .def name, inputs, outputs 指令。反汇编器能将其打印出来。

• VM（fmag_vm.h）是用于测试的 CPU 解释器。它在寄存器文件上运行已编译的流，并将输出留在低的几个寄存器中。

• CLI（tools/cli.c）。cl 驱动程序可以编译着色语言文件（-c）、汇编文本程序（-a）、运行（-r）、反汇编（-d）以及检查（-i）.fmag 文件。

• 超优化器（tools/superopt.c）。寻找计算给定函数的最短 FMAG 序列，并通过参考实现证明其正确性。对于一输入函数穷举搜索，对于两输入函数采样验证。它可以在短程序上通过广度优先搜索寄存器机器状态来寻找真正的最小值，或者在较长程序上使用 STOKE（https://github.com/StanfordPL/stoke）风格的随机搜索。标准库和优化器使用的几个序列就是从它产生的。

• 绑定（odin/fmag.glsl）用于在 Odin 中使用此库。Odin 是我们在 JangaFX 使用的编程语言。

你可以直接构建 lib 源码，但如果你不想这样做，unified.py 会将整个库合并成一个 fmag.h，采用 stb（https://github.com/nothings/stb）风格，这样你只需将这一头文件放入你的项目，并在一个翻译单元中定义 FMAG_IMPLEMENTATION 来引入实现。

指令

如前所述，整个机器围绕一个操作构建：

fmag(a, b, c, d, e) = e > 0 ? fma(a, b, c) : d

当 guard e 为正时，得到 FMA 结果；否则返回 d。

fma(a, b, c) = a*b + c

为什么一条指令就能覆盖这么大范围？这归结于它两个部分各自的能力。FMA 本身就可以求任意多项式，因为串联 FMA 就是霍纳法则（https://en.wikipedia.org/wiki/Horner%27s_method）。有了多项式，再经过几次牛顿迭代（https://en.wikipedia.org/wiki/Newton%27s_method），就能得到数学库的其他部分：除法、平方根、倒数等等。

guard 是指令的另一半，一个条件判断，但本质上零开销，因为 (e > 0) ? x : y 编译为 select 而非跳转。在 GPU 上，guard 的行为就如同 select，相当于用 mix(y, x, e > 0) 或 cndmask 得到的结果。这意味着潜藏在这些程序内部的分支只消耗一次条件移动，warp 保持聚合。正是这一特性使得整个方案在此处可行。硬件永远不需要屏蔽通道或串行化路径，因为根本不存在需要串行的路径。

算术

当你将 guard 固定为一个正常数时，条件判断被取消，剩下的就是普通算术。每个常见数学运算都能自然得出：

fma(a, b, c) = a*b + c = fmag(a, b, c, 0, 1)    # 1 > 0，始终取乘积
add(a, b)    = a + b        = fma(a, 1, b)       # a*1 + b
sub(a, b)    = a - b        = fma(b, -1, a)      # b*-1 + a
mul(a, b)    = a * b        = fma(a, b, 0)       # a*b + 0
neg(a)       = -a           = fma(a, -1, 0)

当 guard 是你实际计算出来的值而非常量时，同一条指令可以在两个值之间选择：

select(t, f, g) = (g > 0) ? t : f = fmag(t, 1, 0, f, g)
relu(x)        = max(x, 0)       = fmag(x, 1, 0, 0, x)
max(a, b)      = (a-b > 0) ? a : b = fmag(a, 1, 0, b, sub(a, b))
min(a, b)      = (a-b > 0) ? b : a = fmag(b, 1, 0, a, sub(a, b))
abs(x)         = (x > 0) ? x : -x = fmag(x, 1, 0, neg(x), x)

这已经用一条指令覆盖了大部分数学库。比较和 select 提供掩码、钳位以及范围缩减所需的分段行为，而乘加提供霍纳和牛顿。在这两部分的共同作用下，几乎没有算术操作不被覆盖。

编码

每条指令由四个 32 位字组成，大小刚好适合放入一个 uvec4，可以发送给着色器。

packet
0-31:   "b operand"
32-63:  "c operand"
64-95:  "guard operand"
96-111: "d destination register"
112-127:"a multiply register"

b、c 和 guard 采用 NaN-boxing（https://craftinginterpreters.com/optimization.html#nan-boxing）。当值为静默 NaN 时，载荷是一个寄存器索引；否则就是浮点字面量本身。这带来的唯一限制是 NaN 和无穷大不能用作字面量，因为标签会变得歧义。不过实际上也没人会写一个字面量 NaN。

最后一个字包含两个裸寄存器索引：目标寄存器 d 和第一个乘数操作数 a，各 16 位，没有留下空间给 NaN 标签。目标寄存器同时也作为回退值 —— 这是一个刻意的设计选择，而非编码限制。该指令读取 r[d]，可能用 FMA 结果覆盖它，在 false 分支下则保持不变，因此结果和回退共享同一槽位。实际上，FMAG 是对寄存器的一次条件写入。这种双重角色使得每条指令的编码恰好为 16 字节，这对于通过缓冲区流式传输程序至关重要。

选择哪个槽位是寄存器

目标寄存器始终必须是寄存器，因为结果必须写回某个地方。第一个乘数操作数 a 也被强制为寄存器，理由很简单：支持 a*b 中的两个字面量因子是无意义的。如果两者都是常量，乘积会在编译时被常量折叠掉，因此最多只有一个因子是字面量，而 a 始终可以是寄存器，没有任何损失。这样一来，b、c 和 guard 可以自由选择是寄存器还是立即数。将 a 挤入剩下的唯一寄存器槽位，也使得整个 16 字节编码能够整洁地适配。

解释器

这个设计最令人满意的特性之一是解释器本身有多小。整个执行引擎只有几行 GLSL。它使用 isnan 解码三个 NaN-boxed 操作数，将最后一个字拆分成两个寄存器索引，由于 a 始终是寄存器，直接读取 r[a]，并使用无分支的 mix。

float r[16]; // 大小根据 FMAG 报告的寄存器数量调整，见下文

vec3 loadops(uvec3 o) {
    vec3 f = uintBitsToFloat(o);
    uvec3 i = o & 15u; // 应掩码至寄存器大小
    return mix(f, vec3(r[i.x], r[i.y], r[i.z]), isnan(f));
}

void exec(uvec4 op) {
    vec3 v = loadops(op.xyz);
    uint d = op.w & 15u;          // 应掩码至寄存器大小
    uint a = (op.w >> 16) & 15u; // 应掩码至寄存器大小
    r[d] = mix(r[d], fma(r[a], v.x, v.y), v.z > 0.0);
}

void main() {
    for (uint i = 0u; i < u_length; ++i)
        exec(b_code[i]);
    // 输入已加载到 r[0..]，输出从 r[...] 读回
}

这里的 mix 是布尔选择重载，而不是线性插值。它丢弃未选中的操作数。

寄存器文件 r 的大小必须与程序所需的标量寄存器数量一致，即 FMAG 写入头部的 regs 字段（每个向量就是那么多标量槽位）。由于 r 是局部数组，它直接影响着色器内部的寄存器压力，因此不能过大。IR 优化器和降级遍努力最小化这个数量，通常很小，但你仍然必须在着色器编译时确定一个大小。一种方法是为几个 2 的幂次大小专门化解释器，并在运行时选择能容纳程序的最小那个。另一种方法是选择一个够大的固定大小并接受压力。对大多数程序来说，16 是一个不错的折中。

RDNA 上的内循环

为了验证这一切是否真的兑现了承诺，我通过 RGA 对 RDNA3（gfx1100）运行了上述解释器，检查 GPU 将要执行的 ISA。内循环紧凑得超出预期：只有一个循环分支，没有别的。没有发散，没有复杂控制流，寄存器读写表现为普通的索引移动。

_L1:
    s_buffer_load_b128 s[12:15], s[4:7], s3 ; 统一取指令
    s_and_b32 s10, s15, 15              ; d = da & 15
    s_bfe_u32 s0, s15, 0x40010          ; a = (da >> 16) & 15
    s_and_b32 s1, s12, 15               ; b 的索引
    s_and_b32 s11, s13, 15              ; c 的索引
    s_and_b32 s2, s14, 15               ; guard 的索引
    s_mov_b32 m0, ... / v_movrels_b32 v.., v0 ; r[idx] 读取
    v_cmp_u_f32 ...                     ; isnan 判断 b, c, guard
    v_cndmask_b32 ...                   ; 操作数选择
    v_fmac_f32 v19, v17, v18           ; r[a]*b + c
    v_cmp_lt_f32 vcc, 0, v20           ; guard > 0
    v_cndmask_b32 v16, v16, v19, vcc   ; guard 选择
    v_movreld_b32 v0, v16              ; r[d] = 结果
    s_cbranch_scc1 _L1                 ; 循环

指令取指是标量加载 s_buffer_load_b128，因为整个 warp 的程序计数器相同，所以整个 warp 只取一条指令。