Kina.ai

Specification

Phase 2: Reverse-Mode Automatic Differentiation (AD) as an Algebraic Effect

1. Overview

Kina treats Automatic Differentiation not as a macro or a runtime tape, but as an algebraic effect handler. During Phase 2, we introduce ReverseModeAD as a first-class effect in the compiler's IR, allowing the compiler to perform a Continuation-Passing Style (CPS) transformation to produce explicit forward and backward passes.

By formalizing AD as an effect, Kina natively bridges side effects and gradient tracking.

2. The ReverseModeAD Effect

The core intuition: whenever a differentiable operation (e.g., Matmul, Add, Relu) executes, it performs an effect that yields its Vector-Jacobian Product (VJP) continuation to the handler.

2.1 Effect Definition

effect ReverseModeAD : 
  | RecordVJP : (Tensor -> Tensor) -> unit

2.2 Forward Pass (Primal)

During the forward execution, an operation calculates the primal value and constructs a closure (vjp_fn) that computes the gradient with respect to its inputs. It then performs RecordVJP(vjp_fn) before returning the primal result.

3. The grad Handler

The grad operator is a syntax sugar for an effect handler that intercepts RecordVJP effects.

let grad = fn(f, x) {
  let tape = ref [];
  let y = handle f(x) with {
    return(res) => res,
    RecordVJP(vjp_fn), k => {
      tape := vjp_fn :: !tape;
      k(()) // Resume forward execution
    }
  };
  
  // Backward pass
  let grad_y = ones_like(y);
  let grad_x = fold_left (fun g vjp -> vjp g) grad_y !tape;
  grad_x
}

Note: The actual implementation in Kina avoids runtime ref mutations by statically lowering the CPS control flow into pure MLIR during the ad_pass.ml compilation phase.

4. Compile-Time Lowering (CPS Transformation)

The middle-end ad_pass.ml lowers the handler construct by CPS-transforming the AST.

  1. Forward Pass Generation: Differentiable IR nodes are replaced with their primal counterparts.
  2. Reverse Pass Generation: The compiler extracts the VJP continuations and threads the gradient variables backward in reverse topological order.
  3. MLIR Emission: The fully static, tape-free forward and backward passes are emitted to the linalg and tensor MLIR dialects.

5. Type System & Placements

The requires_grad property must respect placements and refinements: