Stax 编译栈（二）：自动微分与降级

上一节：第一节　引言、特化、图捕获与图变换　｜　目录　｜　下一节：第三节　调度、融合与 CPU 代码生成

文中参考文献编号 [n] 见第六节；章节号、定理号与习题号全章统一编号，跨节引用照常有效。

5 自动微分与激活重物化

图变换完成后，计算图还要经过自动微分层，才交给 Stax。这一层有两项任务：把 Python 层的调用改写为算子级的图，以及在训练时构造反向计算并决定前向须保存哪些中间值。

5.1 算子级的图

前端记录的是 Python 层的调用，例如 operator.add 或 F.gelu。仅凭被调用的函数，后端无法确定它最终分派到哪个算子重载，也无法得知组合算子的内部构成。Stax 只接受算子级的图，因此即使区域不需要求导，也要在分派器层重新追踪一次，记录每次调用实际到达的算子重载，只是不构造反向部分。推理与训练因而使用同一种图，后端只需面对一种输入。代码清单 5.1 为例 1.1 经这一层后的图。

代码清单 5.1　例 1.1 的算子级图

graph(
    %primals_1 : [num_users=1] = placeholder[target=primals_1]
    %primals_2 : [num_users=1] = placeholder[target=primals_2]
    %add shape=(64, 128) : [num_users=1] = call_function[target=tensorplay.ops.tp.add.Tensor](args = (%primals_1, %primals_2), kwargs = {alpha: 1})
    %relu shape=(64, 128) : [num_users=1] = call_function[target=tensorplay.ops.tp.relu.default](args = (%add,), kwargs = {})
    %mul shape=(64, 128) : [num_users=1] = call_function[target=tensorplay.ops.tp.mul.Scalar](args = (%relu, 2), kwargs = {})
    %sum_0 shape=(64,) : [num_users=1] = call_function[target=tensorplay.ops.tp.sum.dim_IntList](args = (%mul, [-1], False), kwargs = {dtype: tensorplay.undefined})
    return (sum_0,)
)

与代码清单 3.1 相比，operator.add 成为算子重载 add.Tensor 并补全了默认参数 alpha=1，张量与标量的乘法被分派到 mul.Scalar，张量方法 sum(-1) 成为 sum.dim_IntList。每个节点都指向唯一确定的算子重载，后端可据此查找降级规则。重新追踪的过程中还完成两项变换：不属于后端基的算子按分解表展开（4.8 节），原地修改被函数化（4.9 节）。

5.2 反向模式自动微分

设前向计算是一串函数的复合

记 $J_k = \partial \phi_k / \partial z_{k-1}$ 为第 $k$ 步的雅可比矩阵。由链式法则，

反向模式（reverse mode）自动微分从右向左计算这一乘积：

其中 $\bar z_k = \partial L / \partial z_k$ 称为伴随（adjoint）。每一步只需计算一个向量与雅可比矩阵转置的乘积（vector-Jacobian product），不必构造雅可比矩阵本身。对一般的有向无环图，伴随按拓扑逆序累加：一个值的伴随等于其各使用者传回的贡献之和。

上面的链式写法便于理解，但一般的计算是有向无环图。下面在直线程序的模型下精确地叙述反向模式的代价。设 $L$ 由如下程序计算：

其中 $a_k, b_k < k$，$\varphi_k$ 为 $+$、$-$、$\times$、$\div$ 之一，或是只用到 $v_{a_k}$ 的一元初等函数 $\psi$。称 $W = N - d$ 为程序的运算次数（$d$ 为输入维数，不要与上面链的长度 $n$ 混淆）。

定理 5.1（梯度的计算代价 [9, 38]）　设每个一元函数 $\psi$ 的导数 $\psi’(a)$ 可由 $a$ 与 $\psi(a)$ 以至多 $c_\psi$ 次运算求出，$c_\Omega = \max_\psi c_\psi$。则反向模式计算全部 $\nabla_x L$（连同 $L$ 本身）至多需要

$$\bigl(1 + \max(4,\ c_\Omega + 2)\bigr)\, W$$

次运算。对 $\exp$、$\log$、$\tanh$、$\mathrm{sigmoid}$、$\sqrt{\cdot}$，有 $c_\Omega \le 2$，常数不超过 5，与输入维数 $d$ 无关。

证明　算法分两趟。前向趟按 $k = d+1, \dots, N$ 计算并保存每个 $v_k$，共 $W$ 次运算。反向趟令 $\bar v_N = 1$，其余 $\bar v_k = 0$，再按 $k = N, \dots, d+1$ 的次序执行

正确性：断言处理完所有 $k > j$ 之后 $\bar v_j = \partial L / \partial v_j$（把 $v_j$ 视为自变量、其后的计算不变）。$v_j$ 只通过以它为操作数的那些 $v_k$（$k > j$）影响 $L$，由多元链式法则，$\partial L/\partial v_j = \sum_{k:, j \in {a_k, b_k}} (\partial L/\partial v_k)(\partial \varphi_k/\partial v_j)$。按 $k$ 递减的次序处理时，每个 $\bar v_k$ 在被使用之前已经接收了它的全部贡献（归纳假设），右端的每一项恰好被累加一次。取 $j \le d$ 即得 $\nabla_x L$。

代价：逐类计算反向趟中第 $k$ 步的运算次数。加减法：$\bar v_a \mathrel{+}= \bar v_k$，$\bar v_b \mathrel{\pm}= \bar v_k$，共 2 次。乘法 $v_k = v_a v_b$：$\bar v_a \mathrel{+}= \bar v_k v_b$，$\bar v_b \mathrel{+}= \bar v_k v_a$，共 4 次。除法 $v_k = v_a / v_b$：先算 $t = \bar v_k / v_b$，再 $\bar v_a \mathrel{+}= t$，$\bar v_b \mathrel{-}= t, v_k$（因为 $\partial v_k/\partial v_b = -v_k/v_b$），共 4 次。一元函数：求 $\psi’(v_a)$ 用 $c_\psi$ 次（可以用到已保存的 $v_k = \psi(v_a)$），再乘、再加，共 $c_\psi + 2$ 次。每一步至多 $\max(4, c_\Omega + 2)$ 次，反向趟合计至多 $\max(4, c_\Omega+2),W$，加上前向趟即得所述的界。对所列的函数，$\exp’ = \exp$（0 次），$\log’(a) = 1/a$（1 次），$\tanh’ = 1 - v^2$、$\mathrm{sigmoid}’ = v(1 - v)$、$(\sqrt{a})’ = 1/(2v)$（各 2 次）。$\square$

证明中的两趟结构同时说明了反向模式的存储代价：反向趟要用到前向趟保存的全部 $v_k$。张量级的计算图中，每个节点是一个张量算子，上述论证逐算子成立：只要每个算子的向量–雅可比积的代价不超过其前向代价的常数倍（逐元素算子如上；矩阵乘 $C = AB$ 的反向是两次同样规模的矩阵乘），常数便与输入维数无关。Baur 与 Strassen [38] 在只计乘法与除法（非标量运算）的代数复杂性计价模型下给出了常数 3 的版本。

定理 5.1 说明，梯度的计算代价与输入维数无关，这是反向模式适用于参数众多的神经网络的原因。但反向模式另有一项代价：$J_k$ 一般依赖于 $z_{k-1}$，例如 $\phi_k = \tanh$ 时 $J_k = \mathrm{diag}(1 - \tanh^2 z_{k-1})$，因此反向计算须能取得前向的中间值。若全部保存，存储量与前向的步数成正比，这正是训练中激活值占用大量显存的原因。

5.3 联合图

TensorPlay 不另行推导反向图，而是由 TPX 引擎实际执行一次反向，并记录其分派的算子，过程见算法 5.1。

算法 5.1　联合图的追踪

输入：捕获图 G，参数与缓冲区 P，样例输入 X
输出：联合图 J
1  将 P 提升为显式输入，使 G 成为纯函数 F(P, X)
2  在分派器层的追踪模式下执行 F，记录各算子重载
3  为每个可微输出 y_i 创建上游梯度的占位符 tangent_i
4  调用 autograd.grad(y, primals, tangents)，记录引擎执行反向时分派的全部算子
5  为前向与反向记录同一随机数状态，使随机算子在重算时得到相同的随机数
6  返回 J

所得的联合图（joint graph）同时包含前向与反向，反向节点直接引用前向节点的值。联合图把"保存哪些值"这一问题从运行时提前到了编译时：此时可以看到整个反向计算，从而对每个前向值作出保存或重算的决定。

例 5.1　考虑

import tensorplay.nn.functional as F

def f(x, W):
    return F.gelu(x @ W).sum()

x = tp.randn(32, 64)
W = tp.randn(64, 64, requires_grad=True)
cf = tp.compile(f, strict_native=True)
cf(x, W).backward()

其联合图如代码清单 5.2 所示，行首的 [FW] 与 [BW] 表示节点属于前向或反向。

