CUDA Autograd Tape

Shape-erased CUDA tape machinery for dtype-carrying LibTorch.Buffer values. This is the GPU/runtime analogue of the CPU autograd tape: it records node ids, runtime shapes, parent links, and backward callbacks, then performs reverse-mode accumulation over buffers.

@[reducible, inline]

abbrev Runtime.Autograd.LibTorch.Result (α : Type) : Type

Pure error monad for the CUDA tape. Mirrors Engine/Core.

structure Runtime.Autograd.LibTorch.AnyBuffer : Type

Runtime, shape-erased CUDA buffer.

This plays the same role as Spec.SomeTensor in the CPU tape: it pairs runtime Shape metadata with an opaque LibTorch.Buffer handle.

• s : Spec.Shape

  Runtime shape metadata for the buffer.

• buf : Buffer

  Device buffer. Length is expected to be Spec.Shape.size s; dtype is retained natively.

def Runtime.Autograd.LibTorch.AnyBuffer.natToU32Checked (n : ℕ) : Result UInt32

Checked conversion of a Nat length to UInt32, erroring on overflow.

def Runtime.Autograd.LibTorch.AnyBuffer.numelU32 (s : Spec.Shape) : Result UInt32

Number of scalar elements as a UInt32, with every individual shape dimension checked too.

def Runtime.Autograd.LibTorch.AnyBuffer.zeros (s : Spec.Shape) (dtype : Dtype := Dtype.float32) : Result AnyBuffer

Allocate a zero-filled buffer for a given shape.

def Runtime.Autograd.LibTorch.AnyBuffer.checkDtype (value reference : AnyBuffer) : Result Unit

Reject an implicit promotion at a value or cotangent boundary.

def Runtime.Autograd.LibTorch.AnyBuffer.validate (value : AnyBuffer) : Result AnyBuffer

Check that runtime shape metadata agrees with the native CUDA buffer length.

AnyBuffer is intentionally shape-erased, so its structure alone cannot enforce this invariant. Call this at boundaries that accept buffers assembled outside the typed tensor conversion path, before any kernel uses shape-derived indexing.

def Runtime.Autograd.LibTorch.AnyBuffer.add (a b : AnyBuffer) : Result AnyBuffer

Accumulate two AnyBuffer values by elementwise addition, with a dynamic shape check.

This is used by backprop to sum gradient contributions in DAGs. The addition runs only after the shape metadata agree and both native buffer lengths match that shape.

structure Runtime.Autograd.LibTorch.Node : Type

CUDA tape node representing one recorded computation step.

Fields mirror the CPU Engine/Core node:

• value holds the forward buffer.
• parents are tape node ids.
• backward is a local VJP rule producing parent gradient contributions.

• name : Option String

  Optional node name for debugging/pretty-printing.

• value : AnyBuffer

  Forward value computed at this node.

• ownsValue : Bool

  Whether this node owns its forward allocation.

  Shape views and detached references borrow a parent's buffer. Releasing a borrowed value would also invalidate the parent, including parameter mirrors that must survive a tape reset.

• requiresGrad : Bool

  Whether reverse-mode propagation should visit this node.

• parents : Array ℕ

  Parent node ids (dependencies) in the tape.

• cleanup : Array Buffer

  Forward workspace buffers retained only because this node's backward closure may need them.

  The eager runtime releases these buffers explicitly after backprop consumes the tape, so long CUDA training loops do not wait on Lean external-object finalizers for large intermediate allocations.

• backward : AnyBuffer → Result (Array (ℕ × AnyBuffer))

  Local VJP rule for this node.

  Given an upstream cotangent for value, return an array of (parentId, parentCotangent) contributions (one per parent, usually).

structure Runtime.Autograd.LibTorch.Tape : Type

CUDA autograd tape: a grow-only array of nodes. Node ids are array indices.

• nodes : Array Node

  Tape nodes in evaluation order (id = index).

def Runtime.Autograd.LibTorch.Tape.empty : Tape

Empty tape (no nodes).

def Runtime.Autograd.LibTorch.Tape.size (t : Tape) : ℕ

