Module Ir.Prefix_scan

The prefix-scan rewrite (gh-ocannl-1758): a member of the algebraic-rewrite tier over lowered code (Rewrites), gated by the config key prefix_scan (default off, on in the approximate profile, pinned off in reproducible).

A cumulative sum written as the triangular-selector contraction -- a lower-triangular 0/1 mask times the input, a +* "i; i o => o" tri (tensor puzzle 7) -- lowers to an O(n^2) reduction nest y[.., o] += v(.., i) * tri[i, o]. This pass replaces such a nest with one Ir.Low_level.t.Scan_loop per row over o, carrying the running reduction, O(n):

Forward for a prefix (r <= q, r < q), Backward for a suffix (r >= q, r > q), where r is the nest's one reduced loop and q an output loop over the same range. Writing y ⊕ s keeps what the target held: the zero (or -inf) fill of a forward, the earlier contributions of a gradient. The cumsum's backward a.grad[i] += g[o] * tri[i, o] is the same shape with the roles exchanged -- a suffix -- so it becomes a Backward scan with no gradient written by hand.

Recognition, on one perfect serial nest y[..] := y[..] ⊕ rhs (either operand order), ⊕ one of Add, Max, Min, with rhs not reading y and y a float (an integer target declines: the state would round its sums):

Everything else -- a band, a shifted or anti-diagonal triangle, a mask varying with a row, two reduced loops, ranges that differ, a guard or clamp, a mask held in a host-initialized constant or computed by an earlier routine -- is left alone. The mask's and the masked operand's definitions are not removed: they stay as their nodes' definitions, and the optimizer drops them when nothing else reads them.

Numerics. The carried state is at the target's precision widened to f32 (f64 stays f64). A forward cumsum adds in the composed form's order, so it agrees bitwise on finite inputs; cummax and cummin agree exactly. A suffix scan regroups the sum (the composed form adds ascending from q, the scan descending from the end). The composed product form also multiplies the rejected payloads by 0, so an infinite payload outside the triangle makes it NaN where the scan -- the cumulative sum the selector encodes -- stays finite. Hence an approximate-tier gate.

The cummax backward is not this shape (its operand depends on the output position through the gradient gate) and stays composed.

val enabled : unit -> bool

Whether the rewrite applies: the programmatic override if one is set, else the config key prefix_scan.

val set_enabled : bool option -> unit

Programmatic override of the config key, for experiments and tests; None restores the key. Takes effect at the next lowering.

val state_owner : Tnode.t -> int option

The uid of the target node a carried state accumulates into, when the node is a state this module minted.

val reset : unit -> unit

Drops the memoized state nodes (the tier's session-reset hook; also run ahead of an accessibility snapshot). Sibling lowerings of one program mint the same nodes, keyed by the matched site -- the target, the match's ordinal among the scans into it, the precision -- which keeps them on one analysis-cache key.

val rewrite : Low_level.t -> Low_level.t

The pass over raw lowered code (no gate: the tier consults enabled). Every eligible nest is rewritten in place; a routine without one is returned as is. Idempotent: a scan is not a nest.