代码清单 5.2　例 5.1 的联合图（省略了部分元数据）

[FW] %primals_1 : [num_users=2] = placeholder[target=primals_1]
[FW] %primals_2 : [num_users=1] = placeholder[target=primals_2]
[BW] %tangents_1 : [num_users=1] = placeholder[target=tangents_1]
[FW] %matmul shape=(32, 64) = call_function[target=tp.matmul.default](args = (%primals_1, %primals_2))
[FW] %gelu shape=(32, 64) = call_function[target=tp.gelu.default](args = (%matmul, 'none'))
[FW] %sum_0 shape=() = call_function[target=tp.sum.default](args = (%gelu,))
[BW] %expand shape=(32, 64) = call_function[target=tp.expand.default](args = (%tangents_1, [32, 64]))
[BW] %gelu_backward shape=(32, 64) = call_function[target=tp.gelu_backward.default](args = (%expand, %matmul, 'none'))
[BW] %t shape=(64, 64) = call_function[target=tp.t.default](args = (%primals_1,))
[BW] %matmul_0 shape=(64, 64) = call_function[target=tp.matmul.default](args = (%t, %gelu_backward))
[FW] return (sum_0, matmul_0)

primals_1 为 $x$，primals_2 为 $W$，tangents_1 为 $\bar L$。反向部分与 5.2 节的递推一一对应：$\bar L$ 经 sum 的伴随扩展为 $32\times64$；gelu_backward 计算 $\bar h = \bar g \odot \mathrm{GELU}’(h)$，其中 $h = xW$；最后 $\bar W = x^{\mathsf T} \bar h$。$x$ 不需要梯度，图中没有计算 $\bar x$ 的节点。注意 gelu_backward 读取的是 $h$ 而不是 gelu 的输出，这决定了 5.5 节中的划分结果。

5.4 保存与重算的最小割模型

联合图须划分为前向图与反向图。反向所需的前向值有两种取得方式：在前向中保存，或在反向中由已保存的值重新计算。全部保存的计算量最小，但存储最大；只保存输入、其余全部重算的存储最小，但计算量可能翻倍。以重算换存储的思想见于梯度检查点 [10]；在一般的代价模型下，求最优的保存与重算方案是 NP 难问题，Checkmate [11] 将其表述为整数线性规划求解。Stax 采用一个可以在多项式时间内精确求解的简化模型：只决定保存哪些值，并以保存的字节数为代价。

先刻画可行的保存方案。设前向图为 $G_F$，反向读取的前向值集合为 $B$，不宜在反向中重算的值集合为 $U$（定义见下文）。

定义 5.1（可行保存集）　前向值的集合 $C$ 称为可行的，若对每个 $v \in B$，从 $U$ 中的任何值到 $v$ 的每条路径都经过 $C$ 中的某个值（路径的端点也算在内）。

直观地说，反向从 $C$ 出发、只执行廉价的运算，就能重新得到 $B$ 中的每个值，而不必触及 $U$ 中的任何值。下面把这一条件转化为流网络中的割。

定义 5.2（重物化网络）　对前向的每个值 $v$，引入顶点 $v_{\mathrm{in}}$、$v_{\mathrm{out}}$ 与一条容量为 $c(v)$ 的边 $v_{\mathrm{in}} \to v_{\mathrm{out}}$；前向的数据依赖 $u \to v$ 对应容量为无穷的边 $u_{\mathrm{out}} \to v_{\mathrm{in}}$。再引入源点 $s$ 与汇点 $t$：对 $v \in U$，加入无穷容量的边 $s \to v_{\mathrm{in}}$；对 $v \in B$，加入无穷容量的边 $v_{\mathrm{out}} \to t$。

把每个值拆成两个顶点，是把顶点上的代价转化为边上的容量的标准做法：割断边 $v_{\mathrm{in}} \to v_{\mathrm{out}}$ 即表示保存 $v$。

命题 5.1　重物化网络中容量有限的 $s$–$t$ 割与可行保存集一一对应，且割的容量等于 $\sum_{v \in C} c(v)$。

证明　容量有限的割只能割断形如 $v_{\mathrm{in}} \to v_{\mathrm{out}}$ 的边，记被割断的边对应的值集合为 $C$。网络中的每条 $s$–$t$ 路径都具有形式 $s \to u_{\mathrm{in}} \to \cdots \to v_{\mathrm{out}} \to t$，其中 $u \in U$、$v \in B$，中间依次经过前向中一条从 $u$ 到 $v$ 的路径上各值的拆分边。割断全部 $s$–$t$ 路径，等价于前向中每条从 $U$ 到 $B$ 的路径上都有一个值属于 $C$，即 $C$ 可行。反之，对可行的 $C$ 割断相应的边即得一个有限割。割的容量为所割断各边的容量之和。$\square$

于是，求代价最小的可行保存集，就是求网络的最小 $s$–$t$ 割。

定理 5.2（最大流最小割定理 [12]）　在有向网络中，从 $s$ 到 $t$ 的最大流的流量等于 $s$–$t$ 最小割的容量。设 $f$ 为最大流，则在残量网络中自 $s$ 可达的顶点集合 $S$ 与其补集构成一个最小割。

Stax 以 Edmonds–Karp 算法 [13] 求最大流：每次以广度优先搜索在残量网络中寻找最短的增广路径，沿路径推送其瓶颈容量，直到不存在增广路径。该算法的时间复杂度为 $O(|V||E|^2)$。求得最大流后，再以一次广度优先搜索得到 $S$，跨越 $S$ 与其补集的拆分边即为割边。

余下的问题是 $U$ 与容量 $c(v)$ 的取法。一个值被列入 $U$，有以下三种情形：

1. 它不是廉价运算的结果，例如矩阵乘、卷积，重算的计算代价过高；
2. 它在反向中无论如何都须写入内存，即它的某个反向使用者无法与它融合，此时重算不能节省访存；
3. 它是一次规约的结果，判据为输出字节数的 4 倍小于输入字节数之和。重算规约需要重读全部输入，代价与规约本身相同。

容量取

其中 $b(v)$ 为张量的字节数；$d(v)$ 为 $v$ 与其第一个反向读取者之间相隔的前向步数，截断到 $[1, 100]$；若 $v$ 在前向中本就须写入内存，$\beta(v) = 1$，否则 $\beta(v) = 2$。字节数是保存的直接代价；距离因子使存活时间较长的值更倾向于重算，因为它在整个存活期间都占用存储，近似地反映了它对峰值内存的贡献；$\beta$ 反映的是，保存一个原本可在前向中融合掉的值，需要额外写出一次内存。后两个因子是启发式的，它们把最小割模型未能直接表达的时间维度折算进了容量。

此外有两条补充规则：同时被近处与远处读取的值，以及位于很长的可融合链末端的值，强制保存，以免重算把整条链引入反向。若不存在有限割，即某个反向所需的值既不能保存也不能重算，则退回全部保存的划分。

5.5 例 5.1 的划分

例 5.1 的前向有三个值：matmul（即 $h$）、gelu 与 sum_0。按定义 5.2 构造的网络如表 5.1 所示。

表 5.1　例 5.1 的重物化网络

matmul 的容量为 $8192 \times 1.1 = 9011.2$，取整为 9011。gelu 不被任何反向节点读取，$d$ 取上限 100；它可以与后续的求和融合，$\beta = 2$，故容量为 $2 \times 8192 \times 1.1^{100} \approx 2.26\times10^8$。从 $s$ 到 $t$ 只有一条路径 $s \to \text{matmul_in} \to \text{matmul_out} \to t$，其瓶颈为 9011，因此最大流与最小割均为 9011，割边为 matmul_in $\to$ matmul_out，即只保存 $h$。

这一结果符合直觉。$h$ 是矩阵乘的输出，无论是否保存都须写入内存，保存它几乎不增加代价；gelu 的输出与其内部的 erf 一个也没有保存，反向所需的 $\mathrm{GELU}’(h)$ 由 $h$ 重新计算。GELU [14] 定义为 $\mathrm{GELU}(h) = h,\Phi(h)$，其中 $\Phi(h) = \tfrac12\bigl(1 + \mathrm{erf}(h/\sqrt2)\bigr)$ 为标准正态分布函数，其导数为

反向的逐元素内核如代码清单 5.3（CPU）与代码清单 5.4（GPU）所示。

代码清单 5.3　例 5.1 的反向逐元素内核（CPU，变量已重命名）

auto h    = Vectorized<float>::loadu(in_ptr1 + x0, 16);
auto cdf  = 0.5f * (1.0f + (h * 0.7071067811865476f).erf());   // Φ(h)，在反向中重算
auto pdf  = 0.3989422804014327f * (-0.5f * h * h).exp();        // φ(h)
auto grad = upstream * (cdf + h * pdf);                          // GELU'(h) 乘以上游梯度
grad.store(out_ptr0 + x0);

