netcl wiki
architecture

Architecture: Autograd & Tape

Architecture: Autograd & Tape

netcl's automatic differentiation engine is a classical reverse-mode tape-based autograd system built around a per-thread Tape. When you wrap execution inside with ag.Tape():, every differentiable operation creates a Node that records: 1. Its output Tensor buffer 2. A grad_fn closure capable of calculating vector-Jacobian products (local gradients) 3. Its parent nodes in the computational graph

Calling tape.backward(loss) walks the computation graph in reverse topological order, evaluating grad_fn closures and accumulating gradients directly onto leaf parameter tensors.

[!TIP] Theory & Curriculum Link: For the first-principles mathematical derivation of the chain rule, reverse-mode automatic differentiation, and computational graphs, see Curriculum: Calculus & Gradients.

The Computational DAG Structure

The forward computation and the reverse gradient propagation are illustrated below for $y = (x + 1)^2, \; z = 3y, \; \text{loss} = \text{sum}(z)$:

flowchart TD subgraph Forward["1. Forward Computation Pass (Operations)"] direction LR x["x (Leaf Tensor)"] -->|"+ 1"| N1["Node (x + 1)"] N1 -->|"pow(2)"| N2["Node (y = (x+1)²)"] N2 -->|"· 3"| N3["Node (z = 3y)"] N3 -->|"sum"| Loss["Node (Loss)"] end Loss ==>|"backward() trigger"| Seed["Seed Gradient: ∂L/∂L = 1.0"] subgraph Backward["2. Reverse Gradient Propagation (Backward Pass)"] direction RL Seed -->|"∂L/∂z = 1.0"| G3["dL/dz = 1.0"] G3 -->|"· 3"| G2["dL/dy = 3.0"] G2 -->|"· 2(x+1)"| G1["dL/dx = 6(x+1)"] G1 -->|"accumulate"| XGrad["x.grad"] end

Forward edges connect parent nodes to child results. During backward(), netcl traverses reverse edges starting from seed gradient $\frac{\partial \mathcal{L}}{\partial \mathcal{L}} = 1.0$, evaluating each node's grad_fn and accumulating into .grad.

Node: The Graph Record

@dataclass
class Node:
    value: Tensor
    grad_fn: Optional[GradFn]  # Callable[[Tensor], List[Optional[Tensor]]]
    parents: List[Node]
    grad: Optional[Tensor]      # Accumulated gradient
    requires_grad: bool
    op_name: Optional[str]      # For debug prints
    creation_trace: Optional[List[str]]

A Node is distinct from a Tensor: - The Tensor is dumb device memory (cl_mem) with shape and strides. - The Node is the autograd graph record.

Leaf tensors that do not require gradients have no associated Node. Intermediate tensors whose parents do not require gradients are skipped during graph construction.

Node Fields

  • value: The Tensor produced by the forward operation.
  • grad_fn: A closure that takes the upstream gradient and returns local gradients for each parent node.
  • parents: List of input nodes that generated this node.
  • grad: The accumulated gradient tensor flowing back into this node. Leaf parameters store their final gradients here for the optimizer.
  • requires_grad: Boolean flag indicating whether the node participates in backward graph traversal.
  • op_name: Human-readable identifier (e.g. "matmul", "relu", "conv2d").

Tape: The Per-Thread Recorder

class Tape:
    def __enter__(self): ...
    def __exit__(self, exc_type, exc_val, exc_tb): ...
    def backward(self, loss: Node) -> None: ...

The active Tape is bound to the current thread via thread-local storage (threading.local). This provides several guarantees: - Concurrent worker threads run independent autograd tapes without cross-thread contamination. - The default context manager provides an implicit active tape consumed by apply_op when no explicit tape= argument is passed. - Nested tape contexts stack cleanly.

The Tape.backward() Execution Lifecycle

When tape.backward(loss) is invoked: 1. Seed Gradient: Seeds the loss node with $\frac{\partial \mathcal{L}}{\partial \mathcal{L}} = 1.0$ using an OpenCL buffer filled with ones matching loss.value.shape. 2. Reverse Topological Ordering: Builds a reverse topological sort of the graph rooted at loss using Kahn-style depth-first search (build_topo). 3. Kernel Dispatch & Accumulation: Iterates through nodes in reverse order, calling node.grad_fn(node.grad). Local gradients are accumulated into parents using in-place addition (add_inplace), ensuring multi-consumer branch nodes sum incoming gradients correctly. 4. Queue Flush Synchronization: Captures the OpenCL command queue into Tape._pending_flush_queue. Device-to-Host transfers (to_host()) serve as the natural hardware synchronization barrier, avoiding expensive explicit queue.finish() stalls.

Operator Overloading on Nodes

To enable intuitive mathematical syntax, Node implements standard Python dunder methods (__add__, __sub__, __mul__, __matmul__, __pow__):

def __add__(self, other):
    from netcl import autograd as ag
    return ag.add(self, other)

This allows clean, idiomatic expressions where operations are recorded automatically:

y = (x + 1.0) ** 2
z = 3.0 * y
loss = z.sum()

Memory Management in Autograd

  1. Retained Intermediate Activations: Every recorded Node holds a reference to its output Tensor. For deep architectures, this dominates VRAM during training. Wrap evaluation and inference routines in with ag.no_grad(): to discard graph nodes immediately.
  2. Anomaly Detection: Wrapping code in with ag.detect_anomaly(): captures Python stack traces at node creation and scans gradient buffers for NaNs and Infs. Use this strictly during debugging, as host roundtrips degrade throughput.