Specification

Phase 2, Section 4: Control Flow and Iteration

Phase 2: AD as an Effect — algebraic effect and handler infrastructure. Section 4 of 4: Control Flow and Iteration.

1. What this section commits to

This section specifies the syntax, typing rules, and differentiation mechanics for structured control flow in kina-ai. Specifically, it establishes:

Syntactic structures for conditionals (if) and loop iterations (while).
Linear ownership guidelines for loop-carried variables to prevent tensor alias leaks.
Dynamic control-flow taping and activation checkpointing mechanisms during automatic differentiation backpropagation.

2. Structured Control Flow

To support non-linear mathematical operations and iteration, kina-ai provides structured control flow primitives: conditionals and loops. Both are expressions that yield values.

2.1 Conditionals (`if`)

The condition must be a 0D boolean tensor (or boolean scalar). The two branches must evaluate to the same type.

fn absolute_val(x: tensor<F32, [D]>) -> tensor<F32, [D]> {
  // pointwise comparison returns boolean tensor
  let mask = x >= 0.0;
  if (mask) {
    x
  } else {
    -x
  }
}

2.2 Loop Iteration (`while`)

Loops utilize loop-carried variables to maintain linear state tracking. A loop must explicitly declare its initial state bindings, and the body must return the updated states of identical types.

fn power_of_two(exponent: I32) -> tensor<F32, []> {
  let initial_val = tensor<F32, [], @host>.broadcast([]); // 1.0
  
  let final_val = while (i < exponent) [val = initial_val, i = 0] {
    let next_val = val * 2.0;
    let next_i = i + 1;
    (next_val, next_i)
  };
  
  final_val
}

2.3 Linear/Owned Tensor Rules

Since tensors are linear/owned by default, a loop-carried tensor is consumed by the loop body. The body must yield a new tensor of the same type. If a tensor is not yielded, it must be explicitly deallocated or consumed by a terminal operation (e.g. reduction or host copy) to prevent leak errors.

3. AD for Control Flow

Automatic differentiation through control flow requires recording the exact execution path taken during the forward pass.

3.1 Dynamic Branch Tape

When executing under the reverseMode handler:

Conditionals: An execution tag (e.g. 0 for then, 1 for else) is recorded onto the local thread tape. During the backward pass, this tag is read to route the adjoints to the correct branch.
Loops: The iteration count $N$ and the sequence of loop conditions are pushed to the tape. The backward pass executes exactly $N$ iterations in reverse.

4. Activation Checkpointing

For deep iterations or long loops, saving all forward activations for backpropagation consumes $O(N)$ memory, which is prohibited for large tensor models. We specify Activation Checkpointing:

Forward Loop:
[Init] ──> [Step 1] ──> [Step 2] ──> [Step 3] ──> [Step 4] ──> [Target]
            Checkpoint               Checkpoint

Backward Loop:
[Target] ──> Recompute (Step 3 to 4) ──> Backprop Step 4 to 3
         ──> Recompute (Step 1 to 2) ──> Backprop Step 2 to 1

4.1 Checkpoint Interface

The compiler supports an explicit @checkpoint decorator on loops, instructing the code-generation pass to only store loop-carried values every $K$ iterations.

fn dense_loop(init: tensor<F32, [D]>, steps: I32) -> tensor<F32, [D]> {
  // Checkpoint every 4 iterations to balance memory vs compute overhead
  @checkpoint(k=4)
  let result = while (i < steps) [x = init, i = 0] {
    let next_x = relu(x @ weight + bias);
    (next_x, i + 1)
  };
  result
}

This reduces peak memory consumption from $O(N)$ to $O(K + N/K)$. Setting $K = \sqrt{N}$ achieves optimal $O(\sqrt{N})$ space complexity.