Phase 2, Section 1: Scope, Milestones, and Exit Criteria

Phase 2: AD as an Effect — algebraic effect and handler infrastructure. Section 1 of 4: Scope, Milestones, and Exit Criteria.

1. What this section commits to

This section defines the scope boundaries, intermediate milestones, risk management metrics, and final exit criteria for Phase 2 of the Kina research program. Specifically, it establishes:

• The boundaries of the Phase 2 implementation.
• The three engineering milestones, including the load-bearing AD-as-handler Proof of Concept (PoC).
• Specific, measurable abandonment triggers to fail early if the design vision is untenable.
• Consolidated exit criteria that define the phase's completion.

2. Scope Boundaries

Phase 2 shifts kina from a pure functional language to a staged effectful language. To prevent scope creep, we enforce strict boundaries:

In Scope

• Algebraic Effect Handlers: Support for row-polymorphic effect typing in the OCaml compiler, and runtime stack/unwind execution in the Rust runtime.
• Reverse-Mode AD Handler: A language-level reverseMode handler implementing reverse-mode automatic differentiation.
• Differentiable Operators: Gradient computations for all Phase 0 operations and Phase 1 extensions.
• Control Flow: Structured conditionals (if) and loops (while) with tape-based checkpointing for backpropagation.
• Linear/Owned Tensors: Enforcing linear ownership in the type-checker to ensure sound mutation-aware AD.
• Implicit Differentiation: An Implicit/fixpoint effect utilizing the Implicit Function Theorem (IFT) for solver loops.
• Stochastic Estimator Selection: Composed handlers for Random effects (e.g., pathwise vs. score-function).

Deferred (Out of Scope)

• Distributed AD / Parallel Sharding: Sharding of gradients and communication collectives (Phase 3).
• Sparse Gradients: Sparsity-aware AD and MoE expert partitioning (Phase 4).
• GPU Kernel Generation (Tile IR): Compilation of handlers to custom GPU kernels (Phase 5).
• Full Formal Proofs: Lean 4 verification of the AD transform (Phase 6).
• User Tooling & Surface Syntax: Final compiler UX, LSP, and custom error formats (Phase 7).

3. Milestones

Phase 2 is executed in three sequential milestones:

Milestone 1 (Infrastructure) ──> Milestone 2 (AD PoC) ──> Milestone 3 (Control Flow & Stdlib)

Milestone 1: Effect Handler Infrastructure

Goal: Implement row-polymorphic effect typing in the compiler and fiber-like execution in the runtime.

• Compiler: Extend the S-expression surface AST and typed IR to support effect row annotations (e.g., (fn [x] !{Random, State} ...)).
• Type Checker: Implement row subtyping and effect unification.
• Runtime: Add fiber-like stack suspension and continuation capture to support algebraic handler resumption.
• Deliverable: An end-to-end test verifying a simple user-defined effect (e.g., a State cell or a yield generator) compiles, executes, and resumes correctly on CPU.

Milestone 2: The AD-as-Handler PoC (Load-Bearing)

Goal: Verify that reverse-mode AD is expressible as a language-level handler with correct numeric outputs.

• Implementation: Write the reverseMode handler in kina using the effect primitives from Milestone 1.
• Numeric Verification: Differentiate a two-layer perceptron (MLP) forward pass: $$\text{Loss} = \text{mse}(\text{linear}(x, W, b), y)$$
• FFI Verification: Verify that FFI entry points remain stateless and that the handler successfully intercepts and records tape operations.
• Deliverable: The PoC executes end-to-end on CPU, matching PyTorch reference loss and gradients ($\nabla_W, \nabla_b$) within $1e-3$ FP tolerance.

Milestone 3: Control Flow & Standard Library Integration

Goal: Integrate structured loop constructs and expand the standard library to be fully AD-aware.

• Control Flow: Implement if and while constructs in the parser, type-checker, and MLIR emitter.
• Checkpointing: Implement memory-optimized activation checkpointing for loops during backpropagation.
• Stdlib Expansion: Update the standard library to expose differentiable blocks, including RMSNorm, Softmax, and self-attention layer steps.
• Deliverable: A complete, differentiable transformer layer forward and backward pass executing on GPU via library kernel dispatches.

4. Risks and Abandonment Criteria

To keep the research program falsifiable, we commit to the following abandonment thresholds:

1. Handler Performance Overhead

• Risk: Capture/resume continuation overhead in the Rust runtime degrades execution speed.
• Trigger: Running a differentiable transformer block with the AD handler is $> 10\times$ slower than the equivalent hand-taped OCaml-level transformation on CPU.
• Action: Pause; optimize runtime continuation allocation or redesign handlers to compile down to static SSA passes (monomorphization of handlers).

2. SMT / Z3 Unification Complexity

• Risk: Adding effect rows to the bidirectional type-checker causes SMT solver timeouts during constraint discharge.
• Trigger: SMT constraint verification for the proto-stdlib exceeds 5 seconds per function, or timeouts occur on $> 10%$ of standard test cases.
• Action: Fall back to a simplified effect system (e.g., monomorphic effects without row-polymorphic subtyping).

3. PoC Numeric Correctness Failure

• Risk: The AD-as-effect handler produces incorrect gradients due to precision drift, alias leaks, or singular points.
• Trigger: The Milestone 2 PoC fails to match PyTorch gradients within $1e-3$ tolerance on representative inputs after 4 weeks of debugging.
• Action: Abandon in-language AD handlers; fall back to an OCaml-level source-to-source compilation pass.

5. Exit Criteria

Phase 2 is complete when a release candidate satisfies all the following criteria:

Compiler & Type-Checker

• Compiler parses, type-checks, and emits code for row-polymorphic effect signatures.
• Linear/owned tensor rules are strictly enforced (no alias leaks allowed in differentiable functions).
• Support for structured control flow (if, while) implemented and validated.

Runtime & Execution

• Fiber-like runtime supports fast continuation capture, resuming, and unwinding.
• Memory allocator does not leak memory during nested handler executions (verified by Valgrind/Miri).

Automatic Differentiation

• Language-level reverseMode handler successfully computes loss and gradients.
• Gradients for standard MLP and transformer block match PyTorch references within $1e-3$ FP tolerance.
• Stochastic estimators (reparameterize vs. score_function) correctly dispatch based on the composed handler.
• Implicit function theorem (IFT) handler successfully differentiates solver loops without unrolling.

Documentation & Packaging

• Phase 2 specification sections (02, 03, 04) fully written and complete.
• Standard library documentation updated to include effect signatures and differentiable primitives.