Docs / nablatensor-core / com.nablatensor.engine
final class
AadCheckpointPlan
Segment-checkpoint layout for a long tape, shared by the CUDA and Vulkan adjoint code generators.
A fully unrolled reverse sweep keeps every forward value v_i live until the reverse consumes it — for a 1500-node tape that is a live set no GPU register file holds, so the shader spills to scratch memory and occupancy collapses. This plan cuts the tape into segments contiguous pieces of ~sqrt(nodes) nodes each. The forward runs piece by piece and, at each internal boundary, saves only the handful of values that later nodes still read (the carry set). The reverse then walks the pieces backwards: for each it restores the boundary carry, recomputes that piece's forward into a small local set, and runs the piece's reverse. Peak live set drops from O(nodes) to O(sqrt(nodes)) at the cost of recomputing the forward once.
The carry set at boundary b is C[b] = { i < b : some node j >= b reads i as an argument }. A carry node whose last use is at or beyond the final internal boundary is global: it stays in a register for the whole kernel body and is never written to the checkpoint buffer.
Fields
bound[0] == 0, bound[segments] == nodes, strictly increasing.
lastUse[i] = greatest node index that reads i; nodes for the output; -1 if unread.
carry[s] = sorted C[bound[s]] for s in 1..segments-1; carry[0] empty.
Nodes kept in a register for the whole kernel (carried across every internal boundary).
Float offset into a path's checkpoint row where carry[s] \ global starts.
carry[s] \ global, i.e. the nodes segment s actually writes to / reads from the checkpoint buffer.
Total floats per path in the checkpoint buffer.
Methods
Builds a plan, or returns null when checkpointing does not apply: value-only kernels, tapes at or below minNodes, tapes that would split into fewer than two segments, or tapes whose output node is not the final node / is not differentiable (the unrolled path handles those, and they are never the large tapes checkpointing exists for).
Union of every carry[s]: the nodes that need a persistent reverse accumulator.
The checkpoint slot for node i at boundary s, or -1 if it is global / not in the slice.
Value reference for node i: a whole-kernel register for a global carry node, a segment local otherwise.
Segment index containing node i.