Automatic Differentiation (Autodiff)

Autodiff is distinct from both symbolic differentiation and numerical differentiation.

• Symbolic differentiation manipulates mathematical expressions to produce derivative formulas, which can lead to expression swell and is impractical for complex, iterative code.
• Numerical differentiation approximates derivatives using finite differences (e.g., (f(x+ε) - f(x))/ε), which is prone to rounding errors and truncation errors.
• Autodiff computes exact derivatives (to machine precision) by decomposing the program into elementary operations and applying the chain rule systematically, avoiding both expression complexity and numerical inaccuracy.

Autodiff implements the chain rule through two primary modes:

• Forward-Mode Autodiff: Propagates derivatives from inputs to outputs. For a function f: ℝⁿ → ℝᵐ, computing the Jacobian one column at a time is efficient when n << m. It is often used in scientific computing and for computing Hessian-vector products.
• Reverse-Mode Autodiff: Propagates derivatives from outputs back to inputs. For f: ℝⁿ → ℝᵐ, it computes the gradient (a row of the Jacobian) in a single pass when m=1, making it dramatically more efficient for the loss: ℝⁿ → ℝ¹ functions ubiquitous in neural network training (n is the number of parameters). This is the engine of backpropagation. The choice of mode is a fundamental algorithmic decision based on the function's input/output dimensions.

Autodiff systems are built using one of two primary implementation strategies:

• Source Code Transformation (SCT): A compiler (e.g., Tapenade, ADIFOR) analyzes the source code's abstract syntax tree (AST) and generates new source code that computes both the primal function and its derivatives. This allows for whole-program optimization but requires language-specific toolchains.
• Operator Overloading: Leverages a language's ability to redefine operators (common in Python/C++). Libraries like PyTorch's autograd, JAX, and Stan overload tensor operations to build a dynamic computational graph at runtime. Each operation records a functor (the operation) and its input tensors, constructing a trace used for the backward pass. This is more flexible and easier to integrate but may incur runtime overhead for graph construction.

Autodiff operates on an explicit or implicit computational graph (also called a dataflow graph or Wengert list).

• Nodes represent operations (e.g., matrix multiply, ReLU, softmax).
• Edges represent data dependencies (tensors). During the forward pass, the graph is executed, and intermediate results (activations) are computed and often stored. During the reverse pass, the graph is traversed backwards. For each node, the local gradient (derivative of the node's output w.r.t. its inputs) is computed and multiplied by the upstream gradient from later nodes (the chain rule). This yields gradients for the node's inputs, which are propagated further backward. Modern frameworks like TensorFlow (with tf.function) and JAX use just-in-time (JIT) compilation to optimize and compile this graph for efficient execution.

Efficient autodiff requires managing what is differentiated.

• Gradient Tape: A conceptual recording (e.g., in TensorFlow) of operations performed on differentiable variables. Only operations watched by the tape are included in the gradient graph. This provides fine-grained control over the differentiation scope.
• Edge Pruning / Dead Code Elimination: A critical compiler optimization within the autodiff engine. The system analyzes the graph to identify and remove (prune) subgraphs that do not contribute to the requested gradients. For example, if only the gradient w.r.t. a subset of parameters is needed, operations leading only to other parameters are eliminated. This reduces memory and computation overhead for the backward pass.

For NPU acceleration, autodiff is tightly integrated into the graph compilation pipeline.

1. High-Level Graph Capture: A framework (e.g., PyTorch) captures a forward graph, which includes autodiff metadata for potential backward pass generation.
2. Graph Lowering & Fusion: The compiler lowers this graph, applying operator fusion to merge forward (and corresponding backward) operations into efficient kernels. For example, a matmul + bias_add + ReLU forward cluster will have a corresponding fused backward kernel for its gradient.
3. Memory Planning: The compiler performs unified memory planning for both forward and backward passes, reusing buffers between activations and their gradients where possible to minimize peak memory usage—a critical constraint for training large models.
4. Hardware-Specific Gradients: The compiler may select or generate specialized kernel implementations for gradient computations that are optimized for the NPU's memory hierarchy and parallel architecture.

Graph lowering is the process of transforming a high-level, abstract computational graph representation into a lower-level, more hardware-specific representation. For an autodiff system, this involves a series of legalization and conversion passes that take a graph containing differentiable operations and progressively map it to primitive operations available on the target NPU or GPU.

• Process: High-level ops → vendor-specific ops → hardware intrinsics → machine code.
• Role in Autodiff: The lowering pass must preserve the dataflow dependencies required for correct gradient computation during the backward pass.

Static Single Assignment is a property of an intermediate representation where each variable is assigned exactly once. This form drastically simplifies dataflow analysis, which is critical for autodiff algorithms that must trace dependencies from outputs back to inputs to compute gradients.

• Benefit for Autodiff: Creates an unambiguous, acyclic dataflow graph, making it trivial to perform the reverse-mode traversal required for backpropagation.
• Compiler Use: Enables powerful optimizations like constant propagation and dead code elimination on the gradient computation graph generated by autodiff.

Just-In-Time compilation is a dynamic compilation strategy where code is compiled during program execution. In autodiff systems, JIT compilation is often used to generate optimized forward and backward pass kernels on-the-fly, using runtime information like input tensor shapes.

• Advantage: Enables optimizations tailored to the actual execution environment and dynamic graph structures.
• Trade-off: Introduces runtime compilation overhead, which is amortized over repeated executions (common in training loops).
• Contrast: Differs from Ahead-Of-Time (AOT) Compilation, which compiles the entire graph (including gradients) to a static binary before execution.

Graph fusion is a compiler optimization that merges multiple adjacent operators or nodes within a computational graph into a single, compound kernel. This is a critical optimization applied after autodiff has generated the backward pass graph, as it can fuse chains of elementary gradient operations.

• Primary Benefit: Reduces kernel launch overhead and minimizes costly reads/writes of intermediate tensors to high-bandwidth memory.
• Example: Fusing a bias addition, ReLU activation, and their corresponding gradient operations into one custom kernel.
• Impact: Directly improves the performance of the gradient computation step autodiff enables.

Memory planning is a compiler optimization pass that allocates memory buffers for all tensors in a computational graph, aiming to minimize peak memory usage. This is especially vital for the backward pass generated by autodiff, which can temporarily require storage for many intermediate activations and gradients.

• Key Techniques: Buffer reuse (overlapping lifetimes of non-conflicting tensors) and in-place operations.
• Connection to Autodiff: Works in tandem with activation recomputation (checkpointing), a strategy that trades compute for memory by selectively re-calculating activations during the backward pass instead of storing them all.