Number of nodes stored in the tape.

def Runtime.Autograd.LibTorch.Tape.getNode? (t : Tape) (id : ℕ) : Option Node

Read a node by id (returns none if out of bounds).

def Runtime.Autograd.LibTorch.Tape.getValue? (t : Tape) (id : ℕ) : Option AnyBuffer

Read just the stored forward value for a node id.

@[inline]

def Runtime.Autograd.LibTorch.Tape.addNode (t : Tape) (node : Node) : Tape × ℕ

Append a node and return its id.

Invariant: the returned id is t.size, the pre-append size of the tape.

@[simp]

theorem Runtime.Autograd.LibTorch.Tape.addNode_id (t : Tape) (node : Node) : (t.addNode node).2 = t.size

addNode returns the current tape size as the fresh node id.

@[simp]

theorem Runtime.Autograd.LibTorch.Tape.size_addNode (t : Tape) (node : Node) : (t.addNode node).1.size = t.size + 1

Appending a node increases the tape size by one.

@[inline]

def Runtime.Autograd.LibTorch.Tape.leaf (t : Tape) (value : AnyBuffer) (name : Option String := none) (requiresGrad : Bool := true) : Tape × ℕ

Add a leaf node (no parents).

PyTorch comparison: a tensor that enters the graph as a leaf (e.g. input or parameter value).

This low-level constructor records value without validating its external buffer. Code that wraps an externally supplied Buffer must call AnyBuffer.validate first; tape operations validate their operands again when they retrieve values.

def Runtime.Autograd.LibTorch.Tape.requireValue (t : Tape) (id : ℕ) (s : Spec.Shape) : Result Buffer

Read a buffer value from a tape node id, requiring a specific runtime shape.

Fails if:

• the id is invalid, or
• the stored shape does not match s,
• the shape cannot be represented by the UInt32 CUDA ABI, or
• the native buffer length differs from Shape.size s.

def Runtime.Autograd.LibTorch.Tape.requireGrad (dLdyAny : AnyBuffer) (expected : Spec.Shape) : Result AnyBuffer

Require that an upstream gradient matches an expected runtime shape.

This is used inside backward closures to validate/cast the incoming cotangent. Both the shape tag and native buffer length are checked before a VJP reads the cotangent.

@[inline]

def Runtime.Autograd.LibTorch.Tape.unary (t : Tape) (opName : String) (xId : ℕ) (σ τ : Spec.Shape) (forward : Buffer → Buffer) (backward : Buffer → Buffer → Buffer) (ownsValue : Bool := true) : Result (Tape × ℕ)

Generic constructor for unary ops.

You provide:

• forward : Buffer → Buffer
• backward : Buffer → Buffer → Buffer (VJP; given input x and upstream dLdy, return dLdx)

Shapes are explicit and checked dynamically. Reverse mode visits the result only when the input requires gradients, so constant subexpressions do not execute their backward closures.

@[inline]

def Runtime.Autograd.LibTorch.Tape.binary (t : Tape) (opName : String) (aId bId : ℕ) (σ₁ σ₂ τ : Spec.Shape) (forward : Buffer → Buffer → Buffer) (backward : Buffer → Buffer → Buffer → Buffer × Buffer) : Result (Tape × ℕ)

Generic constructor for binary ops.

You provide:

• forward : Buffer → Buffer → Buffer
• backward : Buffer → Buffer → Buffer → (Buffer × Buffer) (VJP; given inputs a, b, and upstream dLdy, return (dLda, dLdb))

Shapes are explicit and checked dynamically. The result requires gradients if either input does.

def Runtime.Autograd.LibTorch.Tape.addGradAll (t : Tape) (grads : Array AnyBuffer) (id : ℕ) (g : AnyBuffer) : Result (Array AnyBuffer)

Internal helper: add a gradient contribution g into the dense gradient array at id.

This checks:

• id is a valid node id,
• the parent requires gradients,
• the contribution shape matches the parent's value shape, then accumulates via AnyBuffer.add.

