Phase 1, Section 12: `reverseMode` result ABI

Compiler transform_ad only. Not Phase 2 handler AD. Not a Phase 1 exit-criteria replacement. No new operator. This section only says how an already-typed reverse-mode result leaves the compiler.

1. What this section commits to

Section 10 builds a primal and one cotangent per reverseMode parameter. Section 11 can apply those cotangents as explicit g tensors. Neither section returned the pair from MLIR, so a .so could not load Kina-executed gradients.

This section locks the result ABI:

The language type of handle f(args…) with reverseMode is a flat tuple of tensors: primal first, then parameter gradients in parameter order. Discrete (non-floating) parameters are omitted; see Section 14 §3.
That tuple lowers to an MLIR multi-result func.func, not an MLIR tuple<> value.
A fn with that result type gets the existing llvm.emit_c_interface wrapper, so a host can inject arguments and read primal plus gradients from a .so.

There is no new surface keyword, IR constructor, or FFI.

2. What this section does not commit to

Tuple projection (now Section 13). This section only returns the flat tuple.
Changing @main to take arguments. Closed programs still bind literals and call functions; host injection uses a fn wrapper, same stance as Section 11’s SGD check.
Changing the seed cotangent. It remains ones of the primal result shape (ConstantTensor (1.0, ret_tensor_ty)).
Device take_along_axis / ScatterAlongAxis.
AdamW, argmax, seeded sampling, KV-cache, or Phase 2 handler AD.
Differentiating non-tensor parameters or non-tensor results.
An MLIR tuple dialect lowering.

3. Language type (amends Section 10)

Section 10’s implementation used a nested pair

(primal, (g₀, …, gₙ₋₁))

The bootstrap surface cannot write a one-element tuple type (T), and the existing .so pipeline (one-shot-bufferize + LLVM C wrappers) does not bufferize MLIR tuple<> types. The result type is therefore flat:

Γ ⊢ f : (T₀, …, Tₙ₋₁) → R
Γ ⊢ argsᵢ : Tᵢ
R = Tensor[…]
each Tᵢ = Tensor[…]                 (this slice; discrete I64 stays an inner let)
─────────────────────────────────────────
Γ ⊢ handle f(args…) with reverseMode
      : (R, T₀, …, Tₙ₋₁)

Constructor order is primal, then each parameter’s cotangent. A single-parameter program has type (R, T₀), written in .kina as

(tensor<F32, [2, 2], @host>, tensor<F32, [2, 2], @host>)

Zero parameters collapse to just R (a one-element tuple unwraps to that tensor). This slice’s examples have at least one differentiable parameter.

The AD pass materializes a flat Tuple of those tensors. It does not wrap the gradients in a second Tuple.

4. Flattening

Any remaining nested TTuple of tensors (legacy IR, multi-carry while) flattens left-to-right, depth-first, before emission:

flatten(Tensor[T, σ]) = [Tensor[T, σ]]
flatten((τ₀, …, τₖ))  = flatten(τ₀) ++ … ++ flatten(τₖ)

A non-tensor element in that spine is unsupported in this slice (fail emit). After flattening, Section 10’s nested pair and §3’s flat tuple produce the same MLIR result list.

5. MLIR function ABI

Let ρ = [τ₀, …, τₘ₋₁] be the flattened tensor types.

`m`	Function result	`return` types
0	(no `->`)	`return`
1	`-> tensor<…>`	`return %v : tensor<…>`
≥ 2	`-> (tensor<…>, …)`	`return %v0, … : tensor<…>, …`

Do not emit tuple<tensor<…>, …> as a function result or operand type. Parentheses appear only on the func.func result list when m ≥ 2. The return op lists types without an extra tuple wrapper.

This applies to:

@main when the program’s result type is a (possibly nested) tensor tuple.
fn results and func.call / func.call_indirect destinations with the same type.

@main stays nullary. It does not gain llvm.emit_c_interface in this section; host .so loads use a fn that already has that attribute.

A call that returns m ≥ 2 values binds m SSA names:

%y, %g = func.call @ad(%x) : (tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>)

6. LLVM C wrapper (host `.so`)

After --llvm-request-c-wrappers and function-boundary bufferization, the existing LLVM C wrapper is:

One tensor result (Section 11 SGD): dest-passing. First argument is MemRefDescriptor *out, then each input descriptor.
Two or more tensor results (this section): one pointer to a packed struct of ranked memref descriptors in flatten order, then each input descriptor.

func.func @ad(%x: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>)
  attributes {llvm.emit_c_interface}

void _mlir_ciface_ad(
  struct { MemRefDescriptor primal; MemRefDescriptor g0; } *outs,
  MemRefDescriptor *in_x);

The callee allocates the result buffers and writes both descriptors into *outs. Ranked memref layout matches Section 11’s SGD loader. This packing is the LLVM C-wrapper rule for a multi-result func.func; it is not a new Kina FFI.

7. Surface

No new tokens. The existing tuple type

(type, type[, type…])

annotates a wrapper fn whose body is handle f(x) with reverseMode.

The StoryLM-0 numeric instance (neg, seed ones):

let f = fn(x: tensor<F32, [2, 2], @host>) -> tensor<F32, [2, 2], @host> {
  let y = -x in
  y
} in
let ad = fn(x: tensor<F32, [2, 2], @host>)
  -> (tensor<F32, [2, 2], @host>, tensor<F32, [2, 2], @host>) {
  handle f(x) with reverseMode
} in
let x = tensor<F32, [2, 2], @host> in
ad(x)

A program that is only handle f(x) with reverseMode (no wrapper) is still legal. @main then returns the flattened tensors computed from the bound literals.

8. Tests

Existing ad_neg.kina / ad_log_softmax.kina / ad_nll.kina / ad_basic.kina goldens: @main has the flattened multi-result signature and a matching return.
ad_neg_grad.kina — §7 wrapper; golden accepted by mlir-opt; the AD fn result is (tensor<2x2xf32>, tensor<2x2xf32>).
Annotated wrapper whose return type is not (R, T₀, …) is rejected.
scripts/validate_reverse_mode_abi.py compiles ad_neg_grad.kina to a .so, injects a patterned x, and checks:
- primal equals -x;
- gradient equals -ones (Section 10 neg VJP with seed ones). Optional PyTorch check when present. This is Kina-executed .so grads for that instance.

9. Relation to SGD

Section 11 examples still take w and g as explicit arguments. Naming g inside .kina is Section 13.

10. Open issues

Tuple projection. Specified in Section 13.
@main C wrapper / argument ABI. Not required while wrapper fns exist; revisit if a closed program must be the .so entry.
Non-tensor / device results in the flatten spine.
Custom seed cotangents (not ones).
Discrete parameters. Specified in Section 14: I64 (and other non-floating) arguments do not appear in the result tuple.