代码清单 5.4　例 5.1 的反向逐元素内核（GPU，Triton，节选；完整内核见 9.3 节）

tmp0  = tl.load(in_ptr0 + (0))                          # 上游梯度：loss 是标量，只有一个元素
tmp1  = tl.broadcast_to(tmp0, [XBLOCK])
tmp2  = tl.load(in_ptr1 + (x0), xmask)                  # 前向保存的 h
tmp5  = libdevice.erf(tmp2 * 0.7071067811865476)
tmp9  = 0.5 * (1.0 + tmp5)                              # Φ(h)，在反向中重算
tmp15 = 0.3989422804014327 * libdevice.exp(-0.5 * (tmp2 * tmp2))   # φ(h)
tmp17 = tmp9 + tmp2 * tmp15                             # GELU'(h) = Φ(h) + h·φ(h)
tl.store(out_ptr0 + (x0), tmp1 * tmp17, xmask)

两种设备上的内核结构相同：只读入 $h$ 与一个标量上游梯度，写出 $ar h$。

erf 在反向内核中由 $h$ 重新计算，并与高斯密度、上游梯度的乘法融合于同一循环。重算的计算量是每个元素一次 erf 与一次 exp，而换得的是前向少写、反向少读两个 $32\times64$ 的张量。

图 5.1　例 5.1 的联合图与最小割。只保存 $h$，erf 在反向的循环中重算。

5.6 运行时的组合

划分所得的前向图与反向图分别交由 Stax 编译，再以一个 autograd.Function 的子类组合，接入 TPX 的计算图：其 forward 调用编译后的前向，将保存的值存入上下文；其 backward 取出保存的值，与上游梯度一起交给编译后的反向。反向图在第一次调用 backward 时才编译，只做推理的调用因此不付出编译反向的代价。

这一做法与 PyTorch 2 [15] 中 AOTAutograd 的最小割划分相同。其成立的前提是后端能够把重算的逐元素链融入反向原本就要执行的循环；若每一步重算都是一个独立的内核，重算的中间值仍要写回内存，以计算换存储的交换将得不偿失。第 6 节说明 Stax 如何满足这一前提。

5.7 前向模式与高阶导数

5.2 节只讨论了反向模式。对 $f: \mathbb R^n \to \mathbb R^m$，雅可比矩阵 $J \in \mathbb R^{m\times n}$ 有两种按列或按行提取的方式。

前向模式（forward mode）沿计算的方向传播切向量：对每个中间值 $z_k$ 维护 $\dot z_k = \partial z_k / \partial x \cdot v$，递推为 $\dot z_k = J_k \dot z_{k-1}$，一次前向计算同时得到 $f(x)$ 与雅可比–向量积 $Jv$。实现上，每个张量携带一个切向量，组成对偶数 $(z, \dot z)$，每个算子按其前向导数公式更新切向量，不需要记录计算图，也不需要保存中间值。

命题 5.2（两种模式的代价）　设计算 $f$ 的代价为 $W$。前向模式求一个 $Jv$、反向模式求一个 $u^{\mathsf T}J$ 的代价都是 $O(W)$。因此求完整的 $J$，前向模式需要 $n$ 次、反向模式需要 $m$ 次，代价分别为 $O(nW)$ 与 $O(mW)$；此外，反向模式需要保存前向的中间值，前向模式不需要。

证明　前向模式中，每个算子的切向量更新是其前向导数公式在一个向量上的求值，代价与算子本身同阶，总和为 $O(W)$；反向模式由定理 5.1。$J$ 的第 $i$ 列为 $J e_i$，第 $j$ 行为 $e_j^{\mathsf T} J$，分别需要 $n$ 次与 $m$ 次。存储的差别来自递推的方向：前向模式的 $\dot z_k$ 只依赖 $\dot z_{k-1}$ 与 $z_{k-1}$，后者此时恰好可用；反向模式从 $\bar z_n$ 开始，须回溯到 $z_0$。$\square$

训练中的损失是标量，$m = 1$，这是反向模式占主导的原因。前向模式适用于输出多、输入少的情形，例如对少数几个超参数求灵敏度，或计算方向导数。

高阶导数由两种模式的组合得到。黑塞矩阵与向量的乘积 $Hv = \nabla(\nabla f \cdot v)$ 可以"前向套反向"地计算：在反向模式求梯度的计算上再做一次前向模式，代价为 $O(W)$，与 Pearlmutter 的结论 [34] 一致；也可以"反向套反向"：以 create_graph=True 求出梯度，使梯度本身带有计算图，再对它求一次梯度。

TensorPlay 的即时执行支持上述各种形式：tensorplay.autograd.forward_ad 提供对偶数（make_dual、unpack_dual），tensorplay.autograd.functional 提供 jvp、vjp、hvp、jacobian 与 hessian；jvp 默认以前向模式计算。

编译路径只支持一阶反向。5.6 节的组合方式把编译后的反向图包装为一个 autograd.Function，反向图本身是编译产物，没有可供再次求导的计算图。这一限制如何呈现给用户，本身就是一个语义契约的问题。

一种做法是什么都不做：一阶梯度不接入计算图，二阶求导便在一个看似无关的地方失败，报错为梯度张量"不需要梯度也没有 grad_fn"，用户无从得知原因是编译。TensorPlay 选择在需要时显式拒绝：

代码清单 5.5　编译区域的反向在需要记录计算图时的包装（节选）

if recording:                                    # 调用方以 create_graph=True 求梯度
    class CompiledFunctionBackward(Function):
        @staticmethod
        def forward(double_ctx, *args):
            return _call(compiled_bw_box[0], list(args[:num_backward_inputs]))

        @staticmethod
        def backward(double_ctx, *args):
            raise RuntimeError(
                "Compiled regions do not support second-order gradients "
                "(double backward). Run the function without compilation "
                "to compute higher-order derivatives."
            )

    grads = iter(CompiledFunctionBackward.apply(*inputs, *primals_for_grad))
else:
    grads = iter(_call(compiled_bw_box[0], inputs))

当梯度模式开启、即调用方要求一阶梯度本身可微时，编译后的反向被再包一层 Function：它的前向就是编译后的反向，结果与原来逐位相同；它的反向不计算任何东西，只抛出一条说明原因与出路的错误。包装的输入除了反向图的实参，还有前向的原始输入 primals_for_grad，这使一阶梯度在计算图中连接到原始输入上，二阶求导沿图回溯时必然到达这个反向并触发错误，而不是因为"没有 grad_fn"在别处失败。于是一阶梯度的值正确且 requires_grad 为真，对它再求导时得到明确的 RuntimeError。

这与 11.3 节严格模式的取向相同：编译器做不到的事，应当在做不到的地方明确失败，而不是给出一个看似成功、实则不完整的结果，或在远处以无关的错误失败。

要真正支持二阶导数，联合图须再做一次追踪，把反向图本身作为前向，求出它的反向；这在 AOT 的框架内是可行的，代价是编译时间与图规模的成倍增长。在此之前，计算黑塞矩阵、黑塞–向量积等高阶导数时，应对未编译的函数调用 tensorplay.autograd.functional 中的接口。

5.8 存储预算下的检查点

5.4 节的最小割模型回答的是：不限存储时，保存哪些值能使反向的重算代价最小。另一个同样基本的问题是反过来的：存储有一个硬上限时，最少要重算多少。本节在最简单也最典型的结构，即一条链上回答它。

设前向是 $l$ 步的链 $z_k = \phi_k(z_{k-1})$，每个状态占一个单位的存储，每一步的代价为 1。反向的第 $k$ 步（计算 $\bar z_{k-1}$）需要 $z_{k-1}$。全部保存需要 $l$ 个单位的存储而不重算；只保存 $z_0$ 时存储为常数，但为取得每个 $z_{k-1}$ 都要从头重算，总代价为 $\Theta(l^2)$。检查点（checkpoint）介于两者之间：只保存少数状态，反向时从最近的检查点重算出所需的状态。

定理 5.3（均匀分段 [10]）　把链等分为 $c$ 段，在每段的起点保存检查点。反向时逐段处理：从该段的检查点重算并保存段内全部状态，再完成这一段的反向。则存储峰值不超过 $c + l/c$，每一步至多被计算两次。取 $c = \sqrt l$ 时存储不超过 $2\sqrt l$，前向的总代价至多为 $2l$。

证明　前向趟只保存 $c$ 个检查点。反向处理某一段时，额外保存该段的 $l/c$ 个状态，处理完即释放，故峰值不超过 $c + l/c$（段的起点本身是检查点，实际还少 1）。每一步在前向趟中计算一次，在其所在段被处理时重算一次。$c + l/c$ 在 $c = \sqrt l$ 时取最小值。$\square$

tensorplay.utils.checkpoint_sequential 实现的正是这一方案：用户把顺序模型分为若干段，每段只保存输入。存储从 $\Theta(l)$ 降为 $\Theta(\sqrt l)$，代价是多做一次前向。允许重算多于一次时，可以做得更好。

