Kina.ai

Specification

Phase 2, Section 2: Effect Handler Infrastructure

Phase 2: AD as an Effect — algebraic effect and handler infrastructure. Section 2 of 4: Effect Handler Infrastructure.

1. What this section commits to

This section specifies the syntax, typing rules, and runtime execution model for algebraic effects and handlers in kina-ai. Specifically, it establishes:

2. Grammar Extensions

We extend the Kina surface syntax (.kina) to support first-class algebraic effects and handlers.

2.1 Lexical & Syntactic Grammar

Type        ::= Tensor[Dtype, Shape, Refinements]
              | fn(Type, ...) !EffectRow -> Type
              | Scalar[Dtype]
              | Unit

EffectRow   ::= { EffectName, ... }                         -- closed row
              | { EffectName, ... | ρ }                     -- open row (row variable ρ)

Decl        ::= ...
              | effect EffectName { EffectOp ... }
              | handler HandlerName(Param, ...) { HandlerBranch, ... }

EffectOp    ::= fn OpName(Param, ...) -> Type

HandlerBranch ::= return(Ident) => Expr
                | EffectName.OpName(Ident, ...), Ident => Expr

Expr        ::= ...
              | perform EffectName.OpName(Expr, ...)
              | handle Expr with HandlerName(Expr, ...)
              | handle Expr with { HandlerBranch, ... }

2.2 Key Commitments

  1. Effects are explicitly declared at top level. Anonymous or inline operations are rejected to ensure clear compilation paths and SMT-provable effect safety.
  2. Row variables ($\rho$) enable row polymorphism. Functions that call functions with arbitrary effects preserve polymorphism: e.g., fn map(f: fn(A) !ρ -> B, xs: List[A]) !ρ -> List[B].
  3. Explicit perform construct. Performing an effect is visually demarcated in code, preventing unexpected control-flow jumps.

3. Type System for Effects

The type checker tracks the set of active effects alongside the value context. The typing judgment is extended to: $$\Gamma; \Sigma \vdash e : T ! E$$ Which reads: "in value context $\Gamma$ and shape/refinement context $\Sigma$, expression $e$ has type $T$ under active effect row $E$."

3.1 Well-Formedness of Effect Rows

An effect row $E$ is well-formed ($\Gamma \vdash E \text{ row}$) if all named effects are declared in the global signature, and any row variable $\rho$ is bound in $\Gamma$.

3.2 Typing Rules

  (OpName : fn(T₁, ..., Tₙ) -> T_ret) ∈ EffectName
  Γ; Σ ⊢ e₁ : T₁ ! E    ...    Γ; Σ ⊢ eₙ : Tₙ ! E
  EffectName ∈ E
──────────────────────────────────────────────────────── (perform)
  Γ; Σ ⊢ perform EffectName.OpName(e₁, ..., eₙ) : T_ret ! E


  Γ; Σ ⊢ e_body : A ! ({EffectName} ∪ E)
  Γ; Σ ⊢ return(x) => e_pure : A -> B ! E
  ∀ Op ∈ EffectName.
    Γ; Σ ⊢ Op(args), k => e_op : (Args_Op × (Res_Op -> B ! E)) -> B ! E
─────────────────────────────────────────────────────────────────────── (handle)
  Γ; Σ ⊢ handle e_body with { return(x) => e_pure, ... } : B ! E

3.3 Row Polymorphism & Subtyping

Effect rows support subtyping via row extension. An effect row $E_1$ is a subtype of $E_2$ ($E_1 \sqsubseteq E_2$) if the set of effects in $E_1$ is a subset of the effects in $E_2$, and their row variables (if present) match.

Coercion from $A ! E_1$ to $A ! E_2$ is implicit when $E_1 \sqsubseteq E_2$.

4. Detailed Syntax Examples

4.1 State Effect Declaration and Usage

Defining a simple State effect that carries a 0D tensor:

effect State {
  fn get() -> tensor<F32, []>
  fn set(val: tensor<F32, []>) -> Unit
}

// A function performing state operations
fn increment() !{State} -> Unit {
  let current = perform State.get();
  perform State.set(current + 1.0)
}

4.2 Pure State Handler

The functional state-passing handler runs a computation f under State, producing a state-transforming function tensor<F32, []> -> A:

fn run_state(init: tensor<F32, []>, f: fn() !{State} -> A) -> A {
  let state_runner = handle f() with {
    return(x) => fn(s: tensor<F32, []>) -> A { 
      x 
    },
    State.get(), k => fn(s: tensor<F32, []>) -> A { 
      k(s)(s) 
    },
    State.set(v), k => fn(s: tensor<F32, []>) -> A { 
      k(())(v) 
    }
  };
  state_runner(init)
}

5. Runtime Execution Model: Fibers and Continuations

The Rust runtime implementing algebraic effects avoids copying heavy system threads. Instead, it utilizes Monomorphized Fibers on the stack.

+─────────────────────────────────────────────────────────────+
| Fiber Stack Frame (Rust Allocation)                         |
|                                                             |
|   +───────────────────+                                     |
|   | Caller context    |                                     |
|   +───────────────────+                                     |
|   | Saved registers   |                                     |
|   +───────────────────+                                     |
|   | Captured yield pt | (Suspension state target)           |
|   +───────────────────+                                     |
|                                                             |
+─────────────────────────────────────────────────────────────+

5.1 Continuation Capture (resume)

  1. Suspension: When perform Effect.Op(...) is called, the current fiber registers the operation code and arguments on the active handler's boundary, saves local register states, and suspends execution.
  2. Continuation Binding: The compiler matches the operation against the surrounding handler block. The handler branch is invoked, receiving the operation arguments and a stack-allocated continuation object k.
  3. Resumption: The continuation k is a callable entity. Executing k(val) restores the saved register context, puts val into the resumption register, and jumps back to the instruction immediately following the perform boundary.
  4. Single-shot Enforcement: Continuations in kina-ai are linear/single-shot by default. Calling the same continuation k twice triggers a runtime safety panic. This matches the single-pass nature of automatic differentiation.

6. Verification and Invariant Auditing

Every compiler backend pass generating LLVM IR or MLIR containing effect row transformations must satisfy: