TensorPlay Stax (2): Autodiff and Lowering

Previous: Part One　Introduction, Specialization, Graph Capture and Graph Transformation　｜　Table of Contents　｜　Next: Part Two　Scheduling, Fusion and CPU Code Generation

Reference numbers [n] in the text are listed in Section 6; section numbers, theorem numbers, and exercise numbers are numbered uniformly across the chapter, and cross-section references remain valid as usual.

5 Automatic Differentiation and Activation Rematerialization

After graph transformation, the computation graph must still pass through the automatic differentiation layer before it is handed to Stax. This layer has two tasks: rewriting Python-level calls into an operator-level graph, and, during training, constructing the backward computation and deciding which intermediate values the forward pass must save.

5.1 The Operator-Level Graph

What the frontend records are Python-level calls, such as operator.add or F.gelu. From the called function alone, the backend cannot determine which operator overload it ultimately dispatches to, nor learn the internal composition of composite operators. Stax accepts only operator-level graphs, so even for regions that need no differentiation, it must be re-traced once at the dispatcher level, recording the operator overload each call actually reaches, only without constructing the backward part. Inference and training therefore use the same kind of graph, and the backend faces only one kind of input. Listing 5.1 is the graph of Example 1.1 after this layer.

Listing 5.1　The operator-level graph of Example 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,)
)

Compared with Listing 3.1, operator.add has become the operator overload add.Tensor with the default argument alpha=1 filled in, tensor–scalar multiplication is dispatched to mul.Scalar, and the tensor method sum(-1) has become sum.dim_IntList. Every node points to a uniquely determined operator overload, and the backend can look up the lowering rule accordingly. The re-tracing also performs two transformations: operators that are not part of the backend base are expanded according to the decomposition table (Section 4.8), and in-place mutations are functionalized (Section 4.9).

5.2 Reverse-Mode Automatic Differentiation

Let the forward computation be a composition of a chain of functions

Write $J_k = \partial \phi_k / \partial z_{k-1}$ for the Jacobian matrix of step $k$. By the chain rule,

Reverse-mode automatic differentiation (reverse mode) evaluates this product from right to left:

where $\bar z_k = \partial L / \partial z_k$ is called the adjoint. Each step needs only a vector times the transpose of the Jacobian matrix (a vector-Jacobian product); the Jacobian matrix itself never has to be constructed. For a general directed acyclic graph, the adjoints accumulate in reverse topological order: the adjoint of a value is the sum of the contributions its users send back.

The chain form above is convenient for understanding, but a general computation is a directed acyclic graph. The following states the cost of reverse mode precisely under the model of a straight-line program. Suppose $L$ is computed by the following program:

where $a_k, b_k < k$, and $\varphi_k$ is one of $+$, $-$, $\times$, $\div$, or a unary elementary function $\psi$ that uses only $v_{a_k}$. Call $W = N - d$ the operation count of the program ($d$ is the input dimension; do not confuse it with the chain length $n$ above).

Theorem 5.1 (Cost of Computing the Gradient [9, 38])　Suppose the derivative $\psi’(a)$ of each unary function $\psi$ can be computed from $a$ and $\psi(a)$ using at most $c_\psi$ operations, and $c_\Omega = \max_\psi c_\psi$. Then reverse mode computes the whole of $\nabla_x L$ (together with $L$ itself) using at most

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

operations. For $\exp$, $\log$, $\tanh$, $\mathrm{sigmoid}$, $\sqrt{\cdot}$ one has $c_\Omega \le 2$, so the constant is at most 5 and is independent of the input dimension $d$.

Proof　The algorithm has two passes. The forward pass computes and saves each $v_k$ in the order $k = d+1, \dots, N$, for $W$ operations in total. The backward pass sets $\bar v_N = 1$ and all other $\bar v_k = 0$, and then performs, in the order $k = N, \dots, d+1$,

Correctness: the invariant is that after all $k > j$ have been processed, $\bar v_j = \partial L / \partial v_j$ (regarding $v_j$ as the independent variable, with the computation after it unchanged). $v_j$ influences $L$ only through those $v_k$ ($k > j$) that have it as an operand, so by the multivariate chain rule, $\partial L/\partial v_j = \sum_{k:, j \in {a_k, b_k}} (\partial L/\partial v_k)(\partial \varphi_k/\partial v_j)$. When processing in decreasing $k$, every $\bar v_k$ has already received all of its contributions before it is used (inductive hypothesis), and each term on the right-hand side is accumulated exactly once. Taking $j \le d$ gives $\nabla_x L$.

Cost: count the operations of step $k$ of the backward pass operation class by operation class. Addition and subtraction: $\bar v_a \mathrel{+}= \bar v_k$, $\bar v_b \mathrel{\pm}= \bar v_k$, 2 operations in total. Multiplication $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 in total. Division $v_k = v_a / v_b$: first compute $t = \bar v_k / v_b$, then $\bar v_a \mathrel{+}= t$ and $\bar v_b \mathrel{-}= t, v_k$ (because $\partial v_k/\partial v_b = -v_k/v_b$), 4 in total. Unary functions: evaluating $\psi’(v_a)$ takes $c_\psi$ operations (it may use the saved $v_k = \psi(v_a)$), then a multiply and an add, for $c_\psi + 2$ in total. Each step is at most $\max(4, c_\Omega + 2)$ operations, the backward pass is at most $\max(4, c_\Omega+2),W$, and adding the forward pass gives the stated bound. For the listed functions, $\exp’ = \exp$ (0 operations), $\log’(a) = 1/a$ (1 operation), $\tanh’ = 1 - v^2$, $\mathrm{sigmoid}’ = v(1 - v)$, and $(\sqrt{a})’ = 1/(2v)$ (2 operations each). $\square$

The two-pass structure in the proof also explains the storage cost of reverse mode: the backward pass needs every $v_k$ saved by the forward pass. In a tensor-level computation graph, each node is a tensor operator, and the argument above holds operator by operator: as long as the cost of each operator’s vector–Jacobian product does not exceed a constant multiple of its forward cost (pointwise operators as above; the reverse of a matrix product $C = AB$ is two matrix products of the same size), the constant is independent of the input dimension. Baur and Strassen [38] give a version with constant 3 under an algebraic complexity cost model that counts only multiplications and divisions (non-scalar operations).

Theorem 5.1 says that the cost of computing the gradient is independent of the input dimension, which is why reverse mode is suitable for neural networks with many parameters. But reverse mode has another cost: $J_k$ generally depends on $z_{k-1}$ — for instance, when $\phi_k = \tanh$ we have $J_k = \mathrm{diag}(1 - \tanh^2 z_{k-1})$ — so the backward computation must be able to obtain the forward intermediate values. If all of them are saved, the storage is proportional to the number of forward steps, which is exactly why activation values occupy so much device memory during training.

5.3 The Joint Graph

TensorPlay does not derive the backward graph separately; instead the TPX engine actually executes one backward pass and records the operators it dispatches, as described in Algorithm 5.1.

Algorithm 5.1　Tracing the Joint Graph

输入：捕获图 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

The resulting joint graph (joint graph) contains both the forward and the backward, and the backward nodes refer directly to the values of the forward nodes. The joint graph moves the question of “which values to save” from run time to compile time: at that point the entire backward computation is visible, so a save-or-recompute decision can be made for every forward value.

Example 5.1　Consider

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()

Its joint graph is shown in Listing 5.2; the [FW] and [BW] prefixes indicate whether a node belongs to the forward or the backward pass.

Listing 5.2　The joint graph of Example 5.1 (some metadata omitted)

[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 is $x$, primals_2 is $W$, and tangents_1 is $\bar L$. The backward part corresponds one-to-one to the recursion of Section 5.2: $\bar L$ is expanded by the adjoint of sum to $32\times64$; gelu_backward computes $\bar h = \bar g \odot \mathrm{GELU}’(h)$, where $h = xW$; and finally $\bar W = x^{\mathsf T} \bar h$. $x$ needs no gradient, and the graph contains no node that computes $\bar x$. Note that gelu_backward reads $h$ rather than the output of gelu, and this is what determines the partition in Section 5.5.

5.4 A Minimum-Cut Model for Save and Recompute

The joint graph must be partitioned into a forward graph and a backward graph. There are two ways to obtain the forward values the backward pass needs: saving them in the forward pass, or recomputing them in the backward pass from already saved values. Saving everything minimizes the amount of computation but maximizes storage; saving only the inputs and recomputing everything else minimizes storage but can double the amount of computation. The idea of trading computation for storage appears in gradient checkpointing [10]; under a general cost model, finding an optimal save-and-recompute scheme is an NP-hard problem, and Checkmate [11] formulates it as an integer linear program. Stax uses a simplified model that can be solved exactly in polynomial time: it decides only which values to save, and uses the number of saved bytes as the cost.

First, characterize the feasible save schemes. Let the forward graph be $G_F$, let $B$ be the set of forward values read by the backward pass, and let $U$ be the set of values that should not be recomputed in the backward pass (defined below).

Definition 5.1 (Feasible Save Set)　A set $C$ of forward values is called feasible if, for every $v \in B$, every path from any value in $U$ to $v$ passes through some value in $C$ (the endpoints of the path count as well).

Intuitively, starting from $C$ and executing only cheap operations, the backward pass can recompute every value in $B$ without touching any value in $U$. The next step turns this condition into a cut in a flow network.

Definition 5.2 (Rematerialization Network)　For each forward value $v$, introduce vertices $v_{\mathrm{in}}$, $v_{\mathrm{out}}$, and an edge $v_{\mathrm{in}} \to v_{\mathrm{out}}$ of capacity $c(v)$; a forward data dependence $u \to v$ corresponds to an infinite-capacity edge $u_{\mathrm{out}} \to v_{\mathrm{in}}$. Then introduce a source $s$ and a sink $t$: for $v \in U$, add an infinite-capacity edge $s \to v_{\mathrm{in}}$; for $v \in B$, add an infinite-capacity edge $v_{\mathrm{out}} \to t$.

Splitting each value into two vertices is the standard way of turning a cost on vertices into capacities on edges: cutting the edge $v_{\mathrm{in}} \to v_{\mathrm{out}}$ means saving $v$.

Proposition 5.1　Finite-capacity $s$–$t$ cuts of the rematerialization network are in bijection with feasible save sets, and the capacity of the cut equals $\sum_{v \in C} c(v)$.

Proof　A finite-capacity cut can only cut edges of the form $v_{\mathrm{in}} \to v_{\mathrm{out}}$; let $C$ be the set of values corresponding to the cut edges. Every $s$–$t$ path in the network has the form $s \to u_{\mathrm{in}} \to \cdots \to v_{\mathrm{out}} \to t$, where $u \in U$ and $v \in B$, passing in between through the split edges of the values along a forward path from $u$ to $v$. Cutting all $s$–$t$ paths is equivalent to every forward path from $U$ to $B$ containing a value in $C$, that is, to $C$ being feasible. Conversely, cutting the corresponding edges for a feasible $C$ yields a finite cut. The capacity of the cut is the sum of the capacities of the cut edges. $\square$

Finding the feasible save set of least cost is therefore exactly finding a minimum $s$–$t$ cut of the network.

Theorem 5.2 (Max-Flow Min-Cut Theorem [12])　In a directed network, the value of a maximum flow from $s$ to $t$ equals the capacity of a minimum $s$–$t$ cut. If $f$ is a maximum flow, then the set $S$ of vertices reachable from $s$ in the residual network, together with its complement, forms a minimum cut.

Stax computes the maximum flow with the Edmonds–Karp algorithm [13]: each iteration uses a breadth-first search to find a shortest augmenting path in the residual network and pushes its bottleneck capacity along the path, until no augmenting path remains. The time complexity of this algorithm is $O(|V||E|^2)$. After the maximum flow has been found, one more breadth-first search yields $S$; the split edges that cross between $S$ and its complement are the cut edges.

The remaining question is how to choose $U$ and the capacities $c(v)$. A value is put into $U$ in the following three cases:

1. It is not the result of a cheap operation, such as a matrix product or a convolution, for which the cost of recomputation is too high;
2. It must be written to memory in the backward pass in any case, that is, one of its backward users cannot be fused with it, in which case recomputation saves no memory traffic;
3. It is the result of a reduction, with the criterion that four times the output byte count is smaller than the sum of the input byte counts. Recomputing a reduction requires rereading all of its inputs and costs the same as the reduction itself.

The capacities are

where $b(v)$ is the number of bytes of the tensor; $d(v)$ is the number of forward steps between $v$ and its first backward reader, clamped to $[1, 100]$; and if $v$ already has to be written to memory in the forward pass, $\beta(v) = 1$, otherwise $\beta(v) = 2$. The byte count is the direct cost of saving; the distance factor biases values with long lifetimes toward recomputation, because such a value occupies storage throughout its entire lifetime, which approximately reflects its contribution to peak memory; and $\beta$ reflects the fact that saving a value that could have been fused away in the forward pass requires one extra write to memory. The latter two factors are heuristic: they fold the temporal dimension that the minimum-cut model does not directly express into the capacity.

There are two additional rules as well: a value read both nearby and far away, and a value at the end of a long fusable chain, are forced to be saved, so that recomputation does not drag the entire chain into the backward pass. If no finite cut exists — that is, if some value needed by the backward pass can neither be saved nor recomputed — the partition falls back to saving everything.

5.5 The Partition of Example 5.1

The forward pass of Example 5.1 has three values: matmul (that is, $h$), gelu, and sum_0. The network built according to Definition 5.2 is shown in Table 5.1.

Table 5.1　The rematerialization network of Example 5.1

The capacity of matmul is $8192 \times 1.1 = 9011.2$, truncated to 9011. gelu is read by no backward node, so $d$ takes the upper bound 100; it can be fused with the subsequent summation, so $\beta = 2$, and the capacity is $2 \times 8192 \times 1.1^{100} \approx 2.26\times10^8$. There is exactly one path from $s$ to $t$, namely $s \to \text{matmul_in} \to \text{matmul_out} \to t$, whose bottleneck is 9011, so both the maximum flow and the minimum cut are 9011, and the cut edge is matmul_in $\to$ matmul_out, that is, only $h$ is saved.

This result is as intuition suggests. $h$ is the output of a matrix product and has to be written to memory whether or not it is saved, so saving it costs almost nothing; neither the output of gelu nor the erf inside it is saved, and the $\mathrm{GELU}’(h)$ needed by the backward pass is recomputed from $h$. GELU [14] is defined as $\mathrm{GELU}(h) = h,\Phi(h)$, where $\Phi(h) = \tfrac12\bigl(1 + \mathrm{erf}(h/\sqrt2)\bigr)$ is the standard normal distribution function, and its derivative is

The backward pointwise kernel is shown in Listing 5.3 (CPU) and Listing 5.4 (GPU).

Listing 5.3　The backward pointwise kernel of Example 5.1 (CPU, variables renamed)

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);

Listing 5.4　The backward pointwise kernel of Example 5.1 (GPU, Triton, abridged; the complete kernel appears in Section 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)

The kernels on the two devices have the same structure: they read only $h$ and a scalar upstream gradient, and write out $ar h$.

In the backward kernel, erf is recomputed from $h$ and fused into the same loop as the multiply by the Gaussian density and the upstream gradient. The recomputation costs one erf and one exp per element, and in exchange the forward pass writes and the backward pass reads two $32\times64$ tensors fewer.

Figure 5.1　The joint graph and minimum cut of Example 5.1. Only $h$ is saved; erf is recomputed in the backward loop.

5.6 Composition at Run Time

The forward graph and the backward graph produced by the partition are handed to Stax for compilation separately, and are then composed through a subclass of autograd.Function that plugs into TPX’s computation graph: its forward calls the compiled forward and stores the saved values into the context; its backward retrieves the saved values and hands them, together with the upstream gradient, to the compiled backward. The backward graph is compiled only when backward is called for the first time, so inference-only calls never pay the cost of compiling the backward.

This is the same minimum-cut partition that AOTAutograd uses in PyTorch 2 [15]. It rests on the premise that the backend can fuse the recomputed pointwise chain into a loop the backward pass would execute anyway; if every recomputed step were a separate kernel, the recomputed intermediate values would still have to be written back to memory and the trade of computation for storage would not pay off. How Stax meets this premise is explained in 第 Section 6.

5.7 Forward Mode and Higher-Order Derivatives

Section 5.2discussed only reverse mode. For $f: \mathbb R^n \to \mathbb R^m$, the Jacobian matrix $J \in \mathbb R^{m\times n}$ can be extracted either by columns or by rows.

Forward mode propagates tangent vectors along the direction of the computation: it maintains $\dot z_k = \partial z_k / \partial x \cdot v$ for every intermediate value $z_k$, with the recursion $\dot z_k = J_k \dot z_{k-1}$, so that a single forward computation yields both $f(x)$ and the Jacobian–vector product $Jv$. In the implementation, each tensor carries a tangent vector, forming a dual number $(z, \dot z)$; each operator updates the tangent vector according to its forward derivative formula, with no need to record a computation graph and no need to save intermediate values.

Proposition 5.2 (Cost of the Two Modes)　Suppose computing $f$ costs $W$. In forward mode, computing one $Jv$ and in reverse mode computing one $u^{\mathsf T}J$ both cost $O(W)$. Hence, to obtain the complete $J$, forward mode needs $n$ applications and reverse mode needs $m$, costing $O(nW)$ and $O(mW)$ respectively; moreover, reverse mode must save the forward intermediate values, while forward mode need not.

Proof　In forward mode, each operator’s tangent-vector update is the evaluation of its forward derivative formula on a vector and costs the same order as the operator itself, so the total is $O(W)$; reverse mode follows from Theorem 5.1. The $i$-th column of $J$ is $J e_i$ and the $j$-th row is $e_j^{\mathsf T} J$, which take $n$ and $m$ applications respectively. The difference in storage comes from the direction of the recursion: $\dot z_k$ in forward mode depends only on $\dot z_{k-1}$ and $z_{k-1}$, and the latter is available exactly at that point; reverse mode starts from $\bar z_n$ and must work its way back to $z_0$. $\square$

The loss in training is a scalar, so $m = 1$, which is why reverse mode dominates. Forward mode suits the case of many outputs and few inputs, such as computing sensitivities with respect to a few hyperparameters or computing directional derivatives.

Higher-order derivatives are obtained by composing the two modes. The product of the Hessian with a vector, $Hv = \nabla(\nabla f \cdot v)$, can be computed by “forward over reverse”: apply forward mode once on top of the reverse-mode gradient computation, at a cost of $O(W)$, consistent with Pearlmutter’s result [34]; it can also be done by “reverse over reverse”: compute the gradient with create_graph=True so that the gradient itself carries a computation graph, and take its gradient once more.

TensorPlay’s eager execution supports all of these forms: tensorplay.autograd.forward_ad provides dual numbers (make_dual, unpack_dual), and tensorplay.autograd.functional provides jvp, vjp, hvp, jacobian, and hessian; jvp computes in forward mode by default.

The compiled path supports only first-order reverse mode. The composition of Section 5.6wraps the compiled backward graph into an autograd.Function; the backward graph is itself a compiled artifact and has no computation graph on which differentiation could be applied again. How this limitation is surfaced to the user is itself a question of semantic contract.

One option is to do nothing: the first-order gradient does not enter a computation graph, and the second-order derivative then fails somewhere that seems unrelated, with an error saying that the gradient tensor “does not require grad and does not have a grad_fn”; the user has no way of learning that compilation is the cause. TensorPlay chooses instead to reject explicitly when the need arises:

Listing 5.5　Wrapper for the backward of a compiled region when a computation graph must be recorded (abridged)

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))

When gradient mode is on, that is, when the caller requires the first-order gradient itself to be differentiable, the compiled backward is wrapped in one more Function: its forward is the compiled backward, whose results are bitwise identical to before; its backward computes nothing and merely raises an error explaining the cause and the way out. Besides the actual arguments of the backward graph, the wrapper takes as inputs the original inputs of the forward pass, primals_for_grad, which makes the first-order gradient connected to the original inputs in the computation graph, so that a second-order derivative walking back along the graph must reach this backward and trigger the error, rather than failing elsewhere with “no grad_fn”. Thus the value of the first-order gradient is correct and requires_grad is true, and differentiating it again yields an explicit RuntimeError.

This is the same attitude that strict mode takes in Section 11.3: what the compiler cannot do should fail explicitly where it cannot be done, rather than yielding a result that appears successful but is in fact incomplete, or failing far away with an unrelated error.

To genuinely support second-order derivatives, the joint graph must be traced once more, taking the backward graph itself as the forward and computing its backward; this is feasible within the AOT framework, at the cost of a multiplicative increase in compile time and graph size. Until then, when computing higher-order derivatives such as the Hessian or Hessian–vector products, call the interfaces in tensorplay.autograd.functional on uncompiled functions.

5.8 Checkpointing under a Storage Budget

The minimum-cut model of Section 5.4answers the question of which values to save so that the recomputation cost of the backward pass is minimized when storage is unconstrained. Another equally basic question is its mirror image: when storage has a hard upper bound, how little must be recomputed? This section answers it for the simplest and most typical structure, a single chain.

Let the forward pass be a chain of $l$ steps, $z_k = \phi_k(z_{k-1})$, where each state occupies one unit of storage and each step costs 1. Step $k$ of the backward pass (computing $\bar z_{k-1}$) needs $z_{k-1}$. Saving everything needs $l$ units of storage and no recomputation; saving only $z_0$ makes the storage constant, but obtaining each $z_{k-1}$ requires recomputing from the start, for a total cost of $\Theta(l^2)$. Checkpointing lies in between: save only a few states, and during the backward pass recompute the needed states from the nearest checkpoint.

Theorem 5.3 (Uniform Partitioning [10])　Split the chain into $c$ equal segments and save a checkpoint at the beginning of each segment. Process the segments one at a time in the backward pass: recompute from that segment’s checkpoint and save all states within the segment, then finish the backward pass for that segment. Then the peak storage does not exceed $c + l/c$, and each step is computed at most twice. Taking $c = \sqrt l$, the storage does not exceed $2\sqrt l$ and the total forward cost is at most $2l$.

Proof　The forward pass saves only $c$ checkpoints. While the backward pass processes a given segment, the $l/c$ states of that segment are additionally saved and released as soon as it is done, so the peak does not exceed $c + l/c$ (the beginning of the segment is itself a checkpoint, so in reality it is 1 less). Each step is computed once in the forward pass and recomputed once when its segment is processed. $c + l/c$ is minimized at $c = \sqrt l$. $\square$

tensorplay.utils.checkpoint_sequential implements exactly this scheme: the user divides a sequential model into several segments, and each segment saves only its input. Storage drops from $\Theta(l)$ to $\Theta(\sqrt l)$, at the cost of one extra forward pass. When more than one recomputation is allowed, one can do better.

Definition 5.3 (Reversal Schedule)　Suppose there are $s$ checkpoint slots, with one slot holding $z_0$ at the start. A reversal schedule is a sequence of the following operations: computing forward a number of steps from some checkpoint (without saving intermediate states, keeping only the current state), writing the current state into a free slot, releasing a slot, and performing step $k$ of the backward pass when $z_{k-1}$ is the current state or is in a slot and $\bar z_k$ has already been computed. If all backward steps are executed in the order $k = l, \dots, 1$ and each forward step $\phi_k$ is computed at most $t$ times, the schedule is said to have repetition number $t$. Write $\beta(s, t)$ for the largest chain length for which such a reversal schedule exists.

Theorem 5.4 (Binomial Checkpointing [44, 45])　$\beta(s, t) \ge \dbinom{s + t}{s}$, and there is a recursive construction that attains this length. Griewank further proved that equality holds [44].

Proof　First prove the recurrence $\beta(s, t) \ge \beta(s, t-1) + \beta(s-1, t)$, then check the boundaries. Let the chain length be $l = m + n$, with $m = \beta(s, t-1)$ and $n = \beta(s-1, t)$. Advance $m$ steps from $z_0$ to obtain $z_m$ and write it into a free slot; the first $m$ steps have then been computed once each. The remaining $n$ steps form a chain starting at $z_m$; using the other $s - 1$ slots (the one occupied by $z_0$ is left alone) complete its reversal with repetition number $t$, which is guaranteed by $n = \beta(s-1, t)$; when the reversal is complete, $\bar z_m$ has been obtained. Release the slot of $z_m$, then reverse the first $m$ steps with $s$ slots: they have already been computed once, so the remaining repetition number is $t - 1$, which is guaranteed by $m = \beta(s, t-1)$.

Boundaries: when $t = 0$ no forward computation can be done, only $z_0$ is available, and only step 1 of the backward pass can be executed, so $\beta(s, 0) \ge 1$. When $s = 1$ there is a single checkpoint, namely $z_0$; advancing $k - 1$ steps from $z_0$ suffices to execute step $k$ of the backward pass, and taking $k = l, \dots, 1$ in turn, step 1 is computed $l - 1$ times, so $\beta(1, t) \ge t + 1$. Both equal the values of $\binom{s+t}{s}$ on the corresponding boundaries, and $\binom{s+t}{s} = \binom{s+t-1}{s} + \binom{s+t-1}{s-1}$ has the same form as the recurrence, so the claim follows by induction. $\square$

The proof of equality — that is, of the fact that no longer chain can be reversed with $s$ slots and repetition number $t$ — is fairly long, and its core is an exchange argument: any reversal schedule can be rewritten into the recursive nesting above without increasing the number of slots or the repetition number [44]. Since $\binom{2s}{s} \approx 4^s / \sqrt{\pi s}$, taking $s = t = O(\log l)$ suffices to reverse a chain of length $l$: both the storage and the recomputation factor grow only as $\log l$. This contrasts with the $\sqrt l$ storage of uniform partitioning.

Two limitations are worth noting. First, the results above are stated for chains; on a general directed acyclic graph, minimizing recomputation under a storage limit is hard: the version with recomputation allowed, formulated with the pebble game, is PSPACE-complete [57], which is why Checkmate uses an integer linear program [11]. Second, Stax’s partitioner only performs the minimum cut of Section 5.4and does not consider a storage budget: partition_min_cut accepts a memory_budget argument but does not currently use it. Checkpointing under a storage budget is specified by the user at the model level through tensorplay.utils.checkpoint.

6 Lowering to the Loop-Level Intermediate Representation

Once the operator-level graph delivered by the automatic differentiation layer enters Stax, it is first lowered to the loop-level intermediate representation. This layer is the core of Stax: it decides in what form each value exists — as a computation that can be inlined into its users, or as a buffer written to memory.

6.1 Interpreted Execution of the Graph

The lowering pass GraphLowering is a graph interpreter that “executes” the nodes one at a time in topological order, except that the result of an execution is an intermediate representation node rather than a tensor. Shape inference and layout decisions are carried out on pseudo-tensors that carry only metadata and never access real data. The three kinds of nodes are handled as follows:

1. placeholder: build an input buffer with the shape, strides, and device of the example input, whose layout cannot be changed;
2. call_function: look up the lowering rule of that operator and call it, obtaining an intermediate representation node. An operator with no lowering rule becomes a fallback kernel and is handed back to the framework for execution unchanged; operators such as matrix products and convolutions are implemented by templates or external libraries (Section 9.8);
3. output: the outputs of the graph must be written to memory; materialization is invoked for each output (Section 6.6), and its layout is determined.

6.2 Lowering Rules

Most lowering rules for pointwise operators are implemented by a common helper function, pointwise; Listing 6.1 is an abridged version of it.

Listing 6.1　The helper function for pointwise lowering (abridged, comments removed)

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,
    )

The logic of the function falls into three steps: derive the iteration range of the output from the shapes of the inputs according to the broadcast rules; obtain a loader for each input, with the loader of a tensor input rewriting its indices according to the broadcast rules and a scalar input becoming a constant; and finally construct a pointwise node whose computation function calls the loaders in turn and applies fn to their results. The whole process allocates no memory and generates no loops.

The rules for unary operators are registered in bulk from a mapping of names to ops methods, such as relu mapping to ops.relu.

There are nearly three hundred lowering rules in total, and by the intermediate representation they produce they fall into the categories of Table 6.1.

Table 6.1　Classification of Lowering Rules

The first four categories are direct instances of Definition 6.1; the later ones extend it.

The indices of an indirect read are determined by data. The $i$-th element of x[idx] is $x(\mathit{idx}(i))$; the read address is no longer an affine function of $i$, but is instead obtained by first reading idx and then using the value read as the index. In the computation function this step is expressed with ops.indirect_indexing, which converts a run-time value into an index variable. Because this value is unknown at compile time, the generated code must handle negative indices (by Python semantics, $-1$ denotes the last element) and check for out-of-bounds access, unless it can prove that the indices fall within the legal range (Section 6.8).

A scatter write is the dual of an indirect read: the write address of y[idx] = v is determined by data. If two positions have the same idx, they write the same element and the result depends on the order of execution. For overwrite semantics, the result of a repeated index has no defined convention in eager execution in the first place, and none is guaranteed after compilation either; for accumulate semantics (scatter_add, index_put(accumulate=True)) the implementation uses atomic addition, and since addition is commutative and associative (under exact arithmetic), concurrent atomic additions give the same result as any serial order.

Concatenation is a piecewise function. Suppose $x_1, \dots, x_m$ are concatenated along dimension 0 with offsets $o_1 = 0 < o_2 < \cdots < o_{m+1}$; then

Stax lowers it to a single pointwise node whose computation function performs, for each $k$, one read masked by $o_k \le i < o_{k+1}$ and merges the results by conditional selection. Concatenation therefore copies no data and can be fused with its users; the cost is that $m$ conditions must be evaluated for every element.

Type promotion is declared together with the lowering rule. Every pointwise operator carries one promotion mode: the default takes the common type of the inputs; division and sqrt, among others, are declared to convert integers to floating point first; and comparison operations are declared to always produce a Boolean. In GPU code, half-precision operations (float16, bfloat16) are carried out in float32, being promoted on load and demoted on store, which is consistent with the error analysis of Section 6.5: if an intermediate result were rounded to half precision, the $u$ in the error bound would grow from $2^{-24}$ to $2^{-11}$ or $2^{-8}$. Every lowered result is finally converted to the data type recorded at tracing time, so that it agrees with the type produced by eager execution.

6.3 Tensors as Functions

Definition 6.1 (Pointwise Nodes and Reduction Nodes)　A pointwise node consists of an iteration range $R$ and a function $\mathit{inner_fn}: R \to \mathbb V$, denoting the computation of one value for every index $i \in R$. A reduction node consists of an iteration range $R$, a reduction range $R_r$, a merge operation $\oplus$, and a function $\mathit{inner_fn}: R \times R_r \to \mathbb V$, denoting

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

where $\mathbb V$ is the set of scalar values and the operations in $\mathit{inner_fn}$ are expressed through the ops interface.

This representation views a tensor as a function from index space to values, which is the same view taken by the function (Func) in Halide [16]: it describes how each element is computed without prescribing the order of computation or the storage. Under this view, fusion between pointwise operations requires no special algorithm at all.

Property 6.1 (Pointwise Fusion Is Function Composition; an Implementation Property)　Let $u$ be a pointwise node whose computation function is $f_u$, and let $v = g(u)$ be a pointwise operation where $g$ acts on a single element. If $u$ is not written to memory, then the computation function of $v$ is

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

where $\pi$ is the index transformation used when $v$ reads $u$ (the identity when there is no view). The computation of $v$ never accesses the storage of $u$.

Explanation　This is a direct consequence of how the lowering pass is constructed, not a mathematical statement that holds independently of the implementation, which is why it is called a property. When $u$ is not written to memory, it exists in the intermediate representation only as the loader $i \mapsto f_u(i)$. The lowering rule for $v$ obtains this loader, calls it after the index transformation $\pi$, and applies $g$ to the result, which gives the formula above. $\square$

By Property 6.1, when several pointwise operations are chained, their computation functions compose layer by layer and end up performing all the computation in a single loop, with intermediate values living only in registers. The computation function of a reduction node can likewise call the loader of its upstream node, so “a pointwise chain followed by a reduction” also fuses into a single unit during lowering. Listing 6.2 gives all the buffers of Example 1.1 after lowering.

Listing 6.2　The buffers of Example 1.1 after lowering (CPU, the variable-name prefix GraphLowering_ omitted)

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,
  ...
))

Apart from the two inputs, only one buffer remains after the whole graph is lowered. add is lowered to a pointwise node that carries only the computation function; the rule for relu wraps an ops.relu around it, and multiplying by 2 does the same; the rule for sum returns a reduction node whose computation function calls the loader composed layer by layer upstream. Before the fusion algorithm even begins, Example 1.1 is already a single kernel.

Figure 6.1　The lowering of Example 1.1. Three pointwise operators compose their computation functions layer by layer and sum calls it directly; only the output is materialized as a buffer.

6.4 Layouts and Views

Definition 6.2 (Layout)　A layout consists of a shape $(s_0, \dots, s_{k-1})$, strides $(t_0, \dots, t_{k-1})$, and an offset $o$; the address of the element $i = (i_0, \dots, i_{k-1})$ is

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

Layouts fall into two categories: fixed layouts whose strides are already determined, and flexible layouts whose strides may still be adjusted by users.

An address is an affine function of the indices. View operations such as transposition, slicing, and broadcasting change no data; they only change the correspondence between indices and addresses, and can therefore all be expressed uniformly as transformations on index space.

Proposition 6.2 (Affine Representation of Views)　Dimension reordering, slicing, broadcast expansion, and the insertion or removal of dimensions of length 1 can all be represented as affine index transformations $\pi(i) = A i + c$, where $A$ is an integer matrix. If the observed tensor has layout $(t, o)$, then the view has layout

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

A composition of affine transformations is again an affine transformation, so any composition of the views above can still be represented by a set of strides and an offset.

Proof　Dimension reordering corresponds to a permutation matrix; the slice $i_j \mapsto a + k, i_j$ corresponds to the transformation whose diagonal entry is the stride $k$ and whose translation is $a$; broadcast expansion corresponds to setting the coefficient of the expanded dimension to 0; and inserting or removing a dimension of length 1 corresponds to adding or removing one row or column of $A$. Substituting into $\mathrm{addr}(\pi(i)) = o + t^{\mathsf T}(A i + c) = (o + t^{\mathsf T} c) + (A^{\mathsf T} t)^{\mathsf T} i$ gives the new strides and offset. Closure under composition follows from $A_2(A_1 i + c_1) + c_2 = (A_2 A_1) i + (A_2 c_1 + c_2)$. $\square$

Broadcast expansion corresponds to a stride of 0: moving along the expanded dimension leaves the address unchanged. A general reshape, on the other hand, need not be affine. Merging two adjacent dimensions $j$ and $j+1$ into a single dimension, the new index $m$ is related to the original indices by $i_j = \lfloor m / s_{j+1} \rfloor$ and $i_{j+1} = m \bmod s_{j+1}$, and the address

is an affine function of $m$ if and only if $t_j = s_{j+1} t_{j+1}$, that is, if the two dimensions are contiguous with each other in memory. When this condition does not hold, the reshape cannot be represented by a set of strides: the integer division and the modulo must be kept in the index expression, or the data must first be copied into a contiguous layout.

Stax distinguishes two cases accordingly. If the observed object already has a fixed layout (an input, or a buffer already written to memory), the view is represented directly as new strides and a new offset on the same storage; otherwise the view stores the index transformation $\pi$, and on a read the indices are transformed first and then the loader of the observed node is called. In both cases the view copies no data, and its effect ultimately shows up as the index expression at read time. Index expressions are represented symbolically and simplified when they are constructed.

The point of flexible layouts is this: if an intermediate buffer is only ever read after transposition, writing it out directly in the transposed order makes both the write and the read contiguous. The layouts of inputs and outputs are decided by the caller and cannot be adjusted.

Flexible layouts are fixed when lowering ends. For each buffer that is still flexible, Stax examines the strides with which its computation function reads its inputs and chooses a loop order that makes the reads as contiguous as possible under that order, then assigns strides to the buffer in the same order so that the writes are contiguous as well. When the read patterns give no preference, the default row-contiguous layout is used. This decision concerns only intermediate buffers and therefore does not change the semantics of the program.

Example 6.1　Suppose f(x, y) = x.t() + y, where $x$ has shape $512\times256$ and $y$ has shape $256\times512$, both row-contiguous. The only buffer after lowering is shown in Listing 6.3.

Listing 6.3　The buffer of Example 6.1 after lowering (targeting CUDA, prefixes omitted)

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])
))

The transposition produced no node of its own; it shows up only as the index i0 + 256*i1 of $x$. By Proposition 6.2, the strides of $x$ are $(256, 1)$ and the transposition corresponds to a permutation matrix, giving new strides $(1, 256)$ and the address $i_0 + 256 i_1$, in agreement with the listing. In the indices of the two inputs, the coefficient of i1 is 256 for one and 1 for the other, that is, along the innermost dimension one input is non-contiguous and the other is contiguous. This difference determines the tiling at code generation time (Section 8.1、Section 9.3).

6.5 Algebraic Properties of Reductions and Their Splitting

How a reduction is implemented depends on the algebraic properties of its merge operation.

Definition 6.3 (Monoid)　A set $M$ together with a binary operation $\oplus$ and an element $e$ is called a monoid if $\oplus$ is associative and $e$ is the identity element: $e \oplus a = a \oplus e = a$. If $\oplus$ is also commutative, it is called a commutative monoid.

Sum, product, maximum, and minimum are all commutative monoids, with identity elements $0$, $1$, $-\infty$, and $+\infty$ respectively. The merge object of argmax is the pair (value, index), compared by value and, when the values are equal, resolved by the smaller index; this too forms a commutative monoid. The one-pass algorithm for mean and variance (the Welford algorithm) can likewise be written as a merge operation over triples (count, mean, second central moment) [17].

Associativity is the basis for parallelizing every reduction: partitioning the reduction range into segments $R_r = R^{(1)} \cup \cdots \cup R^{(k)}$ gives

and the partial results of the segments can be computed independently and then merged. Commutativity, in turn, allows the elements to be accessed in any order, for example when the channels of a vectorized reduction are interleaved.

Floating-point addition is not associative, so the formula above holds only in an approximate sense. Different summation orders give different rounding errors and therefore different error bounds. Let the unit roundoff be $u$; for $n$ numbers [18]:

1. For sequential summation the error satisfies $|\hat s - s| \le (n-1),u \sum_i |a_i| + O(u^2)$;
2. For $k$ segments, summed sequentially within each segment and then merged sequentially across the segments, the coefficient of the error bound is $(n/k - 1) + (k - 1)$;
3. For pairwise tree summation, the coefficient of the error bound is $\lceil \log_2 n \rceil$.

By the second item, taking $k \approx \sqrt n$ makes the coefficient about $2\sqrt n$, significantly smaller than $n - 1$. Listing 1.1 accumulates over 16 vector channels separately and finally reduces to a scalar, which is equivalent to a segmented summation with $k = 16$, and its error bound is no worse than that of sequential summation; the $1.5\times10^{-5}$ difference between it and the result of eager execution comes from the rounding of the two summation orders, not from an error in either of them.

The motivation for splitting is parallelism. If there are few outputs and a long reduction — for example, summing a $64\times16384$ tensor by rows — there are only 64 outputs; on a GPU, if each output is handled by one thread block, there are at most 64 thread blocks, far from filling the device. Splitting cuts each row into $k$ segments, the first level of reduction produces $64k$ partial sums, and the second level merges them, raising the number of parallel work items to $64k$. The cost is writing out and reading back $64k$ partial sums. A CPU has far fewer threads than a GPU has simultaneously resident thread blocks, so parallelizing over outputs is usually enough there; splitting is mainly for the GPU.

Choosing the number of segments must balance two considerations: the number of segments must be enough to fill the device, but each segment must not be too short, otherwise each thread receives too few elements and the share of launch and merge overheads grows. The rule in Stax on the GPU is this: if the number of outputs is at least $2 \cdot 32 \cdot S$ ($S$ is the number of streaming multiprocessors), do not split; otherwise determine the number of segments with Algorithm 6.1. Reductions that carry indices, such as scan, dot product, and argmax, are never split.

Algorithm 6.1　Number of segments for an inner reduction on the 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)⌉

Example 6.2　Suppose f(x) = (x * x).sum(-1), $x$ has shape $64\times16384$, and the target device has $S = 132$ streaming multiprocessors, $w = 32$, and $T = 2048$. Then $m = 256$ and $D_{\min} = 32 \cdot 132 \cdot 2048 = 8,650,752$. Since $n = 64 < 2S = 264$ and $n \cdot n_r = 1,048,576 \le D_{\min}$, take $s = 32$, giving $\lceil 16384 / (32 \cdot 256) \rceil = 2$ segments. The lowered result is shown in Listing 6.4.

Listing 6.4　The two-level reduction of Example 6.2 after lowering (targeting CUDA, prefixes and some fields omitted)

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,
))

The iteration range of the first level of reduction is widened from $[64]$ to $[64, 2]$, and the reduction range shrinks from $[16384]$ to $[8192]$; the index 16384*i0 + 8192*i1 + r0_0 maps segment $i_1$ to the first or the second half of the original reduction dimension, which is exactly the inner sum on the right-hand side of the associativity identity. The second level of reduction merges the partial sums along the dimension of length 2. The 128 partial sums that are written out and read back are negligible compared with the $2^{20}$ input elements.

The computation function of the first level reads the same location twice, because the two operands of x * x are the same loader. The loop-level intermediate representation performs no common subexpression elimination; this redundancy is left to be removed by the in-kernel CSE at code generation time.

6.6 Materialization

The premise of Property 6.1 is that “$u$ is not written to memory”. If $u$ has several users, inlining means that every user recomputes $u$ in its own loop. Whether $u$ is written to memory is called the materialization (realization) decision.

A simple cost model illustrates this decision. Suppose $u$ has $N$ elements, computing each element takes $c$ operations, $q$ input bytes are read, it has $k$ users, the processor has arithmetic throughput $P$ and bandwidth $B$, and each element of $u$ occupies $b$ bytes. With inlining, $u$ is computed $k$ times and its inputs are read $k$ times:

With materialization, $u$ is computed once and written once, and then read $k$ times by its users:

The difference between the two is

When $k = 1$ the difference is negative, so inlining is always better, which is exactly why values with a single user are always inlined. When $k > 1$, the larger $c$ is, the more input $q$ is read, and the more users $k$ there are, the more materialization is favoured.

Stax approximates this criterion with a set of thresholds, checked in turn for nodes that have several users; if any one holds, the value is materialized:

1. On the CPU it includes one of exp, log, log10, log1p, log2, sigmoid, tanh. These transcendental functions are expensive on the CPU, corresponding to a large $c / P$; on the GPU they have dedicated hardware units and their cost is comparatively low;
2. On the CPU, there are more than 5 users and the number of operations in the computation function is close to the limit, corresponding to a large $k$;
3. The number of operations in the computation function exceeds the limit (30 on the GPU, 50 on the CPU), corresponding to a large $c$;
4. More than 4 buffers are read, corresponding to a large $q$; likewise when the accumulated number of reads after layer-by-layer inlining is too large (more than 8 on the GPU, more than 12 on the CPU).

Moreover, the outputs of the graph and the inputs that must be handed to external kernels are necessarily materialized. If some user inevitably requires a value to be materialized, then the value is materialized as soon as it is produced, so that the users placed before it do not each recompute it.

The CPU and the GPU use different thresholds, reflecting the differences in $P/B$ and in the relative cost of transcendental functions between the two kinds of processors. The same program may obtain different materialization results on the two kinds of devices; Example 7.2 is a case in point.

6.7 One Description, Many Interpretations

The ops inside a computation function are not a fixed implementation but a replaceable processor. The same computation function, executed under different processors, yields different results:

1. Executed under the processor that records reads and writes, it yields the buffers and indices it reads and writes, which are the basis of dependence analysis and fusion scoring (Section 7.1);
2. Executed under the counting processor, it yields its number of operations and the set of operations used, which is the basis of the materialization criterion of Section 6.6;
3. Executed under the formatting processor, it yields the text of Listing 6.2;
4. At the code generation stage, executed under a Triton or C++ processor, its output is the statements in the kernel source. For example, the Triton processor prints ops.exp as libdevice.exp(...), and the C++ vector processor prints it as an operation on Vectorized<float>.

The theoretical background of this design is abstract interpretation [19]: the same program, interpreted in different semantic domains, yields different properties of the program. The domain of concrete values gives the results of the computation, the domain of read/write sets gives the data dependences, the domain of operation counts gives the cost estimates, and the domain of source text gives the code. Whoever writes a lowering rule writes the computation function once, and the various analyses and generation steps of the backend are all derived from it.

6.8 Normalization of Loop Bodies and Range Analysis

The computation function is a Python closure, which makes composition easy but analysis awkward. Before scheduling and code generation, the computation function of each buffer is executed once on symbolic indices and recorded as a loop body: a small computation graph in which the operations come from ops calls, while the index expressions are stored separately as symbolic expressions. This is still the approach of Section 6.7, with a processor that records the operations and collects the indices. Two kinds of simplifications are carried out on the loop body, both of which are independent of semantics and affect only efficiency.

The first kind is loop normalization: removing dimensions of length 1, merging adjacent loops that can be merged, and reordering loops according to the strides of the memory accesses. The condition for merging is the same as the condition for reshape in Section 6.4.

Proposition 6.3 (Merging Loops)　Suppose the loop body contains two adjacent nested loops $i_j \in [0, s_j)$ and $i_{j+1} \in [0, s_{j+1})$, and the index of every memory access is an affine function of the loop variables, in which the coefficients of $i_j$ and $i_{j+1}$ are $a_j$ and $a_{j+1}$ respectively. If $a_j = s_{j+1} a_{j+1}$ holds for every memory access, then the two loops can be replaced by $m = s_{j+1} i_j + i_{j+1} \in [0, s_j s_{j+1})$ with every address unchanged.

Proof　$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$, and $(i_j, i_{j+1}) \mapsto m$ is a bijection from $[0, s_j) \times [0, s_{j+1})$ to $[0, s_j s_{j+1})$. $\square$

The first level of reduction of Example 6.2 (Listing 6.4) reads primals_1 at the index 16384*i0 + 8192*i1 + r0_0. After loop normalization, this read appears in the generated Triton kernel as

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

16384*i0 + 8192*i1 has been merged into 8192*x0, where x0 ranges over the dimension of length 128 after merging. This is an instance of Proposition 6.3: $s_1 = 2$ and $a_0 = 16384 = 2 \cdot 8192$. Merging reduces the number of loop levels and makes the innermost loop longer, which favours vectorization and parallel decomposition.

The second kind is range analysis. Code generation has to answer several questions about the ranges of integer values: whether the maximum value of an index expression is less than $2^{31}$, so that addresses can be computed with 32-bit integers; and whether an indirect index is guaranteed to lie in $[0, s)$, so that the correction for negative indices and the bounds check can be dropped. These questions generally cannot be answered exactly at compile time, but they can be approximated conservatively with intervals.

Definition 6.4 (Interval Domain)　A closed interval $[l, u]$ ($l, u \in \mathbb Z \cup {\pm\infty}$) abstracts all the possible values of an integer variable. A loop variable $i \in [0, s)$ is abstracted as $[0, s-1]$, a constant $c$ as $[c, c]$, and a value read from memory as $[-\infty, +\infty]$. Operations are propagated by interval arithmetic, for example

$$[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).$$

Proposition 6.4 (Soundness)　If the interval propagation rule of every operation contains the results of that operation on all concrete inputs, then for every variable in the loop body the interval obtained by the analysis contains all of its values at run time.

Proof　Induction over the topological order of the loop body. The intervals of the inputs (loop variables, constants, values read) contain their values by definition. If the values of the operands of an operation all lie within their intervals, then by the assumption on the propagation rule, the values of the result lie within the interval of the result. $\square$

This is a special case of abstract interpretation [19] on the interval domain. What Proposition 6.4 guarantees is a one-directional conclusion: when the interval says “certainly within $[0, s)$”, the value at run time really is so; when the interval is wider, that does not mean the run-time value goes out of bounds, only that the analysis cannot prove it. Hence the checks are dropped only when the interval is narrow enough, and are otherwise retained, at the cost of a small amount of run-time overhead but with no change in semantics.

Example 6.3　Compare two indirect reads:

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

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

The idx of g1 is read from memory, and its interval is $[-\infty, +\infty]$, so the generated CPU kernel retains the correction for negative indices and the bounds check, as shown in Listing 6.5.

Listing 6.5　The kernel of g1 (CPU, abridged)

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

The index of g2 is computed by arange, and the computation function of arange is the loop variable itself. Index propagation turns the whole index computation into the symbolic expression $(3 i) \bmod 128$, and the indirect read thus becomes a direct read; by Definition 6.4, the interval of this expression is $[0, 127]$, so both the correction and the check are dropped, as shown in Listing 6.6.

Listing 6.6　The kernel of g2 (CPU, abridged)

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))];
}

The kernel of g2 no longer contains the buffer idx: arange, the multiplication, and the modulo have all been absorbed into the index expression. This is another form of Property 6.1: a value that depends only on the loop variables need not be written to memory, and its computation function can be substituted directly into the index of the user.

The conclusions of the range analysis are also used to choose the bit width of indices. If the upper bounds of all the index intervals in a kernel are less than $2^{31}$, address computation can be done with 32-bit integers; on the GPU this reduces the number of registers occupied by each thread and, according to the analysis of Section 9.1, can raise the number of resident warps.

Materialization only decides which values are written to memory. Whether several buffers can be merged into one kernel is decided by the scheduler.

6.9 Semantics of the Core Language and Correctness of the Transformations

The preceding sections have repeatedly used the phrase “preserving semantics”: common subexpression elimination preserves semantics (Section 4.5), pointwise fusion is function composition (Property 6.1), and materializing or not does not change the result (Section 6.6). All of these statements have so far been informal. This section gives a denotational semantics [52] for a core subset of the loop-level intermediate representation, states them as theorems on that basis and proves them, and at the same time explains which transformations hold only under real-number semantics.

Definition 6.5 (Tensor Core Language)　Each tensor variable $y$ carries a domain $D_y \subset \mathbb Z^{d}$ (a finite rectangular box). Expressions are generated by the following grammar:

$$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),$$

where $x$ is an input or a bound variable; $g : \mathbb V^r \to \mathbb V$ is a scalar function; $\pi : D \to D’$ is an affine map taking the index of the result to the index of the observed tensor (Proposition 6.2); and $\mathrm{red}$ reduces along the reduction range $R_r$ with the operation $\oplus$ and the initial value $e_0$. A program is a sequence of bindings $\mathbf{let}\ y_1 = e_1; \dots; \mathbf{let}\ y_m = e_m;\ \mathbf{return}\ y_m$, where each $y_k$ is bound exactly once.

This is precisely the two kinds of nodes of Section 6.3together with the views of Section 6.4: $\mathrm{map}$ corresponds to a pointwise node and $\mathrm{red}$ to a reduction node. The language has no in-place mutation, which is guaranteed by the functionalization of Section 4.9; it also has no operations not determined by the inputs, such as random numbers, and this class of operations is excluded from common subexpression elimination in Section 4.5— the reason will become clear below.

Semantics. Take $\mathbb V$ to be the set of floating-point values (including infinities and NaN); the value of a tensor with domain $D$ is a function $D \to \mathbb V$. An environment $\rho$ maps variables to tensor values. The denotation $[![ e ]!]\rho$ of an expression is defined inductively as follows:

where $r_1 < r_2 < \dots < r_q$ is the lexicographic ordering of the elements of $R_r$. The denotation of a program is the value of $y_m$ after successively extending the environment as $\rho_k = \rho_{k-1}[y_k \mapsto [![ e_k ]!]\rho_{k-1}]$. The semantics of a reduction deliberately prescribes a definite evaluation order, because the floating-point operation $\oplus$ is in general not associative (Section 6.5); only once the order is fixed does “preserving semantics” have a precise meaning.

Definition 6.6 (Semantic Equivalence)　Two expressions (or programs) $e$ and $e’$ are called equivalent, written $e \equiv e’$, if for every well-formed environment $\rho$ (which assigns a tensor value with the correct domain to every free variable) one has $[![ e ]!]\rho = [![ e’ ]!]\rho$, that is, they are pointwise equal as functions $D \to \mathbb V$. Every construct of the core language is a total function under floating-point semantics (an invalid operation yields NaN rather than raising an error), so the denotations are defined everywhere and $\equiv$ is an equivalence relation.

Equality here is bitwise equality of floating-point values, which is far stricter than being numerically close.

Lemma 6.1 (Substitution Lemma)　Suppose the variable $y$ does not occur freely in $e$. Then for every expression $b$ and environment $\rho$,

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

where $b[e/y]$ denotes $b$ with every occurrence of $y$ replaced by $e$.