定义 5.3（反演次序）　设有 $s$ 个检查点槽位，开始时一个槽位存有 $z_0$。一个反演次序是如下操作的序列：从某个检查点出发向前计算若干步（不保存中间状态，只保留当前状态），把当前状态写入一个空闲槽位，释放一个槽位，以及当 $z_{k-1}$ 为当前状态或在槽位中、且 $\bar z_k$ 已经求出时执行反向的第 $k$ 步。若全部反向步骤按 $k = l, \dots, 1$ 执行完毕，且每个前向步骤 $\phi_k$ 至多被计算 $t$ 次，称它是重复数为 $t$ 的反演次序。记 $\beta(s, t)$ 为存在这样的反演次序的最大链长。

定理 5.4（二项式检查点 [44, 45]）　$\beta(s, t) \ge \dbinom{s + t}{s}$，且存在达到这一长度的递归构造。Griewank 进一步证明了等号成立 [44]。

证明　先证递推式 $\beta(s, t) \ge \beta(s, t-1) + \beta(s-1, t)$，再验证边界。设链长 $l = m + n$，$m = \beta(s, t-1)$，$n = \beta(s-1, t)$。从 $z_0$ 前进 $m$ 步得到 $z_m$，写入一个空闲槽位；此时前 $m$ 步各被计算了一次。后 $n$ 步构成一条以 $z_m$ 为起点的链，用其余 $s - 1$ 个槽位（$z_0$ 占用的那个保留不动）按重复数 $t$ 完成它的反演，这由 $n = \beta(s-1, t)$ 保证；反演完成时已求出 $\bar z_m$。释放 $z_m$ 的槽位，再以 $s$ 个槽位反演前 $m$ 步：它们已被计算过一次，剩余的重复数为 $t - 1$，由 $m = \beta(s, t-1)$ 保证。

边界：$t = 0$ 时不能做任何前向计算，只有 $z_0$ 可用，只能执行反向的第 1 步，故 $\beta(s, 0) \ge 1$。$s = 1$ 时只有 $z_0$ 一个检查点，从 $z_0$ 出发前进 $k - 1$ 步即可执行反向的第 $k$ 步，依次取 $k = l, \dots, 1$，第 1 步被计算 $l - 1$ 次，故 $\beta(1, t) \ge t + 1$。两者都等于 $\binom{s+t}{s}$ 在相应边界上的值，而 $\binom{s+t}{s} = \binom{s+t-1}{s} + \binom{s+t-1}{s-1}$ 与递推式同形，由归纳法得证。$\square$

等号（即没有更长的链能以 $s$ 个槽位、重复数 $t$ 反演）的证明较长，其核心是一个交换论证：任一反演次序都可以改写为上面这种递归嵌套的形式，而不增加槽位数与重复数 [44]。由 $\binom{2s}{s} \approx 4^s / \sqrt{\pi s}$，取 $s = t = O(\log l)$ 即可反演长度为 $l$ 的链：存储与重算的倍数都只按 $\log l$ 增长。这与均匀分段的 $\sqrt l$ 存储形成对照。

两点限制值得注意。第一，以上结论针对链；在一般的有向无环图上，给定存储上限时最小化重算是困难的：以卵石博弈表述的允许重算的版本是 PSPACE 完全的 [57]，Checkmate 因而以整数线性规划求解 [11]。第二，Stax 的分区器只做 5.4 节的最小割，不考虑存储预算：partition_min_cut 接受 memory_budget 参数，但目前不使用它。存储预算下的检查点由用户在模型层面通过 tensorplay.utils.checkpoint 指定。

6 降级与循环级中间表示

自动微分层交付的算子级图进入 Stax 后，首先被降级（lowering）为循环级中间表示。这一层是 Stax 的核心：它决定每个值以什么形式存在，是一段可以内联到使用者中的计算，还是一块写入内存的缓冲区。

6.1 图的解释执行

降级器 GraphLowering 是一个图解释器，按拓扑序逐个"执行"节点，只是执行的结果是中间表示节点而不是张量。形状推导与布局决策在只携带元数据的伪张量上进行，不访问真实数据。三类节点的处理方式如下：

1. placeholder：按样例输入的形状、步长与设备建立一个输入缓冲区，其布局不可更改；
2. call_function：查找该算子的降级规则并调用，得到一个中间表示节点。没有降级规则的算子成为回退内核，原样交还框架执行；矩阵乘、卷积等算子由模板或外部库实现（9.8 节）；
3. output：图的输出必须写入内存，对每个输出调用实体化（6.6 节），并确定其布局。

6.2 降级规则

逐元素算子的降级规则大多由一个公共的辅助函数 pointwise 实现，代码清单 6.1 为其节选。

代码清单 6.1　逐元素降级的辅助函数（节选，删去了注释）

def pointwise(fn, *inputs, val=None, out_dtype=None):
    if val is None:
        val = next((x for x in inputs if hasattr(x, "get_size")), None)
    ...
    size = tuple(_extent(s) for s in val.get_size())
    dtype = val.get_dtype()
    device = val.get_device()
    tensor_inputs = [x for x in inputs if hasattr(x, "get_size")]
    if len(tensor_inputs) > 1:
        target = functools.reduce(
            broadcast_symbolic_shapes, (x.get_size() for x in tensor_inputs), ()
        )
        size = tuple(_extent(s) for s in target)
    loaders = [
        _broadcast_loader(x, size) if hasattr(x, "get_size")
        else as_value_node(x, dtype, device).make_loader()
        for x in inputs
    ]

    def inner(index):
        return fn(*[load(index) for load in loaders])

    return Pointwise.create(
        device=device,
        dtype=dtype if out_dtype is None else out_dtype,
        inner_fn=inner,
        ranges=size,
    )

函数的逻辑分为三步：由各输入的形状按广播规则求出输出的迭代范围；为每个输入取得一个读取函数（loader），张量输入的读取函数按广播规则改写下标，标量输入成为常量；最后构造一个逐元素节点，其计算函数依次调用各读取函数并对结果施加 fn。整个过程既不分配内存，也不生成循环。

单目算子的规则由一张名称到 ops 方法的映射批量注册，如 relu 映射为 ops.relu。

降级规则共近三百条，按所产生的中间表示可以分为表 6.1 中的几类。

表 6.1　降级规则的分类

前四类都是定义 6.1 的直接实例，后几类则扩展了它。

间接读取的下标由数据决定。x[idx] 的第 $i$ 个元素是 $x(\mathit{idx}(i))$，读取地址不再是 $i$ 的仿射函数，而是先读取 idx、再以读得的值为下标。计算函数中以 ops.indirect_indexing 表示这一步：它把一个运行时的值转换为下标变量。由于这个值在编译时未知，生成的代码须处理负下标（按 Python 语义，$-1$ 表示最后一个元素）并检查越界，除非能证明下标落在合法范围内（6.8 节）。

散射写入是间接读取的对偶：y[idx] = v 的写入地址由数据决定。若两个位置的 idx 相同，它们写入同一个元素，结果依赖于执行顺序。对于覆盖语义，重复下标的结果在即时执行中本就没有确定的约定，编译后同样不作保证；对于累加语义（scatter_add、index_put(accumulate=True)），以原子加实现，由于加法满足交换律与结合律（在精确算术下），并发的原子加与任何串行顺序给出相同的结果。

拼接是一个分段函数。设沿第 0 维拼接 $x_1, \dots, x_m$，偏移为 $o_1 = 0 < o_2 < \cdots < o_{m+1}$，则

Stax 把它降级为一个逐元素节点，其计算函数对每个 $k$ 做一次以 $o_k \le i < o_{k+1}$ 为掩码的读取，再以条件选择合并。拼接因此不复制数据，可以与其使用者融合；代价是每个元素都要计算 $m$ 个条件。

类型提升与降级规则在同一处声明。每个逐元素算子带有一种提升方式：默认方式取各输入的公共类型；除法与 sqrt 等声明为整数先转为浮点；比较运算声明为结果恒为布尔。在 GPU 代码中，半精度（float16、bfloat16）的运算以 float32 进行，只在读取时升格、写回时降格，这与 6.5 节的误差分析一致：中间结果若以半精度舍入，误差界中的 $u$ 将由 $2^{-24}$ 增大到 $2^{-11}$ 或 $2^{-8}$。每个降级结果最后被转换为追踪时记录的数据类型，以保证与即时执行的类型一致。

6.3 张量作为函数

定义 6.1（逐元素节点与规约节点）　逐元素节点由迭代范围 $R$ 与函数 $\mathit{inner_fn}: R \to \mathbb V$ 组成，表示对每个下标 $i \in R$ 计算一个值。规约节点由迭代范围 $R$、规约范围 $R_r$、合并运算 $\oplus$ 与函数 $\mathit{inner_fn}: R \times R_r \to \mathbb V$ 组成，表示