When the parent does not require gradients, the contribution is not stored anywhere. We still have to release its device buffer. The release is threaded through the existing zero slot with releaseThen, so the cleanup is part of the returned gradient array instead of a dead pure call.

def Runtime.Autograd.LibTorch.Tape.backwardDenseFromStep (t : Tape) (acc : Array AnyBuffer) (id : ℕ) : Result (Array AnyBuffer)

One reverse-mode backprop step at a single node id, updating the dense gradient array.

The incoming array is total: every tape node has a gradient buffer, initialized to zero unless a later node has already contributed to it. That total representation is convenient for the CUDA runtime because every slot has a concrete device buffer that can be released deterministically.

def Runtime.Autograd.LibTorch.Tape.backwardDenseFromLoop (t : Tape) : ℕ → Array AnyBuffer → Result (Array AnyBuffer)

Reverse-mode accumulation over the first n nodes in reverse order.

The recursion visits n-1, n-2, ..., 0; using n = t.size runs the full tape. We keep this as a structural loop rather than a list fold so proof layer callers can reason about one node step at a time.

def Runtime.Autograd.LibTorch.Tape.backwardDenseFrom (t : Tape) (grads0 : Array AnyBuffer) : Result (Array AnyBuffer)

Reverse-mode accumulation starting from an explicit dense gradient array.

This expects grads0.size = t.size. The function is useful for callers that already seeded multiple outputs or want to run a custom cotangent initialization. Before traversal, every slot is checked against its node's shape and native buffer length.

def Runtime.Autograd.LibTorch.Tape.backwardDenseReachableStep (t : Tape) (reachable : Array Bool) (grads : Array AnyBuffer) (id : ℕ) : Result (Array Bool × Array AnyBuffer)

Run one dense reverse step only when id has been reached from the selected output, and mark the parents that receive concrete VJP contributions. The gradient array remains total so consumed CUDA buffers continue to follow the ownership discipline in addGradAll.

def Runtime.Autograd.LibTorch.Tape.backwardDenseReachableLoop (t : Tape) : ℕ → Array Bool → Array AnyBuffer → Result (Array AnyBuffer)

Reverse traversal used by backwardDenseAll, carrying the set of reached node ids.

def Runtime.Autograd.LibTorch.Tape.backwardDenseAll (t : Tape) (outId : ℕ) (seed : AnyBuffer) : Result (Array AnyBuffer)

Reverse-mode accumulation that returns a dense gradient buffer for every node id.

All gradients are initialized to zeros (using each node's runtime Shape), then the output node is seeded and only nodes reached by VJP propagation are executed. Disconnected slots stay explicit zero buffers; their backward closures are not called.

Sparse gradients

Training only needs cotangents for parameter leaves. The sparse traversal below releases each activation cotangent immediately after its local VJP has propagated it and retains owned copies only for node ids selected by the caller.

@[reducible, inline]

abbrev Runtime.Autograd.LibTorch.Tape.SparseGradMap : Type

Device gradients retained for selected tape node ids.

def Runtime.Autograd.LibTorch.Tape.releaseSparseBuffer (b : Buffer) : IO Unit

Sequence a native buffer release inside IO.

def Runtime.Autograd.LibTorch.Tape.releaseSparseGrads (grads : SparseGradMap) : IO Unit

Release every buffer owned by a sparse gradient map.

def Runtime.Autograd.LibTorch.Tape.addSparseGrad (t : Tape) (gradsRef : IO.Ref SparseGradMap) (id : ℕ) (g : AnyBuffer) : IO Unit

Insert or accumulate an owned VJP contribution into the sparse gradient map.

CUDA backward rules return fresh contribution buffers. The sparse map copies a first contribution so its retained values have uniform ownership, then retires the consumed contribution immediately.

def Runtime.Autograd.LibTorch.Tape.backwardSparse (t : Tape) (outId : ℕ) (seed : AnyBuffer) (retain : ℕ → Bool) : IO SparseGradMap

Run reverse mode while retaining gradients only for node ids accepted by retain.

This function consumes seed. The returned buffers are owned by the map and must be released with releaseSparseGrads after the optimizer has consumed them. All other activation gradients are retired during the traversal.