Proof　Induction on the structure of $b$, writing $\rho’ = \rho[y \mapsto [![ e ]!]\rho]$. When $b = y$, both sides are $[![ e ]!]\rho$; when $b$ is another variable $x$, both sides are $\rho(x)$. When $b = \mathrm{map}_g(b_1, \dots, b_r)$, by the semantics definition $[![ b ]!]\rho’$ depends on $\rho’$ only through the various $[![ b_k ]!]\rho’$, which by the inductive hypothesis equal $[![ b_k[e/y] ]!]\rho$, and $\mathrm{map}_g(b_1[e/y], \dots) = b[e/y]$. The cases of $\mathrm{view}$ and $\mathrm{red}$ are analogous. When $b = (\mathbf{let}\ z = e’;\ b’)$, by single assignment we may assume $z \ne y$ and that $z$ does not occur freely in $e$. By the semantics, the left-hand side equals $[![ b’ ]!]\bigl(\rho’[z \mapsto w]\bigr)$, where by the inductive hypothesis $w = [![ e’ ]!]\rho’ = [![ e’[e/y] ]!]\rho$. Since $z$ does not occur freely in $e$, $[![ e ]!]\rho = [![ e ]!]\bigl(\rho[z \mapsto w]\bigr)$; and since $z \ne y$, the two extensions commute, so $\rho’[z \mapsto w] = \rho[z \mapsto w]\bigl[y \mapsto [![ e ]!](\rho[z \mapsto w])\bigr]$. Applying the inductive hypothesis to $b’$ and the environment $\rho[z \mapsto w]$, the left-hand side equals $[![ b’[e/y] ]!]\bigl(\rho[z \mapsto w]\bigr) = [![ \mathbf{let}\ z = e’[e/y];\ b’[e/y] ]!]\rho = [![ b[e/y] ]!]\rho$. $\square$

The proof of Lemma 6.1 uses only one point: the denotation of each construct is determined solely by the denotations of its subexpressions (compositionality), and the denotation is a function of the environment (determinism). Random number generation does not satisfy the latter: two calls to rand give different results, so it is not a function of the environment, and the theorems below do not hold for it.

Theorem 6.1 (Inlining and Materialization Preserve Semantics)　$\mathbf{let}\ y = e;\ b \ \equiv\ b[e/y]$.

Proof　The denotation of the left-hand side is $[![ b ]!]\bigl(\rho[y \mapsto [![ e ]!]\rho]\bigr)$, which by Lemma 6.1 equals that of the right-hand side. $\square$

Read from left to right, Theorem 6.1 is inlining: substituting the producer’s computation function into the consumer, which is exactly the case of a pointwise node not being written to memory during lowering (Section 6.3). Read from right to left, it is materialization: binding the value of a subexpression to a variable and writing it into a buffer (Section 6.6). Both directions preserve the bitwise-equal semantics, so whether to materialize is purely a cost question and not a correctness question, which is precisely why Section 6.6can make the decision using a cost model alone. This conclusion has two implicit premises: the storage type of a buffer is the same as its computation type (if an intermediate value computed in fp32 is stored in bf16, materialization costs one extra rounding), and code generation does not contract multiply-add into fused multiply-add (FMA) instructions, since otherwise the number of roundings would differ before and after inlining.

Corollary 6.1 (The Semantic Form of Property 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)$, and the denotation of both is $\lambda i.\ g(f(x(\pi(i))))$.

Proof　Expanding both sides layer by layer according to the semantic definition gives the result. $\square$

Theorem 6.2 (Common Subexpression Elimination and Dead Code Elimination)　Suppose $y_1$ does not occur freely in $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) If $y$ does not occur freely in $b$, then $\mathbf{let}\ y = e;\ b \ \equiv\ b$.

Proof　(1) The left-hand side binds $y_2$ under $\rho_1 = \rho[y_1 \mapsto [![ e ]!]\rho]$. Since $y_1$ does not occur freely in $e$, $[![ e ]!]\rho_1 = [![ e ]!]\rho = \rho_1(y_1)$. Hence $\rho_2 = \rho_1[y_2 \mapsto \rho_1(y_1)]$, and by Lemma 6.1 (taking the substituted expression to be the variable $y_1$), $[![ b ]!]\rho_2 = [![ b[y_1/y_2] ]!]\rho_1$. (2) Since $y$ does not occur freely in $b$, $[![ b ]!]$ does not read $\rho(y)$. $\square$

The second part of Theorem 6.2 depends on an implicit premise: the evaluation of $e$ always terminates and yields a value. If the language contained operations that raise errors (integer division by zero, an out-of-bounds index, for instance), deleting an unused $e$ would turn “raise an error” into “return normally”. Equivalence would then have to be relaxed to refinement: the transformed program gives the same result on the inputs on which the original is defined, and may give any result on the inputs on which the original raises an error. The verified compiler CompCert states its correctness theorem in terms of exactly this refinement relation [56].

The transformations covered by the first two theorems all preserve bitwise-equal semantics. Splitting and reordering reductions do not: they change the evaluation order prescribed in the definition.

Theorem 6.3 (Real-Number Semantics of Reduction Reordering)　Suppose $\oplus$ is associative and commutative over the reals and $e_0$ is the identity element. Write $[![ \cdot ]!]_{\mathbb R}$ for the semantics in which $\mathbb V$ is replaced by $\mathbb R$ (for $\max$ and $\min$ with initial values $\mp\infty$, by the extended reals $\mathbb R \cup {\pm\infty}$) and operations are interpreted as exact real arithmetic. Then for any permutation of $R_r$, any parenthesization, and any partitioning and merging as in Section 6.5, the real-number semantics of the reduction agrees with the order given in the definition.

Proof　Proved by induction on the number $q$ of elements being reduced: any parenthesization and permutation of the $\oplus$ expression equals the left-associated form of $v_{r_1} \oplus \cdots \oplus v_{r_q}$. When $q \le 2$ this follows immediately from commutativity. When $q > 2$, the outermost level of the expression splits the elements into two nonempty groups $P$ and $Q$; by the inductive hypothesis the two groups equal their own left-associated forms in the defined order; and then, by associativity and commutativity, moving the elements of $Q$ into the left-associated form of $P$ one by one gives the left-associated form of all the elements. A partitioned reduction is a special case of a parenthesization, and since the initial value $e_0$ is the identity element, it can be inserted or deleted at any position. $\square$

The transformations therefore fall into two classes. The first class preserves floating-point semantics: inlining and materialization, common subexpression elimination, dead code elimination, composition of views (Proposition 6.2), and functionalization. For them, the results before and after compilation are bitwise identical, and tests may demand exact equality. The second class preserves only real-number semantics: partitioning and reordering of reductions (Section 6.5), vectorized reductions (Section 8.1), and rescaling in online softmax (Section 9.5). For them, correctness means “the real-number semantics agree and the difference between the floating-point results does not exceed the error bound of Section 6.5”, and tests must leave a tolerance commensurate with the condition number (Section 13.3). Section 4.2rejects algebraic simplifications such as re-association precisely so that graph transformations all fall into the first class; the second class appears only in lowering and code generation, and each occurrence has a corresponding error analysis.

The semantics of this section is denotational and covers only a side-effect-free tensor core. A complete compiler correctness proof must also handle in-place mutation and aliasing (the functionalization of Section 4.9can be regarded as a translation from a semantics with storage to the semantics of this section), guards and specialization (第 Section 2), and run-time fallback. Verified compilers carry out proofs of this kind in a machine proof assistant [56]; Stax performs no formal verification and instead checks the same property empirically through the layered comparison of Section 11.4.

Previous: Part One　Introduction, Specialization, Graph Capture and Graph Transformation　｜　Table of Contents　｜　Next: Part Two　Scheduling, Fusion and CPU Code Generation