$$y(i) = \bigoplus_{r \in R_r} \mathit{inner\_fn}(i, r),\qquad i \in R.$$

其中 $\mathbb V$ 为标量值的集合，$\mathit{inner_fn}$ 中的运算通过 ops 接口表达。

这一表示把张量看作从下标空间到值的函数，与 Halide [16] 中的函数（Func）是同一种观点：它描述每个元素如何计算，而不规定计算的顺序与存储。在这一观点下，逐元素运算之间的融合不需要任何专门的算法。

性质 6.1（逐元素融合即函数复合；实现性质）　设 $u$ 为逐元素节点，其计算函数为 $f_u$；$v = g(u)$ 为逐元素运算，$g$ 作用于单个元素。若 $u$ 不写入内存，则 $v$ 的计算函数为

$$f_v(i) = g\bigl(f_u(\pi(i))\bigr),$$

其中 $\pi$ 为 $v$ 读取 $u$ 时的下标变换（无视图时为恒等）。$v$ 的计算过程中不访问 $u$ 的存储。

说明　这是降级器构造方式的直接后果，而不是可以脱离实现成立的数学命题，因此称为性质。$u$ 未写入内存时，它在中间表示中只以读取函数 $i \mapsto f_u(i)$ 的形式存在。$v$ 的降级规则取得这一读取函数，经下标变换 $\pi$ 后调用，并对结果施加 $g$，即得上式。$\square$

由性质 6.1，多个逐元素运算串联时，计算函数逐层复合，最终在一个循环中完成全部计算，中间值只存在于寄存器中。规约节点的计算函数同样可以调用上游的读取函数，因此"逐元素链后接规约"也在降级时融为一体。代码清单 6.2 给出例 1.1 降级后的全部缓冲区。

代码清单 6.2　例 1.1 降级后的缓冲区（CPU，省略了变量名前缀 GraphLowering_）

InputBuffer(name='primals_1', layout=FixedLayout('cpu', float32, size=[64, 128], stride=[128, 1]))
InputBuffer(name='primals_2', layout=FixedLayout('cpu', float32, size=[64, 128], stride=[128, 1]))
ComputedBuffer(name='buf2', layout=FixedLayout('cpu', float32, size=[64], stride=[1]), data=Reduction(
  'cpu',
  tensorplay.float32,
  def inner_fn(index, rindex):
      i0 = index
      r0_0 = rindex
      tmp0 = ops.load('primals_1', 128*i0 + r0_0)
      tmp1 = ops.load('primals_2', 128*i0 + r0_0)
      tmp2 = ops.add('tmp0', 'tmp1')
      tmp3 = ops.relu('tmp2')
      tmp4 = ops.constant(2, tensorplay.float32)
      tmp5 = ops.mul('tmp3', 'tmp4')
      return tmp5
  ,
  ranges=[64],
  reduction_ranges=[128],
  reduction_type=sum,
  ...
))

除两个输入外，整个图降级后只剩一个缓冲区。add 降级为一个逐元素节点，仅携带计算函数；relu 的规则在其外包一层 ops.relu，乘 2 同理；sum 的规则返回一个规约节点，其计算函数调用上游逐层复合的读取函数。在融合算法开始之前，例 1.1 已经是单一内核。

图 6.1　例 1.1 的降级过程。三个逐元素算子逐层复合计算函数，sum 直接调用之；只有输出被实体化为缓冲区。

6.4 布局与视图

定义 6.2（布局）　布局由形状 $(s_0, \dots, s_{k-1})$、步长 $(t_0, \dots, t_{k-1})$ 与偏移 $o$ 组成，元素 $i = (i_0, \dots, i_{k-1})$ 的地址为

$$\mathrm{addr}(i) = o + \sum_{j} t_j\, i_j = o + t^{\mathsf T} i.$$

布局分为两类：已确定步长的固定布局，与步长尚可由使用者调整的灵活布局。

地址是下标的仿射函数。转置、切片、广播等视图运算都不改变数据，只改变下标与地址的对应关系，因此可以统一表示为下标空间上的变换。

命题 6.2（视图的仿射表示）　维度重排、切片、广播扩展与插入或删除长度为 1 的维度，都可以表示为仿射下标变换 $\pi(i) = A i + c$，其中 $A$ 为整数矩阵。若被观察的张量具有布局 $(t, o)$，则视图的布局为

$$t' = A^{\mathsf T} t,\qquad o' = o + t^{\mathsf T} c.$$

仿射变换的复合仍是仿射变换，因此任意多个上述视图的复合仍可由一组步长与偏移表示。

证明　维度重排对应置换矩阵，切片 $i_j \mapsto a + k, i_j$ 对应对角元为步长 $k$、平移为 $a$ 的变换，广播扩展对应把被扩展维度的系数置为 0，插入或删除长度为 1 的维度对应增加或删除 $A$ 的一行或一列。代入 $\mathrm{addr}(\pi(i)) = o + t^{\mathsf T}(A i + c) = (o + t^{\mathsf T} c) + (A^{\mathsf T} t)^{\mathsf T} i$ 即得新的步长与偏移。复合的封闭性由 $A_2(A_1 i + c_1) + c_2 = (A_2 A_1) i + (A_2 c_1 + c_2)$ 得出。$\square$

广播扩展对应步长为 0：沿被扩展的维度移动时地址不变。一般的 reshape 则不一定是仿射的。把相邻的两维 $j$、$j+1$ 合并为一维，新下标 $m$ 与原下标的关系为 $i_j = \lfloor m / s_{j+1} \rfloor$、$i_{j+1} = m \bmod s_{j+1}$，地址

是 $m$ 的仿射函数，当且仅当 $t_j = s_{j+1} t_{j+1}$，即两维在内存中连续衔接。不满足这一条件时，reshape 无法用一组步长表示，只能在下标表达式中保留整除与取模，或者先把数据复制为连续布局。

Stax 据此区分两种情形。若被观察的对象已有固定布局（例如输入或已写入内存的缓冲区），视图直接表示为同一块存储上的新步长与偏移；否则视图保存下标变换 $\pi$，读取时先变换下标再调用被观察节点的读取函数。两种情形下视图都不复制数据，其效果最终都体现为读取时的下标表达式。下标表达式以符号形式表示并在构造时化简。

灵活布局的意义在于：一个中间缓冲区若只被转置后读取，将它直接按转置后的顺序写出，可以使写入与读取都连续。输入与输出的布局由调用方决定，不可调整。

灵活布局在降级结束时被固定。对每个仍为灵活布局的缓冲区，Stax 考察其计算函数读取各输入时的步长，选出一个循环次序，使按该次序遍历时读取尽量连续，再按同一次序为缓冲区分配步长，使写出也连续。读取方式无法给出偏好时，取按行连续的默认布局。这一决策只涉及中间缓冲区，因而不改变程序的语义。

例 6.1　设 f(x, y) = x.t() + y，其中 $x$ 的形状为 $512\times256$，$y$ 的形状为 $256\times512$，均按行连续。降级后的唯一缓冲区如代码清单 6.3 所示。

代码清单 6.3　例 6.1 降级后的缓冲区（以 CUDA 为目标，省略了前缀）

ComputedBuffer(name='buf2', layout=FixedLayout('cuda:0', float32, size=(256, 512), stride=[512, 1]), data=Pointwise(
  'cuda',
  tensorplay.float32,
  def inner_fn(index):
      i0, i1 = index
      tmp0 = ops.load('primals_1', i0 + 256*i1)
      tmp1 = ops.load('primals_2', 512*i0 + i1)
      tmp2 = ops.add('tmp0', 'tmp1')
      return tmp2
  ,
  ranges=(256, 512),
  origins=OrderedSet([primals_1, primals_2, t, add])
))

转置没有产生任何节点，只体现为 $x$ 的下标 i0 + 256*i1。按命题 6.2，$x$ 的步长为 $(256, 1)$，转置对应置换矩阵，新步长为 $(1, 256)$，地址为 $i_0 + 256 i_1$，与清单一致。两个输入的下标中，i1 的系数分别为 256 与 1，即沿最内层维度一个输入不连续、另一个连续。这一差异决定了代码生成时的分块方式（8.1 节、9.3 节）。

6.5 规约的代数性质与拆分

规约的实现方式依赖于合并运算的代数性质。

定义 6.3（幺半群）　集合 $M$ 连同二元运算 $\oplus$ 与元素 $e$ 称为幺半群，若 $\oplus$ 满足结合律，且 $e$ 为单位元：$e \oplus a = a \oplus e = a$。若 $\oplus$ 还满足交换律，称为交换幺半群。

求和、求积、最大值、最小值都是交换幺半群，单位元分别为 $0$、$1$、$-\infty$、$+\infty$。argmax 的合并对象是（值，下标）对，按值比较、值相等时取较小下标，也构成交换幺半群。均值与方差的一趟算法（Welford 算法）同样可以写为三元组（个数，均值，二阶中心矩）上的合并运算 [17]。

结合律是一切规约并行化的基础：把规约范围划分为若干段 $R_r = R^{(1)} \cup \cdots \cup R^{(k)}$，则

各段的部分结果可以独立计算，再合并。交换律则允许以任意顺序访问元素，例如向量化时各通道交错地分得元素。

浮点加法不满足结合律，上式只在近似意义下成立。不同的求和顺序给出不同的舍入误差，其误差界也不同。设单位舍入误差为 $u$，对 $n$ 个数 [18]：

1. 顺序求和的误差满足 $|\hat s - s| \le (n-1),u \sum_i |a_i| + O(u^2)$；
2. 分为 $k$ 段、段内顺序求和后再顺序合并各段，误差界的系数为 $(n/k - 1) + (k - 1)$；
3. 两两配对的树形求和，误差界的系数为 $\lceil \log_2 n \rceil$。

由第 2 条，取 $k \approx \sqrt n$ 时系数约为 $2\sqrt n$，显著小于 $n - 1$。代码清单 1.1 以 16 个向量通道各自累加、最后归约为标量，相当于 $k = 16$ 的分段求和，其误差界不劣于顺序求和；它与即时执行结果之间 $1.5\times10^{-5}$ 的差异，来自两种求和顺序各自的舍入，而不是其中一种的错误。

拆分的动机是并行度。若输出很少而规约很长，例如对 $64\times16384$ 的张量按行求和，只有 64 个输出；GPU 上若每个输出由一个线程块负责，至多 64 个线程块，远不能占满设备。拆分把每行切为 $k$ 段，第一层规约求出 $64k$ 个部分和，第二层规约合并之，并行的单位增至 $64k$ 个。代价是多写出并读回 $64k$ 个部分和。CPU 的线程数远少于 GPU 可同时驻留的线程块数，按输出并行通常已足够，拆分主要用于 GPU。

拆分段数的选取须兼顾两点：段数要足以占满设备，每段的长度又不能太短，否则每个线程分得的元素过少，启动与合并的开销占比上升。Stax 在 GPU 上的规则是：若输出数不少于 $2 \cdot 32 \cdot S$（$S$ 为流多处理器数），不拆分；否则按算法 6.1 确定段数。扫描、点积与 argmax 一类携带下标的规约不拆分。

算法 6.1　GPU 上内层规约的拆分段数

输入：规约长度 n_r，输出数 n，设备属性（S 个流多处理器，线程束宽 w，每 SM 线程数 T）
输出：段数 k
1  e_min ← w；e_max ← 512；线程数 m ← 8w
2  D_min ← e_min · S · T；D_max ← e_max · S · T
3  if n ≥ 2S: return 1
4  if n_r 不超过免拆分阈值: return 1
5  if n · n_r ≤ D_min:       每线程元素数 s ← e_min
6  elif n · n_r < D_max:     取目标块数 b ← S·T / (2m)，每个输出分得 ⌈b/n⌉ 块，
                             s ← 使 n_r 被均分且接近 ⌈n_r / (m·⌈b/n⌉)⌉ 的段长
7  else:                     s ← n_r 的因子中最接近 e_max 者（相差 50 以内），否则 e_max
8  return ⌈n_r / (s · m)⌉

例 6.2　设 f(x) = (x * x).sum(-1)，$x$ 的形状为 $64\times16384$，目标设备有 $S = 132$ 个流多处理器，$w = 32$，$T = 2048$。则 $m = 256$，$D_{\min} = 32 \cdot 132 \cdot 2048 = 8,650,752$。由于 $n = 64 < 2S = 264$，且 $n \cdot n_r = 1,048,576 \le D_{\min}$，取 $s = 32$，段数为 $\lceil 16384 / (32 \cdot 256) \rceil = 2$。降级结果如代码清单 6.4 所示。

代码清单 6.4　例 6.2 降级后的两层规约（以 CUDA 为目标，省略了前缀与部分字段）

ComputedBuffer(name='buf1', layout=FixedLayout('cuda:0', float32, size=[64, 2], stride=[2, 1]), data=Reduction(
  def inner_fn(index, rindex):
      i0, i1 = index
      r0_0 = rindex
      tmp0 = ops.load('primals_1', 16384*i0 + 8192*i1 + r0_0)
      tmp1 = ops.load('primals_1', 16384*i0 + 8192*i1 + r0_0)
      tmp2 = ops.mul('tmp0', 'tmp1')
      return tmp2
  ,
  ranges=[64, 2], reduction_ranges=[8192], reduction_type=sum,
))
ComputedBuffer(name='buf2', layout=FixedLayout('cuda:0', float32, size=[64], stride=[1]), data=Reduction(
  def inner_fn(index, rindex):
      i0 = index
      r0_0 = rindex
      tmp0 = ops.load('buf1', 2*i0 + r0_0)
      return tmp0
  ,
  ranges=[64], reduction_ranges=[2], reduction_type=sum,
))

第一层规约的迭代范围由 $[64]$ 扩为 $[64, 2]$，规约范围由 $[16384]$ 缩为 $[8192]$，下标 16384*i0 + 8192*i1 + r0_0 将第 $i_1$ 段映射到原规约维的前半或后半，这正是结合律等式右端的内层求和。第二层规约沿长度为 2 的维度合并部分和。多写出并读回的 128 个部分和，相对于 $2^{20}$ 个输入元素可以忽略。

第一层的计算函数对同一位置读取了两次，因为 x * x 的两个操作数是同一个读取函数。循环级中间表示不做公共子表达式消除，这一冗余留到代码生成时由内核内的 CSE 消去。

6.6 实体化

性质 6.1 的前提是"$u$ 不写入内存"。若 $u$ 有多个使用者，内联意味着每个使用者都在自己的循环中重算 $u$。是否把 $u$ 写入内存，称为实体化（realization）决策。

可以用一个简单的代价模型说明这一决策。设 $u$ 有 $N$ 个元素，每个元素的计算需要 $c$ 次运算，读取 $q$ 个输入字节，有 $k$ 个使用者，处理器的算力为 $P$，带宽为 $B$，$u$ 的每个元素占 $b$ 字节。内联时，$u$ 被计算 $k$ 次，其输入被读取 $k$ 次：

实体化时，$u$ 被计算一次、写出一次，再被各使用者读取 $k$ 次：

二者之差为

当 $k = 1$ 时差为负，内联总是更优，这正是单一使用者的值一律内联的原因。$k > 1$ 时，计算量 $c$ 越大、读取的输入 $q$ 越多、使用者 $k$ 越多，越倾向于实体化。

Stax 以一组阈值近似这一判据，对有多个使用者的节点依次检查，任一成立即实体化：

1. 在 CPU 上包含 exp、log、log10、log1p、log2、sigmoid、tanh 之一。这些超越函数在 CPU 上代价较高，相当于 $c / P$ 较大；在 GPU 上有专用的硬件单元，代价相对较低；
2. 在 CPU 上，使用者超过 5 个，且计算函数中的运算数接近上限，对应 $k$ 较大；
3. 计算函数中的运算数超过上限（GPU 上为 30，CPU 上为 50），对应 $c$ 较大；
4. 读取的缓冲区超过 4 个，对应 $q$ 较大；逐层内联后累计读取次数过多（GPU 上超过 8，CPU 上超过 12）时同样实体化。

此外，图的输出以及须交给外部内核的输入必然实体化。若某个使用者必然要求一个值实体化，则在该值产生时立即实体化，以免排在前面的使用者各自重算。

CPU 与 GPU 采用不同的阈值，反映的是两类处理器上 $P/B$ 与超越函数相对代价的差异。同一个程序在两类设备上可能得到不同的实体化结果，例 7.2 即为一例。

6.7 一个描述，多种解释

计算函数中的 ops 不是固定的实现，而是一个可替换的处理器。同一个计算函数在不同的处理器下执行，得到不同的结果：

1. 以记录读写的处理器执行，得到其读取与写入的缓冲区及下标，这是依赖分析与融合打分的依据（7.1 节）；
2. 以计数处理器执行，得到其运算数与所用运算的集合，这是 6.6 节实体化判据的依据；
3. 以格式化处理器执行，得到代码清单 6.2 中的文本；
4. 在代码生成阶段换用 Triton 或 C++ 的处理器执行，输出即为内核源码中的语句。例如 Triton 处理器将 ops.exp 打印为 libdevice.exp(...)，C++ 向量处理器将其打印为 Vectorized<float> 上的运算。

这一设计的理论背景是抽象解释（abstract interpretation）[19]：同一个程序在不同的语义域中解释，得到程序的不同性质。具体值的域给出计算结果，读写集合的域给出数据依赖，运算计数的域给出代价估计，源码文本的域给出代码。编写降级规则的人只需写一次计算函数，后端的各个分析与生成步骤都从中导出。

6.8 循环体的规范化与值域分析

计算函数是 Python 闭包，便于复合，却不便于分析。调度与代码生成之前，每个缓冲区的计算函数在符号下标上执行一次，被记录为一个循环体：一张小的计算图，其中的运算来自 ops 调用，下标表达式则单独保存为符号表达式。这仍是 6.7 节的做法，所用的处理器记录运算并收集下标。在循环体上进行两类与语义无关、只影响效率的化简。

第一类是循环的规范化：删除长度为 1 的维度，合并可以合并的相邻循环，并按访存的步长重排循环。合并的条件与 6.4 节 reshape 的条件相同。

命题 6.3（循环的合并）　设循环体中有相邻的两层循环 $i_j \in [0, s_j)$、$i_{j+1} \in [0, s_{j+1})$，每个访存的下标是循环变量的仿射函数，其中 $i_j$、$i_{j+1}$ 的系数分别为 $a_j$、$a_{j+1}$。若对每个访存都有 $a_j = s_{j+1} a_{j+1}$，则以 $m = s_{j+1} i_j + i_{j+1} \in [0, s_j s_{j+1})$ 代替两层循环，每个访存的地址不变。

证明　$a_j i_j + a_{j+1} i_{j+1} = a_{j+1}(s_{j+1} i_j + i_{j+1}) = a_{j+1} m$，且 $(i_j, i_{j+1}) \mapsto m$ 是 $[0, s_j) \times [0, s_{j+1})$ 到 $[0, s_j s_{j+1})$ 的双射。$\square$

例 6.2 的第一层规约（代码清单 6.4）读取 primals_1 的下标为 16384*i0 + 8192*i1 + r0_0。循环规范化之后，生成的 Triton 内核中这一读取写为

tmp0 = tl.load(in_ptr0 + (r0_1 + 8192*x0), r0_mask & xmask, ...)

16384*i0 + 8192*i1 被合并为 8192*x0，其中 x0 遍历合并后长度为 128 的维度。这是命题 6.3 的一例：$s_1 = 2$，$a_0 = 16384 = 2 \cdot 8192$。合并减少了循环的层数，使最内层循环更长，有利于向量化与并行的划分。

第二类是值域分析。代码生成需要回答若干关于整数取值范围的问题：下标表达式的最大值是否小于 $2^{31}$，从而可以用 32 位整数计算地址；间接下标是否一定落在 $[0, s)$ 中，从而可以省去负下标的修正与越界检查。这些问题在编译时一般无法精确回答，但可以用区间做保守的近似。

定义 6.4（区间域）　以闭区间 $[l, u]$（$l, u \in \mathbb Z \cup {\pm\infty}$）抽象一个整数变量的全部可能取值。循环变量 $i \in [0, s)$ 抽象为 $[0, s-1]$，常数 $c$ 抽象为 $[c, c]$，从内存读取的值抽象为 $[-\infty, +\infty]$。运算按区间算术传播，例如

$$[l_1, u_1] + [l_2, u_2] = [l_1 + l_2,\ u_1 + u_2],\qquad [l, u] \bmod c = [0, c-1]\ (l \ge 0,\ c > 0).$$

命题 6.4（可靠性）　若每个运算的区间传播规则都包含该运算在所有具体输入上的结果，则对循环体中的每个变量，分析所得的区间包含其在运行时的一切取值。

证明　按循环体的拓扑序归纳。输入（循环变量、常数、读取的值）的区间按定义包含其取值。若一个运算的各操作数的取值都在其区间内，由传播规则的假设，结果的取值在结果的区间内。$\square$

这是抽象解释 [19] 在区间域上的特例。命题 6.4 保证的是单向的结论：区间说"一定在 $[0, s)$ 内"时，运行时的值确实如此；区间较宽时，并不意味着运行时的值会越界，只是分析无法证明。因此只有在区间足够窄时才省去检查，否则保留之，代价是少量运行时开销，而不会改变语义。

例 6.3　比较两种间接读取：

def g1(x, idx):                     # 下标来自输入
    return x[idx] * 2

def g2(x):                          # 下标由循环变量算出
    idx = (tp.arange(128) * 3) % 128
    return x[idx] * 2

g1 的 idx 从内存读取，区间为 $[-\infty, +\infty]$，生成的 CPU 内核保留了负下标的修正与越界检查，如代码清单 6.5 所示。

代码清单 6.5　g1 的内核（CPU，节选）

auto tmp0 = VectorizedN<int64_t,2>::loadu(in_ptr0 + x0, 16);    // 读取 idx
auto tmp4 = tmp0 + VectorizedN<int64_t,2>(128);
auto tmp7 = VecMask<int64_t,2>(tmp0 < VectorizedN<int64_t,2>(0));
auto tmp8 = decltype(tmp4)::blendv(tmp0, tmp4, tmp7.template cast<int64_t,2>());  // 负下标加 128
...
TP_CHECK((VecMask<int64_t,2>((VectorizedN<int64_t,2>(0) <= tmp10) & (tmp10 < VectorizedN<int64_t,2>(128)))).all_masked(),
         "index out of bounds: 0 <= tmp10 < 128");
...
tmpbuf[x0_inner] = in_ptr1[tmp9[x0_inner]];                       // 按下标逐个读取 x

g2 的下标由 arange 计算而来，arange 的计算函数就是循环变量本身。下标传播把整个下标计算化为符号表达式 $(3 i) \bmod 128$，间接读取于是变为直接读取；由定义 6.4，该表达式的区间为 $[0, 127]$，修正与检查都被省去，如代码清单 6.6 所示。

代码清单 6.6　g2 的内核（CPU，节选）

for (long x0_inner = 0; x0_inner < 16; x0_inner++)
{
    tmpbuf[x0_inner] = in_ptr0[(static_cast<int64_t>(3*x0 + 3*x0_inner) % static_cast<int64_t>(128))];
}

g2 的内核中不再有 idx 这块缓冲区：arange、乘法与取模都被吸收进了下标表达式。这是性质 6.1 的另一种形式：一个只依赖于循环变量的值不必写入内存，它的计算函数可以直接代入使用者的下标中。

值域分析的结论还用于选择下标的位宽。若一个内核中全部下标表达式的区间上界都小于 $2^{31}$，地址计算可以用 32 位整数进行；在 GPU 上，这减少了每个线程占用的寄存器，按 9.1 节的分析，可以提高驻留的线程束数。

实体化只决定哪些值写入内存。多个缓冲区能否合并为同一个内核，由调度器决定。

6.9 核心语言的语义与变换的正确性

前面各节反复使用"保持语义"一词：公共子表达式消除保持语义（4.5 节），逐元素融合即函数复合（性质 6.1），实体化与否不改变结果（6.6 节）。这些说法到此为止都是非形式的。本节为循环级中间表示的一个核心子集给出指称语义 [52]，在此基础上把它们陈述为定理并加以证明，同时说明哪些变换只在实数语义下成立。

定义 6.5（张量核心语言）　每个张量变量 $y$ 带有一个定义域 $D_y \subset \mathbb Z^{d}$（有限的长方体）。表达式由下面的文法生成：

$$e ::= x \ \mid\ \mathrm{map}_g(e_1, \dots, e_r) \ \mid\ \mathrm{view}_\pi(e) \ \mid\ \mathrm{red}_{\oplus, e_0, R_r}(e),$$

其中 $x$ 为输入或已绑定的变量；$g : \mathbb V^r \to \mathbb V$ 为标量函数；$\pi : D \to D’$ 为把结果的下标映到被观察张量下标的仿射映射（命题 6.2）；$\mathrm{red}$ 沿规约范围 $R_r$ 以运算 $\oplus$ 与初值 $e_0$ 规约。程序是一串绑定 $\mathbf{let}\ y_1 = e_1; \dots; \mathbf{let}\ y_m = e_m;\ \mathbf{return}\ y_m$，每个 $y_k$ 只被绑定一次。

这正是 6.3 节的两类节点加上 6.4 节的视图：$\mathrm{map}$ 对应逐元素节点，$\mathrm{red}$ 对应规约节点。语言中没有原地修改，这由 4.9 节的函数化保证；也没有随机数之类不由输入决定的运算，这类运算在 4.5 节被排除在公共子表达式消除之外，下文会看到原因。

**语义。**取 $\mathbb V$ 为浮点数的集合（含无穷与 NaN），定义域为 $D$ 的张量的值是一个函数 $D \to \mathbb V$。环境 $\rho$ 把变量映为张量值。表达式的指称 $[![ e ]!]\rho$ 归纳定义如下：

其中 $r_1 < r_2 < \dots < r_q$ 是 $R_r$ 的元素按字典序的排列。程序的指称为依次扩展环境 $\rho_k = \rho_{k-1}[y_k \mapsto [![ e_k ]!]\rho_{k-1}]$ 之后 $y_m$ 的值。规约的语义特意规定了一个确定的求值次序，因为浮点运算 $\oplus$ 一般不满足结合律（6.5 节）；只有规定了次序，“保持语义"才有确切的含义。

定义 6.6（语义等价）　称两个表达式（或程序）$e$、$e’$ 等价，记为 $e \equiv e’$，若对一切良构的环境 $\rho$（为每个自由变量指定一个定义域正确的张量值），都有 $[![ e ]!]\rho = [![ e’ ]!]\rho$，即作为 $D \to \mathbb V$ 的函数逐点相等。核心语言中的每种构造在浮点语义下都是全函数（无效运算给出 NaN 而不是报错），故指称处处有定义，$\equiv$ 是等价关系。

这里的相等是浮点值的逐位相等，比"数值上接近"严格得多。

引理 6.1（替换引理）　设变量 $y$ 不在 $e$ 中自由出现，则对一切表达式 $b$ 与环境 $\rho$，

$$[\![ b ]\!]\bigl(\rho[y \mapsto [\![ e ]\!]\rho]\bigr) = [\![ b[e / y] ]\!]\rho,$$

其中 $b[e/y]$ 表示把 $b$ 中 $y$ 的每次出现替换为 $e$。

证明　对 $b$ 的结构归纳，记 $\rho’ = \rho[y \mapsto [![ e ]!]\rho]$。$b = y$ 时两端都是 $[![ e ]!]\rho$；$b$ 为其他变量 $x$ 时两端都是 $\rho(x)$。$b = \mathrm{map}_g(b_1, \dots, b_r)$ 时，由语义定义，$[![ b ]!]\rho’$ 只通过各 $[![ b_k ]!]\rho’$ 依赖于 $\rho’$，由归纳假设它们分别等于 $[![ b_k[e/y] ]!]\rho$，而 $\mathrm{map}_g(b_1[e/y], \dots) = b[e/y]$。$\mathrm{view}$ 与 $\mathrm{red}$ 同理。$b = (\mathbf{let}\ z = e’;\ b’)$ 时，由单次赋值可设 $z \ne y$ 且 $z$ 不在 $e$ 中自由出现。按语义，左端等于 $[![ b’ ]!]\bigl(\rho’[z \mapsto w]\bigr)$，其中由归纳假设 $w = [![ e’ ]!]\rho’ = [![ e’[e/y] ]!]\rho$。因 $z$ 不在 $e$ 中自由出现，$[![ e ]!]\rho = [![ e ]!]\bigl(\rho[z \mapsto w]\bigr)$；又 $z \ne y$，两次扩展可以交换，故 $\rho’[z \mapsto w] = \rho[z \mapsto w]\bigl[y \mapsto [![ e ]!](\rho[z \mapsto w])\bigr]$。对 $b’$ 与环境 $\rho[z \mapsto w]$ 用归纳假设，左端等于 $[![ b’[e/y] ]!]\bigl(\rho[z \mapsto w]\bigr) = [![ \mathbf{let}\ z = e’[e/y];\ b’[e/y] ]!]\rho = [![ b[e/y] ]!]\rho$。$\square$

引理 6.1 的证明只用到一点：每种构造的指称只由其子表达式的指称决定（组合性），且指称是环境的函数（确定性）。随机数生成不满足后一条：两次调用 rand 的结果不同，它不是环境的函数，下面的定理对它不成立。

定理 6.1（内联与实体化保持语义）　$\mathbf{let}\ y = e;\ b \ \equiv\ b[e/y]$。

证明　左端的指称是 $[![ b ]!]\bigl(\rho[y \mapsto [![ e ]!]\rho]\bigr)$，由引理 6.1 等于右端。$\square$

从左向右读，定理 6.1 是内联：把生产者的计算函数代入使用者，这正是降级时逐元素节点不写入内存的情形（6.3 节）。从右向左读，它是实体化：把一个子表达式的值绑定到变量并写入缓冲区（6.6 节）。两个方向都保持逐位相等的语义，所以是否实体化纯粹是代价问题，不是正确性问题，这正是 6.6 节可以只用代价模型作决策的依据。这一结论有两个隐含的前提：缓冲区的存储类型与计算类型相同（若以 bf16 存储一个以 fp32 计算的中间值，实体化会多一次舍入），且代码生成不把乘加收缩为融合乘加（FMA）指令，否则内联前后的舍入次数不同。

推论 6.1（性质 6.1 的语义形式）　$\mathrm{map}g\bigl(\mathrm{view}\pi(\mathrm{map}f(x))\bigr) \equiv \mathrm{map}{g \circ f}\bigl(\mathrm{view}_\pi(x)\bigr)$，二者的指称都是 $\lambda i.\ g(f(x(\pi(i))))$。

证明　按语义定义逐层展开两端即得。$\square$

定理 6.2（公共子表达式消除与死代码消除）　设 $y_1$ 不在 $e$ 中自由出现。（1）$\mathbf{let}\ y_1 = e;\ \mathbf{let}\ y_2 = e;\ b \ \equiv\ \mathbf{let}\ y_1 = e;\ b[y_1 / y_2]$。（2）若 $y$ 不在 $b$ 中自由出现，则 $\mathbf{let}\ y = e;\ b \ \equiv\ b$。

证明　（1）左端在 $\rho_1 = \rho[y_1 \mapsto [![ e ]!]\rho]$ 下绑定 $y_2$。$y_1$ 不在 $e$ 中自由出现，故 $[![ e ]!]\rho_1 = [![ e ]!]\rho = \rho_1(y_1)$。于是 $\rho_2 = \rho_1[y_2 \mapsto \rho_1(y_1)]$，由引理 6.1（取被替换的表达式为变量 $y_1$），$[![ b ]!]\rho_2 = [![ b[y_1/y_2] ]!]\rho_1$。（2）$y$ 不在 $b$ 中自由出现，$[![ b ]!]$ 不读取 $\rho(y)$。$\square$

定理 6.2 的第 2 条依赖于一个隐含的前提：$e$ 的求值总会终止并给出一个值。若语言中有会报错的运算（例如整数除以零、越界的索引），删去一个不被使用的 $e$ 会把"报错"变为"正常返回”。这时等价须放宽为精化：变换后的程序在原程序有定义的输入上给出相同的结果，在原程序报错的输入上可以给出任意结果。经过验证的编译器 CompCert 正是以这种精化关系陈述其正确性定理的 [56]。

前两个定理覆盖的变换都保持逐位相等的语义。规约的拆分与重排则不然：它们改变了定义中规定的求值次序。

定理 6.3（规约重排的实数语义）　设 $\oplus$ 在实数上满足结合律与交换律，$e_0$ 为单位元。以 $[![ \cdot ]!]_{\mathbb R}$ 记把 $\mathbb V$ 换成 $\mathbb R$（对以 $\mp\infty$ 为初值的 $\max$、$\min$，换成扩展实数 $\mathbb R \cup {\pm\infty}$）、运算按精确的实数运算解释的语义。则对 $R_r$ 的任意排列、任意加括号的方式，以及按 6.5 节的任意分段与合并，规约的实数语义都与定义中的次序相同。

证明　对规约元素的个数 $q$ 归纳证明：任意加括号与排列的 $\oplus$ 表达式都等于 $v_{r_1} \oplus \cdots \oplus v_{r_q}$ 的左结合形式。$q \le 2$ 时由交换律直接成立。$q > 2$ 时，表达式的最外层把元素分为非空的两组 $P$ 与 $Q$，由归纳假设，两组分别等于各自元素按定义次序的左结合式；再由结合律与交换律逐个把 $Q$ 中的元素移入 $P$ 的左结合式中，即得全体元素的左结合式。分段规约是加括号的特例，初值 $e_0$ 为单位元，可在任意位置插入或删去。$\square$

于是变换分为两类。第一类保持浮点语义，包括内联与实体化、公共子表达式消除、死代码消除、视图的复合（命题 6.2）以及函数化：对它们，编译前后的结果逐位相同，测试可以要求精确相等。第二类只保持实数语义，包括规约的分段与重排（6.5 节）、向量化的规约（8.1 节）与在线 softmax 的重新缩放（9.5 节）：对它们，正确性的含义是"实数语义相同，且浮点结果之差不超过 6.5 节的误差界"，测试必须留有与条件数相称的容差（13.3 节）。4.2 节拒绝重结合之类的代数化简，正是为了让图变换全部落在第一类中；第二类只出现在降级与代码生成中，并且每一处都有对应的误差分析。

本节的语义是指称式的，只覆盖了无副作用的张量核心。一个完整的编译器正确性证明还须处理原地修改与别名（4.9 节的函数化可以视为从带存储的语义到本节语义的翻译）、守卫与特化（第 2 节）以及运行时的回退。经过验证的编译器在机器证明助手中完成了这类证明 [56]；Stax 没有做形式验证，它以 11.4 节的分层比对在经验上检验同一个性质。

上一节：第一节　引言、特化、图捕获与图变换　｜　目录　｜　下一节：第三节　调度、融合与 CPU 代码生成