Skip to main content

logicaffeine_compile/spatial/passes/
from_qwen_ops.rs

1//! `from_qwen_ops` — every op of an authored transformer forward, lowered as a COMPOSITION of the closed
2//! compute algebra: chains of [`ComputeOp::Generic`] and [`ComputeOp::MatmulTiles`], and nothing else.
3//!
4//! ★ NOT ONE NEW `ComputeOp`. Rope, silu, sigmoid, gelu, decay, l2norm, transpose, a causal conv ring, a
5//! Gated-DeltaNet step and GQA attention all lower through this file, and `spatial_algebra_lock`'s frozen
6//! per-printer arm counts do not move for any of them. That is the campaign's whole claim, and it holds
7//! because B-P3 made the pieces general: a reduction is an ITERATOR KIND, a broadcast is an ABSENT TERM, a
8//! transform is an `SfpuKind`, a piecewise region is a `Select` on an AXIS, and an access map is a signed
9//! FLAT affine offset rather than one term per axis.
10//!
11//! ── The layout contract of a composed program ────────────────────────────────────────────────────────
12//!
13//! ★ EVERY ACTIVATION IS FLAT: `1 × n`. This is not a convenience, it is what makes composition possible
14//! at all. An `ElemAccess` is a flat offset into a page read as `tile.rows × (page_tiles · tile.cols)`, so
15//! element `p` of a `1 × n` tensor is at offset `p` — the SAME number the language's flat `item p` uses.
16//! Give the same tensor a `2 × 4` shape in one section and a `1 × 8` shape in the next and the two
17//! sections disagree about where element 4 lives, with no affine map able to reconcile them (the
18//! conversion needs `p / cols`, which is not affine in `p`). Flat is the one layout every op can both
19//! produce and consume.
20//!
21//! ★ A 2-D VIEW IS A BAND, NOT A RESHAPE. Where an op is per-head — rope's rotations, l2norm across
22//! heads, a GDN head's memory — the heads are contiguous RANGES of the flat column axis, and a range is
23//! a `Pred::AxisLt` on a CONSTANT. So one `Generic` computes every head by nesting one `Select` per head
24//! (`banded`), and inside head `h`'s band an operand that is indexed per-head is read through that
25//! head's own base (`c - h·width`), while a per-head SCALAR is read through a CONSTANT map: the
26//! reduction that produced it broadcast it across the whole of its tile row, and inside one band the
27//! row is fixed. No divisibility, no reshape, no gather.
28
29use std::collections::BTreeSet;
30
31use super::from_tensor_op::{
32    const_map, eps_nano, ident_map, items_per_tile, page_width, per_col_map, per_row_map, MatmulNames,
33    Template,
34};
35use super::{SpatialReject, LV_GEN_COL, LV_GEN_K, LV_GEN_ROW, LV_ITEM_X, LV_ITEM_Y, LV_KBLOCK};
36use crate::codegen_sva::fabric_target::Precision;
37use crate::spatial::ir::*;
38use crate::spatial::pipeline::StageNames;
39use crate::spatial::target::{precision_int_bits, SpatialTarget};
40use crate::tensor_recognize::qwen::Window;
41use crate::tensor_recognize::qwen::{
42    ArgMaxSpec, AttentionSpec, CausalConvSpec, DecaySpec, FullAttentionSpec, GateAct, GatedDeltaSpec,
43    GeluSpec, HeadSliceSpec, IntClampSpec, IntRequantSpec, JoinSpec, L2NormSpec, LutAttentionSpec, MoeGroups, MoeScore, MoeSpec,
44    ReluSpec,
45    RopeSpec, RowGatherSpec, RowIndex, RowsGatherSpec, ScaleSpec, SetRowSpec, SigmoidSpec, SiluSpec, SoftcapSpec,
46    TransposeSpec,
47};
48use crate::tensor_recognize::{Fill, FusedRowSpec, HeadMatmulSpec};
49
50// ── The flat vocabulary ───────────────────────────────────────────────────────────────────────────
51
52/// The iteration space of a FLAT stage: ONE row, `n` columns. The evaluator writes the result at
53/// `plane[c / tile][(r % tile) · tile + c % tile]`, so a single-row space lands the whole answer in tile
54/// row 0 — which IS the `1 × n` layout every activation has.
55fn flat_space(n: usize) -> Vec<IterSpec> {
56    vec![
57        IterSpec {
58            var: LV_GEN_ROW,
59            extent: 1,
60            kind: IterKind::Parallel,
61        },
62        IterSpec {
63            var: LV_GEN_COL,
64            extent: n.max(1),
65            kind: IterKind::Parallel,
66        },
67    ]
68}
69
70/// A 2-D stage's space: `rows` parallel rows by `cols` columns, the columns either walked or folded.
71fn head_space(rows: usize, cols: usize, reduce: Option<ReducePool>) -> Vec<IterSpec> {
72    vec![
73        IterSpec {
74            var: LV_GEN_ROW,
75            extent: rows.max(1),
76            kind: IterKind::Parallel,
77        },
78        IterSpec {
79            var: LV_GEN_COL,
80            extent: cols.max(1),
81            kind: match reduce {
82                Some(pool) => IterKind::Reduce { pool },
83                None => IterKind::Parallel,
84            },
85        },
86    ]
87}
88
89/// A flat map along the column axis with a signed base: element `c + base` of the operand's page.
90fn at(base: i64) -> ElemAccess {
91    ElemAccess {
92        base,
93        terms: vec![(LV_GEN_COL, 1)],
94        indirect: None,
95    }
96}
97
98/// A map that reads ONE fixed element of the operand's page — the shape a per-head scalar takes inside
99/// a band, where the head (and therefore the tile row the reduction broadcast it across) is constant.
100fn fixed(base: i64) -> ElemAccess {
101    ElemAccess {
102        base,
103        terms: Vec::new(),
104        indirect: None,
105    }
106}
107
108/// A 2-D map over a page `width` elements wide: `base + row·row_stride + col·col_stride`.
109fn map2(base: i64, row_stride: i64, col_stride: i64) -> ElemAccess {
110    ElemAccess {
111        base,
112        terms: vec![(LV_GEN_ROW, row_stride), (LV_GEN_COL, col_stride)],
113        indirect: None,
114    }
115}
116
117fn op(i: usize) -> Box<TensorBody> {
118    Box::new(TensorBody::Operand(i))
119}
120
121fn konst(v: f64) -> Box<TensorBody> {
122    Box::new(TensorBody::Const(v.to_bits()))
123}
124
125fn add(a: Box<TensorBody>, b: Box<TensorBody>) -> Box<TensorBody> {
126    Box::new(TensorBody::Add(a, b))
127}
128
129fn sub(a: Box<TensorBody>, b: Box<TensorBody>) -> Box<TensorBody> {
130    Box::new(TensorBody::Sub(a, b))
131}
132
133fn mul(a: Box<TensorBody>, b: Box<TensorBody>) -> Box<TensorBody> {
134    Box::new(TensorBody::Mul(a, b))
135}
136
137fn unary(kind: SfpuKind, a: Box<TensorBody>) -> Box<TensorBody> {
138    Box::new(TensorBody::Unary { kind, arg: a })
139}
140
141/// `x` soft-capped at `c`, `c·tanh(x/c)`, as `2c·σ(2x/c) − c`: one exponential and one reciprocal, reaching `±c` in both
142/// tails.
143fn capped(x: Box<TensorBody>, c: f64) -> Box<TensorBody> {
144    let sigmoid = unary(SfpuKind::Recip, add(konst(1.0), unary(SfpuKind::Exp, mul(x, konst(-2.0 / c)))));
145    sub(mul(konst(2.0 * c), sigmoid), konst(c))
146}
147
148/// `x / (1 + e^-x)` written out of the algebra rather than as `SfpuKind::Silu`, for the places a gate is
149/// one factor of a larger body (Gated DeltaNet's `silu(z)`); the standalone op uses the fused kind.
150fn silu_expr(x: Box<TensorBody>) -> Box<TensorBody> {
151    Box::new(TensorBody::Unary {
152        kind: SfpuKind::Silu,
153        arg: x,
154    })
155}
156
157/// A Gated DeltaNet's output gate on `z`: `silu(z)`, or `σ(z)` as `1/(1 + e^(−z))`.
158fn delta_gate(gate: GateAct, z: Box<TensorBody>) -> Box<TensorBody> {
159    match gate {
160        GateAct::Silu => silu_expr(z),
161        GateAct::Sigmoid => unary(SfpuKind::Recip, add(konst(1.0), unary(SfpuKind::Exp, mul(z, konst(-1.0))))),
162    }
163}
164
165/// Nested `Select`s over the flat column axis: `bands[k] = (hi, body)` computes `body` for every column
166/// below `hi` that the earlier bands did not claim. The bands must be given in increasing `hi`, and the
167/// last one is the fall-through — so a partition of `[0, n)` costs exactly `bands.len() - 1` predicates
168/// and every column is covered.
169fn banded(bands: Vec<(usize, TensorBody)>) -> TensorBody {
170    let mut it = bands.into_iter().rev();
171    let (_, mut acc) = it.next().expect("a banded body has at least one band");
172    for (hi, body) in it {
173        acc = TensorBody::Select {
174            cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Const(hi)),
175            then: Box::new(body),
176            otherwise: Box::new(acc),
177        };
178    }
179    acc
180}
181
182/// Which of the target's elementwise datapaths a stage's body runs on: the float one every float op takes, or the
183/// integer one an int8 model's requant, saturation and ReLU take — named by the op, never guessed from its body (a
184/// float ReLU's `max(x, 0.0)` is integral arithmetic over float data).
185#[derive(Clone, Copy, Debug, PartialEq, Eq)]
186enum Datapath {
187    Float,
188    Int,
189}
190
191impl Datapath {
192    /// The datapath a value of the walk's element kind rides.
193    fn of(float: bool) -> Datapath {
194        if float {
195            Datapath::Float
196        } else {
197            Datapath::Int
198        }
199    }
200
201    fn precision(self, target: &SpatialTarget) -> Result<Precision, SpatialReject> {
202        match self {
203            Datapath::Float => elem_precision(target),
204            Datapath::Int => super::from_tensor_op::elementwise_int(target).ok_or(
205                SpatialReject::MissingCompute {
206                    needed: "an integer elementwise datapath",
207                    target: target.name,
208                },
209            ),
210        }
211    }
212}
213
214/// The float precision the target's elementwise datapath offers (the RDU's witnessed one is bf16).
215pub(crate) fn elem_precision(target: &SpatialTarget) -> Result<Precision, SpatialReject> {
216    super::from_tensor_op::elementwise_float(target).ok_or(SpatialReject::MissingCompute {
217        needed: "a float elementwise datapath",
218        target: target.name,
219    })
220}
221
222/// The width in elements of a buffer holding `page_tiles` tiles.
223fn wide(page_tiles: usize, target: &SpatialTarget) -> usize {
224    page_width(page_tiles, target)
225}
226
227/// Tiles a flat `1 × n` row occupies.
228fn tiles_of(n: usize, target: &SpatialTarget) -> usize {
229    n.div_ceil(target.tile.cols).max(1)
230}
231
232/// The ONE work item of a flat section: the whole `1 × cols_t` tile row.
233///
234/// ★ ONE ITEM, NOT ONE PER TILE. `place` clones the WHOLE program per site, so a vector split into one
235/// item per tile becomes that many copies of every stage — which is how a 64-lane decision program ended
236/// up with two identical `Scalar` stages carrying the same 75-cell body. A stage whose compute covers
237/// the entire row in one pass has exactly one item, and the site count follows the work rather than the
238/// tile count.
239fn whole_row(cols_t: usize) -> Vec<WorkItem> {
240    vec![WorkItem {
241        id: WorkItemId::new(0),
242        coord: (0, 0),
243        out: TileRect {
244            row0_t: 0,
245            col0_t: 0,
246            h_t: 1,
247            w_t: cols_t,
248        },
249    }]
250}
251
252/// A flat operand slot: which source tensor it reads and through which map.
253struct Slot {
254    buf: BufferId,
255    map: ElemAccess,
256}
257
258/// A flat section under construction: the loads of its operands, then ONE `Generic` over the whole row — or, for a
259/// result wider than an accumulator page, over one block of it an item ([`FlatBuild::column_blocked`]).
260struct FlatBuild {
261    t: Template,
262    precision: Precision,
263    loads: Vec<(StageId, BufferId)>,
264    blocks: Option<ColumnBlocks>,
265}
266
267/// A flat result blocked by columns: `items` items of `blk` result tiles each.
268#[derive(Clone, Copy, Debug)]
269struct ColumnBlocks {
270    blk: usize,
271    items: usize,
272}
273
274impl FlatBuild {
275    fn new(target: &SpatialTarget) -> Result<FlatBuild, SpatialReject> {
276        FlatBuild::on(Datapath::Float, target)
277    }
278
279    /// A flat section whose activations ride `datapath`.
280    fn on(datapath: Datapath, target: &SpatialTarget) -> Result<FlatBuild, SpatialReject> {
281        FlatBuild::at(datapath.precision(target)?, target)
282    }
283
284    /// A flat section whose activations are `precision`, chosen by its caller.
285    fn at(precision: Precision, target: &SpatialTarget) -> Result<FlatBuild, SpatialReject> {
286        Ok(FlatBuild {
287            t: Template::new(target)?,
288            precision,
289            loads: Vec::new(),
290            blocks: None,
291        })
292    }
293
294    /// A float section whose `n`-column result is wider than an accumulator page: its pointwise stage
295    /// ([`Self::pointwise`]) runs an item per block of `blk` result tiles, every operand loaded whole, so the page an item
296    /// computes is one block. A result a page holds is the one whole-row item, as [`Self::new`] makes it.
297    fn columns(n: usize, target: &SpatialTarget) -> Result<FlatBuild, SpatialReject> {
298        let mut f = FlatBuild::new(target)?;
299        if row_needs_blocking(n, target) {
300            let cols_t = tiles_of(n, target);
301            let blk = row_block(cols_t, target);
302            f.blocks = Some(ColumnBlocks { blk, items: cols_t / blk });
303        }
304        Ok(f)
305    }
306
307    /// The buffer a section's `n`-column result is computed into: the whole row, or one block of it.
308    fn result_page(&mut self, name: &str, n: usize, precision: Precision, target: &SpatialTarget) -> BufferId {
309        let tiles = self.blocks.map_or(tiles_of(n, target), |b| b.blk);
310        self.t.buffer(name, tiles, precision)
311    }
312
313    /// `dst[c] = body(c)` over the result's `n` columns, every map and predicate written in RESULT columns. In a
314    /// column-blocked section the stage walks its item's block, and a column `c` of the block is result column
315    /// `c + o` for the item's first column `o`, read from a compile-time table one page an item: a map that walks the
316    /// columns gathers through `o`; a predicate's constant bound `k` becomes the table's `k − o`; a bound or a gather
317    /// index the section computes is moved by `o` in a stage of its own first. Returns the stages, in order.
318    fn pointwise(
319        &mut self,
320        name: &str,
321        n: usize,
322        slots: Vec<Slot>,
323        body: TensorBody,
324        dst: BufferId,
325        target: &SpatialTarget,
326    ) -> Result<Vec<StageId>, SpatialReject> {
327        let Some(ColumnBlocks { blk, items }) = self.blocks else {
328            let w = wide(self.t.page_tiles(dst), target);
329            return Ok(vec![self.compute(name, flat_space(n), slots, ident_map(w), body, dst)]);
330        };
331        let tw = target.tile.cols;
332        let span = items * blk * tw;
333        // The constant column bounds, each a table entry.
334        let mut bounds: Vec<usize> = Vec::new();
335        fn constant_bounds(b: &TensorBody, out: &mut Vec<usize>) {
336            match b {
337                TensorBody::Operand(_) | TensorBody::Const(_) => {}
338                TensorBody::Add(x, y)
339                | TensorBody::Sub(x, y)
340                | TensorBody::Mul(x, y)
341                | TensorBody::Min(x, y)
342                | TensorBody::Max(x, y)
343                | TensorBody::FloorDiv(x, y) => {
344                    constant_bounds(x, out);
345                    constant_bounds(y, out);
346                }
347                TensorBody::Unary { arg, .. } | TensorBody::Shr { arg, .. } => constant_bounds(arg, out),
348                TensorBody::Select { cond, then, otherwise } => {
349                    if let (LV_GEN_COL, PredRhs::Const(k)) = (cond.var(), cond.rhs()) {
350                        if !out.contains(&k) {
351                            out.push(k);
352                        }
353                    }
354                    constant_bounds(then, out);
355                    constant_bounds(otherwise, out);
356                }
357            }
358        }
359        constant_bounds(&body, &mut bounds);
360        // Item `i`'s page: its first column `o = i·blk·tile`, then each bound less `o`.
361        let per = (1 + bounds.len()).div_ceil(tw);
362        let magnitude = span.max(bounds.iter().copied().max().unwrap_or(0)) + span;
363        let ix = index_type(target, magnitude as u64)?;
364        let values: Vec<i64> = (0..items)
365            .flat_map(|i| {
366                let o = (i * blk * tw) as i64;
367                let mut page: Vec<i64> = std::iter::once(o).chain(bounds.iter().map(|k| *k as i64 - o)).collect();
368                page.resize(per * tw, 0);
369                page
370            })
371            .collect();
372        let table = self.t.tensor(&format!("{name}__columns"), TensorRole::Param, 1, items * per * tw, ix, TensorInit::Values(values));
373        let at_table = self.t.buffer_deep(&format!("in_{name}_columns"), 1, per, ix);
374        let l_table = self.t.load_block_at(&format!("load_{name}_columns"), table, 0, vec![(LV_ITEM_X, per)], 1, per, items * per, at_table);
375        self.loads.push((l_table, at_table));
376        let mut stages = Vec::new();
377        // The block's slots: the first column, each bound, then the section's own, then what is moved by `o`.
378        let first = Slot { buf: at_table, map: fixed(0) };
379        let mut out: Vec<Slot> = std::iter::once(first).chain((0..bounds.len()).map(|m| Slot { buf: at_table, map: fixed(1 + m as i64) })).collect();
380        let shift = out.len();
381        let one = wide(1, target);
382        let mut moved = |f: &mut FlatBuild, what: &str, from: &Slot, body: Box<TensorBody>| {
383            let page = f.t.scratch_deep(&format!("{name}_{what}"), 1, 1, ix);
384            stages.push(f.compute(&format!("{name}_{what}"), flat_space(1), vec![Slot { buf: at_table, map: fixed(0) }, Slot { buf: from.buf, map: from.map.clone() }], ident_map(one), *body, page));
385            page
386        };
387        let mut own = Vec::with_capacity(slots.len());
388        let mut extra: Vec<Slot> = Vec::new();
389        for (j, slot) in slots.iter().enumerate() {
390            let walks = slot.map.terms.iter().find(|(v, _)| *v == LV_GEN_COL).map_or(0, |(_, s)| *s);
391            let map = match (walks, slot.map.indirect) {
392                (0, _) => slot.map.clone(),
393                (s, None) => ElemAccess {
394                    indirect: Some(IndirectTerm { operand: 0, stride: s, modulus: span + 1 }),
395                    ..slot.map.clone()
396                },
397                // A gather already indexed by data: its index moves by `o` too, shifted to stay non-negative.
398                (s, Some(ind)) => {
399                    let lift = ind.stride.unsigned_abs() as usize * ind.modulus;
400                    let index = &slots[ind.operand];
401                    let page = moved(self, &format!("index{j}"), index, add(add(mul(op(0), konst(s as f64)), mul(op(1), konst(ind.stride as f64))), konst(lift as f64)));
402                    extra.push(Slot { buf: page, map: fixed(0) });
403                    ElemAccess {
404                        base: slot.map.base - lift as i64,
405                        terms: slot.map.terms.clone(),
406                        indirect: Some(IndirectTerm { operand: usize::MAX - (extra.len() - 1), stride: 1, modulus: s.unsigned_abs() as usize * span + 2 * lift + 1 }),
407                    }
408                }
409            };
410            own.push(Slot { buf: slot.buf, map });
411        }
412        // A bound the section computes moves by `o`.
413        let mut computed: Vec<(usize, usize)> = Vec::new();
414        fn computed_bounds(b: &TensorBody, out: &mut Vec<usize>) {
415            match b {
416                TensorBody::Operand(_) | TensorBody::Const(_) => {}
417                TensorBody::Add(x, y)
418                | TensorBody::Sub(x, y)
419                | TensorBody::Mul(x, y)
420                | TensorBody::Min(x, y)
421                | TensorBody::Max(x, y)
422                | TensorBody::FloorDiv(x, y) => {
423                    computed_bounds(x, out);
424                    computed_bounds(y, out);
425                }
426                TensorBody::Unary { arg, .. } | TensorBody::Shr { arg, .. } => computed_bounds(arg, out),
427                TensorBody::Select { cond, then, otherwise } => {
428                    if let (LV_GEN_COL, PredRhs::Operand(o)) = (cond.var(), cond.rhs()) {
429                        if !out.contains(&o) {
430                            out.push(o);
431                        }
432                    }
433                    computed_bounds(then, out);
434                    computed_bounds(otherwise, out);
435                }
436            }
437        }
438        let mut operands = Vec::new();
439        computed_bounds(&body, &mut operands);
440        for o in operands {
441            let page = moved(self, &format!("bound{o}"), &slots[o], sub(op(1), op(0)));
442            extra.push(Slot { buf: page, map: fixed(0) });
443            computed.push((o, extra.len() - 1));
444        }
445        let base = shift + own.len();
446        // The placeholders for the moved indices name their slot past the section's own.
447        for slot in &mut own {
448            if let Some(ind) = slot.map.indirect.as_mut() {
449                if ind.operand > usize::MAX / 2 {
450                    ind.operand = base + (usize::MAX - ind.operand);
451                }
452            }
453        }
454        out.extend(own);
455        out.extend(extra);
456        fn rewrite(b: &TensorBody, shift: usize, bounds: &[usize], computed: &[(usize, usize)], base: usize) -> TensorBody {
457            let r = |x: &TensorBody| Box::new(rewrite(x, shift, bounds, computed, base));
458            match b {
459                TensorBody::Operand(i) => TensorBody::Operand(i + shift),
460                TensorBody::Const(c) => TensorBody::Const(*c),
461                TensorBody::Add(x, y) => TensorBody::Add(r(x), r(y)),
462                TensorBody::Sub(x, y) => TensorBody::Sub(r(x), r(y)),
463                TensorBody::Mul(x, y) => TensorBody::Mul(r(x), r(y)),
464                TensorBody::Min(x, y) => TensorBody::Min(r(x), r(y)),
465                TensorBody::Max(x, y) => TensorBody::Max(r(x), r(y)),
466                TensorBody::FloorDiv(x, y) => TensorBody::FloorDiv(r(x), r(y)),
467                TensorBody::Shr { arg, by } => TensorBody::Shr { arg: r(arg), by: *by },
468                TensorBody::Unary { kind, arg } => TensorBody::Unary { kind: *kind, arg: r(arg) },
469                TensorBody::Select { cond, then, otherwise } => {
470                    let rhs = match (cond.var(), cond.rhs()) {
471                        (LV_GEN_COL, PredRhs::Const(k)) => PredRhs::Operand(1 + bounds.iter().position(|b| *b == k).expect("every bound is in the table")),
472                        (LV_GEN_COL, PredRhs::Operand(o)) => {
473                            PredRhs::Operand(base + computed.iter().find(|(from, _)| *from == o).expect("every computed bound is moved").1)
474                        }
475                        (_, PredRhs::Operand(o)) => PredRhs::Operand(o + shift),
476                        (_, rhs) => rhs,
477                    };
478                    let cond = match cond {
479                        Pred::AxisLt(v, _) => Pred::AxisLt(*v, rhs),
480                        Pred::AxisGe(v, _) => Pred::AxisGe(*v, rhs),
481                    };
482                    TensorBody::Select { cond, then: r(then), otherwise: r(otherwise) }
483                }
484            }
485        }
486        let body = rewrite(&body, shift, &bounds, &computed, base);
487        let w = wide(blk, target);
488        stages.push(self.compute(name, flat_space(blk * tw), out, ident_map(w), body, dst));
489        Ok(stages)
490    }
491
492    /// Declare a flat `1 × n` operand tensor and load the whole row into a buffer.
493    fn input(&mut self, name: &str, n: usize, fill: &Fill, target: &SpatialTarget) -> BufferId {
494        let p = self.precision;
495        self.input_as(name, n, fill, p, target)
496    }
497
498    /// [`Self::input`] at an explicit element precision — an integer table, or an INDEX (the target's address
499    /// integer, [`index_type`]: a float datapath holds integers exactly only up to its mantissa).
500    fn input_as(
501        &mut self,
502        name: &str,
503        n: usize,
504        fill: &Fill,
505        precision: Precision,
506        target: &SpatialTarget,
507    ) -> BufferId {
508        self.input_paged(name, n, fill, precision, 2, target)
509    }
510
511    /// [`Self::input_as`] into a buffer of `depth` pages: one for an operand a one-item section loads once.
512    fn input_paged(
513        &mut self,
514        name: &str,
515        n: usize,
516        fill: &Fill,
517        precision: Precision,
518        depth: usize,
519        target: &SpatialTarget,
520    ) -> BufferId {
521        let cols_t = tiles_of(n, target);
522        let id = self.t.tensor(
523            name,
524            TensorRole::Input,
525            1,
526            n,
527            precision,
528            super::from_tensor_op::init_of(fill, n),
529        );
530        let buf = self.t.buffer_deep(&format!("in_{name}"), depth, cols_t, precision);
531        let l = self
532            .t
533            .load(&format!("load_{name}"), id, vec![], cols_t, buf);
534        self.loads.push((l, buf));
535        buf
536    }
537
538    /// Tiles `first..first + tiles_of(n)` of a `1 × total` input, `n` elements from tile `first`, into a page of their
539    /// own loaded once: one block of a row wider than a page.
540    #[allow(clippy::too_many_arguments)]
541    fn input_slice(
542        &mut self,
543        name: &str,
544        total: usize,
545        first: usize,
546        n: usize,
547        fill: &Fill,
548        precision: Precision,
549        target: &SpatialTarget,
550    ) -> BufferId {
551        let id = self.t.tensor(name, TensorRole::Input, 1, total, precision, super::from_tensor_op::init_of(fill, total));
552        let tiles = tiles_of(n, target);
553        let buf = self.t.buffer_deep(&format!("in_{name}"), 1, tiles, precision);
554        let l = self.t.load_block_at(&format!("load_{name}"), id, first, vec![], 1, tiles, tiles_of(total, target), buf);
555        self.loads.push((l, buf));
556        buf
557    }
558
559    /// The accumulation-path scratch between two compute stages, in the accumulator the section's precision folds
560    /// in: fp32 for a float section, an integer as wide as the datapath for an integer one.
561    fn scratch(&mut self, name: &str, page_tiles: usize) -> BufferId {
562        let p = AccPolicy::for_result(self.precision).interm_precision();
563        self.t.scratch(name, page_tiles, p)
564    }
565
566    fn compute(
567        &mut self,
568        name: &str,
569        iters: Vec<IterSpec>,
570        slots: Vec<Slot>,
571        out_map: ElemAccess,
572        body: TensorBody,
573        dst: BufferId,
574    ) -> StageId {
575        let ins: Vec<BufferId> = slots.iter().map(|s| s.buf).collect();
576        let maps: Vec<ElemAccess> = slots.into_iter().map(|s| s.map).collect();
577        self.t.stage(
578            name,
579            StageKind::Compute(ComputeOp::Generic {
580                iters,
581                ins: maps,
582                out: out_map,
583                body,
584            }),
585            ins,
586            vec![dst],
587        )
588    }
589
590    /// Wire every load to the first chain stage that reads it, then chain the stages, then the store.
591    /// ★ A SECTION ALWAYS CALLS ITS RESULT AN `Output`. Whether it is really an activation handed to
592    /// the next section, or the carried state of the step, is not knowable from inside one section —
593    /// it is a property of the WHOLE program, so `compose` decides it and a section lowered in
594    /// isolation stays a program that validates on its own.
595    fn finish(
596        self,
597        out_name: &str,
598        n: usize,
599        chain: Vec<StageId>,
600        out_buf: BufferId,
601        name: String,
602        claim: NumericsClaim,
603        target: &SpatialTarget,
604    ) -> Result<SpatialProgram, SpatialReject> {
605        let p = self.precision;
606        self.finish_as(out_name, n, chain, out_buf, name, claim, p, target)
607    }
608
609    /// [`Self::finish`] with the result tensor at an explicit precision (an index, an integer row).
610    #[allow(clippy::too_many_arguments)]
611    fn finish_as(
612        mut self,
613        out_name: &str,
614        n: usize,
615        chain: Vec<StageId>,
616        out_buf: BufferId,
617        name: String,
618        claim: NumericsClaim,
619        out_precision: Precision,
620        target: &SpatialTarget,
621    ) -> Result<SpatialProgram, SpatialReject> {
622        let cols_t = tiles_of(n, target);
623        let y = self.t.tensor(
624            out_name,
625            TensorRole::Output,
626            1,
627            n,
628            out_precision,
629            TensorInit::Zero,
630        );
631        let st = match self.blocks {
632            None => self.t.store(&format!("store_{out_name}"), y, vec![], cols_t, out_buf),
633            Some(ColumnBlocks { blk, .. }) => self.t.store(&format!("store_{out_name}"), y, vec![(LV_ITEM_X, blk)], blk, out_buf),
634        };
635        self.wire_and_finish(chain, out_buf, st, cols_t, name, claim, target)
636    }
637
638    /// Wire every load to the first chain stage that reads it, chain the stages, drain the last into the
639    /// store `st`, and finish over one whole-row work item `cols_t` tiles wide — the part of [`Self::finish`]
640    /// shared by a section whose store is not the plain whole-row one (a scatter).
641    #[allow(clippy::too_many_arguments)]
642    fn wire_and_finish(
643        self,
644        chain: Vec<StageId>,
645        out_buf: BufferId,
646        st: StageId,
647        cols_t: usize,
648        name: String,
649        claim: NumericsClaim,
650        target: &SpatialTarget,
651    ) -> Result<SpatialProgram, SpatialReject> {
652        self.wire_and_finish_many(chain, &[(out_buf, st)], cols_t, name, claim, target)
653    }
654
655    /// [`Self::wire_and_finish`] for a section with SEVERAL results, each drained by its own store — every
656    /// store fed by the chain stage that produces its buffer.
657    fn wire_and_finish_many(
658        mut self,
659        chain: Vec<StageId>,
660        stores: &[(BufferId, StageId)],
661        cols_t: usize,
662        name: String,
663        claim: NumericsClaim,
664        target: &SpatialTarget,
665    ) -> Result<SpatialProgram, SpatialReject> {
666        for (l, b) in std::mem::take(&mut self.loads) {
667            let consumer = *chain.iter().find(|c| self.t.stage_reads(**c, b)).ok_or(
668                SpatialReject::NotLowered("an operand nothing in the chain reads"),
669            )?;
670            self.t.edge(l, b, consumer);
671        }
672        for i in 0..chain.len() {
673            let b = self.t.stage_output(chain[i]);
674            if let Some(next) = chain
675                .iter()
676                .skip(i + 1)
677                .find(|c| self.t.stage_reads(**c, b))
678            {
679                self.t.edge(chain[i], b, *next);
680            }
681        }
682        for (b, st) in stores {
683            let producer = *chain
684                .iter()
685                .rev()
686                .find(|c| self.t.stage_output(**c) == *b)
687                .ok_or(SpatialReject::NotLowered(
688                    "a stored buffer no stage of the chain produces",
689                ))?;
690            self.t.edge(producer, *b, *st);
691        }
692        let acc = AccPolicy::for_result(self.precision);
693        let items = match self.blocks {
694            None => whole_row(cols_t),
695            Some(ColumnBlocks { blk, .. }) => row_items(cols_t, blk),
696        };
697        self.t.finish(name, items, acc, claim, target)
698    }
699}
700
701/// The claim every float composition earns: the chain rounds at each stage boundary exactly where the
702/// datapath does, and the reference is the VM's f64 arithmetic.
703fn approx(what: &str, p: Precision) -> NumericsClaim {
704    NumericsClaim::Approx {
705        reason: format!("{what} in {p:?}"),
706    }
707}
708
709/// What a map's numbers promise: exact when it computes integers with integer arithmetic — an int8 model's
710/// requant, saturation and ReLU, evaluated in `i64` — and [`approx`] in the float precision otherwise.
711fn claim_of(what: &str, body: &TensorBody, p: Precision) -> NumericsClaim {
712    if precision_int_bits(p).is_some() && body.integral() {
713        NumericsClaim::Exact
714    } else {
715        approx(what, p)
716    }
717}
718
719// ── The element-wise family: one `Generic`, one body ──────────────────────────────────────────────
720
721/// A unary map `y = f(x)` over a `rows × cols` tensor: a row in the row form below, a matrix tile by tile.
722fn unary_map(
723    name: &str,
724    kind: &str,
725    datapath: Datapath,
726    body: impl Fn(Box<TensorBody>) -> Box<TensorBody>,
727    x: (&str, &Fill),
728    out: &str,
729    shape: (usize, usize),
730    target: &SpatialTarget,
731) -> Result<SpatialProgram, SpatialReject> {
732    let (rows, n) = shape;
733    if rows > 1 {
734        return matrix_unary_map(name, kind, datapath, body, x, out, shape, target);
735    }
736    let cols_t = tiles_of(n, target);
737    if cols_t > row_block_tiles(target) {
738        return blocked_unary_map(name, kind, datapath, body, x, out, n, target);
739    }
740    let mut b = FlatBuild::on(datapath, target)?;
741    let xin = b.input(x.0, n, x.1, target);
742    let out_buf = b.t.buffer("out", cols_t, b.precision);
743    let w = wide(cols_t, target);
744    let s = b.compute(
745        kind,
746        flat_space(n),
747        vec![Slot {
748            buf: xin,
749            map: at(0),
750        }],
751        ident_map(w),
752        *body(op(0)),
753        out_buf,
754    );
755    let claim = claim_of(kind, &body(op(0)), b.precision);
756    b.finish(
757        out,
758        n,
759        vec![s],
760        out_buf,
761        format!("{name}_{n}"),
762        claim,
763        target,
764    )
765}
766
767/// `y = body(a, b)` over a `1 × n` row a site holds whole: one item over the row, both operands loaded whole and mapped
768/// together — the form [`unary_map`] gives a row, for an element-wise op of two rows. The caller has chosen the
769/// precision and the claim (the element-wise lowering, from its operands' fills).
770#[allow(clippy::too_many_arguments)]
771pub(crate) fn row_binary_map(
772    name: String,
773    kind: &str,
774    body: TensorBody,
775    a: (&str, &Fill),
776    b: (&str, &Fill),
777    out: &str,
778    n: usize,
779    precision: Precision,
780    claim: NumericsClaim,
781    target: &SpatialTarget,
782) -> Result<SpatialProgram, SpatialReject> {
783    let cols_t = tiles_of(n, target);
784    let mut fb = FlatBuild::at(precision, target)?;
785    let ain = fb.input(a.0, n, a.1, target);
786    let bin = fb.input(b.0, n, b.1, target);
787    let out_buf = fb.t.buffer("out", cols_t, precision);
788    let w = wide(cols_t, target);
789    let s = fb.compute(
790        kind,
791        flat_space(n),
792        vec![Slot { buf: ain, map: at(0) }, Slot { buf: bin, map: at(0) }],
793        ident_map(w),
794        body,
795        out_buf,
796    );
797    fb.finish(out, n, vec![s], out_buf, name, claim, target)
798}
799
800/// How many tiles of one row a single site may hold resident. A flat section keeps its whole row in
801/// local memory, which is free while the row is a handful of tiles and impossible once it is not: at
802/// Qwen3.8's real `d = 5120` a row is 160 tiles against an accumulator budget of eight, and the
803/// whole-row form is refused as `AccTilesExceeded` long before local memory is consulted.
804pub fn row_needs_blocking(cols: usize, target: &SpatialTarget) -> bool {
805    tiles_of(cols, target) > row_block_tiles(target)
806}
807
808pub(crate) fn row_block_tiles(target: &SpatialTarget) -> usize {
809    target.effective_acc_tiles(AccPolicy::F32, true).max(1)
810}
811
812/// ★ A POINTWISE MAP HAS NO CROSS-ELEMENT DEPENDENCE, so a row wider than a site can hold is simply
813/// WORK, not a problem: one item per tile, a one-tile page, and the row streams. This is the same shape
814/// `from_tensor_op::from_elementwise` has always had.
815///
816/// The whole-row form above stays for narrow rows on purpose — see `whole_row`: `place` clones the
817/// program per site, so splitting a two-tile vector into two items buys nothing and doubles the stages.
818/// The rule is therefore about the ROW, not about the operation: one item while the row fits the site's
819/// budget, one item per tile once it does not.
820#[allow(clippy::too_many_arguments)]
821fn blocked_unary_map(
822    name: &str,
823    kind: &str,
824    datapath: Datapath,
825    body: impl Fn(Box<TensorBody>) -> Box<TensorBody>,
826    x: (&str, &Fill),
827    out: &str,
828    n: usize,
829    target: &SpatialTarget,
830) -> Result<SpatialProgram, SpatialReject> {
831    let precision = datapath.precision(target)?;
832    let cols_t = tiles_of(n, target);
833    let mut t = Template::new(target)?;
834    let xt = t.tensor(
835        x.0,
836        TensorRole::Input,
837        1,
838        n,
839        precision,
840        super::from_tensor_op::init_of(x.1, n),
841    );
842    let yt = t.tensor(out, TensorRole::Output, 1, n, precision, TensorInit::Zero);
843    let xin = t.buffer("in_x", 1, precision);
844    let out_buf = t.buffer("out", 1, precision);
845    let lx = t.load("load_x", xt, vec![(LV_ITEM_X, 1)], 1, xin);
846    let w = wide(1, target);
847    let s = t.stage(
848        kind,
849        StageKind::Compute(ComputeOp::Generic {
850            iters: flat_space(w),
851            ins: vec![ident_map(w)],
852            out: ident_map(w),
853            body: *body(op(0)),
854        }),
855        vec![xin],
856        vec![out_buf],
857    );
858    let st = t.store(
859        &format!("store_{out}"),
860        yt,
861        vec![(LV_ITEM_X, 1)],
862        1,
863        out_buf,
864    );
865    t.edge(lx, xin, s);
866    t.edge(s, out_buf, st);
867    t.finish(
868        format!("{name}_{n}"),
869        items_per_tile(cols_t),
870        AccPolicy::for_result(precision),
871        claim_of(kind, &body(op(0)), precision),
872        target,
873    )
874}
875
876/// A unary map over a MATRIX: the tensors keep the program's `rows × cols`, and every tile of the 2-D tile grid
877/// is one work item mapped over its full `tile × tile` page — the tiling a matmul reads the same tensor in, so a
878/// composition joins the two without a layout conversion.
879#[allow(clippy::too_many_arguments)]
880fn matrix_unary_map(
881    name: &str,
882    kind: &str,
883    datapath: Datapath,
884    body: impl Fn(Box<TensorBody>) -> Box<TensorBody>,
885    x: (&str, &Fill),
886    out: &str,
887    shape: (usize, usize),
888    target: &SpatialTarget,
889) -> Result<SpatialProgram, SpatialReject> {
890    let (rows, cols) = shape;
891    let precision = datapath.precision(target)?;
892    let mut t = Template::new(target)?;
893    let xt = t.tensor(
894        x.0,
895        TensorRole::Input,
896        rows,
897        cols,
898        precision,
899        super::from_tensor_op::init_of(x.1, rows * cols),
900    );
901    let yt = t.tensor(
902        out,
903        TensorRole::Output,
904        rows,
905        cols,
906        precision,
907        TensorInit::Zero,
908    );
909    let (rows_t, cols_t) = (
910        rows.div_ceil(target.tile.rows),
911        cols.div_ceil(target.tile.cols),
912    );
913    let steps = super::from_tensor_op::tile_grid_steps(rows_t, cols_t);
914    let xin = t.buffer("in_x", 1, precision);
915    let out_buf = t.buffer("out", 1, precision);
916    let lx = t.load("load_x", xt, steps.clone(), 1, xin);
917    let w = target.tile.cols;
918    let s = t.stage(
919        kind,
920        StageKind::Compute(ComputeOp::Generic {
921            iters: super::from_tensor_op::space(target, w, None),
922            ins: vec![ident_map(w)],
923            out: ident_map(w),
924            body: *body(op(0)),
925        }),
926        vec![xin],
927        vec![out_buf],
928    );
929    let st = t.store(&format!("store_{out}"), yt, steps, 1, out_buf);
930    t.edge(lx, xin, s);
931    t.edge(s, out_buf, st);
932    t.finish(
933        format!("{name}_{rows}x{cols}"),
934        super::from_tensor_op::items_per_tile_grid(rows_t, cols_t),
935        AccPolicy::for_result(precision),
936        claim_of(kind, &body(op(0)), precision),
937        target,
938    )
939}
940
941/// `y = silu(x)` — `x / (1 + e^-x)`, one fused SFPU kind.
942pub fn from_silu(
943    spec: &SiluSpec,
944    names: &StageNames,
945    shape: (usize, usize),
946    target: &SpatialTarget,
947) -> Result<SpatialProgram, SpatialReject> {
948    unary_map(
949        "silu",
950        "silu",
951        Datapath::Float,
952        |x| silu_expr(x),
953        names.operand(0),
954        &names.out,
955        pointwise_shape(shape, spec.n),
956        target,
957    )
958}
959
960/// `y = (x · M) >> S` — an int8 model's requant: the exact product, then an arithmetic shift.
961pub fn from_int_requant(
962    spec: &IntRequantSpec,
963    names: &StageNames,
964    shape: (usize, usize),
965    target: &SpatialTarget,
966) -> Result<SpatialProgram, SpatialReject> {
967    let (m, s) = (spec.m as f64, spec.s);
968    unary_map(
969        "int_requant",
970        "int_requant",
971        Datapath::Int,
972        |x| Box::new(TensorBody::Shr { arg: mul(x, konst(m)), by: s }),
973        names.operand(0),
974        &names.out,
975        pointwise_shape(shape, spec.n),
976        target,
977    )
978}
979
980/// `y = min(max(x, LO), HI)` — saturation into `[LO, HI]`.
981pub fn from_int_clamp(
982    spec: &IntClampSpec,
983    names: &StageNames,
984    shape: (usize, usize),
985    target: &SpatialTarget,
986) -> Result<SpatialProgram, SpatialReject> {
987    let (lo, hi) = (spec.lo as f64, spec.hi as f64);
988    unary_map(
989        "int_clamp",
990        "int_clamp",
991        Datapath::Int,
992        |x| Box::new(TensorBody::Min(Box::new(TensorBody::Max(x, konst(lo))), konst(hi))),
993        names.operand(0),
994        &names.out,
995        pointwise_shape(shape, spec.n),
996        target,
997    )
998}
999
1000/// `y = max(x, 0)`, in the tensor's element kind.
1001pub fn from_relu(
1002    spec: &ReluSpec,
1003    names: &StageNames,
1004    shape: (usize, usize),
1005    target: &SpatialTarget,
1006) -> Result<SpatialProgram, SpatialReject> {
1007    unary_map(
1008        "relu",
1009        "relu",
1010        if spec.int { Datapath::Int } else { Datapath::Float },
1011        |x| Box::new(TensorBody::Max(x, konst(0.0))),
1012        names.operand(0),
1013        &names.out,
1014        pointwise_shape(shape, spec.n),
1015        target,
1016    )
1017}
1018
1019/// The shape a pointwise op maps over: the program's `rows × cols` when it accounts for the `n` elements the
1020/// recognizer counted, a `1 × n` row otherwise.
1021fn pointwise_shape(shape: (usize, usize), n: usize) -> (usize, usize) {
1022    if shape.0 * shape.1 == n {
1023        shape
1024    } else {
1025        (1, n)
1026    }
1027}
1028
1029/// `y = sigmoid(x)` — `1/(1 + e^-x)`, the reciprocal of the silu denominator. No new SFPU kind: `Exp`
1030/// and `Recip` already exist and the composition of the two IS the sigmoid.
1031pub fn from_sigmoid(
1032    spec: &SigmoidSpec,
1033    names: &StageNames,
1034    shape: (usize, usize),
1035    target: &SpatialTarget,
1036) -> Result<SpatialProgram, SpatialReject> {
1037    unary_map(
1038        "sigmoid",
1039        "sigmoid",
1040        Datapath::Float,
1041        |x| {
1042            unary(
1043                SfpuKind::Recip,
1044                add(konst(1.0), unary(SfpuKind::Exp, sub(konst(0.0), x))),
1045            )
1046        },
1047        names.operand(0),
1048        &names.out,
1049        pointwise_shape(shape, spec.n),
1050        target,
1051    )
1052}
1053
1054/// `y = gelu(x)` — ggml's tanh approximation in its one-exponential form `x / (1 + e^-2z)` with
1055/// `z = √(2/π)·(x + 0.044715·x³)`, exactly as the parser desugars it (same arrangement, same tails).
1056pub fn from_gelu(
1057    spec: &GeluSpec,
1058    names: &StageNames,
1059    shape: (usize, usize),
1060    target: &SpatialTarget,
1061) -> Result<SpatialProgram, SpatialReject> {
1062    const SQRT_2_OVER_PI: f64 = 0.797_884_560_802_865_4;
1063    const K: f64 = 0.044_715;
1064    unary_map(
1065        "gelu",
1066        "gelu",
1067        Datapath::Float,
1068        |x| {
1069            let cube = mul(x.clone(), mul(x.clone(), x.clone()));
1070            let z = mul(konst(SQRT_2_OVER_PI), add(x.clone(), mul(konst(K), cube)));
1071            let two_z = mul(konst(2.0), z);
1072            mul(
1073                x,
1074                unary(
1075                    SfpuKind::Recip,
1076                    add(konst(1.0), unary(SfpuKind::Exp, sub(konst(0.0), two_z))),
1077                ),
1078            )
1079        },
1080        names.operand(0),
1081        &names.out,
1082        pointwise_shape(shape, spec.n),
1083        target,
1084    )
1085}
1086
1087/// `y = x·C` — one multiply by the program's constant.
1088pub fn from_scale(
1089    spec: &ScaleSpec,
1090    names: &StageNames,
1091    shape: (usize, usize),
1092    target: &SpatialTarget,
1093) -> Result<SpatialProgram, SpatialReject> {
1094    let c = spec.factor.get();
1095    unary_map(
1096        "scale",
1097        "scale",
1098        Datapath::Float,
1099        |x| mul(x, konst(c)),
1100        names.operand(0),
1101        &names.out,
1102        pointwise_shape(shape, spec.n),
1103        target,
1104    )
1105}
1106
1107/// `y = C·tanh(x/C)` — each element soft-capped (`capped`'s form: one exponential and one reciprocal).
1108pub fn from_softcap(
1109    spec: &SoftcapSpec,
1110    names: &StageNames,
1111    shape: (usize, usize),
1112    target: &SpatialTarget,
1113) -> Result<SpatialProgram, SpatialReject> {
1114    let c = spec.cap.get();
1115    unary_map(
1116        "softcap",
1117        "softcap",
1118        Datapath::Float,
1119        |x| capped(x, c),
1120        names.operand(0),
1121        &names.out,
1122        pointwise_shape(shape, spec.n),
1123        target,
1124    )
1125}
1126
1127/// `y[i] = (1 + e^(ar[i] + dt[i]))^(alog[i])` — Gated DeltaNet's forget gate.
1128///
1129/// ★ A REAL POWER NEEDS A REAL LOG, AND THAT IS WHY `SfpuKind::Ln` EXISTS. `base^rate` with a
1130/// data-dependent, non-integer, negative `rate` is `exp(rate · ln base)` and nothing else; no
1131/// composition of exp, reciprocal and rsqrt reaches it. The base is `1 + e^u > 1`, so the log is
1132/// always taken of a positive number.
1133pub fn from_decay(
1134    spec: &DecaySpec,
1135    names: &StageNames,
1136    target: &SpatialTarget,
1137) -> Result<SpatialProgram, SpatialReject> {
1138    if spec.rows > 1 {
1139        return dense_decay(spec, names, target);
1140    }
1141    let n = spec.n;
1142    let mut b = FlatBuild::new(target)?;
1143    let ar = b.input(names.operand(0).0, n, names.operand(0).1, target);
1144    let dt = b.input(names.operand(1).0, n, names.operand(1).1, target);
1145    let al = b.input(names.operand(2).0, n, names.operand(2).1, target);
1146    let cols_t = tiles_of(n, target);
1147    let out_buf = b.t.buffer("out", cols_t, b.precision);
1148    let w = wide(cols_t, target);
1149    let base = add(konst(1.0), unary(SfpuKind::Exp, add(op(0), op(1))));
1150    let body = *unary(SfpuKind::Exp, mul(op(2), unary(SfpuKind::Ln, base)));
1151    let slots = vec![
1152        Slot {
1153            buf: ar,
1154            map: at(0),
1155        },
1156        Slot {
1157            buf: dt,
1158            map: at(0),
1159        },
1160        Slot {
1161            buf: al,
1162            map: at(0),
1163        },
1164    ];
1165    let s = b.compute("decay", flat_space(n), slots, ident_map(w), body, out_buf);
1166    let claim = approx("decay", b.precision);
1167    b.finish(
1168        &names.out,
1169        n,
1170        vec![s],
1171        out_buf,
1172        format!("decay_{n}"),
1173        claim,
1174        target,
1175    )
1176}
1177
1178/// [`from_decay`] over a block of rows — the `rows × heads` matrix a projection wrote, every row decayed by the same
1179/// per-head bias and rate. One item per tile of the matrix; the bias and rate tiles of its columns are read down
1180/// every row of it (a map with no row term — a broadcast is a missing term, not an operator).
1181fn dense_decay(
1182    spec: &DecaySpec,
1183    names: &StageNames,
1184    target: &SpatialTarget,
1185) -> Result<SpatialProgram, SpatialReject> {
1186    let (rows, heads) = (spec.rows, spec.n / spec.rows.max(1));
1187    if rows * heads != spec.n || heads == 0 {
1188        return Err(SpatialReject::NotLowered(
1189            "a block decay whose rows do not divide its elements",
1190        ));
1191    }
1192    let (tr, tw) = (target.tile.rows, target.tile.cols);
1193    let (rows_t, cols_t) = (rows.div_ceil(tr), heads.div_ceil(tw));
1194    let dt = elem_precision(target)?;
1195    let mut t = Template::new(target)?;
1196    let (ar_name, ar_fill) = names.operand(0);
1197    let (dt_name, dt_fill) = names.operand(1);
1198    let (al_name, al_fill) = names.operand(2);
1199    let ar = t.tensor(
1200        ar_name,
1201        TensorRole::Input,
1202        rows,
1203        heads,
1204        dt,
1205        super::from_tensor_op::init_of(ar_fill, spec.n),
1206    );
1207    let bias = t.tensor(
1208        dt_name,
1209        TensorRole::Input,
1210        1,
1211        heads,
1212        dt,
1213        super::from_tensor_op::init_of(dt_fill, heads),
1214    );
1215    let rate = t.tensor(
1216        al_name,
1217        TensorRole::Input,
1218        1,
1219        heads,
1220        dt,
1221        super::from_tensor_op::init_of(al_fill, heads),
1222    );
1223    let out_t = t.tensor(
1224        &names.out,
1225        TensorRole::Output,
1226        rows,
1227        heads,
1228        dt,
1229        TensorInit::Zero,
1230    );
1231    let (ar_in, bias_in, rate_in, out) = (
1232        t.buffer("in_ar", 1, dt),
1233        t.buffer("in_bias", 1, dt),
1234        t.buffer("in_rate", 1, dt),
1235        t.buffer("out", 1, dt),
1236    );
1237    let l_ar = t.load_block_at(
1238        "load_ar",
1239        ar,
1240        0,
1241        vec![(LV_ITEM_Y, cols_t), (LV_ITEM_X, 1)],
1242        1,
1243        1,
1244        cols_t,
1245        ar_in,
1246    );
1247    let l_bias = t.load_block_at("load_bias", bias, 0, vec![(LV_ITEM_X, 1)], 1, 1, 1, bias_in);
1248    let l_rate = t.load_block_at("load_rate", rate, 0, vec![(LV_ITEM_X, 1)], 1, 1, 1, rate_in);
1249    let down_rows = ElemAccess {
1250        base: 0,
1251        terms: vec![(LV_GEN_COL, 1)],
1252        indirect: None,
1253    };
1254    let base = add(konst(1.0), unary(SfpuKind::Exp, add(op(0), op(1))));
1255    let s = t.stage(
1256        "decay",
1257        StageKind::Compute(ComputeOp::Generic {
1258            iters: super::from_tensor_op::space(target, tw, None),
1259            ins: vec![ident_map(tw), down_rows.clone(), down_rows],
1260            out: ident_map(tw),
1261            body: *unary(SfpuKind::Exp, mul(op(2), unary(SfpuKind::Ln, base))),
1262        }),
1263        vec![ar_in, bias_in, rate_in],
1264        vec![out],
1265    );
1266    let st = t.store_block(
1267        "store_decay",
1268        out_t,
1269        0,
1270        vec![(LV_ITEM_Y, cols_t), (LV_ITEM_X, 1)],
1271        1,
1272        1,
1273        cols_t,
1274        out,
1275    );
1276    for (from, buf) in [(l_ar, ar_in), (l_bias, bias_in), (l_rate, rate_in)] {
1277        t.edge(from, buf, s);
1278    }
1279    t.edge(s, out, st);
1280    let items: Vec<WorkItem> = (0..rows_t)
1281        .flat_map(|r| (0..cols_t).map(move |c| (r, c)))
1282        .enumerate()
1283        .map(|(id, (r, c))| WorkItem {
1284            id: WorkItemId::new(id),
1285            coord: (r, c),
1286            out: TileRect {
1287                row0_t: r,
1288                col0_t: c,
1289                h_t: 1,
1290                w_t: 1,
1291            },
1292        })
1293        .collect();
1294    let mut prog = t.finish(
1295        format!("decay_{rows}x{heads}"),
1296        items,
1297        AccPolicy::for_result(dt),
1298        approx("decay", dt),
1299        target,
1300    )?;
1301    prog.work.plan = Some(BlockPlan {
1302        grid: (rows_t, cols_t),
1303        ..tile_row_plan(1, 1, dt)
1304    });
1305    Ok(prog)
1306}
1307
1308/// `Kt[j·rows + i] = X[i·cols + j]` — a TRANSPOSE IS A MAP, not an operator.
1309///
1310/// ★ A STAGE'S RESULT IS PLACED BY ITS ITERATION SPACE, so a flat result must iterate the flat output. Its
1311/// element `c = j·rows + i` (new row `j`, new column `i`) is `X[i·cols + j]`, and inside new row `j`'s band
1312/// `[j·rows, (j+1)·rows)` that offset is `c·cols + j·(1 − rows·cols)` — affine. So a matrix transpose is `cols`
1313/// bands of one read each: a map per band, still no operator.
1314///
1315/// ★ A VECTOR TRANSPOSES TO ITSELF. A `1 × n` row and an `n × 1` column are the same flat sequence, so their
1316/// transpose is one plain read — no bands at all.
1317pub fn from_transpose(
1318    spec: &TransposeSpec,
1319    names: &StageNames,
1320    float: bool,
1321    target: &SpatialTarget,
1322) -> Result<SpatialProgram, SpatialReject> {
1323    let (rows, cols) = (spec.rows, spec.cols);
1324    let n = rows * cols;
1325    if n == 0 {
1326        return Err(SpatialReject::NotLowered("a transpose of an empty tensor"));
1327    }
1328    let mut b = FlatBuild::on(Datapath::of(float), target)?;
1329    let xin = b.input(names.operand(0).0, n, names.operand(0).1, target);
1330    let cols_t = tiles_of(n, target);
1331    let out_buf = b.t.buffer("out", cols_t, b.precision);
1332    let body_and_slots = if rows == 1 || cols == 1 {
1333        (
1334            TensorBody::Operand(0),
1335            vec![Slot {
1336                buf: xin,
1337                map: at(0),
1338            }],
1339        )
1340    } else {
1341        let mut slots = Vec::with_capacity(cols);
1342        let mut bands = Vec::with_capacity(cols);
1343        for j in 0..cols {
1344            let base = j as i64 * (1 - n as i64);
1345            slots.push(Slot {
1346                buf: xin,
1347                map: ElemAccess {
1348                    base,
1349                    terms: vec![(LV_GEN_COL, cols as i64)],
1350                    indirect: None,
1351                },
1352            });
1353            bands.push(((j + 1) * rows, TensorBody::Operand(slots.len() - 1)));
1354        }
1355        (banded(bands), slots)
1356    };
1357    let (body, slots) = body_and_slots;
1358    let s = b.compute(
1359        "transpose",
1360        flat_space(n),
1361        slots,
1362        ident_map(wide(cols_t, target)),
1363        body,
1364        out_buf,
1365    );
1366    b.finish(
1367        &names.out,
1368        n,
1369        vec![s],
1370        out_buf,
1371        format!("transpose_{rows}x{cols}"),
1372        NumericsClaim::Exact,
1373        target,
1374    )
1375}
1376
1377/// `P = the first|last K of every N of X` — group `g`'s `K` elements from `from` land at `[g·K, (g+1)·K)`: one band per
1378/// group, each an affine read `X[c + g·(N − K) + from]`.
1379pub fn from_head_slice(
1380    spec: &HeadSliceSpec,
1381    names: &StageNames,
1382    target: &SpatialTarget,
1383) -> Result<SpatialProgram, SpatialReject> {
1384    let (every, from, width) = (spec.every, spec.from, spec.width);
1385    let groups = if every == 0 { 0 } else { spec.n / every };
1386    let n = groups * width;
1387    if n == 0 || from + width > every {
1388        return Err(SpatialReject::NotLowered("a head slice of nothing, or a part past its group"));
1389    }
1390    if whole_tiles(&[every, from, width], target) {
1391        return tile_head_slice(spec, names, groups, target);
1392    }
1393    let mut b = FlatBuild::columns(n, target)?;
1394    let xin = b.input(names.operand(0).0, spec.n, names.operand(0).1, target);
1395    let precision = b.precision;
1396    let out_buf = b.result_page("out", n, precision, target);
1397    let mut slots = Vec::with_capacity(groups);
1398    let mut bands = Vec::with_capacity(groups);
1399    for g in 0..groups {
1400        slots.push(Slot {
1401            buf: xin,
1402            map: at((g * (every - width) + from) as i64),
1403        });
1404        bands.push(((g + 1) * width, TensorBody::Operand(g)));
1405    }
1406    let chain = b.pointwise("head_slice", n, slots, banded(bands), out_buf, target)?;
1407    b.finish(
1408        &names.out,
1409        n,
1410        chain,
1411        out_buf,
1412        format!("head_slice_{groups}x{width}_of_{every}"),
1413        NumericsClaim::Exact,
1414        target,
1415    )
1416}
1417
1418/// `J = A joined with B across H heads` — unit `u` (row `u / H`, head `u mod H`) holds its `ha` elements of `A` then its
1419/// `hb` of `B` at `[u·(ha + hb), (u+1)·(ha + hb))`: two bands per unit, `A[c − u·hb]` then `B[c − (u+1)·ha]`.
1420pub fn from_join(spec: &JoinSpec, names: &StageNames, target: &SpatialTarget) -> Result<SpatialProgram, SpatialReject> {
1421    let (rows, heads) = (spec.rows, spec.heads);
1422    if heads == 0 || spec.wa % heads != 0 || spec.wb % heads != 0 {
1423        return Err(SpatialReject::NotLowered("a join whose widths do not divide among its heads"));
1424    }
1425    let (ha, hb) = (spec.wa / heads, spec.wb / heads);
1426    let (units, w) = (rows * heads, ha + hb);
1427    let n = units * w;
1428    if n == 0 {
1429        return Err(SpatialReject::NotLowered("a join of nothing"));
1430    }
1431    if ha > 0 && hb > 0 && whole_tiles(&[ha, hb], target) {
1432        return tile_join(names, (rows, heads), (ha, hb), target);
1433    }
1434    let mut b = FlatBuild::columns(n, target)?;
1435    let ain = b.input(names.operand(0).0, rows * spec.wa, names.operand(0).1, target);
1436    let bin = b.input(names.operand(1).0, rows * spec.wb, names.operand(1).1, target);
1437    let precision = b.precision;
1438    let out_buf = b.result_page("out", n, precision, target);
1439    let mut slots = Vec::with_capacity(2 * units);
1440    let mut bands = Vec::with_capacity(2 * units);
1441    for u in 0..units {
1442        let (ui, ha_i, hb_i) = (u as i64, ha as i64, hb as i64);
1443        if ha > 0 {
1444            slots.push(Slot { buf: ain, map: at(-ui * hb_i) });
1445            bands.push((u * w + ha, TensorBody::Operand(slots.len() - 1)));
1446        }
1447        if hb > 0 {
1448            slots.push(Slot { buf: bin, map: at(-(ui + 1) * ha_i) });
1449            bands.push(((u + 1) * w, TensorBody::Operand(slots.len() - 1)));
1450        }
1451    }
1452    let chain = b.pointwise("join", n, slots, banded(bands), out_buf, target)?;
1453    b.finish(
1454        &names.out,
1455        n,
1456        chain,
1457        out_buf,
1458        format!("join_{rows}x{heads}_{ha}+{hb}"),
1459        NumericsClaim::Exact,
1460        target,
1461    )
1462}
1463
1464/// `Y = X matmul W across H heads` — the weighted-sum shape: the folded axis is the ROW (`e` over a head's `a` inputs),
1465/// the flat output column `c = u·b + r` for unit `u` (row `u / H`, head `h = u mod H`), and each unit's band reads
1466/// `X[u·a + e]` and `W[(h·a + e)·b + c − u·b]`.
1467pub fn from_head_matmul(
1468    spec: &HeadMatmulSpec,
1469    names: &StageNames,
1470    target: &SpatialTarget,
1471) -> Result<SpatialProgram, SpatialReject> {
1472    let (m, heads, a, bw) = (spec.m, spec.heads, spec.a, spec.b);
1473    let (units, n) = (m * heads, m * heads * bw);
1474    if n == 0 || a == 0 {
1475        return Err(SpatialReject::NotLowered("a per-head matmul of nothing"));
1476    }
1477    if m == 1 && whole_tiles(&[a, bw], target) {
1478        return tile_head_matmul(spec, names, target);
1479    }
1480    let mut f = FlatBuild::new(target)?;
1481    let x = f.input(names.operand(0).0, m * heads * a, names.operand(0).1, target);
1482    let w = f.input(names.operand(1).0, heads * a * bw, names.operand(1).1, target);
1483    let cols_t = tiles_of(n, target);
1484    let out_buf = f.t.buffer("out", cols_t, f.precision);
1485    let tile_elems = target.tile.rows * target.tile.cols;
1486    let ones = f.t.tensor(
1487        "hm_scaler",
1488        TensorRole::Param,
1489        target.tile.rows,
1490        target.tile.cols,
1491        f.precision,
1492        TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
1493    );
1494    let ones_in = f.t.buffer("in_scaler", 1, f.precision);
1495    let ls = f.t.load("load_scaler", ones, vec![], 1, ones_in);
1496    f.loads.push((ls, ones_in));
1497    let iters = vec![
1498        IterSpec {
1499            var: LV_GEN_ROW,
1500            extent: a,
1501            kind: IterKind::Reduce { pool: ReducePool::Sum },
1502        },
1503        IterSpec {
1504            var: LV_GEN_COL,
1505            extent: n,
1506            kind: IterKind::Parallel,
1507        },
1508    ];
1509    let mut slots = vec![Slot { buf: ones_in, map: const_map() }];
1510    let mut bands = Vec::with_capacity(units);
1511    for u in 0..units {
1512        let h = u % heads;
1513        slots.push(Slot {
1514            buf: x,
1515            map: ElemAccess { base: (u * a) as i64, terms: vec![(LV_GEN_ROW, 1)], indirect: None },
1516        });
1517        let x_slot = slots.len() - 1;
1518        slots.push(Slot {
1519            buf: w,
1520            map: ElemAccess {
1521                base: (h * a * bw) as i64 - (u * bw) as i64,
1522                terms: vec![(LV_GEN_ROW, bw as i64), (LV_GEN_COL, 1)],
1523                indirect: None,
1524            },
1525        });
1526        let w_slot = slots.len() - 1;
1527        bands.push(((u + 1) * bw, *mul(op(x_slot), op(w_slot))));
1528    }
1529    let s = f.compute("head_matmul", iters, slots, per_col_map(), banded(bands), out_buf);
1530    let claim = approx("per-head matmul", f.precision);
1531    f.finish(&names.out, n, vec![s], out_buf, format!("head_matmul_{m}x{heads}x{a}x{bw}"), claim, target)
1532}
1533
1534// ── The latent-attention family at whole-tile widths ─────────────────────────────────────────────────
1535//
1536// ★ AT WHOLE-TILE WIDTHS A SLICE AND A JOIN ARE TILE MOVES, AND A PER-HEAD PRODUCT IS A PER-HEAD GEMV. The flat forms
1537// above hold the whole row in one accumulator page and every weight in one site, which is free at a toy's widths and
1538// impossible at DeepSeek-V2's: sixteen heads of 576 are 288 tiles against eight, and one head's latent expansion is
1539// 128 × 512. Every published latent attention's widths are whole tiles (128 and 64 for the query's halves, 512 for the
1540// latent, 128 for the values), so these forms take them: each item moves or computes one block of the result, and only
1541// that block's operands leave memory.
1542
1543/// Whether every width is a whole number of tiles.
1544fn whole_tiles(widths: &[usize], target: &SpatialTarget) -> bool {
1545    widths.iter().all(|w| w % target.tile.cols == 0)
1546}
1547
1548/// The items of a flat row of `groups` runs of `run_t` tiles, `blk` tiles an item: item `(g, j)` owns tiles
1549/// `[g·run_t + j·blk, g·run_t + (j+1)·blk)`.
1550fn grouped_block_items(groups: usize, run_t: usize, blk: usize) -> Vec<WorkItem> {
1551    (0..groups)
1552        .flat_map(|g| (0..run_t / blk).map(move |j| (g, j)))
1553        .enumerate()
1554        .map(|(id, (g, j))| WorkItem {
1555            id: WorkItemId::new(id),
1556            coord: (g, j),
1557            out: TileRect {
1558                row0_t: 0,
1559                col0_t: g * run_t + j * blk,
1560                h_t: 1,
1561                w_t: blk,
1562            },
1563        })
1564        .collect()
1565}
1566
1567/// `P = the first|last K of every N of X` at whole-tile widths: block `j` of group `g`'s part is the block at tile
1568/// `g·N_t + from_t + j·blk` of `X`, an affine tile copy.
1569fn tile_head_slice(
1570    spec: &HeadSliceSpec,
1571    names: &StageNames,
1572    groups: usize,
1573    target: &SpatialTarget,
1574) -> Result<SpatialProgram, SpatialReject> {
1575    let tw = target.tile.cols;
1576    let (every_t, from_t, width_t) = (spec.every / tw, spec.from / tw, spec.width / tw);
1577    let blk = row_block(width_t, target);
1578    let precision = elem_precision(target)?;
1579    let (x_name, x_fill) = names.operand(0);
1580    let mut t = Template::new(target)?;
1581    let x = t.tensor(x_name, TensorRole::Input, 1, spec.n, precision, super::from_tensor_op::init_of(x_fill, spec.n));
1582    let y = t.tensor(&names.out, TensorRole::Output, 1, groups * spec.width, precision, TensorInit::Zero);
1583    let xin = t.buffer("in_x", blk, precision);
1584    let lx = t.load_block_at("load_x", x, from_t, vec![(LV_ITEM_Y, every_t), (LV_ITEM_X, blk)], 1, blk, spec.n / tw, xin);
1585    let (copy, out) = copy_page(&mut t, xin, blk, precision, target);
1586    let st = t.store_block(
1587        &format!("store_{}", names.out),
1588        y,
1589        0,
1590        vec![(LV_ITEM_Y, width_t), (LV_ITEM_X, blk)],
1591        1,
1592        blk,
1593        groups * width_t,
1594        out,
1595    );
1596    t.edge(lx, xin, copy);
1597    t.edge(copy, out, st);
1598    t.finish(
1599        format!("head_slice_{groups}x{}_of_{}_tiled", spec.width, spec.every),
1600        grouped_block_items(groups, width_t, blk),
1601        AccPolicy::for_result(precision),
1602        NumericsClaim::Exact,
1603        target,
1604    )
1605}
1606
1607/// `J = A joined with B across H heads` at whole-tile parts: unit `u`'s run of `(ha + hb)/32` tiles is its `ha/32` tiles
1608/// of `A` then its `hb/32` of `B`, moved `blk` tiles an item (`blk` dividing both parts).
1609///
1610/// ★ EVERY ITEM RUNS THE SAME STAGES, SO THE SOURCE IS DATA. Block `s` of a unit comes from `A` for `s < ha/blk` and
1611/// from `B` after. Each item reads both: `A`'s block `min(s, ha/blk − 1)` and `B`'s block `max(s − ha/blk, 0)`, each
1612/// an [`IndirectTile`] over the unit's run addressed by a compile-time table one tile per block, and keeps `A`'s where
1613/// a third table holds the block's width and `B`'s where it holds zero. Every address is in its tensor, whatever the
1614/// block.
1615fn tile_join(
1616    names: &StageNames,
1617    (rows, heads): (usize, usize),
1618    (ha, hb): (usize, usize),
1619    target: &SpatialTarget,
1620) -> Result<SpatialProgram, SpatialReject> {
1621    let tw = target.tile.cols;
1622    let (ha_t, hb_t) = (ha / tw, hb / tw);
1623    let units = rows * heads;
1624    let blk = (1..=ha_t.min(hb_t).min(row_block_tiles(target)))
1625        .rev()
1626        .find(|d| ha_t % d == 0 && hb_t % d == 0)
1627        .unwrap_or(1);
1628    let (na, nb) = (ha_t / blk, hb_t / blk);
1629    let blocks = na + nb;
1630    let precision = elem_precision(target)?;
1631    let ix = index_type(target, (blk * tw).max(blocks) as u64 + 1)?;
1632    let (a_name, a_fill) = names.operand(0);
1633    let (b_name, b_fill) = names.operand(1);
1634    let mut t = Template::new(target)?;
1635    let a = t.tensor(a_name, TensorRole::Input, 1, units * ha, precision, super::from_tensor_op::init_of(a_fill, units * ha));
1636    let b = t.tensor(b_name, TensorRole::Input, 1, units * hb, precision, super::from_tensor_op::init_of(b_fill, units * hb));
1637    let y = t.tensor(&names.out, TensorRole::Output, 1, units * (ha + hb), precision, TensorInit::Zero);
1638    // One value a block, at element 0 of its tile.
1639    let per_block = |t: &mut Template, what: &str, value: &dyn Fn(usize) -> usize| {
1640        let tiles = (0..blocks).flat_map(|s| std::iter::once(value(s) as i64).chain(std::iter::repeat_n(0, tw - 1))).collect();
1641        t.tensor(&format!("{}__join_{what}", names.out), TensorRole::Param, 1, blocks * tw, ix, TensorInit::Values(tiles))
1642    };
1643    let from_a = per_block(&mut t, "from_a", &|s| s.min(na - 1));
1644    let from_b = per_block(&mut t, "from_b", &|s| s.saturating_sub(na).min(nb - 1));
1645    let keep_a = per_block(&mut t, "keep_a", &|s| if s < na { blk * tw } else { 0 });
1646    let gather = |t: &mut Template, part: &str, tensor: TensorId, table: TensorId, run_t: usize, bound: usize| {
1647        let pick = t.buffer_deep(&format!("in_{part}_block"), 1, 1, ix);
1648        let l_pick = t.load_block_at(&format!("load_{part}_block"), table, 0, vec![(LV_ITEM_X, 1)], 1, 1, 1, pick);
1649        let page = t.buffer(&format!("in_{part}"), blk, precision);
1650        let l_page = t.stage(
1651            &format!("load_{part}"),
1652            StageKind::Load {
1653                tensor,
1654                pattern: AccessPattern {
1655                    origin: TileAddr {
1656                        base: 0,
1657                        steps: vec![(LV_ITEM_Y, run_t)],
1658                        indirect: Some(IndirectTile { index: pick, tile_stride: blk, elems_per_tile: 1, bound: (0, bound) }),
1659                    },
1660                    block_h: 1,
1661                    block_w: blk,
1662                    stride_w: 1,
1663                    stride_h: 1,
1664                },
1665                num_blocks: 1,
1666                multicast: None,
1667            },
1668            vec![pick],
1669            vec![page],
1670        );
1671        t.edge(l_pick, pick, l_page);
1672        (l_page, page)
1673    };
1674    let (l_a, a_in) = gather(&mut t, "a", a, from_a, ha_t, na - 1);
1675    let (l_b, b_in) = gather(&mut t, "b", b, from_b, hb_t, nb - 1);
1676    let keep_in = t.buffer_deep("in_keep_a", 1, 1, ix);
1677    let l_keep = t.load_block_at("load_keep_a", keep_a, 0, vec![(LV_ITEM_X, 1)], 1, 1, 1, keep_in);
1678    let w = wide(blk, target);
1679    let out = t.buffer("out", blk, precision);
1680    let s = t.stage(
1681        "join",
1682        StageKind::Compute(ComputeOp::Generic {
1683            iters: flat_space(w),
1684            ins: vec![ident_map(w), ident_map(w), fixed(0)],
1685            out: ident_map(w),
1686            body: TensorBody::Select {
1687                cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(2)),
1688                then: op(0),
1689                otherwise: op(1),
1690            },
1691        }),
1692        vec![a_in, b_in, keep_in],
1693        vec![out],
1694    );
1695    let st = t.store_block(
1696        &format!("store_{}", names.out),
1697        y,
1698        0,
1699        vec![(LV_ITEM_Y, ha_t + hb_t), (LV_ITEM_X, blk)],
1700        1,
1701        blk,
1702        units * (ha_t + hb_t),
1703        out,
1704    );
1705    for (from, buf, to) in [(l_a, a_in, s), (l_b, b_in, s), (l_keep, keep_in, s), (s, out, st)] {
1706        t.edge(from, buf, to);
1707    }
1708    t.finish(
1709        format!("join_{rows}x{heads}_{ha}+{hb}_tiled"),
1710        grouped_block_items(units, ha_t + hb_t, blk),
1711        AccPolicy::for_result(precision),
1712        NumericsClaim::Exact,
1713        target,
1714    )
1715}
1716
1717/// `Y = X matmul W across H heads` for one row at whole-tile widths: item `(h, o)` is output tile `o` of head `h`,
1718/// `Σ_e X[h·a + e]·W[h·a + e, o·32 + c]` — head `h`'s `a/32` input tiles against the `a/32 × 1` tile column `o` of
1719/// its rows of `W`, which `W`'s declaration as the `(H·a) × b` matrix it is lays out as tiles. Needs a square tile: a
1720/// weight tile's rows are the input's elements. Only the head's rows of
1721/// one tile column move an item, so `W` is never on chip whole.
1722fn tile_head_matmul(
1723    spec: &HeadMatmulSpec,
1724    names: &StageNames,
1725    target: &SpatialTarget,
1726) -> Result<SpatialProgram, SpatialReject> {
1727    let (heads, a, bw) = (spec.heads, spec.a, spec.b);
1728    let (tr, tw) = (target.tile.rows, target.tile.cols);
1729    if tr != tw {
1730        return Err(SpatialReject::NotLowered("a per-head matmul in tiles on a tile that is not square — a weight tile's rows are the input's elements"));
1731    }
1732    let (a_t, b_t) = (a / tw, bw / tw);
1733    let precision = elem_precision(target)?;
1734    let (x_name, x_fill) = names.operand(0);
1735    let (w_name, w_fill) = names.operand(1);
1736    let mut t = Template::new(target)?;
1737    let scaler = t.tensor("hm_scaler", TensorRole::Param, tr, tw, precision, TensorInit::FloatValues(vec![1.0f64.to_bits(); tr * tw]));
1738    let x = t.tensor(x_name, TensorRole::Input, 1, heads * a, precision, super::from_tensor_op::init_of(x_fill, heads * a));
1739    let w = t.tensor(w_name, TensorRole::Param, heads * a, bw, precision, super::from_tensor_op::init_of(w_fill, heads * a * bw));
1740    let y = t.tensor(&names.out, TensorRole::Output, 1, heads * bw, precision, TensorInit::Zero);
1741    let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
1742    let x_in = t.buffer("in_x", a_t, precision);
1743    let w_in = t.buffer("in_w", a_t, precision);
1744    let out = t.buffer("out", 1, precision);
1745    let l_sc = t.load_block_at("load_scaler", scaler, 0, Vec::new(), 1, 1, 1, sc_in);
1746    let l_x = t.load_block_at("load_x", x, 0, vec![(LV_ITEM_Y, a_t)], 1, a_t, heads * a_t, x_in);
1747    let l_w = t.load_block_at("load_w", w, 0, vec![(LV_ITEM_Y, a_t * b_t), (LV_ITEM_X, 1)], a_t, 1, b_t, w_in);
1748    // ★ A PAGE IS ONE PLANE OF ITS TILES SIDE BY SIDE, so the weight's `a/32` stacked tiles lie as tile `k` at columns
1749    // `[32k, 32k + 32)`: input `e = 32k + r`, output `c` is at row `r`, column `32k + c`, and the input row's `e` at
1750    // column `32k + r` of row 0. The fold runs over both: the tile's rows `r` and the tiles `k`.
1751    let s = t.stage(
1752        "head_dot",
1753        StageKind::Compute(ComputeOp::Generic {
1754            iters: vec![
1755                IterSpec { var: LV_GEN_ROW, extent: tr, kind: IterKind::Reduce { pool: ReducePool::Sum } },
1756                IterSpec { var: LV_GEN_COL, extent: tw, kind: IterKind::Parallel },
1757                IterSpec { var: LV_GEN_K, extent: a_t, kind: IterKind::Reduce { pool: ReducePool::Sum } },
1758            ],
1759            ins: vec![
1760                ElemAccess { base: 0, terms: vec![(LV_GEN_ROW, (a_t * tw) as i64), (LV_GEN_K, tw as i64), (LV_GEN_COL, 1)], indirect: None },
1761                ElemAccess { base: 0, terms: vec![(LV_GEN_K, tw as i64), (LV_GEN_ROW, 1)], indirect: None },
1762                const_map(),
1763            ],
1764            out: per_col_map(),
1765            body: *mul(op(0), op(1)),
1766        }),
1767        vec![w_in, x_in, sc_in],
1768        vec![out],
1769    );
1770    let st = t.store_block(&format!("store_{}", names.out), y, 0, vec![(LV_ITEM_Y, b_t), (LV_ITEM_X, 1)], 1, 1, heads * b_t, out);
1771    for (from, buf, to) in [(l_w, w_in, s), (l_x, x_in, s), (l_sc, sc_in, s), (s, out, st)] {
1772        t.edge(from, buf, to);
1773    }
1774    let prog = t.finish(
1775        format!("head_matmul_1x{heads}x{a}x{bw}_tiled"),
1776        grouped_block_items(heads, b_t, 1),
1777        AccPolicy::F32,
1778        approx("per-head matmul", precision),
1779        target,
1780    )?;
1781    if site_footprint(&prog) > target.max_local_mem_bytes() {
1782        return Err(SpatialReject::NotLowered("a per-head matmul whose head's inputs and weight column do not fit a site"));
1783    }
1784    Ok(prog)
1785}
1786
1787// ── The per-head family: a reduction, then a banded rescale ───────────────────────────────────────
1788
1789/// The chain every per-row normalization shares: square → fold the row → transform the row scalar.
1790/// Returns `(chain, scale_buffer, scale_width)`; the caller adds the banded rescale that consumes it.
1791fn row_scale_chain(
1792    b: &mut FlatBuild,
1793    xin: BufferId,
1794    rows: usize,
1795    cols: usize,
1796    xform: SfpuKind,
1797    target: &SpatialTarget,
1798) -> (Vec<StageId>, BufferId, usize) {
1799    let xw = wide(b.t.page_tiles(xin), target);
1800    let tile_elems = target.tile.rows * target.tile.cols;
1801    // The reduce engine's scaler operand, at 1.0: Tensix's `reduce_tile` multiplies every term by a
1802    // constant tile and REQUIRES one. Every divide in these chains happens ONCE, in the transform.
1803    let ones = b.t.tensor(
1804        "row_scaler",
1805        TensorRole::Param,
1806        target.tile.rows,
1807        target.tile.cols,
1808        b.precision,
1809        TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
1810    );
1811    let ones_in = b.t.buffer("in_scaler", 1, b.precision);
1812    let l = b.t.load("load_scaler", ones, vec![], 1, ones_in);
1813    b.loads.push((l, ones_in));
1814
1815    let sq_tiles = tiles_of(cols, target);
1816    let sqw = wide(sq_tiles, target);
1817    let xsq = b.scratch("xsq", sq_tiles);
1818    // The square reads the FLAT input through a 2-D space (`row` is the head) and writes the
1819    // `rows × cols` plane the reduction folds.
1820    let s1 = b.compute(
1821        "square",
1822        head_space(rows, cols, None),
1823        vec![
1824            Slot {
1825                buf: xin,
1826                map: map2(0, cols as i64, 1),
1827            },
1828            Slot {
1829                buf: xin,
1830                map: map2(0, cols as i64, 1),
1831            },
1832        ],
1833        ident_map(sqw),
1834        *mul(op(0), op(1)),
1835        xsq,
1836    );
1837    let ssum = b.scratch("row_sum", 1);
1838    let s2 = b.compute(
1839        "reduce_rows",
1840        head_space(rows, cols, Some(ReducePool::Sum)),
1841        vec![
1842            Slot {
1843                buf: xsq,
1844                map: map2(0, sqw as i64, 1),
1845            },
1846            Slot {
1847                buf: ones_in,
1848                map: const_map(),
1849            },
1850        ],
1851        per_row_map(wide(1, target)),
1852        TensorBody::Operand(0),
1853        ssum,
1854    );
1855    let one_w = wide(1, target);
1856    let scale = b.scratch("row_scale", 1);
1857    let s3 = b.compute(
1858        "row_xform",
1859        head_space(rows, target.tile.cols, None),
1860        vec![Slot {
1861            buf: ssum,
1862            map: per_row_map(one_w),
1863        }],
1864        per_row_map(one_w),
1865        *unary(xform, op(0)),
1866        scale,
1867    );
1868    let _ = xw;
1869    (vec![s1, s2, s3], scale, one_w)
1870}
1871
1872/// `y[r, i] = x[r, i] / sqrt(Σ_i x[r, i]² + 1e-6)` over `rows` heads of `cols` — per-head L2
1873/// normalization on a FLAT `1 × rows·cols` tensor.
1874pub fn from_l2norm(
1875    spec: &L2NormSpec,
1876    names: &StageNames,
1877    target: &SpatialTarget,
1878) -> Result<SpatialProgram, SpatialReject> {
1879    let (rows, cols) = (spec.rows, spec.cols);
1880    let n = rows * cols;
1881    if row_needs_blocking(n, target) && !row_needs_blocking(cols, target) {
1882        return blocked_l2norm(spec, names, target);
1883    }
1884    let mut b = FlatBuild::new(target)?;
1885    let xin = b.input(names.operand(0).0, n, names.operand(0).1, target);
1886    let cols_t = tiles_of(n, target);
1887    let out_buf = b.t.buffer("out", cols_t, b.precision);
1888    let ow = wide(cols_t, target);
1889    // `over: 1` — l2norm divides by the root SUM of squares, not the root mean.
1890    let (mut chain, scale, one_w) = row_scale_chain(
1891        &mut b,
1892        xin,
1893        rows,
1894        cols,
1895        SfpuKind::Rsqrt {
1896            eps_nano: eps_nano(spec.eps)?,
1897            over: 1,
1898        },
1899        target,
1900    );
1901    // ★ THE BANDED RESCALE. Head `h` occupies the flat columns `[h·cols, (h+1)·cols)`, and its scale
1902    // sits broadcast across tile row `h` — so inside the band the scale's map is a CONSTANT offset.
1903    let mut slots = vec![Slot {
1904        buf: xin,
1905        map: at(0),
1906    }];
1907    let mut bands = Vec::new();
1908    for h in 0..rows {
1909        slots.push(Slot {
1910            buf: scale,
1911            map: fixed((h * one_w) as i64),
1912        });
1913        bands.push(((h + 1) * cols, *mul(op(0), op(slots.len() - 1))));
1914    }
1915    let s = b.compute(
1916        "rescale",
1917        flat_space(n),
1918        slots,
1919        ident_map(ow),
1920        banded(bands),
1921        out_buf,
1922    );
1923    chain.push(s);
1924    let claim = approx("per-head l2 normalization", b.precision);
1925    b.finish(
1926        &names.out,
1927        n,
1928        chain,
1929        out_buf,
1930        format!("l2norm_{rows}x{cols}"),
1931        claim,
1932        target,
1933    )
1934}
1935
1936/// The COMPOSITION-LAYOUT twin of [`super::from_tensor_op::from_fused_row`]'s RMSNorm: the same chain
1937/// (square → fold → rsqrt-once → rescale → per-column weight), over a FLAT `1 × rows·cols` tensor
1938/// instead of a `rows × cols` one.
1939///
1940/// ★ WHY A TWIN AND NOT A PARAMETER. `from_fused_row` gives the tensor the shape the SOURCE declared —
1941/// right for a program that is one op, and exactly what a hundred green locks pin. Inside a composition
1942/// the same tensor is a `1 × n` activation whichever way the sugar views it, and the two views put
1943/// element `cols` in different places. Both lower the identical arithmetic in the identical order, so
1944/// they agree to the last bit; `spatial_qwen_ops_lock` checks that against the VM rather than against
1945/// each other.
1946pub fn from_row_norm_flat(
1947    spec: &FusedRowSpec,
1948    names: &StageNames,
1949    target: &SpatialTarget,
1950) -> Result<SpatialProgram, SpatialReject> {
1951    use crate::tensor_recognize::RowReduceKind;
1952    match spec.reduce {
1953        RowReduceKind::RmsNorm => {}
1954        RowReduceKind::Softmax if spec.weight.is_none() && spec.residual.is_none() && spec.bias.is_none() => {
1955            let (x, fill) = names.operand(0);
1956            return flat_softmax_rows((x, fill), &names.out, spec.rows, spec.cols, format!("softmax_flat_{}x{}", spec.rows, spec.cols), target);
1957        }
1958        _ => return Err(SpatialReject::NotLowered("a composed layernorm row op (or a softmax with a weight, residual or bias) — no flat twin yet")),
1959    }
1960    let (rows, cols) = (spec.rows, spec.cols);
1961    let n = rows * cols;
1962    // ★ A PER-HEAD NORM REDUCES PER HEAD, so the head is the item — the same fact `blocked_l2norm` turns
1963    // on. The flat form keeps all `rows·cols` resident: at Qwen3.8's `24 heads x 256` that is 192 tiles
1964    // against eight. One head is `cols = 256`, eight tiles, and the fold is over exactly the elements the
1965    // norm is defined on.
1966    if row_needs_blocking(n, target)
1967        && !row_needs_blocking(cols, target)
1968        && spec.residual.is_none()
1969        && spec.bias.is_none()
1970    {
1971        return blocked_head_rms(spec, names, target);
1972    }
1973    let mut b = FlatBuild::new(target)?;
1974    let xin = b.input(names.operand(0).0, n, names.operand(0).1, target);
1975    // The spec's operand order is `[x, weight?, residual?, bias?]`.
1976    let mut rest = 1usize;
1977    let w_in = spec.weight.map(|_| {
1978        let i = rest;
1979        rest += 1;
1980        b.input(names.operand(i).0, cols, names.operand(i).1, target)
1981    });
1982    let res_in = spec.residual.map(|_| {
1983        let i = rest;
1984        rest += 1;
1985        b.input(names.operand(i).0, n, names.operand(i).1, target)
1986    });
1987    if spec.bias.is_some() {
1988        return Err(SpatialReject::NotLowered(
1989            "a composed rmsnorm with a bias term",
1990        ));
1991    }
1992    let cols_t = tiles_of(n, target);
1993    let out_buf = b.t.buffer("out", cols_t, b.precision);
1994    let ow = wide(cols_t, target);
1995    // A fused residual add is part of the value that is normalized, so it happens first and its result
1996    // is what both the reduction and the rescale read.
1997    let mut chain: Vec<StageId> = Vec::new();
1998    let src = match res_in {
1999        Some(r) => {
2000            let xr = b.scratch("xr", cols_t);
2001            let s = b.compute(
2002                "add_residual",
2003                flat_space(n),
2004                vec![
2005                    Slot {
2006                        buf: xin,
2007                        map: at(0),
2008                    },
2009                    Slot { buf: r, map: at(0) },
2010                ],
2011                ident_map(ow),
2012                *add(op(0), op(1)),
2013                xr,
2014            );
2015            chain.push(s);
2016            xr
2017        }
2018        None => xin,
2019    };
2020    let (mut norm_chain, scale, one_w) = row_scale_chain(
2021        &mut b,
2022        src,
2023        rows,
2024        cols,
2025        SfpuKind::Rsqrt {
2026            eps_nano: eps_nano(spec.eps)?,
2027            over: cols,
2028        },
2029        target,
2030    );
2031    chain.append(&mut norm_chain);
2032    let mut slots = vec![Slot {
2033        buf: src,
2034        map: at(0),
2035    }];
2036    let mut bands = Vec::new();
2037    for h in 0..rows {
2038        let head_base = (h * cols) as i64;
2039        slots.push(Slot {
2040            buf: scale,
2041            map: fixed((h * one_w) as i64),
2042        });
2043        let scale_slot = slots.len() - 1;
2044        let value = match w_in {
2045            Some(wb) => {
2046                slots.push(Slot {
2047                    buf: wb,
2048                    map: at(-head_base),
2049                });
2050                mul(mul(op(0), op(scale_slot)), op(slots.len() - 1))
2051            }
2052            None => mul(op(0), op(scale_slot)),
2053        };
2054        bands.push(((h + 1) * cols, *value));
2055    }
2056    let s = b.compute(
2057        "rescale",
2058        flat_space(n),
2059        slots,
2060        ident_map(ow),
2061        banded(bands),
2062        out_buf,
2063    );
2064    chain.push(s);
2065    let claim = approx("row reduction", b.precision);
2066    b.finish(
2067        &names.out,
2068        n,
2069        chain,
2070        out_buf,
2071        format!("rmsnorm_flat_{rows}x{cols}"),
2072        claim,
2073        target,
2074    )
2075}
2076
2077/// `y[r, j] = exp(x[r, j] − max_r) / Σ_j exp(x[r, j] − max_r)` over `rows` rows of `cols`, on a FLAT
2078/// `1 × rows·cols` activation — the row softmax of the composition layout. Attention's probabilities are this
2079/// section (heads as rows, cache positions as columns), and so is a standalone `softmax of … over rows`: one
2080/// chain — row max, centred exp, row sum, one reciprocal per row, banded rescale — in one place.
2081fn flat_softmax_rows(
2082    input: (&str, &Fill),
2083    out_name: &str,
2084    rows: usize,
2085    cols: usize,
2086    label: String,
2087    target: &SpatialTarget,
2088) -> Result<SpatialProgram, SpatialReject> {
2089    flat_softmax_rows_at(
2090        input,
2091        out_name,
2092        rows,
2093        cols,
2094        label,
2095        elem_precision(target)?,
2096        None,
2097        target,
2098    )
2099}
2100
2101/// [`flat_softmax_rows`] with its input and result at `precision` — the accumulator's width for a softmax whose
2102/// result DECIDES something (a router's expert choice), where a bf16 rounding of a logit moves its weight by
2103/// `e^(δ)` and can flip a top-K pick.
2104///
2105/// `sinks` names a flat tensor of one logit per row that joins the row's softmax and takes no value (gpt-oss's attention
2106/// sinks): the row's maximum starts from it and its term joins the denominator, as ggml's `soft_max_ext` takes them.
2107#[allow(clippy::too_many_arguments)]
2108fn flat_softmax_rows_at(
2109    input: (&str, &Fill),
2110    out_name: &str,
2111    rows: usize,
2112    cols: usize,
2113    label: String,
2114    precision: Precision,
2115    sinks: Option<(&str, &Fill)>,
2116    target: &SpatialTarget,
2117) -> Result<SpatialProgram, SpatialReject> {
2118    let mut f = FlatBuild::new(target)?;
2119    f.precision = precision;
2120    let (h, l) = (rows, cols);
2121    let n = h * l;
2122    let sc = f.input(input.0, n, input.1, target);
2123    let cols_t = tiles_of(n, target);
2124    let out_buf = f.t.buffer("probs", cols_t, f.precision);
2125    let ow = wide(cols_t, target);
2126    let tile_elems = target.tile.rows * target.tile.cols;
2127    // A composition unifies tensors by name, so a shared constant's name says what it holds: a tile of ones at the
2128    // target's element precision is `sm_scaler`, and at another precision (a router's fp32 softmax) it says which.
2129    let scaler = if precision == elem_precision(target)? {
2130        "sm_scaler".to_string()
2131    } else {
2132        format!("sm_scaler_{}", format!("{precision:?}").to_lowercase())
2133    };
2134    let ones = f.t.tensor(
2135        &scaler,
2136        TensorRole::Param,
2137        target.tile.rows,
2138        target.tile.cols,
2139        f.precision,
2140        TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
2141    );
2142    let ones_in = f.t.buffer("in_scaler", 1, f.precision);
2143    let ls = f.t.load("load_scaler", ones, vec![], 1, ones_in);
2144    f.loads.push((ls, ones_in));
2145    let one_w = wide(1, target);
2146    // max over the head's row, then Σ exp(x − max), then exp(x − max)/Σ.
2147    let rmax = f.scratch("row_max", 1);
2148    let s1 = f.compute(
2149        "row_max",
2150        head_space(h, l, Some(ReducePool::Max)),
2151        vec![
2152            Slot {
2153                buf: sc,
2154                map: map2(0, l as i64, 1),
2155            },
2156            Slot {
2157                buf: ones_in,
2158                map: const_map(),
2159            },
2160        ],
2161        per_row_map(one_w),
2162        TensorBody::Operand(0),
2163        rmax,
2164    );
2165    // A row's sink: element `r` of a flat `1 × rows` tensor.
2166    let sink_map = ElemAccess {
2167        base: 0,
2168        terms: vec![(LV_GEN_ROW, 1)],
2169        indirect: None,
2170    };
2171    let sink = sinks.map(|(name, fill)| f.input(name, h, fill, target));
2172    let mut sink_stages = Vec::new();
2173    let rmax = match sink {
2174        Some(sk) => {
2175            let with_sink = f.scratch("row_max_sink", 1);
2176            sink_stages.push(f.compute(
2177                "row_max_sink",
2178                head_space(h, target.tile.cols, None),
2179                vec![
2180                    Slot {
2181                        buf: rmax,
2182                        map: per_row_map(one_w),
2183                    },
2184                    Slot {
2185                        buf: sk,
2186                        map: sink_map.clone(),
2187                    },
2188                ],
2189                per_row_map(one_w),
2190                TensorBody::Max(op(0), op(1)),
2191                with_sink,
2192            ));
2193            with_sink
2194        }
2195        None => rmax,
2196    };
2197    let ex = f.scratch("expx", tiles_of(l, target));
2198    let exw = wide(tiles_of(l, target), target);
2199    let s2 = f.compute(
2200        "exp_centred",
2201        head_space(h, l, None),
2202        vec![
2203            Slot {
2204                buf: sc,
2205                map: map2(0, l as i64, 1),
2206            },
2207            Slot {
2208                buf: rmax,
2209                map: per_row_map(one_w),
2210            },
2211        ],
2212        ident_map(exw),
2213        *unary(SfpuKind::Exp, sub(op(0), op(1))),
2214        ex,
2215    );
2216    let rsum = f.scratch("row_sum", 1);
2217    let s3 = f.compute(
2218        "row_sum",
2219        head_space(h, l, Some(ReducePool::Sum)),
2220        vec![
2221            Slot {
2222                buf: ex,
2223                map: map2(0, exw as i64, 1),
2224            },
2225            Slot {
2226                buf: ones_in,
2227                map: const_map(),
2228            },
2229        ],
2230        per_row_map(one_w),
2231        TensorBody::Operand(0),
2232        rsum,
2233    );
2234    let rsum = match sink {
2235        Some(sk) => {
2236            let with_sink = f.scratch("row_sum_sink", 1);
2237            sink_stages.push(f.compute(
2238                "row_sum_sink",
2239                head_space(h, target.tile.cols, None),
2240                vec![
2241                    Slot {
2242                        buf: rsum,
2243                        map: per_row_map(one_w),
2244                    },
2245                    Slot {
2246                        buf: sk,
2247                        map: sink_map,
2248                    },
2249                    Slot {
2250                        buf: rmax,
2251                        map: per_row_map(one_w),
2252                    },
2253                ],
2254                per_row_map(one_w),
2255                *add(op(0), unary(SfpuKind::Exp, sub(op(1), op(2)))),
2256                with_sink,
2257            ));
2258            with_sink
2259        }
2260        None => rsum,
2261    };
2262    let inv = f.scratch("row_inv", 1);
2263    let s4 = f.compute(
2264        "recip",
2265        head_space(h, target.tile.cols, None),
2266        vec![Slot {
2267            buf: rsum,
2268            map: per_row_map(one_w),
2269        }],
2270        per_row_map(one_w),
2271        *unary(SfpuKind::Recip, op(0)),
2272        inv,
2273    );
2274    // `ex` is the `h × l` plane: head `head`'s row `j` sits at `head·exw + j`, so the flat output
2275    // column `c = head·l + j` reads it at `c + head·(exw − l)` — one map per head, all affine.
2276    let mut slots: Vec<Slot> = Vec::new();
2277    let mut bands = Vec::new();
2278    for head in 0..h {
2279        slots.push(Slot {
2280            buf: ex,
2281            map: at((head * (exw - l)) as i64),
2282        });
2283        let e_slot = slots.len() - 1;
2284        slots.push(Slot {
2285            buf: inv,
2286            map: fixed((head * one_w) as i64),
2287        });
2288        let i_slot = slots.len() - 1;
2289        bands.push(((head + 1) * l, *mul(op(e_slot), op(i_slot))));
2290    }
2291    let s5 = f.compute(
2292        "scale",
2293        flat_space(n),
2294        slots,
2295        ident_map(ow),
2296        banded(bands),
2297        out_buf,
2298    );
2299    let claim = approx("row softmax", f.precision);
2300    let chain = match sink_stages.as_slice() {
2301        [max, sum] => vec![s1, *max, s2, s3, *sum, s4, s5],
2302        _ => vec![s1, s2, s3, s4, s5],
2303    };
2304    f.finish(
2305        out_name,
2306        n,
2307        chain,
2308        out_buf,
2309        label,
2310        claim,
2311        target,
2312    )
2313}
2314
2315/// RoPE-NEOX over `h` heads of `hd` dims, rotating each head's leading `2·half` dims as `(i, i+half)`
2316/// pairs and passing the remaining `hd − 2·half` through — Qwen3.8's partial rotary.
2317///
2318/// ★ THREE REGIONS PER HEAD, ALL BANDS OF ONE FLAT AXIS. Head `h`'s rotated first half is the columns
2319/// `[h·hd, h·hd + half)`, its second half `[h·hd + half, h·hd + 2·half)` and its pass-through tail
2320/// `[h·hd + 2·half, (h+1)·hd)`. Each bound is a CONSTANT, so the whole op is one `Generic` with
2321/// `3·h − 1` predicates and no per-head stages at all. `q` is read through three shifted maps (`c`,
2322/// `c + half`, `c − half`) shared by every head; the cosine/sine tables are read through a per-head
2323/// base, because their index is the position WITHIN the head.
2324///
2325/// ★ A CACHE IS READ WHERE THE POSITION POINTS, IN THE ROPE'S OWN SECTION. With `from position P` the tables are
2326/// `cache` rows of angles and row `r` of the block turns by row `P + r`. Rows of WHOLE TILES are fetched alone: the
2327/// block's `rows·half` angles are a tile run from tile `(P − 1)·half_t` — a literal in the load's base, a named `P` an
2328/// [`IndirectTile`] over the guarded `P − 1` — so a real model's cache of thousands of positions never comes on chip.
2329/// Rows narrower than a tile have no tile of their own: the cache is loaded whole and every table read adds the flat
2330/// offset `(P − 1)·half` — a literal as one more constant, a named one as an [`IndirectTerm`] checked against the
2331/// positions the block admits. Either way a position past the cache faults where every tier's check fails instead of
2332/// reading angles that are not there.
2333pub fn from_rope(
2334    spec: &RopeSpec,
2335    names: &StageNames,
2336    target: &SpatialTarget,
2337) -> Result<SpatialProgram, SpatialReject> {
2338    let (heads, hd, half, rows) = (spec.h, spec.hd, spec.half, spec.rows);
2339    if heads == 0 || hd == 0 || half == 0 || 2 * half > hd {
2340        return Err(SpatialReject::NotLowered(
2341            "a rope whose rotated width does not fit its head dimension",
2342        ));
2343    }
2344    if rows == 0 || heads % rows != 0 {
2345        return Err(SpatialReject::NotLowered(
2346            "a rope whose table rows do not divide its heads",
2347        ));
2348    }
2349    let table_rows = spec.from.map_or(rows, |(_, cache)| cache);
2350    if let Some((RowIndex::Literal(p), cache)) = spec.from {
2351        if p < 1 || p as usize + rows - 1 > cache {
2352            return Err(SpatialReject::NotLowered("a literal rope position whose block leaves the cache — the program fails at run time on every tier, so there is no section to lower"));
2353        }
2354    }
2355    // Head `h` turns by table row `h / per_row`: its reads of the tables start `row·half` in.
2356    let per_row = heads / rows;
2357    let row_shift = |h: usize| ((h / per_row) * half) as i64;
2358    let n = heads * hd;
2359    // ★ A ROTARY EMBEDDING NEVER CROSSES A HEAD, so the head is the natural work item. The flat form
2360    // keeps all `heads·hd` elements resident and tells heads apart with three bands EACH; at Qwen3.8's
2361    // attention width that is 192 tiles against an eight-tile budget and 72 bands. One head is `hd = 256`
2362    // — eight tiles, exactly the budget — and inside it the band structure collapses to the three regions
2363    // the rotation actually has, because the page IS the head.
2364    // A head is the item only when it is whole tiles; a head of 80 shares tiles with its neighbours, so the rotation is
2365    // the one flat stage, blocked by the result's columns past an accumulator page.
2366    if row_needs_blocking(n, target) && !row_needs_blocking(hd, target) && hd % target.tile.cols == 0 {
2367        return blocked_rope(spec, names, target);
2368    }
2369    let mut b = FlatBuild::columns(n, target)?;
2370    let q = b.input(names.operand(0).0, n, names.operand(0).1, target);
2371    let tw = target.tile.cols;
2372    // The tables, and the position's offset into them still to be added to every table read: none (a table row per
2373    // row, one row for all, or a cache whose block is fetched alone), a literal, or data.
2374    let (rc, rs, mut chain, shift, position) = match spec.from {
2375        Some((from, cache)) if half % tw == 0 => {
2376            let (rc, rs) = cache_block_loads(&mut b, names, from, cache, rows, half / tw, target)?;
2377            (rc, rs, Vec::new(), 0i64, None)
2378        }
2379        from => {
2380            let n_table = table_rows * half;
2381            if from.is_some() && tiles_of(n_table, target) > row_block_tiles(target) {
2382                return Err(SpatialReject::NotLowered("a rope reading a cache whose rows are narrower than a tile loads the cache whole, and this cache is wider than a site's page"));
2383            }
2384            let rc = b.input(names.operand(1).0, n_table, names.operand(1).1, target);
2385            let rs = b.input(names.operand(2).0, n_table, names.operand(2).1, target);
2386            match from {
2387                None => (rc, rs, Vec::new(), 0i64, None),
2388                Some((RowIndex::Literal(p), _)) => {
2389                    (rc, rs, Vec::new(), (p - 1) * half as i64, None)
2390                }
2391                Some((RowIndex::Named(_), cache)) => {
2392                    let (p_name, p_fill) = names.operand(3);
2393                    let ix = index_type(
2394                        target,
2395                        index_space((cache * half) as u64, half as u64, p_fill),
2396                    )?;
2397                    let index = b.input_as(p_name, 1, p_fill, ix, target);
2398                    let (stages, lo) = row_offset_stages(&mut b, index, half, ix, target);
2399                    (rc, rs, stages, 0, Some((lo, (cache - rows + 1) * half)))
2400                }
2401            }
2402        }
2403    };
2404    let precision = b.precision;
2405    let out_buf = b.result_page("out", n, precision, target);
2406
2407    let mut slots = vec![
2408        Slot { buf: q, map: at(0) }, // q[c]
2409        Slot {
2410            buf: q,
2411            map: at(half as i64),
2412        }, // q[c + half]
2413        Slot {
2414            buf: q,
2415            map: at(-(half as i64)),
2416        }, // q[c - half]
2417    ];
2418    let (q_at, q_up, q_dn) = (0usize, 1usize, 2usize);
2419    let offset = position.map(|(lo, extent)| {
2420        slots.push(Slot {
2421            buf: lo,
2422            map: fixed(0),
2423        });
2424        IndirectTerm {
2425            operand: slots.len() - 1,
2426            stride: 1,
2427            modulus: extent,
2428        }
2429    });
2430    let table = |off: i64| ElemAccess {
2431        indirect: offset,
2432        ..at(off)
2433    };
2434    let mut bands = Vec::new();
2435    for h in 0..heads {
2436        let base = (h * hd) as i64 - row_shift(h) - shift;
2437        // Region 1: out = q[c]·rc[c − h·hd + row·half] − q[c + half]·rs[c − h·hd + row·half].
2438        slots.push(Slot {
2439            buf: rc,
2440            map: table(-base),
2441        });
2442        let rc1 = slots.len() - 1;
2443        slots.push(Slot {
2444            buf: rs,
2445            map: table(-base),
2446        });
2447        let rs1 = slots.len() - 1;
2448        bands.push((
2449            h * hd + half,
2450            *sub(mul(op(q_at), op(rc1)), mul(op(q_up), op(rs1))),
2451        ));
2452        // Region 2: out = q[c − half]·rs[c − half − h·hd + row·half] + q[c]·rc[c − half − h·hd + row·half].
2453        slots.push(Slot {
2454            buf: rs,
2455            map: table(-base - half as i64),
2456        });
2457        let rs2 = slots.len() - 1;
2458        slots.push(Slot {
2459            buf: rc,
2460            map: table(-base - half as i64),
2461        });
2462        let rc2 = slots.len() - 1;
2463        bands.push((
2464            h * hd + 2 * half,
2465            *add(mul(op(q_dn), op(rs2)), mul(op(q_at), op(rc2))),
2466        ));
2467        // Region 3: the pass-through tail (empty when the rotary is full).
2468        if 2 * half < hd {
2469            bands.push(((h + 1) * hd, TensorBody::Operand(q_at)));
2470        }
2471    }
2472    chain.extend(b.pointwise("rope", n, slots, banded(bands), out_buf, target)?);
2473    let claim = approx("rotary embedding", b.precision);
2474    b.finish(
2475        &names.out,
2476        n,
2477        chain,
2478        out_buf,
2479        rope_name(spec),
2480        claim,
2481        target,
2482    )
2483}
2484
2485/// The `rows` rows a rope block reads of a cache of `cache` positions whose rows are `half_t` whole tiles: ONE tile run
2486/// per table from tile `(P − 1)·half_t`, loaded into buffers the rope's compute reads as a table row per row. A literal
2487/// `P` is the load's base; a named one is loaded, decremented and checked once — a guard of stride zero keeping the
2488/// block's last row among the cache's, as a row index is — and handed to each table's load as an [`IndirectTile`]
2489/// through a FIFO copy of its own (a FIFO has one reader). Returns the cosine and sine buffers.
2490fn cache_block_loads(
2491    b: &mut FlatBuild,
2492    names: &StageNames,
2493    from: RowIndex,
2494    cache: usize,
2495    rows: usize,
2496    half_t: usize,
2497    target: &SpatialTarget,
2498) -> Result<(BufferId, BufferId), SpatialReject> {
2499    let precision = b.precision;
2500    let span = rows * half_t;
2501    let declare = |b: &mut FlatBuild, k: usize| {
2502        let (name, fill) = names.operand(k);
2503        let n = cache * half_t * target.tile.cols;
2504        (
2505            b.t.tensor(
2506                name,
2507                TensorRole::Input,
2508                1,
2509                n,
2510                precision,
2511                super::from_tensor_op::init_of(fill, n),
2512            ),
2513            b.t.buffer(&format!("in_{name}"), span, precision),
2514        )
2515    };
2516    let ((ct, cin), (st, sin)) = (declare(b, 1), declare(b, 2));
2517    let load = |b: &mut FlatBuild,
2518                name: &str,
2519                tensor: TensorId,
2520                out: BufferId,
2521                origin: TileAddr,
2522                index: Option<(StageId, BufferId)>| {
2523        let l = b.t.stage(
2524            name,
2525            StageKind::Load {
2526                tensor,
2527                pattern: AccessPattern {
2528                    origin,
2529                    block_h: 1,
2530                    block_w: span,
2531                    stride_w: 1,
2532                    stride_h: span,
2533                },
2534                num_blocks: 1,
2535                multicast: None,
2536            },
2537            index.map(|(_, i)| vec![i]).unwrap_or_default(),
2538            vec![out],
2539        );
2540        if let Some((s, i)) = index {
2541            b.t.edge(s, i, l);
2542        }
2543        b.loads.push((l, out));
2544    };
2545    match from {
2546        RowIndex::Literal(p) => {
2547            for (name, tensor, out) in [("load_rc", ct, cin), ("load_rs", st, sin)] {
2548                load(
2549                    b,
2550                    name,
2551                    tensor,
2552                    out,
2553                    TileAddr {
2554                        base: (p as usize - 1) * half_t,
2555                        steps: Vec::new(),
2556                        indirect: None,
2557                    },
2558                    None,
2559                );
2560            }
2561        }
2562        RowIndex::Named(_) => {
2563            let (p_name, p_fill) = names.operand(3);
2564            let ix = index_type(target, index_space(cache as u64, 1, p_fill))?;
2565            let w1 = wide(1, target);
2566            let pt = b.t.tensor(
2567                p_name,
2568                TensorRole::Input,
2569                1,
2570                1,
2571                ix,
2572                super::from_tensor_op::init_of(p_fill, 1),
2573            );
2574            let pin = b.t.buffer("in_position", 1, ix);
2575            let lp = b.t.load("load_position", pt, vec![], 1, pin);
2576            let scalar = |b: &mut FlatBuild,
2577                          name: &str,
2578                          from: BufferId,
2579                          map: ElemAccess,
2580                          body: TensorBody,
2581                          out: BufferId| {
2582                b.t.stage(
2583                    name,
2584                    StageKind::Compute(ComputeOp::Generic {
2585                        iters: flat_space(1),
2586                        ins: vec![map],
2587                        out: ident_map(w1),
2588                        body,
2589                    }),
2590                    vec![from],
2591                    vec![out],
2592                )
2593            };
2594            let biased = b.t.scratch_deep("position_biased", 1, 1, ix);
2595            let s_biased = scalar(
2596                b,
2597                "position_biased",
2598                pin,
2599                fixed(0),
2600                *sub(op(0), konst(1.0)),
2601                biased,
2602            );
2603            b.t.edge(lp, pin, s_biased);
2604            let row = b.t.scratch_deep("position_row", 1, 1, ix);
2605            let guard = ElemAccess {
2606                base: 0,
2607                terms: Vec::new(),
2608                indirect: Some(IndirectTerm {
2609                    operand: 0,
2610                    stride: 0,
2611                    modulus: cache - rows + 1,
2612                }),
2613            };
2614            let s_row = scalar(
2615                b,
2616                "position_row",
2617                biased,
2618                guard,
2619                TensorBody::Operand(0),
2620                row,
2621            );
2622            b.t.edge(s_biased, biased, s_row);
2623            for (name, tensor, out) in [("load_rc", ct, cin), ("load_rs", st, sin)] {
2624                let copy = b.t.buffer(&format!("{name}_row"), 1, ix);
2625                let s_copy = scalar(
2626                    b,
2627                    &format!("{name}_row"),
2628                    row,
2629                    fixed(0),
2630                    TensorBody::Operand(0),
2631                    copy,
2632                );
2633                b.t.edge(s_row, row, s_copy);
2634                let indirect = Some(IndirectTile {
2635                    index: copy,
2636                    tile_stride: half_t,
2637                    elems_per_tile: 1,
2638                    bound: (0, cache - rows),
2639                });
2640                load(
2641                    b,
2642                    name,
2643                    tensor,
2644                    out,
2645                    TileAddr {
2646                        base: 0,
2647                        steps: Vec::new(),
2648                        indirect,
2649                    },
2650                    Some((s_copy, copy)),
2651                );
2652            }
2653        }
2654    }
2655    Ok((cin, sin))
2656}
2657
2658/// A rope program's name: its heads, head dimension and rotated half — and its rows, when there is more than one.
2659fn rope_name(spec: &RopeSpec) -> String {
2660    let rows = if spec.rows > 1 {
2661        format!("_rows{}", spec.rows)
2662    } else {
2663        String::new()
2664    };
2665    format!("rope_{}x{}_half{}{rows}", spec.h, spec.hd, spec.half)
2666}
2667
2668// ── The carried-state family ──────────────────────────────────────────────────────────────────────
2669
2670/// `Y[c] = Σ_t CS'[t, c]·WC[t, c]` where `CS' = shift(CS) ← X` — the depthwise causal conv1d over a
2671/// `taps`-deep decode ring, as TWO sections.
2672///
2673/// ★ THE RING IS WRITTEN IN ONE SECTION AND READ IN THE NEXT, AND THAT IS THE WHOLE POINT. The update
2674/// and the reduction are separated by a barrier, so the order between them is a fact the certificate
2675/// checks (`StateOrder`) and the evaluator enforces (`StaleStateRead`) rather than an accident of which
2676/// engine happened to run first. Fold them into one section and a reversed scheduler would compute the
2677/// convolution against LAST token's ring and nothing would notice.
2678pub fn from_causal_conv(
2679    spec: &CausalConvSpec,
2680    names: &StageNames,
2681    target: &SpatialTarget,
2682) -> Result<Vec<SpatialProgram>, SpatialReject> {
2683    let (taps, chan) = (spec.taps, spec.chan);
2684    if taps == 0 || chan == 0 {
2685        return Err(SpatialReject::NotLowered(
2686            "a causal conv with no taps or no channels",
2687        ));
2688    }
2689    if spec.rows > 1 {
2690        return block_causal_conv(spec, names, target);
2691    }
2692    if ring_is_matrix(spec, target) {
2693        return blocked_causal_conv(spec, names, target);
2694    }
2695    flat_causal_conv(spec, names, target)
2696}
2697
2698/// Whether a conv ring of `taps × chan` is held as the matrix it is — a column block of every tap an item — rather than
2699/// as one flat row: when the row is wider than an accumulator page, or the single step's flat form does not fit a site
2700/// (a flat row fills a tile row of the datapath's tile a row of elements, so its pages grow with the tile's height). One
2701/// decision for the single step and a block of rows, so a prefill chunk leaves the ring in the layout a decode reads.
2702fn ring_is_matrix(spec: &CausalConvSpec, target: &SpatialTarget) -> bool {
2703    if row_needs_blocking(spec.taps * spec.chan, target) {
2704        return true;
2705    }
2706    let step = CausalConvSpec { rows: 1, ..spec.clone() };
2707    let names = StageNames { out: "y".into(), ins: vec!["x".into(), "w".into(), "cs".into()], fills: vec![Fill::Unknown; 3] };
2708    flat_causal_conv(&step, &names, target).map_or(true, |sections| sections.iter().any(|p| site_footprint(p) > target.max_local_mem_bytes()))
2709}
2710
2711/// The single step's conv over a flat ring: the ring update, then the reduction against the weight.
2712fn flat_causal_conv(
2713    spec: &CausalConvSpec,
2714    names: &StageNames,
2715    target: &SpatialTarget,
2716) -> Result<Vec<SpatialProgram>, SpatialReject> {
2717    let (taps, chan) = (spec.taps, spec.chan);
2718    let ring_len = taps * chan;
2719    let (x_name, x_fill) = names.operand(0);
2720    let (w_name, w_fill) = names.operand(1);
2721    let (cs_name, cs_fill) = names.operand(2);
2722
2723    // ── Section 1: the ring update. `CS'[t, c] = t < taps-1 ? CS[t+1, c] : X[c]`, over the flat ring.
2724    let shift = {
2725        let mut b = FlatBuild::new(target)?;
2726        let cs = b.input(cs_name, ring_len, cs_fill, target);
2727        let x = b.input(x_name, chan, x_fill, target);
2728        let cols_t = tiles_of(ring_len, target);
2729        let out_buf = b.t.buffer("ring", cols_t, b.precision);
2730        let ow = wide(cols_t, target);
2731        // Two bands: the shifted rows, then the appended newest row.
2732        let slots = vec![
2733            Slot {
2734                buf: cs,
2735                map: at(chan as i64),
2736            }, // CS[c + chan] — one row up
2737            Slot {
2738                buf: x,
2739                map: at(-(((taps - 1) * chan) as i64)),
2740            }, // X[c - (taps-1)·chan]
2741        ];
2742        let body = banded(vec![
2743            ((taps - 1) * chan, TensorBody::Operand(0)),
2744            (ring_len, TensorBody::Operand(1)),
2745        ]);
2746        let s = b.compute(
2747            "ring_shift",
2748            flat_space(ring_len),
2749            slots,
2750            ident_map(ow),
2751            body,
2752            out_buf,
2753        );
2754        let claim = NumericsClaim::Exact;
2755        b.finish(
2756            cs_name,
2757            ring_len,
2758            vec![s],
2759            out_buf,
2760            format!("conv_ring_{taps}x{chan}"),
2761            claim,
2762            target,
2763        )?
2764    };
2765
2766    // ── Section 2: the tap reduction against the UPDATED ring.
2767    let reduce = {
2768        let mut b = FlatBuild::new(target)?;
2769        let cs = b.input(cs_name, ring_len, cs_fill, target);
2770        let wc = b.input(w_name, ring_len, w_fill, target);
2771        let cols_t = tiles_of(chan, target);
2772        let out_buf = b.t.buffer("out", cols_t, b.precision);
2773        let tile_elems = target.tile.rows * target.tile.cols;
2774        let ones = b.t.tensor(
2775            "conv_scaler",
2776            TensorRole::Param,
2777            target.tile.rows,
2778            target.tile.cols,
2779            b.precision,
2780            TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
2781        );
2782        let ones_in = b.t.buffer("in_scaler", 1, b.precision);
2783        let l = b.t.load("load_scaler", ones, vec![], 1, ones_in);
2784        b.loads.push((l, ones_in));
2785        // ★ THE TAP AXIS IS THE REDUCED ONE, AND IT IS THE *ROW*. Folding the ROW axis leaves one value
2786        // per COLUMN, written DOWN the output tile — so the answer is present at row 0 of every column,
2787        // which IS the flat `1 × chan` layout the next section reads. (Folding the column axis instead
2788        // would leave one value per ROW, and a flat reader would see only the first head's.) The channel
2789        // is the column, the tap is the row, and `CS[t, c]` is the flat offset `t·chan + c`.
2790        let iters = vec![
2791            IterSpec {
2792                var: LV_GEN_ROW,
2793                extent: taps,
2794                kind: IterKind::Reduce {
2795                    pool: ReducePool::Sum,
2796                },
2797            },
2798            IterSpec {
2799                var: LV_GEN_COL,
2800                extent: chan,
2801                kind: IterKind::Parallel,
2802            },
2803        ];
2804        let slots = vec![
2805            Slot {
2806                buf: cs,
2807                map: map2(0, chan as i64, 1),
2808            },
2809            Slot {
2810                buf: wc,
2811                map: map2(0, chan as i64, 1),
2812            },
2813            Slot {
2814                buf: ones_in,
2815                map: const_map(),
2816            },
2817        ];
2818        let s = b.compute(
2819            "tap_reduce",
2820            iters,
2821            slots,
2822            per_col_map(),
2823            *mul(op(0), op(1)),
2824            out_buf,
2825        );
2826        let claim = approx("causal convolution", b.precision);
2827        b.finish(
2828            &names.out,
2829            chan,
2830            vec![s],
2831            out_buf,
2832            format!("conv_taps_{taps}x{chan}"),
2833            claim,
2834            target,
2835        )?
2836    };
2837    Ok(vec![shift, reduce])
2838}
2839
2840/// `Y = causal conv(X, WC, CS)` over a block of `rows` rows. The rows enter the ring one after another, so with
2841/// `E = [CS; X]` — the ring's `taps` rows, then the block's — row `r` convolves the window `E[r + 1 … r + taps]` and the
2842/// ring the block leaves is `E[rows … rows + taps − 1]`. TWO sections: every output row against the ring as it came in,
2843/// then the ring update — the barrier between them is the order "read the old ring, then write the new one".
2844///
2845/// `X` and `Y` are the DENSE `rows × chan` matrices a projection writes and a pointwise op reads. The ring and the weight
2846/// keep the layout the single step gives them — a `taps × chan` matrix where the ring is wider than a site's row (every
2847/// tap of a channel tile in one tile), the flat `1 × taps·chan` row otherwise — so a decode step continues from the ring
2848/// a prefill chunk leaves. A column block is whole tiles of a matrix ring, the whole row of a flat one.
2849///
2850/// ★ THE HALO IS THE ONLY NON-UNIFORM READ, AND ITS SOURCE IS DATA. Output row `y` of tile row `R` reads window row
2851/// `j = y + 1 + t − taps` of the block: its own tile row when `j ≥ 0`, else the row `j` before it — the previous tile
2852/// row of `X`, or, for the first tile row, the ring. The choice is a per-item threshold from a compile-time table (the
2853/// tile height for the first tile row, zero after), so every item runs the same stages.
2854fn block_causal_conv(
2855    spec: &CausalConvSpec,
2856    names: &StageNames,
2857    target: &SpatialTarget,
2858) -> Result<Vec<SpatialProgram>, SpatialReject> {
2859    let (rows, taps, chan) = (spec.rows, spec.taps, spec.chan);
2860    let (tr, tw) = (target.tile.rows, target.tile.cols);
2861    if taps > tr {
2862        return Err(SpatialReject::NotLowered(
2863            "a block causal conv whose ring is taller than a tile row",
2864        ));
2865    }
2866    // The single step's ring layout, decided the way `from_causal_conv` decides it.
2867    let matrix_ring = ring_is_matrix(spec, target);
2868    if matrix_ring && chan % tw != 0 {
2869        return Err(SpatialReject::NotLowered(
2870            "a blocked causal conv whose channel count is not a whole number of tiles",
2871        ));
2872    }
2873    let (cols_t, rows_t) = (chan.div_ceil(tw), rows.div_ceil(tr));
2874    let dt = elem_precision(target)?;
2875    let ix = index_type(target, (rows_t * cols_t + tr) as u64)?;
2876    let blk = if matrix_ring {
2877        row_block(cols_t, target)
2878    } else {
2879        cols_t
2880    };
2881    let n_b = cols_t / blk;
2882    let wc_w = wide(blk, target);
2883    let ring_t = tiles_of(taps * chan, target);
2884    // Where ring row `k` sits in its page: plane row `k` of a matrix ring's tile row, or `k·chan` along a flat one.
2885    let seg = if matrix_ring { wc_w } else { chan };
2886    let ring_tiles = if matrix_ring { blk } else { ring_t };
2887    let (x_name, x_fill) = names.operand(0);
2888    let (w_name, w_fill) = names.operand(1);
2889    let (cs_name, cs_fill) = names.operand(2);
2890    let ring_tensor = |t: &mut Template, name: &str, role: TensorRole, fill: &Fill| {
2891        let (r, c) = if matrix_ring {
2892            (taps, chan)
2893        } else {
2894            (1, taps * chan)
2895        };
2896        t.tensor(
2897            name,
2898            role,
2899            r,
2900            c,
2901            dt,
2902            super::from_tensor_op::init_of(fill, taps * chan),
2903        )
2904    };
2905    let ring_page = |t: &mut Template, name: &str, tensor: TensorId, into: BufferId| {
2906        if matrix_ring {
2907            t.load_block_at(
2908                name,
2909                tensor,
2910                0,
2911                vec![(LV_ITEM_X, blk)],
2912                1,
2913                blk,
2914                cols_t,
2915                into,
2916            )
2917        } else {
2918            t.load_block_at(name, tensor, 0, Vec::new(), 1, ring_t, ring_t, into)
2919        }
2920    };
2921    let per_tile_row = |t: &mut Template, what: &str, values: Vec<i64>| {
2922        let tiles = values
2923            .into_iter()
2924            .flat_map(|v| std::iter::once(v).chain(std::iter::repeat_n(0, tw - 1)))
2925            .collect();
2926        t.tensor(
2927            &format!("{}__conv_{what}", names.out),
2928            TensorRole::Param,
2929            1,
2930            rows_t * tw,
2931            ix,
2932            TensorInit::Values(tiles),
2933        )
2934    };
2935    let rows_map = |base: i64, stride: usize| ElemAccess {
2936        base,
2937        terms: vec![(LV_GEN_ROW, stride as i64), (LV_GEN_COL, 1)],
2938        indirect: None,
2939    };
2940
2941    // ── Section 1: every output row, against the ring as it came in.
2942    let outputs = {
2943        let mut t = Template::new(target)?;
2944        let x = t.tensor(
2945            x_name,
2946            TensorRole::Input,
2947            rows,
2948            chan,
2949            dt,
2950            super::from_tensor_op::init_of(x_fill, rows * chan),
2951        );
2952        let w = ring_tensor(&mut t, w_name, TensorRole::Input, w_fill);
2953        let cs = ring_tensor(&mut t, cs_name, TensorRole::Input, cs_fill);
2954        let y = t.tensor(
2955            &names.out,
2956            TensorRole::Output,
2957            rows,
2958            chan,
2959            dt,
2960            TensorInit::Zero,
2961        );
2962        let prev_rows = per_tile_row(
2963            &mut t,
2964            "previous_row",
2965            (0..rows_t).map(|r| r.saturating_sub(1) as i64).collect(),
2966        );
2967        let halo = per_tile_row(
2968            &mut t,
2969            "halo",
2970            (0..rows_t)
2971                .map(|r| if r == 0 { tr as i64 } else { 0 })
2972                .collect(),
2973        );
2974        let (c_in, p_in, ring_in, w_in) = (
2975            t.buffer("in_rows", blk, dt),
2976            t.buffer("in_previous_rows", blk, dt),
2977            t.buffer("in_ring", ring_tiles, dt),
2978            t.buffer("in_weight", ring_tiles, dt),
2979        );
2980        let (prev_in, halo_in) = (
2981            t.buffer("in_previous_row", 1, ix),
2982            t.buffer("in_halo", 1, ix),
2983        );
2984        let l_c = t.load_block_at(
2985            "load_rows",
2986            x,
2987            0,
2988            vec![(LV_ITEM_Y, cols_t), (LV_ITEM_X, blk)],
2989            1,
2990            blk,
2991            cols_t,
2992            c_in,
2993        );
2994        let l_prev = t.load_block_at(
2995            "load_previous_row",
2996            prev_rows,
2997            0,
2998            vec![(LV_ITEM_Y, 1)],
2999            1,
3000            1,
3001            1,
3002            prev_in,
3003        );
3004        let origin = TileAddr {
3005            base: 0,
3006            steps: vec![(LV_ITEM_X, blk)],
3007            indirect: Some(IndirectTile {
3008                index: prev_in,
3009                tile_stride: cols_t,
3010                elems_per_tile: 1,
3011                bound: (0, rows_t - 1),
3012            }),
3013        };
3014        let l_p = t.stage(
3015            "load_previous_rows",
3016            StageKind::Load {
3017                tensor: x,
3018                pattern: AccessPattern {
3019                    origin,
3020                    block_h: 1,
3021                    block_w: blk,
3022                    stride_w: 1,
3023                    stride_h: cols_t,
3024                },
3025                num_blocks: 1,
3026                multicast: None,
3027            },
3028            vec![prev_in],
3029            vec![p_in],
3030        );
3031        t.edge(l_prev, prev_in, l_p);
3032        let l_ring = ring_page(&mut t, "load_ring", cs, ring_in);
3033        let l_w = ring_page(&mut t, "load_weight", w, w_in);
3034        let l_halo = t.load_block_at("load_halo", halo, 0, vec![(LV_ITEM_Y, 1)], 1, 1, 1, halo_in);
3035        let (mut ins, mut maps) = (vec![halo_in], vec![fixed(0)]);
3036        let mut slot = |b: BufferId, m: ElemAccess| {
3037            ins.push(b);
3038            maps.push(m);
3039            op(ins.len() - 1)
3040        };
3041        let mut body: Option<Box<TensorBody>> = None;
3042        for tap in 0..taps {
3043            let shift = 1 + tap as i64 - taps as i64;
3044            let own = slot(c_in, rows_map(shift * wc_w as i64, wc_w));
3045            let before = slot(p_in, rows_map((tr as i64 + shift) * wc_w as i64, wc_w));
3046            let ring = slot(ring_in, rows_map((1 + tap as i64) * seg as i64, seg));
3047            let weight = slot(
3048                w_in,
3049                ElemAccess {
3050                    base: (tap * seg) as i64,
3051                    terms: vec![(LV_GEN_COL, 1)],
3052                    indirect: None,
3053                },
3054            );
3055            let halo_read = Box::new(TensorBody::Select {
3056                cond: Pred::AxisLt(LV_GEN_ROW, PredRhs::Operand(0)),
3057                then: ring,
3058                otherwise: before,
3059            });
3060            let window = Box::new(TensorBody::Select {
3061                cond: Pred::AxisGe(LV_GEN_ROW, PredRhs::Const(taps - 1 - tap)),
3062                then: own,
3063                otherwise: halo_read,
3064            });
3065            let term = mul(weight, window);
3066            body = Some(match body {
3067                None => term,
3068                Some(acc) => add(acc, term),
3069            });
3070        }
3071        let out = t.buffer("out", blk, dt);
3072        let s = t.stage(
3073            "conv_rows",
3074            StageKind::Compute(ComputeOp::Generic {
3075                iters: super::from_tensor_op::space(target, wc_w, None),
3076                ins: maps,
3077                out: ident_map(wc_w),
3078                body: *body.expect("a conv has taps"),
3079            }),
3080            ins.clone(),
3081            vec![out],
3082        );
3083        for (from, buf) in [
3084            (l_c, c_in),
3085            (l_p, p_in),
3086            (l_ring, ring_in),
3087            (l_w, w_in),
3088            (l_halo, halo_in),
3089        ] {
3090            t.edge(from, buf, s);
3091        }
3092        let st = t.store_block(
3093            "store_y",
3094            y,
3095            0,
3096            vec![(LV_ITEM_Y, cols_t), (LV_ITEM_X, blk)],
3097            1,
3098            blk,
3099            cols_t,
3100            out,
3101        );
3102        t.edge(s, out, st);
3103        let items: Vec<WorkItem> = (0..rows_t)
3104            .flat_map(|r| (0..n_b).map(move |b| (r, b)))
3105            .enumerate()
3106            .map(|(id, (r, b))| WorkItem {
3107                id: WorkItemId::new(id),
3108                coord: (r, b),
3109                out: TileRect {
3110                    row0_t: r,
3111                    col0_t: b * blk,
3112                    h_t: 1,
3113                    w_t: blk,
3114                },
3115            })
3116            .collect();
3117        let mut prog = t.finish(
3118            format!("conv_rows_{rows}x{chan}_taps{taps}"),
3119            items,
3120            AccPolicy::F32,
3121            approx("causal convolution", dt),
3122            target,
3123        )?;
3124        prog.work.plan = Some(BlockPlan {
3125            grid: (rows_t, n_b),
3126            ..tile_row_plan(blk, blk, dt)
3127        });
3128        prog
3129    };
3130
3131    // ── Section 2: the ring the block leaves, `CS'[t] = E[rows + t]` — the old ring's rows past the block's start, then
3132    // the block's last rows, each from the tile row of `X` that holds it. Only the sources a band reads are loaded.
3133    let ring = {
3134        let mut t = Template::new(target)?;
3135        let x = t.tensor(
3136            x_name,
3137            TensorRole::Input,
3138            rows,
3139            chan,
3140            dt,
3141            super::from_tensor_op::init_of(x_fill, rows * chan),
3142        );
3143        let cs = ring_tensor(&mut t, cs_name, TensorRole::State, cs_fill);
3144        #[derive(Clone, Copy, PartialEq)]
3145        enum Source {
3146            Ring,
3147            Rows(usize),
3148        }
3149        // Tap `t`'s source and where its row sits: ring row `rows + t`, or row `xr` of `X` at plane row `xr mod tile`.
3150        let sources: Vec<(Source, usize)> = (0..taps)
3151            .map(|tap| {
3152                let e = rows + tap;
3153                if e < taps {
3154                    (Source::Ring, e)
3155                } else {
3156                    (Source::Rows((e - taps) / tr), (e - taps) % tr)
3157                }
3158            })
3159            .collect();
3160        let mut pages: Vec<(Source, StageId, BufferId)> = Vec::new();
3161        for (source, _) in &sources {
3162            if pages.iter().any(|(s, _, _)| s == source) {
3163                continue;
3164            }
3165            let (load, buf) = match *source {
3166                Source::Ring => {
3167                    let b = t.buffer("in_ring", ring_tiles, dt);
3168                    (ring_page(&mut t, "load_ring", cs, b), b)
3169                }
3170                Source::Rows(tile_row) => {
3171                    let b = t.buffer(&format!("in_rows_{tile_row}"), blk, dt);
3172                    (
3173                        t.load_block_at(
3174                            &format!("load_rows_{tile_row}"),
3175                            x,
3176                            tile_row * cols_t,
3177                            vec![(LV_ITEM_X, blk)],
3178                            1,
3179                            blk,
3180                            cols_t,
3181                            b,
3182                        ),
3183                        b,
3184                    )
3185                }
3186            };
3187            pages.push((*source, load, buf));
3188        }
3189        let page_of = |source: Source| {
3190            pages
3191                .iter()
3192                .find(|(s, _, _)| *s == source)
3193                .map(|(_, _, b)| *b)
3194                .expect("every tap's source is loaded")
3195        };
3196        let (mut ins, mut maps, mut bands) = (Vec::new(), Vec::new(), Vec::new());
3197        for (tap, (source, row)) in sources.iter().enumerate() {
3198            ins.push(page_of(*source));
3199            // A matrix ring's tap `t` is plane row `t` of the output: read plane row `row` — `row − t` rows away. A flat
3200            // ring's tap `t` is the band `[t·chan, (t+1)·chan)` along the output row.
3201            let (row_stride, base) = match (matrix_ring, source) {
3202                (true, _) => (wc_w, (*row as i64 - tap as i64) * wc_w as i64),
3203                (false, Source::Ring) => (seg, (*row as i64 - tap as i64) * seg as i64),
3204                (false, Source::Rows(_)) => (wc_w, *row as i64 * wc_w as i64 - (tap * seg) as i64),
3205            };
3206            let map = if matrix_ring {
3207                rows_map(base, row_stride)
3208            } else {
3209                ElemAccess {
3210                    base,
3211                    terms: vec![(LV_GEN_COL, 1)],
3212                    indirect: None,
3213                }
3214            };
3215            maps.push(map);
3216            bands.push((
3217                if matrix_ring {
3218                    tap + 1
3219                } else {
3220                    (tap + 1) * seg
3221                },
3222                TensorBody::Operand(ins.len() - 1),
3223            ));
3224        }
3225        let (iters, body, out_w) = if matrix_ring {
3226            // Rows past the ring's taps are the tile's padding: they carry a zero.
3227            bands.push((tr, *konst(0.0)));
3228            (
3229                super::from_tensor_op::space(target, wc_w, None),
3230                banded_on(LV_GEN_ROW, bands),
3231                wc_w,
3232            )
3233        } else {
3234            (flat_space(taps * seg), banded(bands), wide(ring_t, target))
3235        };
3236        let out = t.buffer("ring", ring_tiles, dt);
3237        let s = t.stage(
3238            "ring_update",
3239            StageKind::Compute(ComputeOp::Generic {
3240                iters,
3241                ins: maps,
3242                out: ident_map(out_w),
3243                body,
3244            }),
3245            ins,
3246            vec![out],
3247        );
3248        for (_, load, buf) in &pages {
3249            t.edge(*load, *buf, s);
3250        }
3251        let st = if matrix_ring {
3252            t.store_block(
3253                "store_ring",
3254                cs,
3255                0,
3256                vec![(LV_ITEM_X, blk)],
3257                1,
3258                blk,
3259                cols_t,
3260                out,
3261            )
3262        } else {
3263            t.store_block("store_ring", cs, 0, Vec::new(), 1, ring_t, ring_t, out)
3264        };
3265        t.edge(s, out, st);
3266        // An item writes its column block of a matrix ring, or — flat — the whole ring.
3267        let ring_w = if matrix_ring { blk } else { ring_t };
3268        let items: Vec<WorkItem> = (0..n_b)
3269            .map(|b| WorkItem {
3270                id: WorkItemId::new(b),
3271                coord: (0, b),
3272                out: TileRect {
3273                    row0_t: 0,
3274                    col0_t: b * ring_w,
3275                    h_t: 1,
3276                    w_t: ring_w,
3277                },
3278            })
3279            .collect();
3280        let mut prog = t.finish(
3281            format!("conv_ring_rows_{rows}x{chan}_taps{taps}"),
3282            items,
3283            AccPolicy::F32,
3284            NumericsClaim::Exact,
3285            target,
3286        )?;
3287        prog.work.plan = Some(BlockPlan {
3288            grid: (1, n_b),
3289            ..tile_row_plan(ring_w, ring_w, dt)
3290        });
3291        prog
3292    };
3293    Ok(vec![outputs, ring])
3294}
3295
3296/// The Gated DeltaNet step over `hv` value heads reading `hk` key heads at head dim `d`:
3297/// `u = S·k`, `S ← α·S + β·(v − α·u)·kᵀ`, `o = S·q / √d`, RMSNorm'd, scaled by `GN`, gated by `silu(z)`.
3298///
3299/// ★ THREE SECTIONS, BECAUSE THE MEMORY IS READ ONCE AND WRITTEN ONCE. Section 1 reads `S` to form the
3300/// prediction `u`; section 2 writes the new `S` from the old one and `u`; section 3 reads the new `S`
3301/// to form the output. Each boundary is a barrier, so "read the old memory, then write, then read the
3302/// new one" is enforced rather than hoped for, and `StaleStateRead` catches any scheduler that tries
3303/// otherwise. Fusing them would make a wrong answer indistinguishable from a right one.
3304///
3305/// ★ A BLOCK OF ROWS IS THE STEP, ROW AFTER ROW. The delta rule is a recurrence — row `r` reads the memory row `r − 1`
3306/// wrote — so a block of tokens is its rows' three sections in turn against the one memory, each reading its row of the
3307/// operands in place and writing its row of the output. `dense` names the block's tables the program lays out as
3308/// matrices; the output is written in its layout.
3309pub fn from_gated_delta(
3310    spec: &GatedDeltaSpec,
3311    names: &StageNames,
3312    dense: &BTreeSet<String>,
3313    target: &SpatialTarget,
3314) -> Result<Vec<SpatialProgram>, SpatialReject> {
3315    let (hv, hk, d, rows) = (spec.hv, spec.hk, spec.d, spec.rows);
3316    if hv == 0 || hk == 0 || d == 0 || rows == 0 {
3317        return Err(SpatialReject::NotLowered(
3318            "a gated delta step with no heads, no head dimension or no rows",
3319        ));
3320    }
3321    // The flat form keeps the whole `hv·d²` memory as one row; a memory past one accumulator page of its row is blocked
3322    // by head: a head of whole tiles as it is, any other widened to whole tiles.
3323    if row_needs_blocking(hv * d * d, target) {
3324        let carried = names.operand(6).0;
3325        return if d % target.tile.cols == 0 {
3326            blocked_gated_delta(spec, names, dense, d, &|_| carried.to_string(), target)
3327        } else {
3328            widened_gated_delta(spec, names, target)
3329        };
3330    }
3331    let mut sections = Vec::new();
3332    for r in 0..rows {
3333        sections.extend(flat_gated_delta_row(spec, names, r, target)?);
3334        // A row of the flat output shares its tiles with the rows beside it, so each row's output is written into
3335        // the table by the row write that merges one row into its tiles.
3336        if rows > 1 {
3337            let at = RowIndex::Literal(r as i64 + 1);
3338            let write = SetRowSpec {
3339                table: spec.out,
3340                index: at,
3341                source: spec.out,
3342                rows,
3343                cols: hv * d,
3344                count: 1,
3345            };
3346            let written = StageNames {
3347                out: names.out.clone(),
3348                ins: vec![names.out.clone(), format!("{}__o{r}", names.out)],
3349                fills: vec![Fill::ConstF(0f64.to_bits()), Fill::Unknown],
3350            };
3351            sections.push(from_set_row(
3352                &write,
3353                &written,
3354                true,
3355                dense.contains(&names.out),
3356                target,
3357            )?);
3358        }
3359    }
3360    Ok(sections)
3361}
3362
3363/// Row `r` of [`from_gated_delta`]'s flat form — its prediction, its memory update and its output — reading the row's
3364/// operands in place: row `r` of `Q` and `K` from element `r·hk·d`, of `V` and `Z` from `r·hv·d`, of `B` and `A` from
3365/// `r·hv`. A one-row step is row 0 of itself.
3366fn flat_gated_delta_row(
3367    spec: &GatedDeltaSpec,
3368    names: &StageNames,
3369    r: usize,
3370    target: &SpatialTarget,
3371) -> Result<Vec<SpatialProgram>, SpatialReject> {
3372    let (hv, hk, d, rows) = (spec.hv, spec.hk, spec.d, spec.rows);
3373    let s_len = hv * d * d;
3374    let out_len = hv * d;
3375    let (qk_at, v_at, h_at) = ((r * hk * d) as i64, (r * out_len) as i64, (r * hv) as i64);
3376    let tag = if rows == 1 {
3377        String::new()
3378    } else {
3379        format!("_row{r}")
3380    };
3381    let o_name = if rows == 1 {
3382        names.out.clone()
3383    } else {
3384        format!("{}__o{r}", names.out)
3385    };
3386    let (q_name, q_fill) = names.operand(0);
3387    let (k_name, k_fill) = names.operand(1);
3388    let (v_name, v_fill) = names.operand(2);
3389    let (z_name, z_fill) = names.operand(3);
3390    let (b_name, b_fill) = names.operand(4);
3391    let (a_name, a_fill) = names.operand(5);
3392    let (st_name, st_fill) = names.operand(6);
3393    let (gn_name, gn_fill) = names.operand(7);
3394    let u_name = if rows == 1 {
3395        format!("{}_u", names.out)
3396    } else {
3397        format!("{}_u{r}", names.out)
3398    };
3399
3400    // ── Section 1: `u[h·d + a] = Σ_b S[h·d² + a·d + b]·k[(h mod hk)·d + b]`, the prediction from the
3401    // memory BEFORE it is updated. The reduced axis is the column `b`; the head and row `a` ride the
3402    // flat output column through per-head bands.
3403    let predict = {
3404        let mut f = FlatBuild::new(target)?;
3405        let st = f.input(st_name, s_len, st_fill, target);
3406        let k = f.input(k_name, rows * hk * d, k_fill, target);
3407        let cols_t = tiles_of(out_len, target);
3408        let out_buf = f.t.buffer("u", cols_t, f.precision);
3409        let tile_elems = target.tile.rows * target.tile.cols;
3410        let ones = f.t.tensor(
3411            "gd_scaler",
3412            TensorRole::Param,
3413            target.tile.rows,
3414            target.tile.cols,
3415            f.precision,
3416            TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
3417        );
3418        let ones_in = f.t.buffer("in_scaler", 1, f.precision);
3419        let l = f.t.load("load_scaler", ones, vec![], 1, ones_in);
3420        f.loads.push((l, ones_in));
3421        // ★ THE FOLD IS ON THE ROW, THE ANSWER IS ON THE COLUMN. Folding the row axis writes the
3422        // result DOWN every row of its column, so row 0 of column `c` holds it — which is exactly the
3423        // flat `1 × n` layout the next section reads, and it spans as many tiles as the answer needs.
3424        // Folding the column axis instead leaves one value per ROW, and a flat reader would see only
3425        // the first of them: the same shape, a completely different answer.
3426        let iters = vec![
3427            IterSpec {
3428                var: LV_GEN_ROW,
3429                extent: d,
3430                kind: IterKind::Reduce {
3431                    pool: ReducePool::Sum,
3432                },
3433            },
3434            IterSpec {
3435                var: LV_GEN_COL,
3436                extent: out_len,
3437                kind: IterKind::Parallel,
3438            },
3439        ];
3440        let mut slots = vec![Slot {
3441            buf: ones_in,
3442            map: const_map(),
3443        }];
3444        let mut bands = Vec::new();
3445        for h in 0..hv {
3446            // `S[h·d² + a·d + b]` with `a = c − h·d` is `c·d + b` — the per-head term cancels, so the
3447            // memory is read by ONE map for every head.
3448            slots.push(Slot {
3449                buf: st,
3450                map: ElemAccess {
3451                    base: 0,
3452                    terms: vec![(LV_GEN_COL, d as i64), (LV_GEN_ROW, 1)],
3453                    indirect: None,
3454                },
3455            });
3456            let s_slot = slots.len() - 1;
3457            // `k[(h mod hk)·d + b]` — constant along the output column, affine in the folded row.
3458            slots.push(Slot {
3459                buf: k,
3460                map: ElemAccess {
3461                    base: ((h % hk) * d) as i64 + qk_at,
3462                    terms: vec![(LV_GEN_ROW, 1)],
3463                    indirect: None,
3464                },
3465            });
3466            let k_slot = slots.len() - 1;
3467            bands.push(((h + 1) * d, *mul(op(s_slot), op(k_slot))));
3468        }
3469        let body = banded(bands);
3470        let s = f.compute("gd_predict", iters, slots, per_col_map(), body, out_buf);
3471        let claim = approx("gated-delta prediction", f.precision);
3472        f.finish(
3473            &u_name,
3474            out_len,
3475            vec![s],
3476            out_buf,
3477            format!("gd_predict_{hv}x{d}{tag}"),
3478            claim,
3479            target,
3480        )?
3481    };
3482
3483    // ── Section 2: the rank-1 memory update `S[h, a, c] = α_h·S[h, a, c] + coef[h, a]·k[h, c]` with
3484    // `coef[h, a] = β_h·(v[h·d + a] − α_h·u[h·d + a])`.
3485    let update = {
3486        let mut f = FlatBuild::new(target)?;
3487        let st = f.input(st_name, s_len, st_fill, target);
3488        let u = f.input(&u_name, out_len, &Fill::Unknown, target);
3489        let v = f.input(v_name, rows * out_len, v_fill, target);
3490        let k = f.input(k_name, rows * hk * d, k_fill, target);
3491        let al = f.input(a_name, rows * hv, a_fill, target);
3492        let be = f.input(b_name, rows * hv, b_fill, target);
3493        let cols_t = tiles_of(s_len, target);
3494        let out_buf = f.t.buffer("state", cols_t, f.precision);
3495        let ow = wide(cols_t, target);
3496        // The flat state column is `c = h·d² + a·d + j`. Inside head `h`'s band, `a = (c − h·d²) / d`
3497        // and `j = (c − h·d²) % d` — NEITHER is affine, so the head's band is split further into `d`
3498        // ROW bands, one per `a`, inside which both are affine again. `d` is small (128 in Qwen3.8's
3499        // real config, 4 in the tiny one) and the bands are constants, so this is a wide body rather
3500        // than a deep one.
3501        let mut slots = vec![Slot {
3502            buf: st,
3503            map: at(0),
3504        }];
3505        let st_slot = 0usize;
3506        let mut bands = Vec::new();
3507        for h in 0..hv {
3508            let hb = (h * d * d) as i64;
3509            slots.push(Slot {
3510                buf: al,
3511                map: fixed(h as i64 + h_at),
3512            });
3513            let a_slot = slots.len() - 1;
3514            slots.push(Slot {
3515                buf: be,
3516                map: fixed(h as i64 + h_at),
3517            });
3518            let b_slot = slots.len() - 1;
3519            for a in 0..d {
3520                let row_base = hb + (a * d) as i64;
3521                // v[h·d + a] and u[h·d + a] are fixed within this row band.
3522                slots.push(Slot {
3523                    buf: v,
3524                    map: fixed((h * d + a) as i64 + v_at),
3525                });
3526                let v_slot = slots.len() - 1;
3527                slots.push(Slot {
3528                    buf: u,
3529                    map: fixed((h * d + a) as i64),
3530                });
3531                let u_slot = slots.len() - 1;
3532                // k[(h mod hk)·d + j] with j = c − row_base.
3533                slots.push(Slot {
3534                    buf: k,
3535                    map: at(((h % hk) * d) as i64 + qk_at - row_base),
3536                });
3537                let k_slot = slots.len() - 1;
3538                let coef = mul(op(b_slot), sub(op(v_slot), mul(op(a_slot), op(u_slot))));
3539                let value = add(mul(op(a_slot), op(st_slot)), mul(coef, op(k_slot)));
3540                bands.push((row_base as usize + d, *value));
3541            }
3542        }
3543        let s = f.compute(
3544            "gd_update",
3545            flat_space(s_len),
3546            slots,
3547            ident_map(ow),
3548            banded(bands),
3549            out_buf,
3550        );
3551        let claim = approx("gated-delta memory update", f.precision);
3552        f.finish(
3553            st_name,
3554            s_len,
3555            vec![s],
3556            out_buf,
3557            format!("gd_update_{hv}x{d}{tag}"),
3558            claim,
3559            target,
3560        )?
3561    };
3562
3563    // ── Section 3: `o = S_new·q / √d`, RMSNorm over the head, scaled by `GN`, gated by `silu(z)`.
3564    let read = {
3565        let mut f = FlatBuild::new(target)?;
3566        let st = f.input(st_name, s_len, st_fill, target);
3567        let q = f.input(q_name, rows * hk * d, q_fill, target);
3568        let z = f.input(z_name, rows * out_len, z_fill, target);
3569        let gn = f.input(gn_name, d, gn_fill, target);
3570        let cols_t = tiles_of(out_len, target);
3571        let out_buf = f.t.buffer("out", cols_t, f.precision);
3572        let ow = wide(cols_t, target);
3573        let tile_elems = target.tile.rows * target.tile.cols;
3574        let ones = f.t.tensor(
3575            "gd_scaler",
3576            TensorRole::Param,
3577            target.tile.rows,
3578            target.tile.cols,
3579            f.precision,
3580            TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
3581        );
3582        let ones_in = f.t.buffer("in_scaler", 1, f.precision);
3583        let l = f.t.load("load_scaler", ones, vec![], 1, ones_in);
3584        f.loads.push((l, ones_in));
3585        let ovec = f.scratch("gd_o", cols_t);
3586        let iters = vec![
3587            IterSpec {
3588                var: LV_GEN_ROW,
3589                extent: d,
3590                kind: IterKind::Reduce {
3591                    pool: ReducePool::Sum,
3592                },
3593            },
3594            IterSpec {
3595                var: LV_GEN_COL,
3596                extent: out_len,
3597                kind: IterKind::Parallel,
3598            },
3599        ];
3600        let inv_sqrt_d = 1.0f64 / (d as f64).sqrt();
3601        let mut slots = vec![Slot {
3602            buf: ones_in,
3603            map: const_map(),
3604        }];
3605        let mut bands = Vec::new();
3606        for h in 0..hv {
3607            slots.push(Slot {
3608                buf: st,
3609                map: ElemAccess {
3610                    base: 0,
3611                    terms: vec![(LV_GEN_COL, d as i64), (LV_GEN_ROW, 1)],
3612                    indirect: None,
3613                },
3614            });
3615            let s_slot = slots.len() - 1;
3616            slots.push(Slot {
3617                buf: q,
3618                map: ElemAccess {
3619                    base: ((h % hk) * d) as i64 + qk_at,
3620                    terms: vec![(LV_GEN_ROW, 1)],
3621                    indirect: None,
3622                },
3623            });
3624            let q_slot = slots.len() - 1;
3625            bands.push((
3626                (h + 1) * d,
3627                *mul(mul(op(s_slot), op(q_slot)), konst(inv_sqrt_d)),
3628            ));
3629        }
3630        let s1 = f.compute("gd_read", iters, slots, per_col_map(), banded(bands), ovec);
3631        // The per-head RMSNorm of the ungated output, then the weight and the gate.
3632        let (mut chain, scale, one_w) = row_scale_chain(
3633            &mut f,
3634            ovec,
3635            hv,
3636            d,
3637            SfpuKind::Rsqrt {
3638                eps_nano: eps_nano(spec.eps)?,
3639                over: d,
3640            },
3641            target,
3642        );
3643        chain.insert(0, s1);
3644        let mut slots = vec![
3645            Slot {
3646                buf: ovec,
3647                map: at(0),
3648            },
3649            Slot {
3650                buf: z,
3651                map: at(v_at),
3652            },
3653        ];
3654        let (o_slot, z_slot) = (0usize, 1usize);
3655        let mut bands = Vec::new();
3656        for h in 0..hv {
3657            let head_base = (h * d) as i64;
3658            slots.push(Slot {
3659                buf: scale,
3660                map: fixed((h * one_w) as i64),
3661            });
3662            let sc = slots.len() - 1;
3663            slots.push(Slot {
3664                buf: gn,
3665                map: at(-head_base),
3666            });
3667            let g = slots.len() - 1;
3668            bands.push((
3669                (h + 1) * d,
3670                *mul(mul(op(g), mul(op(o_slot), op(sc))), delta_gate(spec.gate, op(z_slot))),
3671            ));
3672        }
3673        let s2 = f.compute(
3674            "gd_gate",
3675            flat_space(out_len),
3676            slots,
3677            ident_map(ow),
3678            banded(bands),
3679            out_buf,
3680        );
3681        chain.push(s2);
3682        let claim = approx("gated-delta output", f.precision);
3683        f.finish(
3684            &o_name,
3685            out_len,
3686            chain,
3687            out_buf,
3688            format!("gd_read_{hv}x{d}{tag}"),
3689            claim,
3690            target,
3691        )?
3692    };
3693    Ok(vec![predict, update, read])
3694}
3695
3696/// A section writing `out`, a flat `1 × n` tensor whose element `e` is the sink logit of query head `head_of(e)` (or
3697/// [`MASKED`] for `None`, a row no head fills): the per-row sinks a softmax over the heads' rows reads, from the
3698/// attention's one logit per head (`sinks`, `per_row` of them). Runs of elements with one head share a band.
3699fn spread_sinks(
3700    sinks: (&str, &Fill),
3701    per_row: usize,
3702    out: &str,
3703    n: usize,
3704    head_of: &dyn Fn(usize) -> Option<usize>,
3705    target: &SpatialTarget,
3706) -> Result<SpatialProgram, SpatialReject> {
3707    let mut f = FlatBuild::new(target)?;
3708    let sk = f.input(sinks.0, per_row, sinks.1, target);
3709    let cols_t = tiles_of(n, target);
3710    let out_buf = f.t.buffer("sinks", cols_t, f.precision);
3711    let mut slots: Vec<Slot> = Vec::new();
3712    let mut bands: Vec<(usize, TensorBody)> = Vec::new();
3713    let mut e = 0;
3714    while e < n {
3715        let head = head_of(e);
3716        let end = (e..n).find(|x| head_of(*x) != head).unwrap_or(n);
3717        let body = match head {
3718            Some(h) => {
3719                slots.push(Slot { buf: sk, map: fixed(h as i64) });
3720                op(slots.len() - 1)
3721            }
3722            None => konst(MASKED),
3723        };
3724        bands.push((end, *body));
3725        e = end;
3726    }
3727    if slots.is_empty() {
3728        slots.push(Slot { buf: sk, map: fixed(0) });
3729    }
3730    let stage = f.compute("spread_sinks", flat_space(n), slots, ident_map(wide(cols_t, target)), banded(bands), out_buf);
3731    let claim = NumericsClaim::Exact;
3732    f.finish(out, n, vec![stage], out_buf, format!("spread_sinks_{n}"), claim, target)
3733}
3734
3735/// GQA attention of a block of query rows over a cache of `l` rows: `score[j] = (q·k_j)/√hd`, row-softmax, then
3736/// the weighted V sum — THREE sections, because a softmax cannot be one `Generic` (its maximum and its
3737/// denominator are reductions whose results the next stage broadcasts back). One row is decode's single query; a
3738/// causal block of `T` rows is a prefill chunk's queries, row `r` attending one position past row `r − 1`.
3739pub fn from_attention(
3740    spec: &AttentionSpec,
3741    names: &StageNames,
3742    target: &SpatialTarget,
3743) -> Result<Vec<SpatialProgram>, SpatialReject> {
3744    let (h, kv, hd, l, t) = (spec.h, spec.kv, spec.hd, spec.l, spec.rows);
3745    if h == 0 || kv == 0 || hd == 0 || l == 0 || t == 0 {
3746        return Err(SpatialReject::NotLowered(
3747            "an attention with no heads, no query rows, no head dimension or an empty cache",
3748        ));
3749    }
3750    if spec.window.is_some_and(|w| w.span() == 0) {
3751        return Err(SpatialReject::NotLowered("a window of no positions — every row's softmax would be empty"));
3752    }
3753    // The sinks' operand follows `Q, K, V` and a computed position.
3754    let sinks = spec.sinks.map(|_| 3 + usize::from(matches!(spec.pos, Some(RowIndex::Named(_)))));
3755    let heads = QueryHeads {
3756        rows: t,
3757        per_row: h,
3758        group: (h / kv).max(1),
3759    };
3760    // The positions attended: the whole cache, a literal prefix of it, or — masked — the first `P + r` of query row
3761    // `r`, for a `P` the program computes (its operand follows `Q, K, V`) or a literal one under a block of rows.
3762    let (attended, limit) = match spec.pos {
3763        None if t > 1 => {
3764            return Err(SpatialReject::NotLowered(
3765                "a block of query rows with no causal bound",
3766            ))
3767        }
3768        None => (l, None),
3769        Some(RowIndex::Literal(p)) => {
3770            if p < 1 || p as usize + (t - 1) > l {
3771                return Err(SpatialReject::NotLowered("a literal attention length outside its cache — the program fails at run time on every tier, so there is no section to lower"));
3772            }
3773            let p = p as usize;
3774            if t == 1 {
3775                (p, None)
3776            } else {
3777                (p + (t - 1), Some(Bound::Literal(p)))
3778            }
3779        }
3780        Some(RowIndex::Named(_)) => (l, Some(Bound::Operand(3))),
3781    };
3782    // ★ A WINDOW KEEPS THE LAST `W` OF EACH ROW'S POSITIONS, so its rows are always masked: an unmasked decode row's bound
3783    // is the whole of what it attends.
3784    let limit = match (spec.window, limit) {
3785        (Some(_), None) => Some(Bound::Literal(attended)),
3786        (_, limit) => limit,
3787    };
3788    let tag = if t == 1 { "attn" } else { "causal_attn" };
3789    let scale = spec.scale.map_or(1.0 / (hd as f64).sqrt(), |s| s.get());
3790    // ★ A HEAD IN WHOLE TILES TAKES THE FLASH FORM — every real attention layer does. Query head `v` reads K/V head
3791    // `(v mod h) / group`, which is not affine in `v`, so the item carries both — `(v, K/V head)` — and every address
3792    // is affine in its own coordinate; the cache streams in blocks.
3793    if attention_reads_dense(h, kv, hd, spec.vd, target) {
3794        let g = FlashGeometry {
3795            q: heads,
3796            kvw: kv,
3797            hd,
3798            vd: spec.vd,
3799            v_part: None,
3800            rows: l,
3801            l: attended,
3802            limit,
3803            window: spec.window,
3804            sinks,
3805            grouped: false,
3806            scale,
3807            cap: spec.softcap.map(|c| c.get()),
3808        };
3809        // ★ ONE ROW'S QUERY HEADS THAT SHARE A K/V HEAD STREAM IT ONCE. A decode step's heads fold one item per K/V
3810        // head when the group fits a tile's rows and its regrouping fits a site; otherwise an item per query head.
3811        if t == 1 && heads.group > 1 && heads.group <= target.tile.rows && !row_needs_blocking(hd.max(spec.vd), target) {
3812            if let Ok(sections) = grouped_decode_attention(names, g, tag, target) {
3813                return Ok(sections);
3814            }
3815        }
3816        return flash_attention(names, g, tag, target);
3817    }
3818    attention_sections(
3819        names,
3820        heads,
3821        HeadWidths { kv, hd, vd: spec.vd, scale, cap: spec.softcap.map(|c| c.get()) },
3822        Cache {
3823            rows: l,
3824            attended,
3825            limit,
3826            window: spec.window,
3827        },
3828        sinks,
3829        tag,
3830        target,
3831    )
3832}
3833
3834/// The query heads an attention computes: `rows` query rows of `per_row` heads each, a row's heads reading K/V heads
3835/// in groups of `group`. Query head `v` is head `v mod per_row` of row `v / per_row`: its query and its output sit
3836/// at `v·hd` of the flat `rows × (per_row·hd)` block, and it reads K/V head `(v mod per_row) / group`.
3837///
3838/// ★ ONE SCHEME, EVERY ATTENTION. Decode is one row whose heads share K/V heads (GQA); a prefill chunk's causal
3839/// block is `T` such rows; full attention is `t` rows whose heads read their own K/V head (`group = 1`).
3840#[derive(Clone, Copy, Debug)]
3841struct QueryHeads {
3842    rows: usize,
3843    per_row: usize,
3844    group: usize,
3845}
3846
3847impl QueryHeads {
3848    fn count(&self) -> usize {
3849        self.rows * self.per_row
3850    }
3851    fn row_of(&self, v: usize) -> usize {
3852        v / self.per_row
3853    }
3854    fn kv_of(&self, v: usize) -> usize {
3855        (v % self.per_row) / self.group
3856    }
3857}
3858
3859/// The first row's live length `P`, for a bound the cache's extent does not fix: query row `r` attends the first
3860/// `P + r` positions.
3861#[derive(Clone, Copy, Debug)]
3862enum Bound {
3863    /// Fixed by the program, and checked against the cache when the attention is lowered.
3864    Literal(usize),
3865    /// Held by operand `i`, and checked at run time so every row stays inside the cache: `1 ≤ P ≤ rows − (T − 1)`.
3866    Operand(usize),
3867}
3868
3869/// The cache an attention reads: `rows` positions declared, the first `attended` of them covered, and — masked —
3870/// the bound each query row's live positions end at, and the window of positions before it a row keeps.
3871#[derive(Clone, Copy, Debug)]
3872struct Cache {
3873    rows: usize,
3874    attended: usize,
3875    limit: Option<Bound>,
3876    window: Option<Window>,
3877}
3878
3879/// The largest `P` a block of `q_rows` query rows admits over a cache of `rows`: its last row's positions all cached.
3880fn bound_max(rows: usize, q_rows: usize) -> usize {
3881    rows + 1 - q_rows
3882}
3883
3884/// The widths of an attention's heads: `kv` K/V heads per position, keys `hd` and values `vd` wide, and the factor
3885/// every score takes.
3886#[derive(Clone, Copy, Debug)]
3887struct HeadWidths {
3888    kv: usize,
3889    hd: usize,
3890    vd: usize,
3891    scale: f64,
3892    /// The soft cap every scaled score takes before the mask, `c·tanh(s/c)`.
3893    cap: Option<f64>,
3894}
3895
3896/// The three sections every unmasked attention is: scores, the row softmax, the weighted V sum — over the
3897/// [`QueryHeads`], each attending over `l` positions of a K holding `kv` heads of `hd` per position
3898/// (`k[j·(kv·hd) + kvh·hd + e]`) and a V holding `kv` heads of `vd` (`v[j·(kv·vd) + kvh·vd + e]`).
3899fn attention_sections(
3900    names: &StageNames,
3901    heads: QueryHeads,
3902    widths: HeadWidths,
3903    cache: Cache,
3904    sinks: Option<usize>,
3905    tag: &str,
3906    target: &SpatialTarget,
3907) -> Result<Vec<SpatialProgram>, SpatialReject> {
3908    let HeadWidths { kv, hd, vd, scale, cap } = widths;
3909    let h = heads.count();
3910    let kv_of = |v: usize| heads.kv_of(v);
3911    let l = cache.attended;
3912    let kv_len = cache.rows * kv * hd;
3913    let v_len = cache.rows * kv * vd;
3914    let (q_name, q_fill) = names.operand(0);
3915    let (k_name, k_fill) = names.operand(1);
3916    let (v_name, v_fill) = names.operand(2);
3917    let scores_name = format!("{}_scores", names.out);
3918    let probs_name = format!("{}_probs", names.out);
3919
3920    // ── Section 1: the scores, `h · l` of them, flat as `1 × (h·l)`.
3921    let score = {
3922        let mut f = FlatBuild::new(target)?;
3923        let q = f.input(q_name, h * hd, q_fill, target);
3924        let k = f.input(k_name, kv_len, k_fill, target);
3925        let n = h * l;
3926        let cols_t = tiles_of(n, target);
3927        let out_buf = f.t.buffer("scores", cols_t, f.precision);
3928        let ow = wide(cols_t, target);
3929        let tile_elems = target.tile.rows * target.tile.cols;
3930        let ones = f.t.tensor(
3931            "at_scaler",
3932            TensorRole::Param,
3933            target.tile.rows,
3934            target.tile.cols,
3935            f.precision,
3936            TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
3937        );
3938        let ones_in = f.t.buffer("in_scaler", 1, f.precision);
3939        let ls = f.t.load("load_scaler", ones, vec![], 1, ones_in);
3940        f.loads.push((ls, ones_in));
3941        // The folded axis is the ROW (`e` over the head dimension); the flat output column is
3942        // `c = head·l + j`, so the scores come out as a flat `1 × (h·l)` row.
3943        let iters = vec![
3944            IterSpec {
3945                var: LV_GEN_ROW,
3946                extent: hd,
3947                kind: IterKind::Reduce {
3948                    pool: ReducePool::Sum,
3949                },
3950            },
3951            IterSpec {
3952                var: LV_GEN_COL,
3953                extent: n,
3954                kind: IterKind::Parallel,
3955            },
3956        ];
3957        let mut slots = vec![Slot {
3958            buf: ones_in,
3959            map: const_map(),
3960        }];
3961        let mut bands = Vec::new();
3962        for head in 0..h {
3963            let kvh = kv_of(head);
3964            // `q[head·hd + e]`: constant along the output column, affine in the folded row.
3965            slots.push(Slot {
3966                buf: q,
3967                map: ElemAccess {
3968                    base: (head * hd) as i64,
3969                    terms: vec![(LV_GEN_ROW, 1)],
3970                    indirect: None,
3971                },
3972            });
3973            let q_slot = slots.len() - 1;
3974            // `kc[j·(kv·hd) + kvh·hd + e]` with `j = c − head·l`.
3975            let j_stride = (kv * hd) as i64;
3976            let base = (kvh * hd) as i64 - (head * l) as i64 * j_stride;
3977            slots.push(Slot {
3978                buf: k,
3979                map: ElemAccess {
3980                    base,
3981                    terms: vec![(LV_GEN_COL, j_stride), (LV_GEN_ROW, 1)],
3982                    indirect: None,
3983                },
3984            });
3985            let k_slot = slots.len() - 1;
3986            bands.push((
3987                (head + 1) * l,
3988                *mul(mul(op(q_slot), op(k_slot)), konst(scale)),
3989            ));
3990        }
3991        let mut chain = Vec::new();
3992        let scored = match (cache.limit, cap) {
3993            (None, None) => out_buf,
3994            _ => f.scratch("raw_scores", cols_t),
3995        };
3996        chain.push(f.compute(
3997            "attn_scores",
3998            iters,
3999            slots,
4000            per_col_map(),
4001            banded(bands),
4002            scored,
4003        ));
4004        // ★ A SOFT CAP takes every scaled score before the mask, so a masked score stays masked.
4005        let scored = match cap {
4006            None => scored,
4007            Some(c) => {
4008                let to = if cache.limit.is_some() { f.scratch("capped_scores", cols_t) } else { out_buf };
4009                chain.push(f.compute("attn_cap", flat_space(n), vec![Slot { buf: scored, map: at(0) }], ident_map(ow), *capped(op(0), c), to));
4010                to
4011            }
4012        };
4013        if let Some(bound) = cache.limit {
4014            // ★ THE MASK. Head `v`'s position `j` sits at column `v·l + j`, so it is live iff the column is below
4015            // `v·l + P + r`, `r` the head's query row: one threshold per head — a constant for a literal `P`, else built
4016            // from `P` once `P` is checked to keep every row inside the cache.
4017            let row_base = |v: usize| v * l + heads.row_of(v);
4018            let mut slots = vec![Slot {
4019                buf: scored,
4020                map: at(0),
4021            }];
4022            let below: Vec<PredRhs> = match bound {
4023                Bound::Literal(p) => (0..h).map(|v| PredRhs::Const(row_base(v) + p)).collect(),
4024                Bound::Operand(li) => {
4025                    let (pn, pf) = names.operand(li);
4026                    let ix = index_type(target, cache.rows as u64)?;
4027                    let p_in = f.input_as(pn, 1, pf, ix, target);
4028                    let (mut guard, pc) =
4029                        checked_count(&mut f, p_in, bound_max(cache.rows, heads.rows), ix, target);
4030                    chain.append(&mut guard);
4031                    let thr = f.t.scratch_deep("thresholds", 1, tiles_of(h, target), ix);
4032                    let bands: Vec<(usize, TensorBody)> = (0..h)
4033                        .map(|v| (v + 1, *add(konst(row_base(v) as f64), op(0))))
4034                        .collect();
4035                    chain.push(f.compute(
4036                        "thresholds",
4037                        flat_space(h),
4038                        vec![Slot {
4039                            buf: pc,
4040                            map: fixed(0),
4041                        }],
4042                        ident_map(wide(tiles_of(h, target), target)),
4043                        banded(bands),
4044                        thr,
4045                    ));
4046                    (0..h)
4047                        .map(|v| {
4048                            slots.push(Slot {
4049                                buf: thr,
4050                                map: fixed(v as i64),
4051                            });
4052                            PredRhs::Operand(slots.len() - 1)
4053                        })
4054                        .collect()
4055                }
4056            };
4057            // A window keeps the last `W` positions below each threshold — the column also at or past it less `W` — and a
4058            // chunk of `C` the positions of the row's own chunk: head `v` attending `n` keeps from column
4059            // `v·l + ⌊(n − 1)/C⌋·C`, `n` its threshold less `v·l`.
4060            let floors: Vec<Option<PredRhs>> = match (cache.window, &bound) {
4061                (None, _) => vec![None; h],
4062                (Some(w), Bound::Literal(p)) => (0..h)
4063                    .map(|v| {
4064                        let floor = v * l + w.first(p + heads.row_of(v)) - 1;
4065                        (floor > v * l).then_some(PredRhs::Const(floor))
4066                    })
4067                    .collect(),
4068                (Some(Window::Chunked(c)), Bound::Operand(_)) => {
4069                    let thr = match below[0] {
4070                        PredRhs::Operand(i) => slots[i].buf,
4071                        PredRhs::Const(_) => unreachable!("a computed bound's thresholds are an operand"),
4072                    };
4073                    let ix = index_type(target, cache.rows as u64)?;
4074                    let floor = f.t.scratch_deep("floors", 1, tiles_of(h, target), ix);
4075                    let c = c as f64;
4076                    let bands: Vec<(usize, TensorBody)> = (0..h)
4077                        .map(|v| {
4078                            let base = (v * l) as f64;
4079                            let n = sub(op(0), konst(base));
4080                            (v + 1, *add(konst(base), mul(Box::new(TensorBody::FloorDiv(sub(n, konst(1.0)), konst(c))), konst(c))))
4081                        })
4082                        .collect();
4083                    chain.push(f.compute(
4084                        "floors",
4085                        flat_space(h),
4086                        vec![Slot { buf: thr, map: ident_map(wide(tiles_of(h, target), target)) }],
4087                        ident_map(wide(tiles_of(h, target), target)),
4088                        banded(bands),
4089                        floor,
4090                    ));
4091                    (0..h)
4092                        .map(|v| {
4093                            slots.push(Slot { buf: floor, map: fixed(v as i64) });
4094                            Some(PredRhs::Operand(slots.len() - 1))
4095                        })
4096                        .collect()
4097                }
4098                (Some(Window::Sliding(w)), Bound::Operand(_)) => {
4099                    let thr = match below[0] {
4100                        PredRhs::Operand(i) => slots[i].buf,
4101                        PredRhs::Const(_) => unreachable!("a computed bound's thresholds are an operand"),
4102                    };
4103                    let ix = index_type(target, cache.rows as u64)?;
4104                    let floor = f.t.scratch_deep("floors", 1, tiles_of(h, target), ix);
4105                    chain.push(f.compute(
4106                        "floors",
4107                        flat_space(h),
4108                        vec![Slot { buf: thr, map: ident_map(wide(tiles_of(h, target), target)) }],
4109                        ident_map(wide(tiles_of(h, target), target)),
4110                        *sub(op(0), konst(w as f64)),
4111                        floor,
4112                    ));
4113                    (0..h)
4114                        .map(|v| {
4115                            slots.push(Slot { buf: floor, map: fixed(v as i64) });
4116                            Some(PredRhs::Operand(slots.len() - 1))
4117                        })
4118                        .collect()
4119                }
4120            };
4121            let bands: Vec<(usize, TensorBody)> = below
4122                .into_iter()
4123                .zip(floors)
4124                .enumerate()
4125                .map(|(v, (rhs, floor))| {
4126                    let kept = match floor {
4127                        Some(floor) => Box::new(TensorBody::Select {
4128                            cond: Pred::AxisGe(LV_GEN_COL, floor),
4129                            then: op(0),
4130                            otherwise: konst(MASKED),
4131                        }),
4132                        None => op(0),
4133                    };
4134                    (
4135                        (v + 1) * l,
4136                        TensorBody::Select {
4137                            cond: Pred::AxisLt(LV_GEN_COL, rhs),
4138                            then: kept,
4139                            otherwise: konst(MASKED),
4140                        },
4141                    )
4142                })
4143                .collect();
4144            chain.push(f.compute(
4145                "mask",
4146                flat_space(n),
4147                slots,
4148                ident_map(ow),
4149                banded(bands),
4150                out_buf,
4151            ));
4152        }
4153        let claim = approx("attention scores", f.precision);
4154        f.finish(
4155            &scores_name,
4156            n,
4157            chain,
4158            out_buf,
4159            format!("{tag}_scores_{h}x{l}"),
4160            claim,
4161            target,
4162        )?
4163    };
4164
4165    // ── Section 2: the row softmax over each head's `l` scores, each head's sink in its denominator. A decode row's
4166    // softmax rows are its heads; a causal block's rows repeat them, so the sinks are first spread over the rows.
4167    let sink_rows_name = format!("{}__sink_rows", names.out);
4168    let (sink_rows, spread) = match sinks.map(|i| names.operand(i)) {
4169        None => (None, None),
4170        Some(sk) if heads.rows == 1 => (Some(sk), None),
4171        Some(sk) => {
4172            let spread = spread_sinks(sk, heads.per_row, &sink_rows_name, h, &|v| Some(v % heads.per_row), target)?;
4173            (Some((sink_rows_name.as_str(), &Fill::Unknown)), Some(spread))
4174        }
4175    };
4176    let softmax = flat_softmax_rows_at(
4177        (&scores_name, &Fill::Unknown),
4178        &probs_name,
4179        h,
4180        l,
4181        format!("{tag}_softmax_{h}x{l}"),
4182        elem_precision(target)?,
4183        sink_rows,
4184        target,
4185    )?;
4186
4187    // ── Section 3: the weighted V sum, `ctx[head·vd + e] = Σ_j p[head·l + j]·vc[j·(kv·vd) + kvh·vd + e]`.
4188    let weighted = {
4189        let mut f = FlatBuild::new(target)?;
4190        let p = f.input(&probs_name, h * l, &Fill::Unknown, target);
4191        let v = f.input(v_name, v_len, v_fill, target);
4192        let n = h * vd;
4193        let cols_t = tiles_of(n, target);
4194        let out_buf = f.t.buffer("out", cols_t, f.precision);
4195        let tile_elems = target.tile.rows * target.tile.cols;
4196        let ones = f.t.tensor(
4197            "wv_scaler",
4198            TensorRole::Param,
4199            target.tile.rows,
4200            target.tile.cols,
4201            f.precision,
4202            TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
4203        );
4204        let ones_in = f.t.buffer("in_scaler", 1, f.precision);
4205        let ls = f.t.load("load_scaler", ones, vec![], 1, ones_in);
4206        f.loads.push((ls, ones_in));
4207        // The folded axis is the ROW (`j` over the cache); the flat output column is `c = head·vd + e`.
4208        let iters = vec![
4209            IterSpec {
4210                var: LV_GEN_ROW,
4211                extent: l,
4212                kind: IterKind::Reduce {
4213                    pool: ReducePool::Sum,
4214                },
4215            },
4216            IterSpec {
4217                var: LV_GEN_COL,
4218                extent: n,
4219                kind: IterKind::Parallel,
4220            },
4221        ];
4222        let mut slots = vec![Slot {
4223            buf: ones_in,
4224            map: const_map(),
4225        }];
4226        let mut bands = Vec::new();
4227        for head in 0..h {
4228            let kvh = kv_of(head);
4229            // `p[head·l + j]`: constant along the output column, affine in the folded row.
4230            slots.push(Slot {
4231                buf: p,
4232                map: ElemAccess {
4233                    base: (head * l) as i64,
4234                    terms: vec![(LV_GEN_ROW, 1)],
4235                    indirect: None,
4236                },
4237            });
4238            let p_slot = slots.len() - 1;
4239            // `vc[j·(kv·vd) + kvh·vd + e]` with `e = c − head·vd`.
4240            let base = (kvh * vd) as i64 - (head * vd) as i64;
4241            slots.push(Slot {
4242                buf: v,
4243                map: ElemAccess {
4244                    base,
4245                    terms: vec![(LV_GEN_ROW, (kv * vd) as i64), (LV_GEN_COL, 1)],
4246                    indirect: None,
4247                },
4248            });
4249            let v_slot = slots.len() - 1;
4250            bands.push(((head + 1) * vd, *mul(op(p_slot), op(v_slot))));
4251        }
4252        let s = f.compute(
4253            "attn_weighted",
4254            iters,
4255            slots,
4256            per_col_map(),
4257            banded(bands),
4258            out_buf,
4259        );
4260        let claim = approx("attention weighted sum", f.precision);
4261        f.finish(
4262            &names.out,
4263            n,
4264            vec![s],
4265            out_buf,
4266            format!("{tag}_wv_{h}x{vd}"),
4267            claim,
4268            target,
4269        )?
4270    };
4271    Ok([score].into_iter().chain(spread).chain([softmax, weighted]).collect())
4272}
4273
4274/// `C = lut attention(Q, K, V)` — an int8 model's single-query attention, integer throughout, in THREE sections: the
4275/// scores `s = q·k_j` over the cache; per head, the weights `w_j = E[clamp(((max − s_j)·M) >> S, 0, n − 1)]` for the
4276/// positions attended and 0 past them; and the context `⌊Σ_j w_j·v_j / Σ_j w_j⌋`. A masked position is fenced off
4277/// only where it could change an answer — the maximum and the weights — and the sums need no mask: a weight of 0 adds
4278/// nothing. Every value the masked step computes is one the unmasked step could: a masked position folds into the
4279/// maximum as the head's first score, which every `P ≥ 1` attends, so no sentinel widens what the maximum holds.
4280pub fn from_lut_attention(
4281    spec: &LutAttentionSpec,
4282    names: &StageNames,
4283    target: &SpatialTarget,
4284) -> Result<Vec<SpatialProgram>, SpatialReject> {
4285    let (h, kv, hd, l, n) = (spec.h, spec.kv, spec.hd, spec.l, spec.table_len);
4286    if h == 0 || kv == 0 || hd == 0 || l == 0 || n == 0 || h % kv != 0 {
4287        return Err(SpatialReject::NotLowered(
4288            "a LUT attention with no heads, no head dimension, no cache or no table, or query heads that do not group over its K/V heads",
4289        ));
4290    }
4291    let group = h / kv;
4292    // The positions attended: the whole cache, a literal prefix of it, or — masked — the first `P` for a `P` the
4293    // program computes (its operand follows `Q, K, V, E`).
4294    let (attended, bound) = match spec.pos {
4295        None => (l, None),
4296        Some(RowIndex::Literal(p)) => {
4297            if p < 1 || p as usize > l {
4298                return Err(SpatialReject::NotLowered("a literal attention length outside its cache — the program fails at run time on every tier, so there is no section to lower"));
4299            }
4300            (p as usize, None)
4301        }
4302        Some(RowIndex::Named(_)) => (l, Some(4usize)),
4303    };
4304    let (q_name, q_fill) = names.operand(0);
4305    let (k_name, k_fill) = names.operand(1);
4306    let (v_name, v_fill) = names.operand(2);
4307    let (e_name, e_fill) = names.operand(3);
4308    let scores_name = format!("{}_scores", names.out);
4309    let weights_name = format!("{}_weights", names.out);
4310    let kv_len = l * kv * hd;
4311    let nl = h * attended;
4312    let tag = format!("lut_attn_{h}x{attended}");
4313
4314    // ── Section 1: the scores, `h · attended` of them, flat: `s[head·attended + j] = Σ_e q[head·hd + e]·k[j, kvh, e]`.
4315    let score = {
4316        let mut f = FlatBuild::on(Datapath::Int, target)?;
4317        let q = f.input(q_name, h * hd, q_fill, target);
4318        let k = f.input(k_name, kv_len, k_fill, target);
4319        let out_buf = f.t.buffer("scores", tiles_of(nl, target), f.precision);
4320        let iters = vec![
4321            IterSpec {
4322                var: LV_GEN_ROW,
4323                extent: hd,
4324                kind: IterKind::Reduce {
4325                    pool: ReducePool::Sum,
4326                },
4327            },
4328            IterSpec {
4329                var: LV_GEN_COL,
4330                extent: nl,
4331                kind: IterKind::Parallel,
4332            },
4333        ];
4334        let mut slots = Vec::new();
4335        let mut bands = Vec::new();
4336        let j_stride = (kv * hd) as i64;
4337        for head in 0..h {
4338            let kvh = head / group;
4339            slots.push(Slot {
4340                buf: q,
4341                map: ElemAccess {
4342                    base: (head * hd) as i64,
4343                    terms: vec![(LV_GEN_ROW, 1)],
4344                    indirect: None,
4345                },
4346            });
4347            let q_slot = slots.len() - 1;
4348            slots.push(Slot {
4349                buf: k,
4350                map: ElemAccess {
4351                    base: (kvh * hd) as i64 - (head * attended) as i64 * j_stride,
4352                    terms: vec![(LV_GEN_COL, j_stride), (LV_GEN_ROW, 1)],
4353                    indirect: None,
4354                },
4355            });
4356            let k_slot = slots.len() - 1;
4357            bands.push(((head + 1) * attended, *mul(op(q_slot), op(k_slot))));
4358        }
4359        let s = f.compute("lut_scores", iters, slots, per_col_map(), banded(bands), out_buf);
4360        f.finish(&scores_name, nl, vec![s], out_buf, format!("{tag}_scores"), NumericsClaim::Exact, target)?
4361    };
4362
4363    // ── Section 2: per head, the row maximum over the attended positions, the LUT index, and the weight.
4364    let weights = {
4365        let mut f = FlatBuild::on(Datapath::Int, target)?;
4366        let sc = f.input(&scores_name, nl, &Fill::Unknown, target);
4367        let e = f.input(e_name, n, e_fill, target);
4368        let one_w = wide(1, target);
4369        let flat_w = wide(tiles_of(nl, target), target);
4370        let mut chain = Vec::new();
4371        // When the bound is the program's: `P` checked to `1 ≤ P ≤ L` (or the run stops), and head `head`'s
4372        // threshold `head·attended + P` — its position `j` sits at flat column `head·attended + j`, live iff below it.
4373        let masked = match bound {
4374            Some(li) => {
4375                let (pn, pf) = names.operand(li);
4376                let ix = index_type(target, l as u64)?;
4377                let p_in = f.input_as(pn, 1, pf, ix, target);
4378                let (mut guard, pc) = checked_count(&mut f, p_in, l, ix, target);
4379                chain.append(&mut guard);
4380                let thr = f.t.scratch_deep("thresholds", 1, tiles_of(h, target), ix);
4381                let bands: Vec<(usize, TensorBody)> = (0..h)
4382                    .map(|head| (head + 1, *add(konst((head * attended) as f64), op(0))))
4383                    .collect();
4384                chain.push(f.compute(
4385                    "thresholds",
4386                    flat_space(h),
4387                    vec![Slot { buf: pc, map: fixed(0) }],
4388                    ident_map(wide(tiles_of(h, target), target)),
4389                    banded(bands),
4390                    thr,
4391                ));
4392                Some((pc, thr))
4393            }
4394            None => None,
4395        };
4396        // A body that holds at a live position and is `otherwise` past `P`, against the bound in slot `bound_slot`.
4397        let fenced = |bound_slot: Option<usize>, then: Box<TensorBody>, otherwise: Box<TensorBody>| -> TensorBody {
4398            match bound_slot {
4399                Some(bs) => TensorBody::Select {
4400                    cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(bs)),
4401                    then,
4402                    otherwise,
4403                },
4404                None => *then,
4405            }
4406        };
4407        // The row maximum: ONE FOLD PER HEAD, each a single row over its own slice of the flat scores — a page holds
4408        // `t` rows, and on a scalar-tile part that is one. The fold's column IS `j`, read against `P` itself.
4409        let mut rmax = Vec::with_capacity(h);
4410        for head in 0..h {
4411            let m = f.scratch(&format!("row_max_{head}"), 1);
4412            let mut slots = vec![Slot {
4413                buf: sc,
4414                map: at((head * attended) as i64),
4415            }];
4416            let p_slot = masked.map(|(pc, _)| {
4417                slots.push(Slot { buf: pc, map: fixed(0) });
4418                slots.len() - 1
4419            });
4420            // Past `P` the fold takes the head's first score again: attended at every `P ≥ 1`, never above the maximum.
4421            let first = masked.map(|_| {
4422                slots.push(Slot {
4423                    buf: sc,
4424                    map: fixed((head * attended) as i64),
4425                });
4426                slots.len() - 1
4427            });
4428            chain.push(f.compute(
4429                &format!("row_max_{head}"),
4430                head_space(1, attended, Some(ReducePool::Max)),
4431                slots,
4432                per_row_map(one_w),
4433                fenced(p_slot, op(0), first.map_or_else(|| konst(0.0), op)),
4434                m,
4435            ));
4436            rmax.push(m);
4437        }
4438        // `idx[head·attended + j] = clamp(((max_head − s_j)·M) >> S, 0, n − 1)`, 0 past `P`.
4439        let idx = f.scratch("lut_index", tiles_of(nl, target));
4440        let mut slots = vec![Slot { buf: sc, map: at(0) }];
4441        let mut bands = Vec::new();
4442        for (head, &m) in rmax.iter().enumerate() {
4443            slots.push(Slot { buf: m, map: fixed(0) });
4444            let mx = slots.len() - 1;
4445            let thr_slot = masked.map(|(_, thr)| {
4446                slots.push(Slot {
4447                    buf: thr,
4448                    map: fixed(head as i64),
4449                });
4450                slots.len() - 1
4451            });
4452            let index = Box::new(TensorBody::Min(
4453                Box::new(TensorBody::Max(
4454                    Box::new(TensorBody::Shr {
4455                        arg: mul(sub(op(mx), op(0)), konst(spec.m as f64)),
4456                        by: spec.s,
4457                    }),
4458                    konst(0.0),
4459                )),
4460                konst((n - 1) as f64),
4461            ));
4462            bands.push(((head + 1) * attended, fenced(thr_slot, index, konst(0.0))));
4463        }
4464        chain.push(f.compute(
4465            "lut_index",
4466            flat_space(nl),
4467            slots,
4468            ident_map(flat_w),
4469            banded(bands),
4470            idx,
4471        ));
4472        // `w = E[idx]` at a live position, 0 past `P`: a weight of 0 adds nothing to either sum.
4473        let out_buf = f.t.buffer("weights", tiles_of(nl, target), f.precision);
4474        let mut slots = vec![
4475            Slot { buf: idx, map: at(0) },
4476            Slot {
4477                buf: e,
4478                map: ElemAccess {
4479                    base: 0,
4480                    terms: Vec::new(),
4481                    indirect: Some(IndirectTerm {
4482                        operand: 0,
4483                        stride: 1,
4484                        modulus: n,
4485                    }),
4486                },
4487            },
4488        ];
4489        let mut bands = Vec::new();
4490        for head in 0..h {
4491            let thr_slot = masked.map(|(_, thr)| {
4492                slots.push(Slot {
4493                    buf: thr,
4494                    map: fixed(head as i64),
4495                });
4496                slots.len() - 1
4497            });
4498            bands.push(((head + 1) * attended, fenced(thr_slot, op(1), konst(0.0))));
4499        }
4500        chain.push(f.compute(
4501            "lut_weight",
4502            flat_space(nl),
4503            slots,
4504            ident_map(flat_w),
4505            banded(bands),
4506            out_buf,
4507        ));
4508        f.finish(&weights_name, nl, chain, out_buf, format!("{tag}_weights"), NumericsClaim::Exact, target)?
4509    };
4510
4511    // ── Section 3: the context, `ctx[head·hd + e] = ⌊Σ_j w[head·attended + j]·v[j, kvh, e] / Σ_j w[head·attended + j]⌋`.
4512    let context = {
4513        let mut f = FlatBuild::on(Datapath::Int, target)?;
4514        let w = f.input(&weights_name, nl, &Fill::Unknown, target);
4515        let v = f.input(v_name, kv_len, v_fill, target);
4516        let one_w = wide(1, target);
4517        let nh = h * hd;
4518        let mut chain = Vec::new();
4519        // One denominator fold per head, as the maxima are.
4520        let mut den = Vec::with_capacity(h);
4521        for head in 0..h {
4522            let d = f.scratch(&format!("den_{head}"), 1);
4523            chain.push(f.compute(
4524                &format!("den_{head}"),
4525                head_space(1, attended, Some(ReducePool::Sum)),
4526                vec![Slot {
4527                    buf: w,
4528                    map: at((head * attended) as i64),
4529                }],
4530                per_row_map(one_w),
4531                TensorBody::Operand(0),
4532                d,
4533            ));
4534            den.push(d);
4535        }
4536        let num = f.scratch("num", tiles_of(nh, target));
4537        let iters = vec![
4538            IterSpec {
4539                var: LV_GEN_ROW,
4540                extent: attended,
4541                kind: IterKind::Reduce {
4542                    pool: ReducePool::Sum,
4543                },
4544            },
4545            IterSpec {
4546                var: LV_GEN_COL,
4547                extent: nh,
4548                kind: IterKind::Parallel,
4549            },
4550        ];
4551        let mut slots = Vec::new();
4552        let mut bands = Vec::new();
4553        for head in 0..h {
4554            let kvh = head / group;
4555            slots.push(Slot {
4556                buf: w,
4557                map: ElemAccess {
4558                    base: (head * attended) as i64,
4559                    terms: vec![(LV_GEN_ROW, 1)],
4560                    indirect: None,
4561                },
4562            });
4563            let w_slot = slots.len() - 1;
4564            slots.push(Slot {
4565                buf: v,
4566                map: ElemAccess {
4567                    base: (kvh * hd) as i64 - (head * hd) as i64,
4568                    terms: vec![(LV_GEN_ROW, (kv * hd) as i64), (LV_GEN_COL, 1)],
4569                    indirect: None,
4570                },
4571            });
4572            let v_slot = slots.len() - 1;
4573            bands.push(((head + 1) * hd, *mul(op(w_slot), op(v_slot))));
4574        }
4575        chain.push(f.compute("num", iters, slots, per_col_map(), banded(bands), num));
4576        let out_buf = f.t.buffer("out", tiles_of(nh, target), f.precision);
4577        let mut slots = vec![Slot { buf: num, map: at(0) }];
4578        let mut bands = Vec::new();
4579        for (head, &d) in den.iter().enumerate() {
4580            slots.push(Slot { buf: d, map: fixed(0) });
4581            let ds = slots.len() - 1;
4582            bands.push(((head + 1) * hd, TensorBody::FloorDiv(op(0), op(ds))));
4583        }
4584        chain.push(f.compute(
4585            "floor_div",
4586            flat_space(nh),
4587            slots,
4588            ident_map(wide(tiles_of(nh, target), target)),
4589            banded(bands),
4590            out_buf,
4591        ));
4592        f.finish(&names.out, nh, chain, out_buf, format!("{tag}_context"), NumericsClaim::Exact, target)?
4593    };
4594    Ok(vec![score, weights, context])
4595}
4596
4597/// `C = full attention(Q, K, V)` over `t` positions and `h` heads — `t` query rows of `h` heads, each reading its
4598/// own K/V head and attending over all `t` positions. Unmasked, as the vision tower's block is: every row sees
4599/// every position.
4600pub fn from_full_attention(
4601    spec: &FullAttentionSpec,
4602    names: &StageNames,
4603    target: &SpatialTarget,
4604) -> Result<Vec<SpatialProgram>, SpatialReject> {
4605    let (t, h, hd) = (spec.t, spec.h, spec.hd);
4606    if t == 0 || h == 0 || hd == 0 {
4607        return Err(SpatialReject::NotLowered(
4608            "a full attention with no positions, no heads or no head dimension",
4609        ));
4610    }
4611    let heads = QueryHeads {
4612        rows: t,
4613        per_row: h,
4614        group: 1,
4615    };
4616    if full_attention_reads_dense(hd, target) {
4617        return flash_attention(
4618            names,
4619            FlashGeometry {
4620                q: heads,
4621                kvw: h,
4622                hd,
4623                vd: hd,
4624                v_part: None,
4625                rows: t,
4626                l: t,
4627                limit: None,
4628                window: None,
4629                sinks: None,
4630                grouped: false,
4631                scale: 1.0 / (hd as f64).sqrt(),
4632                cap: None,
4633            },
4634            "full_attn",
4635            target,
4636        );
4637    }
4638    attention_sections(
4639        names,
4640        heads,
4641        HeadWidths { kv: h, hd, vd: hd, scale: 1.0 / (hd as f64).sqrt(), cap: None },
4642        Cache {
4643            rows: t,
4644            attended: t,
4645            limit: None,
4646            window: None,
4647        },
4648        None,
4649        "full_attn",
4650        target,
4651    )
4652}
4653
4654/// The plan of a relayout: one item per `blk`-tile block of each of `tiles_r` dense tile rows.
4655fn relayout_plan(tiles_r: usize, row_t: usize, blk: usize, dt: Precision) -> BlockPlan {
4656    BlockPlan {
4657        grid: (tiles_r, row_t / blk),
4658        ..tile_row_plan(row_t, blk, dt)
4659    }
4660}
4661
4662fn relayout_items(tiles_r: usize, row_t: usize, blk: usize) -> Vec<WorkItem> {
4663    (0..tiles_r)
4664        .flat_map(|r| (0..row_t / blk).map(move |b| (r, b)))
4665        .enumerate()
4666        .map(|(id, (r, b))| WorkItem {
4667            id: WorkItemId::new(id),
4668            coord: (r, b),
4669            out: TileRect {
4670                row0_t: r,
4671                col0_t: b * blk,
4672                h_t: 1,
4673                w_t: blk,
4674            },
4675        })
4676        .collect()
4677}
4678
4679/// Whether a `rows × cols` block moves between flat rows and a dense matrix a tile row at a time: whole tiles across.
4680fn relayout_tiles(cols: usize, target: &SpatialTarget) -> bool {
4681    cols % target.tile.cols == 0
4682}
4683
4684/// `to` — a DENSE `rows × cols` matrix — from `from`, the same matrix as flat rows.
4685///
4686/// ★ A PAGE IS A FLAT ROW-MAJOR PLANE, SO A RELAYOUT IS AN AFFINE MAP. Item `(R, b)` loads the flat tiles dense tile
4687/// row `R` holds — its `h` rows' `blk`-tile segments, `row_t` tiles apart — as ONE page of `h·blk` tiles side by
4688/// side, and writes `out(y, c) = in[y·(blk·tile) + c]`: row `y` of the tile row is the page's `y`-th segment. Every
4689/// item of a section loads the same `h` rows, so rows that are not whole tile rows take TWO sections — the whole tile
4690/// rows, then the last tile row's `rows mod tile` rows — each writing its own tiles of `to` (a part of it, so each
4691/// declares it updated in part). A row that is not whole tiles has no segments of its own; then the block must be one
4692/// small tile row, loaded whole — `out(y, c) = in[y·cols + c]` below `rows` and `cols`, zero past them.
4693pub(crate) fn relayout_rows_to_dense(
4694    from: (&str, &Fill),
4695    to: &str,
4696    rows: usize,
4697    cols: usize,
4698    dt: Precision,
4699    target: &SpatialTarget,
4700) -> Result<Vec<SpatialProgram>, SpatialReject> {
4701    let (tr, tw) = (target.tile.rows, target.tile.cols);
4702    let row_t = cols.div_ceil(tw);
4703    let below = |axis: LoopVar, bound: usize, then: Box<TensorBody>| TensorBody::Select {
4704        cond: Pred::AxisLt(axis, PredRhs::Const(bound)),
4705        then,
4706        otherwise: konst(0.0),
4707    };
4708    let name = format!("relayout_dense_{rows}x{cols}");
4709    let declare = |t: &mut Template, role: TensorRole| {
4710        let src = t.tensor(
4711            from.0,
4712            TensorRole::Input,
4713            1,
4714            rows * cols,
4715            dt,
4716            super::from_tensor_op::init_of(from.1, rows * cols),
4717        );
4718        (src, t.tensor(to, role, rows, cols, dt, TensorInit::Zero))
4719    };
4720    if relayout_tiles(cols, target) {
4721        // Rows `first … first + n − 1`, a whole number of tile rows or the last tile row's rest, as one section; `part`
4722        // when the section writes only part of `to`.
4723        let rows_section =
4724            |first: usize, n: usize, part: bool| -> Result<SpatialProgram, SpatialReject> {
4725                let h = n.min(tr);
4726                let tiles_r = n.div_ceil(tr);
4727                let r0 = first / tr;
4728                // A page holds `h` rows' segments of `blk` tiles each: the widest block of columns the site holds — built
4729                // and measured, so the sizing cannot disagree with the section it sizes.
4730                let tiled = |blk: usize| -> Result<SpatialProgram, SpatialReject> {
4731                    let mut t = Template::new(target)?;
4732                    let (src, dst) = declare(
4733                        &mut t,
4734                        if part {
4735                            TensorRole::State
4736                        } else {
4737                            TensorRole::Output
4738                        },
4739                    );
4740                    let w = wide(blk, target);
4741                    let inp = t.buffer("in_rows", h * blk, dt);
4742                    let out = t.buffer("out", blk, dt);
4743                    let load = t.load_block_at(
4744                        "load_rows",
4745                        src,
4746                        first * row_t,
4747                        vec![(LV_ITEM_Y, tr * row_t), (LV_ITEM_X, blk)],
4748                        h,
4749                        blk,
4750                        row_t,
4751                        inp,
4752                    );
4753                    let s = t.stage(
4754                        "relayout",
4755                        StageKind::Compute(ComputeOp::Generic {
4756                            iters: super::from_tensor_op::space(target, w, None),
4757                            ins: vec![ident_map(w)],
4758                            out: ident_map(w),
4759                            body: below(LV_GEN_ROW, h, op(0)),
4760                        }),
4761                        vec![inp],
4762                        vec![out],
4763                    );
4764                    let st = t.store_block(
4765                        "store_dense",
4766                        dst,
4767                        r0 * row_t,
4768                        vec![(LV_ITEM_Y, row_t), (LV_ITEM_X, blk)],
4769                        1,
4770                        blk,
4771                        row_t,
4772                        out,
4773                    );
4774                    t.edge(load, inp, s);
4775                    t.edge(s, out, st);
4776                    let items: Vec<WorkItem> = relayout_items(tiles_r, row_t, blk)
4777                        .into_iter()
4778                        .map(|mut it| {
4779                            it.out.row0_t += r0;
4780                            it
4781                        })
4782                        .collect();
4783                    let n_items = items.len();
4784                    let mut prog = t.finish(
4785                        format!("{name}_rows{first}"),
4786                        items,
4787                        AccPolicy::for_result(dt),
4788                        NumericsClaim::Exact,
4789                        target,
4790                    )?;
4791                    let plan = relayout_plan(tiles_r, row_t, blk, dt);
4792                    prog.work.plan = Some(if part {
4793                        BlockPlan {
4794                            grid: (1, n_items),
4795                            ..plan
4796                        }
4797                    } else {
4798                        plan
4799                    });
4800                    Ok(prog)
4801                };
4802                let budget = target.max_local_mem_bytes();
4803                (1..=row_block(row_t, target))
4804                    .rev()
4805                    .filter(|d| row_t % d == 0)
4806                    .map(tiled)
4807                    .find(|p| p.as_ref().is_ok_and(|p| site_footprint(p) <= budget))
4808                    .unwrap_or(Err(SpatialReject::NotLowered(
4809                        "a relayout of one tile row of flat rows does not fit a site",
4810                    )))
4811            };
4812        let whole = if rows <= tr { rows } else { rows - rows % tr };
4813        return if whole == rows {
4814            Ok(vec![rows_section(0, rows, false)?])
4815        } else {
4816            Ok(vec![
4817                rows_section(0, whole, true)?,
4818                rows_section(whole, rows - whole, true)?,
4819            ])
4820        };
4821    }
4822    // Rows narrower than a tile: a dense tile row's `tile` rows are ONE contiguous run of `tile·cols` flat elements,
4823    // tile-aligned when `tile·cols` is whole tiles, so each tile row is an item that loads its run as one page and
4824    // reads `out(y, c) = in[y·cols + c]`. The last tile row's shorter run is a section of its own.
4825    let run_t = (tr * cols).div_ceil(tw);
4826    let page_t = if rows <= tr {
4827        tiles_of(rows * cols, target)
4828    } else {
4829        run_t
4830    };
4831    if (rows > tr && (tr * cols) % tw != 0) || page_t > row_block_tiles(target) {
4832        return Err(SpatialReject::NotLowered("a block of rows whose relayout into a matrix is neither whole tiles across nor a tile row of narrow rows one page holds"));
4833    }
4834    let narrow_section =
4835        |first: usize, n: usize, part: bool| -> Result<SpatialProgram, SpatialReject> {
4836            let (h, tiles_r, r0) = (n.min(tr), n.div_ceil(tr), first / tr);
4837            let n_t = tiles_of(h * cols, target);
4838            let mut t = Template::new(target)?;
4839            let (src, dst) = declare(
4840                &mut t,
4841                if part {
4842                    TensorRole::State
4843                } else {
4844                    TensorRole::Output
4845                },
4846            );
4847            let inp = t.buffer("in_rows", n_t, dt);
4848            let out = t.buffer("out", row_t, dt);
4849            let load = t.load_block_at(
4850                "load_rows",
4851                src,
4852                r0 * run_t,
4853                vec![(LV_ITEM_Y, run_t)],
4854                1,
4855                n_t,
4856                n_t,
4857                inp,
4858            );
4859            let flat = ElemAccess {
4860                base: 0,
4861                terms: vec![(LV_GEN_ROW, cols as i64), (LV_GEN_COL, 1)],
4862                indirect: None,
4863            };
4864            let w = wide(row_t, target);
4865            let s = t.stage(
4866                "relayout",
4867                StageKind::Compute(ComputeOp::Generic {
4868                    iters: super::from_tensor_op::space(target, w, None),
4869                    ins: vec![flat],
4870                    out: ident_map(w),
4871                    body: below(LV_GEN_ROW, h, Box::new(below(LV_GEN_COL, cols, op(0)))),
4872                }),
4873                vec![inp],
4874                vec![out],
4875            );
4876            let st = t.store_block(
4877                "store_dense",
4878                dst,
4879                r0 * row_t,
4880                vec![(LV_ITEM_Y, row_t)],
4881                1,
4882                row_t,
4883                row_t,
4884                out,
4885            );
4886            t.edge(load, inp, s);
4887            t.edge(s, out, st);
4888            let items: Vec<WorkItem> = relayout_items(tiles_r, row_t, row_t)
4889                .into_iter()
4890                .map(|mut it| {
4891                    it.out.row0_t += r0;
4892                    it
4893                })
4894                .collect();
4895            let n_items = items.len();
4896            let mut prog = t.finish(
4897                format!("{name}_rows{first}"),
4898                items,
4899                AccPolicy::for_result(dt),
4900                NumericsClaim::Exact,
4901                target,
4902            )?;
4903            let plan = relayout_plan(tiles_r, row_t, row_t, dt);
4904            prog.work.plan = Some(if part {
4905                BlockPlan {
4906                    grid: (1, n_items),
4907                    ..plan
4908                }
4909            } else {
4910                plan
4911            });
4912            Ok(prog)
4913        };
4914    let whole = if rows <= tr { rows } else { rows - rows % tr };
4915    if whole == rows {
4916        Ok(vec![narrow_section(0, rows, false)?])
4917    } else {
4918        Ok(vec![
4919            narrow_section(0, whole, true)?,
4920            narrow_section(whole, rows - whole, true)?,
4921        ])
4922    }
4923}
4924
4925/// `to` — the flat rows of `from`, a DENSE `rows × cols` matrix whose rows are whole tiles: item `(y, b)` is row `y`'s
4926/// `b`-th `blk`-tile segment. It fetches dense tile row `⌊y / tile⌋` through the row itself — an [`IndirectTile`] with
4927/// `elems_per_tile = tile`, as a dense row read does — lifts in-tile row `y mod tile` out of the page, and writes the
4928/// segment where flat row `y` holds it. Both numbers come from a compile-time table, two tiles per row, so the section
4929/// is the same for every row count.
4930pub(crate) fn relayout_dense_to_rows(
4931    from: (&str, &Fill),
4932    to: &str,
4933    rows: usize,
4934    cols: usize,
4935    dt: Precision,
4936    target: &SpatialTarget,
4937) -> Result<SpatialProgram, SpatialReject> {
4938    if cols % target.tile.cols != 0 {
4939        return relayout_small_dense_to_rows(from, to, rows, cols, dt, target);
4940    }
4941    let picks: Vec<usize> = (0..rows).collect();
4942    relayout_dense_picks_to_rows(from, to, rows, &picks, cols, dt, target)
4943}
4944
4945/// [`relayout_dense_to_rows`] of the rows `picks` names: flat row `y` of `to` is row `picks[y]` of `from`, a dense
4946/// `rows × cols` matrix whose rows are whole tiles.
4947fn relayout_dense_picks_to_rows(
4948    from: (&str, &Fill),
4949    to: &str,
4950    rows: usize,
4951    picks: &[usize],
4952    cols: usize,
4953    dt: Precision,
4954    target: &SpatialTarget,
4955) -> Result<SpatialProgram, SpatialReject> {
4956    let (tr, tw) = (target.tile.rows, target.tile.cols);
4957    let row_t = cols / tw;
4958    let blk = row_block(row_t, target);
4959    let w = wide(blk, target);
4960    let ix = index_type(target, rows as u64)?;
4961    let mut t = Template::new(target)?;
4962    let src = t.tensor(
4963        from.0,
4964        TensorRole::Input,
4965        rows,
4966        cols,
4967        dt,
4968        super::from_tensor_op::init_of(from.1, rows * cols),
4969    );
4970    let flat_rows = picks.len();
4971    let dst = t.tensor(
4972        to,
4973        TensorRole::Output,
4974        1,
4975        flat_rows * cols,
4976        dt,
4977        TensorInit::Zero,
4978    );
4979    let at: Vec<i64> = picks
4980        .iter()
4981        .flat_map(|y| [*y as i64, (*y % tr) as i64])
4982        .flat_map(|v| std::iter::once(v).chain(std::iter::repeat_n(0, tw - 1)))
4983        .collect();
4984    let table = t.tensor(
4985        &format!("{to}__row_at"),
4986        TensorRole::Param,
4987        1,
4988        2 * flat_rows * tw,
4989        ix,
4990        TensorInit::Values(at),
4991    );
4992    let row_in = t.buffer("in_row", 1, ix);
4993    let in_tile = t.buffer("in_tile_row", 1, ix);
4994    let l_row = t.load_block_at("load_row", table, 0, vec![(LV_ITEM_Y, 2)], 1, 1, 1, row_in);
4995    let l_in = t.load_block_at(
4996        "load_in_tile_row",
4997        table,
4998        1,
4999        vec![(LV_ITEM_Y, 2)],
5000        1,
5001        1,
5002        1,
5003        in_tile,
5004    );
5005    let page = t.buffer("in_tile_rows", blk, dt);
5006    let origin = TileAddr {
5007        base: 0,
5008        steps: vec![(LV_ITEM_X, blk)],
5009        indirect: Some(IndirectTile {
5010            index: row_in,
5011            tile_stride: row_t,
5012            elems_per_tile: tr,
5013            bound: (0, rows - 1),
5014        }),
5015    };
5016    let load = t.stage(
5017        "gather_tile_row",
5018        StageKind::Load {
5019            tensor: src,
5020            pattern: AccessPattern {
5021                origin,
5022                block_h: 1,
5023                block_w: blk,
5024                stride_w: 1,
5025                stride_h: row_t,
5026            },
5027            num_blocks: 1,
5028            multicast: None,
5029        },
5030        vec![row_in],
5031        vec![page],
5032    );
5033    t.edge(l_row, row_in, load);
5034    let out = t.buffer("out", blk, dt);
5035    let lift = ElemAccess {
5036        base: 0,
5037        terms: vec![(LV_GEN_COL, 1)],
5038        indirect: Some(IndirectTerm {
5039            operand: 1,
5040            stride: w as i64,
5041            modulus: tr,
5042        }),
5043    };
5044    let s = t.stage(
5045        "lift_row",
5046        StageKind::Compute(ComputeOp::Generic {
5047            iters: flat_space(w),
5048            ins: vec![lift, fixed(0)],
5049            out: ident_map(w),
5050            body: TensorBody::Operand(0),
5051        }),
5052        vec![page, in_tile],
5053        vec![out],
5054    );
5055    t.edge(load, page, s);
5056    t.edge(l_in, in_tile, s);
5057    let st = t.store_block(
5058        "store_row",
5059        dst,
5060        0,
5061        vec![(LV_ITEM_Y, row_t), (LV_ITEM_X, blk)],
5062        1,
5063        blk,
5064        row_t,
5065        out,
5066    );
5067    t.edge(s, out, st);
5068    let items: Vec<WorkItem> = (0..flat_rows)
5069        .flat_map(|y| (0..row_t / blk).map(move |b| (y, b)))
5070        .enumerate()
5071        .map(|(id, (y, b))| WorkItem {
5072            id: WorkItemId::new(id),
5073            coord: (y, b),
5074            out: TileRect {
5075                row0_t: 0,
5076                col0_t: y * row_t + b * blk,
5077                h_t: 1,
5078                w_t: blk,
5079            },
5080        })
5081        .collect();
5082    let name = if flat_rows == rows && picks.iter().enumerate().all(|(y, r)| y == *r) {
5083        format!("relayout_rows_{rows}x{cols}")
5084    } else {
5085        format!("relayout_rows_{flat_rows}_of_{rows}x{cols}")
5086    };
5087    let mut prog = t.finish(
5088        name,
5089        items,
5090        AccPolicy::for_result(dt),
5091        NumericsClaim::Exact,
5092        target,
5093    )?;
5094    // The grid counts the flat result's blocks: one tile row of every row's segments.
5095    prog.work.plan = Some(BlockPlan {
5096        grid: (1, flat_rows * (row_t / blk)),
5097        ..tile_row_plan(row_t, blk, dt)
5098    });
5099    Ok(prog)
5100}
5101
5102/// [`relayout_dense_to_rows`] for a matrix whose rows are not whole tiles: then it must be one small tile row, loaded
5103/// whole, and flat element `C = y·cols + c` is dense page element `(y, c)` — at plane offset `C + y·(W − cols)`, `W` the
5104/// page's width — a map affine within row `y`'s band of `cols` columns.
5105fn relayout_small_dense_to_rows(
5106    from: (&str, &Fill),
5107    to: &str,
5108    rows: usize,
5109    cols: usize,
5110    dt: Precision,
5111    target: &SpatialTarget,
5112) -> Result<SpatialProgram, SpatialReject> {
5113    let row_t = cols.div_ceil(target.tile.cols);
5114    let n_t = tiles_of(rows * cols, target);
5115    if rows > target.tile.rows || row_t > row_block_tiles(target) || n_t > row_block_tiles(target) {
5116        return Err(SpatialReject::NotLowered("a matrix whose relayout into flat rows is neither whole tiles across nor one small tile row"));
5117    }
5118    let w = wide(row_t, target);
5119    let mut t = Template::new(target)?;
5120    let src = t.tensor(
5121        from.0,
5122        TensorRole::Input,
5123        rows,
5124        cols,
5125        dt,
5126        super::from_tensor_op::init_of(from.1, rows * cols),
5127    );
5128    let dst = t.tensor(to, TensorRole::Output, 1, rows * cols, dt, TensorInit::Zero);
5129    let inp = t.buffer("in_dense", row_t, dt);
5130    let out = t.buffer("out", n_t, dt);
5131    let load = t.load_block_at("load_dense", src, 0, Vec::new(), 1, row_t, row_t, inp);
5132    let maps: Vec<ElemAccess> = (0..rows)
5133        .map(|y| ElemAccess {
5134            base: (y * (w - cols)) as i64,
5135            terms: vec![(LV_GEN_COL, 1)],
5136            indirect: None,
5137        })
5138        .collect();
5139    let bands: Vec<(usize, TensorBody)> = (0..rows)
5140        .map(|y| ((y + 1) * cols, TensorBody::Operand(y)))
5141        .collect();
5142    let s = t.stage(
5143        "relayout",
5144        StageKind::Compute(ComputeOp::Generic {
5145            iters: flat_space(rows * cols),
5146            ins: maps,
5147            out: ident_map(wide(n_t, target)),
5148            body: banded(bands),
5149        }),
5150        vec![inp; rows],
5151        vec![out],
5152    );
5153    let st = t.store_block("store_rows", dst, 0, Vec::new(), 1, n_t, n_t, out);
5154    t.edge(load, inp, s);
5155    t.edge(s, out, st);
5156    let items = vec![WorkItem {
5157        id: WorkItemId::new(0),
5158        coord: (0, 0),
5159        out: TileRect {
5160            row0_t: 0,
5161            col0_t: 0,
5162            h_t: 1,
5163            w_t: n_t,
5164        },
5165    }];
5166    let mut prog = t.finish(
5167        format!("relayout_rows_{rows}x{cols}"),
5168        items,
5169        AccPolicy::for_result(dt),
5170        NumericsClaim::Exact,
5171        target,
5172    )?;
5173    prog.work.plan = Some(BlockPlan {
5174        grid: (1, 1),
5175        ..tile_row_plan(n_t, n_t, dt)
5176    });
5177    Ok(prog)
5178}
5179
5180// ── Flash attention: the real-width form ──────────────────────────────────────────────────────────────
5181//
5182// ★ AN ATTENTION HEAD NEVER FITS A SITE AT REAL LENGTH, SO IT STREAMS. The cache streams past each item in BLOCKS —
5183// one tile row of positions of the item's K/V head — as the K blocks of a FOLD ([`FoldLevel`]): each block's scores
5184// `S = Q·Kᵀ/√hd` fold into a running softmax that never leaves the site, `m' = max(m, rowmax S)`, `α = e^(m − m')`,
5185// `P = e^(S − m')`, `s' = α·s + rowsum P`, `O' = α·O + P·V`. The rescale is the identity `e^(x−m) = e^(x−m')·
5186// e^(m−m')`, so nothing is approximated, and the epilogue's `O / s` is the attention. An item is a TILE of query
5187// rows of one head — FlashAttention-2's outer loop; one row, for a decode step — so the rows of a tile share every
5188// block they read, and nothing but the attention itself is ever stored.
5189//
5190// ★ A BLOCK IS ONE TILE ROW OF THE CACHE. The cache is dense tiles, `rows × (kvw·hd)`: block `b` of K/V head `x` is
5191// tile row `b`'s `hd / tile` tiles from `x·hd`, an address affine in the block. Row `r` of query tile `q` sees the
5192// positions below its limit — `P + q·tile + r` for a causal block, `P` for a decode row, `l` for a full attention —
5193// so block `b` masks its positions from `limit − start` on, `start = b·tile` carried from block to block as a
5194// counter. A tail block's positions past `l` fall to the same limit.
5195//
5196// ★ THE RUNNING STATE KEEPS THE ACCUMULATOR'S WIDTH. `m` is an exponent shift: rounded to bf16 it would move every
5197// weight of the row by up to 1%. `m`, `s` and `O` are fp32 carries; only the attention leaves at the datapath's
5198// precision.
5199
5200/// Where a flash attention's query heads sit. Item `(v, x)` owns query head `v` of [`QueryHeads`] `q` and reads its
5201/// K/V head `x` of a cache of `rows` positions holding `kvw` heads each; a head is `hd` wide. The blocks cover the
5202/// first `l` positions, and `limit` masks query row `r`'s positions from `P + r` on.
5203///
5204/// `grouped`: the queries are ONE row regrouped by K/V head ([`grouped_decode_attention`]) — a `(kvw · tile) × hd`
5205/// matrix whose tile row `x` holds K/V head `x`'s query heads, one per row. Item `x` is then K/V head `x`, and every
5206/// row is limited by the one query row's `P`.
5207#[derive(Clone, Copy, Debug)]
5208struct FlashGeometry {
5209    q: QueryHeads,
5210    kvw: usize,
5211    hd: usize,
5212    /// The values' width a head: the keys' `hd`, or narrower (a latent attention's values are its keys' latent part).
5213    vd: usize,
5214    /// The value pages this fold computes, `(first, count)`, when not all of them: a fold too large for a site is one
5215    /// fold per page ([`flash_attention`]).
5216    v_part: Option<(usize, usize)>,
5217    rows: usize,
5218    l: usize,
5219    limit: Option<Bound>,
5220    /// A window: row `r` keeps only the last `W` of the positions below its limit, or those of its own chunk of `C`.
5221    window: Option<Window>,
5222    /// The operand holding one sink logit per query head, which joins each head's softmax and takes no value.
5223    sinks: Option<usize>,
5224    grouped: bool,
5225    /// The factor every score takes.
5226    scale: f64,
5227    /// The soft cap every scaled score takes before its block's maximum, `c·tanh(s/c)`.
5228    cap: Option<f64>,
5229}
5230
5231/// Whether keys `hd` and values `vd` wide take the flash form on `target`: whole tiles. A key of any width is one
5232/// contraction into a tile of scores; values wider than an accumulator page fold a page at a time.
5233fn flash_applies(hd: usize, vd: usize, target: &SpatialTarget) -> bool {
5234    hd % target.tile.cols == 0 && vd % target.tile.cols == 0 && hd > 0 && vd > 0
5235}
5236
5237/// Whether an attention of `h` query heads over `kv` K/V heads, keys `hd` and values `vd` wide, reads its K and V as
5238/// DENSE tiles on `target` — the flash form's cache layout. A table an attention reads dense is written dense too, so
5239/// this is the one place the two agree.
5240pub fn attention_reads_dense(h: usize, kv: usize, hd: usize, vd: usize, target: &SpatialTarget) -> bool {
5241    flash_applies(hd, vd, target) && kv > 0 && h % kv == 0
5242}
5243
5244/// [`attention_reads_dense`] for a full attention, whose query heads read K/V heads one to one.
5245pub fn full_attention_reads_dense(hd: usize, target: &SpatialTarget) -> bool {
5246    flash_applies(hd, hd, target)
5247}
5248
5249/// The mask value a masked score takes: far below any real score, and still finite in every float format.
5250const MASKED: f64 = -1.0e30;
5251
5252/// The fan-ins of a merge tree over `nb` partials, bottom level first: the fewest levels any fan-in up to `fmax`
5253/// allows, and at each level the SMALLEST fan-in that still reaches one partial in the levels left — so the tree
5254/// pads as little as its depth allows. Their product is the partial slot count, `≥ nb`.
5255fn merge_fans(nb: usize, fmax: usize) -> Vec<usize> {
5256    if nb <= 1 {
5257        return Vec::new();
5258    }
5259    let reaches = |f: usize, k: u32, r: usize| f.checked_pow(k).is_none_or(|p| p >= r);
5260    let mut levels = 1u32;
5261    while !reaches(fmax, levels, nb) {
5262        levels += 1;
5263    }
5264    let mut fans = Vec::new();
5265    let mut r = nb;
5266    for k in (1..=levels).rev() {
5267        let f = (2..=fmax).find(|f| reaches(*f, k, r)).unwrap_or(fmax);
5268        fans.push(f);
5269        r = r.div_ceil(f);
5270    }
5271    fans
5272}
5273
5274/// `P`, checked to be in `1..=rows`, from a one-element count operand — a guard of stride zero on `P − 1`, the
5275/// same range check a row index gets. Returns the stages and the buffer holding the checked `P`.
5276fn checked_count(
5277    f: &mut FlatBuild,
5278    count: BufferId,
5279    rows: usize,
5280    ix: Precision,
5281    target: &SpatialTarget,
5282) -> (Vec<StageId>, BufferId) {
5283    let one = wide(1, target);
5284    let biased = f.t.scratch_deep("count_biased", 1, 1, ix);
5285    let s0 = f.compute(
5286        "count_biased",
5287        flat_space(1),
5288        vec![Slot {
5289            buf: count,
5290            map: fixed(0),
5291        }],
5292        ident_map(one),
5293        *sub(op(0), konst(1.0)),
5294        biased,
5295    );
5296    let checked = f.t.scratch_deep("count_checked", 1, 1, ix);
5297    let guard = ElemAccess {
5298        base: 0,
5299        terms: Vec::new(),
5300        indirect: Some(IndirectTerm {
5301            operand: 0,
5302            stride: 0,
5303            modulus: rows,
5304        }),
5305    };
5306    let s1 = f.compute(
5307        "count_checked",
5308        flat_space(1),
5309        vec![Slot {
5310            buf: biased,
5311            map: guard,
5312        }],
5313        ident_map(one),
5314        *add(op(0), konst(1.0)),
5315        checked,
5316    );
5317    (vec![s0, s1], checked)
5318}
5319
5320fn flash_scaler_name(names: &StageNames) -> String {
5321    format!("{}__flash_scaler", names.out)
5322}
5323
5324/// Bytes of local memory one site of a section holds: every buffer, at its declared depth.
5325fn site_footprint(p: &SpatialProgram) -> usize {
5326    p.buffers.iter().map(|b| p.buffer_bytes(b)).sum()
5327}
5328
5329/// ★ TWO PAGES WHERE THEY FIT, ONE WHERE THEY DO NOT. A buffer an item fills or drains once holds a second page, so the
5330/// next item begins while this one's pages drain; where the second pages leave no room in a site (fp32 at a real
5331/// width), one page computes the same values, an item's pages draining before the next item's begin. `build(pages)`
5332/// is the section with that many pages in each such buffer.
5333fn paged(
5334    build: impl Fn(usize) -> Result<SpatialProgram, SpatialReject>,
5335    target: &SpatialTarget,
5336) -> Result<SpatialProgram, SpatialReject> {
5337    let doubled = build(2)?;
5338    if site_footprint(&doubled) <= target.max_local_mem_bytes() {
5339        Ok(doubled)
5340    } else {
5341        build(1)
5342    }
5343}
5344
5345/// The flash fold of an attention: ONE section, an item per query tile and head, the cache streaming past as its
5346/// K blocks.
5347/// A flash fold's query layout: the query rows it holds, the heads a row carries, the query tile rows, and its items.
5348fn fold_items(g: &FlashGeometry, target: &SpatialTarget) -> (usize, usize, usize, usize) {
5349    let tr = target.tile.rows;
5350    let (rows, per_row) = if g.grouped { (g.kvw * tr, 1) } else { (g.q.rows, g.q.per_row) };
5351    let tiles_q = rows.div_ceil(tr);
5352    (rows, per_row, tiles_q, tiles_q * per_row)
5353}
5354
5355/// The tensor a flash fold reads its rows' sinks from: one tile per item, row `r`'s sink in element `r`.
5356fn flash_sinks_name(names: &StageNames) -> String {
5357    format!("{}__flash_sinks", names.out)
5358}
5359
5360fn flash_fold(
5361    names: &StageNames,
5362    g: FlashGeometry,
5363    tag: &str,
5364    target: &SpatialTarget,
5365) -> Result<SpatialProgram, SpatialReject> {
5366    paged(|pages| flash_fold_paged(names, g, tag, target, pages), target)
5367}
5368
5369/// [`flash_fold`] with `pages` pages in each streamed key and value page and in the output page.
5370fn flash_fold_paged(
5371    names: &StageNames,
5372    g: FlashGeometry,
5373    tag: &str,
5374    target: &SpatialTarget,
5375    pages: usize,
5376) -> Result<SpatialProgram, SpatialReject> {
5377    let (tr, tw) = (target.tile.rows, target.tile.cols);
5378    let head_t = g.hd / tw;
5379    let (rows, per_row, tiles_q, items) = fold_items(&g, target);
5380    let width = per_row * g.hd;
5381    let width_t = per_row * head_t;
5382    // ★ THE VALUES FOLD AN ACCUMULATOR PAGE AT A TIME. A row's running output is `vd` wide; past a page it is `nv`
5383    // blocks of `vt` tiles, each its own carry, block product, update and normalization, all sharing one row's scores,
5384    // maximum and sum.
5385    let v_t = g.vd / tw;
5386    let vt = row_block(v_t, target);
5387    let (j0, nv) = g.v_part.unwrap_or((0, v_t / vt));
5388    // The fold's own part of a head's values: `nv` pages from page `j0`.
5389    let part_t = nv * vt;
5390    let out_width = per_row * part_t * tw;
5391    let out_width_t = per_row * part_t;
5392    let part = |name: &str, j: usize| if nv == 1 { name.to_string() } else { format!("{name}{j}") };
5393    let nb = g.l.div_ceil(tr);
5394    let precision = elem_precision(target)?;
5395    let acc = AccPolicy::F32.interm_precision();
5396    let ix = index_type(target, (g.rows + tiles_q * tr) as u64)?;
5397    let (q_name, q_fill) = names.operand(0);
5398    let (k_name, k_fill) = names.operand(1);
5399    let (v_name, v_fill) = names.operand(2);
5400    let kv_len = g.rows * g.kvw * g.hd;
5401    let v_len = g.rows * g.kvw * g.vd;
5402    let scale = g.scale;
5403    let (w1, wd) = (wide(1, target), wide(vt, target));
5404    let mut t = Template::new(target)?;
5405    let scaler = t.tensor(
5406        &flash_scaler_name(names),
5407        TensorRole::Param,
5408        tr,
5409        tw,
5410        precision,
5411        TensorInit::ConstF(1.0f64.to_bits()),
5412    );
5413    // The queries and the result are DENSE — a tile row holds `tile` query rows of every head.
5414    let q = t.tensor(
5415        q_name,
5416        TensorRole::Input,
5417        rows,
5418        width,
5419        precision,
5420        super::from_tensor_op::init_of(q_fill, rows * width),
5421    );
5422    let k = t.tensor(
5423        k_name,
5424        TensorRole::Input,
5425        g.rows,
5426        g.kvw * g.hd,
5427        precision,
5428        super::from_tensor_op::init_of(k_fill, kv_len),
5429    );
5430    let v = t.tensor(
5431        v_name,
5432        TensorRole::Input,
5433        g.rows,
5434        g.kvw * g.vd,
5435        precision,
5436        super::from_tensor_op::init_of(v_fill, v_len),
5437    );
5438    let c = t.tensor(
5439        &names.out,
5440        TensorRole::Output,
5441        rows,
5442        out_width,
5443        precision,
5444        TensorInit::Zero,
5445    );
5446    let generic =
5447        |iters: Vec<IterSpec>, ins: Vec<ElemAccess>, out: ElemAccess, body: TensorBody| {
5448            StageKind::Compute(ComputeOp::Generic {
5449                iters,
5450                ins,
5451                out,
5452                body,
5453            })
5454        };
5455    let m2 = |terms: &[(LoopVar, i64)]| ElemAccess {
5456        base: 0,
5457        terms: terms.to_vec(),
5458        indirect: None,
5459    };
5460    let contract = |r: usize, c: usize, k: usize| {
5461        vec![
5462            IterSpec {
5463                var: LV_GEN_ROW,
5464                extent: r,
5465                kind: IterKind::Parallel,
5466            },
5467            IterSpec {
5468                var: LV_GEN_COL,
5469                extent: c,
5470                kind: IterKind::Parallel,
5471            },
5472            IterSpec {
5473                var: LV_GEN_K,
5474                extent: k,
5475                kind: IterKind::Reduce {
5476                    pool: ReducePool::Sum,
5477                },
5478            },
5479        ]
5480    };
5481
5482    // ── The prologue, once per item: the scaler, the query tile, and each row's limit.
5483    let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
5484    let q_in = t.buffer_deep("in_q", 1, head_t, precision);
5485    let l_sc = t.load_block_at("load_scaler", scaler, 0, Vec::new(), 1, 1, 1, sc_in);
5486    // Item `y` is tile row `y / per_row`, head `y mod per_row` of the dense queries: linear tile `y·head_t`.
5487    let l_q = t.load_block_at(
5488        "load_q",
5489        q,
5490        0,
5491        vec![(LV_ITEM_Y, head_t)],
5492        1,
5493        head_t,
5494        width_t,
5495        q_in,
5496    );
5497    // ★ A MASK ONLY WHERE A POSITION CAN FALL PAST A LIMIT: a causal or decode bound, or a cache that ends inside a
5498    // tile row. Its block start moves only when there is more than one block — one block starts at 0, and its limit is
5499    // the row's own.
5500    let masked = g.limit.is_some() || g.l % tr != 0 || g.window.is_some();
5501    let counted = masked && nb > 1;
5502    let row_lim = masked.then(|| t.scratch_deep("row_limit", 1, 1, ix));
5503    // The input holding `P`, when the program computes it: the fold's blocks then run to the position.
5504    let mut position: Option<TensorId> = None;
5505    // ★ ROW `r` OF QUERY TILE `q` SEES THE POSITIONS BELOW `P + q·tile + r`. `P` is checked first when the program
5506    // computes it; `q·tile` comes from a compile-time table, one tile per item; `r` from a tile holding `r` in row `r`.
5507    // A full attention's rows all see its `l` positions.
5508    if let Some(row_lim) = row_lim {
5509        match g.limit {
5510            None => {
5511                t.stage(
5512                    "row_limit",
5513                    generic(
5514                        head_space(tr, 1, None),
5515                        Vec::new(),
5516                        ident_map(w1),
5517                        *konst(g.l as f64),
5518                    ),
5519                    Vec::new(),
5520                    vec![row_lim],
5521                );
5522            }
5523            Some(bound) => {
5524                let (mut ins, mut maps, mut feeds) = (Vec::new(), Vec::new(), Vec::new());
5525                let p = match bound {
5526                    Bound::Literal(p) => konst(p as f64),
5527                    Bound::Operand(li) => {
5528                        let (pn, pf) = names.operand(li);
5529                        let pt = t.tensor(
5530                            pn,
5531                            TensorRole::Input,
5532                            1,
5533                            1,
5534                            ix,
5535                            super::from_tensor_op::init_of(pf, 1),
5536                        );
5537                        position = Some(pt);
5538                        let p_in = t.buffer_deep(&format!("in_{pn}"), 1, 1, ix);
5539                        let l_p = t.load_block_at(
5540                            &format!("load_{pn}"),
5541                            pt,
5542                            0,
5543                            Vec::new(),
5544                            1,
5545                            1,
5546                            1,
5547                            p_in,
5548                        );
5549                        let biased = t.scratch_deep("count_biased", 1, 1, ix);
5550                        let s_b = t.stage(
5551                            "count_biased",
5552                            generic(
5553                                flat_space(1),
5554                                vec![fixed(0)],
5555                                ident_map(w1),
5556                                *sub(op(0), konst(1.0)),
5557                            ),
5558                            vec![p_in],
5559                            vec![biased],
5560                        );
5561                        t.edge(l_p, p_in, s_b);
5562                        let guard = ElemAccess {
5563                            base: 0,
5564                            terms: Vec::new(),
5565                            indirect: Some(IndirectTerm {
5566                                operand: 0,
5567                                stride: 0,
5568                                modulus: bound_max(g.rows, g.q.rows),
5569                            }),
5570                        };
5571                        ins.push(biased);
5572                        maps.push(guard);
5573                        feeds.push((s_b, biased));
5574                        add(op(0), konst(1.0))
5575                    }
5576                };
5577                // A regrouped tile's rows are heads of the one query row: every row's limit is `P`.
5578                let limit = if g.grouped {
5579                    p
5580                } else {
5581                    let tile_rows: Vec<i64> = (0..items)
5582                        .flat_map(|y| {
5583                            std::iter::once(((y / per_row) * tr) as i64)
5584                                .chain(std::iter::repeat_n(0, tw - 1))
5585                        })
5586                        .collect();
5587                    let tiles_t = t.tensor(
5588                        &format!("{}__flash_tile_rows", names.out),
5589                        TensorRole::Param,
5590                        1,
5591                        items * tw,
5592                        ix,
5593                        TensorInit::Values(tile_rows),
5594                    );
5595                    let q_row = t.buffer_deep("in_tile_row", 1, 1, ix);
5596                    let l_qr = t.load_block_at(
5597                        "load_tile_row",
5598                        tiles_t,
5599                        0,
5600                        vec![(LV_ITEM_Y, 1)],
5601                        1,
5602                        1,
5603                        1,
5604                        q_row,
5605                    );
5606                    ins.push(q_row);
5607                    maps.push(fixed(0));
5608                    feeds.push((l_qr, q_row));
5609                    let q_row_op = op(ins.len() - 1);
5610                    let row_index = t.tensor(
5611                        &format!("{}__flash_row_index", names.out),
5612                        TensorRole::Param,
5613                        tr,
5614                        tw,
5615                        ix,
5616                        TensorInit::Values((0..tr * tw).map(|e| (e / tw) as i64).collect()),
5617                    );
5618                    let r_in = t.buffer_deep("in_row_index", 1, 1, ix);
5619                    let l_r =
5620                        t.load_block_at("load_row_index", row_index, 0, Vec::new(), 1, 1, 1, r_in);
5621                    ins.push(r_in);
5622                    maps.push(per_row_map(w1));
5623                    feeds.push((l_r, r_in));
5624                    let r_op = op(ins.len() - 1);
5625                    add(add(p, q_row_op), r_op)
5626                };
5627                let s = t.stage(
5628                    "row_limit",
5629                    generic(head_space(tr, 1, None), maps, ident_map(w1), *limit),
5630                    ins,
5631                    vec![row_lim],
5632                );
5633                for (from, buf) in feeds {
5634                    t.edge(from, buf, s);
5635                }
5636            }
5637        }
5638    }
5639    // The carried state, and its first page: the empty softmax (`m = −10³⁰`, `s = 0`, `O = 0`) at block start 0 — or,
5640    // with sinks, each row's sink as its maximum and a sum of one.
5641    let m_run = t.scratch_deep("run_max", 2, 1, acc);
5642    let s_run = t.scratch_deep("run_sum", 2, 1, acc);
5643    let o_runs: Vec<BufferId> = (0..nv).map(|j| t.scratch_deep(&part("run_out", j), 2, vt, acc)).collect();
5644    match g.sinks {
5645        None => {
5646            t.stage(
5647                "init_max",
5648                generic(head_space(tr, tw, None), Vec::new(), ident_map(w1), *konst(MASKED)),
5649                Vec::new(),
5650                vec![m_run],
5651            );
5652        }
5653        Some(_) => {
5654            let sinks = t.tensor(
5655                &flash_sinks_name(names),
5656                TensorRole::Input,
5657                1,
5658                items * tw,
5659                precision,
5660                super::from_tensor_op::init_of(&Fill::Unknown, items * tw),
5661            );
5662            let sk_in = t.buffer_deep("in_sinks", 1, 1, precision);
5663            let l_sk = t.load_block_at("load_sinks", sinks, 0, vec![(LV_ITEM_Y, 1)], 1, 1, 1, sk_in);
5664            let row_of_page = ElemAccess {
5665                base: 0,
5666                terms: vec![(LV_GEN_ROW, 1)],
5667                indirect: None,
5668            };
5669            let s = t.stage(
5670                "init_max",
5671                generic(head_space(tr, tw, None), vec![row_of_page], ident_map(w1), *op(0)),
5672                vec![sk_in],
5673                vec![m_run],
5674            );
5675            t.edge(l_sk, sk_in, s);
5676        }
5677    }
5678    t.stage(
5679        "init_sum",
5680        generic(
5681            head_space(tr, tw, None),
5682            Vec::new(),
5683            ident_map(w1),
5684            *konst(if g.sinks.is_some() { 1.0 } else { 0.0 }),
5685        ),
5686        Vec::new(),
5687        vec![s_run],
5688    );
5689    for (j, &o_run) in o_runs.iter().enumerate() {
5690        t.stage(
5691            &part("init_out", j),
5692            generic(
5693                head_space(tr, vt * tw, None),
5694                Vec::new(),
5695                ident_map(wd),
5696                *konst(0.0),
5697            ),
5698            Vec::new(),
5699            vec![o_run],
5700        );
5701    }
5702    let start = counted.then(|| {
5703        let start = t.scratch_deep("block_start", 2, 1, ix);
5704        t.stage(
5705            "init_start",
5706            generic(flat_space(1), Vec::new(), ident_map(w1), *konst(0.0)),
5707            Vec::new(),
5708            vec![start],
5709        );
5710        start
5711    });
5712
5713    // ── The body, once per block: block `b` of K/V head `x` is tile row `b` of the cache from tile `x·head_t`.
5714    let row_stride = g.kvw * head_t;
5715    let v_row_stride = g.kvw * v_t;
5716    let k_in = t.buffer_deep("in_k", pages, head_t, precision);
5717    let v_ins: Vec<BufferId> = (0..nv).map(|j| t.buffer_deep(&part("in_v", j), pages, vt, precision)).collect();
5718    // Block `b` of K/V head `x`, `cols` tiles from tile `base` of the head's run of `head` tiles in a row of `stride`.
5719    let cache_block = |t: &mut Template, name: &str, tensor: TensorId, (base, cols, head, stride): (usize, usize, usize, usize), into: BufferId| {
5720        let origin = TileAddr {
5721            base,
5722            steps: vec![(LV_ITEM_X, head), (LV_KBLOCK, stride)],
5723            indirect: None,
5724        };
5725        t.stage(
5726            name,
5727            StageKind::Load {
5728                tensor,
5729                pattern: AccessPattern {
5730                    origin,
5731                    block_h: 1,
5732                    block_w: cols,
5733                    stride_w: 1,
5734                    stride_h: stride,
5735                },
5736                num_blocks: nb,
5737                multicast: None,
5738            },
5739            Vec::new(),
5740            vec![into],
5741        )
5742    };
5743    let l_k = cache_block(&mut t, "load_k", k, (0, head_t, head_t, row_stride), k_in);
5744    let l_vs: Vec<StageId> = (0..nv)
5745        .map(|j| cache_block(&mut t, &part("load_v", j), v, ((j0 + j) * vt, vt, v_t, v_row_stride), v_ins[j]))
5746        .collect();
5747    // Each row's limit in the block's own positions: its limit less the block's start.
5748    let lim = match (row_lim, start) {
5749        (Some(row_lim), Some(start)) => {
5750            let lim_b = t.scratch_deep("block_limit", 1, 1, ix);
5751            t.stage(
5752                "block_limit",
5753                generic(
5754                    head_space(tr, 1, None),
5755                    vec![per_row_map(w1), fixed(0)],
5756                    ident_map(w1),
5757                    *sub(op(0), op(1)),
5758                ),
5759                vec![row_lim, start],
5760                vec![lim_b],
5761            );
5762            Some(lim_b)
5763        }
5764        (row_lim, _) => row_lim,
5765    };
5766    // A window keeps the last `W` positions below each row's limit: the block's floor is its limit less `W`. A chunk of
5767    // `C` keeps the row's own chunk: its floor, `⌊(limit − 1)/C⌋·C`, less the block's start.
5768    let floor = match (g.window, lim) {
5769        (Some(Window::Chunked(c)), Some(_)) => {
5770            let row_lim = row_lim.expect("a chunked row is masked");
5771            let floor_b = t.scratch_deep("block_floor", 1, 1, ix);
5772            let c = c as f64;
5773            let chunk = mul(Box::new(TensorBody::FloorDiv(sub(op(0), konst(1.0)), konst(c))), konst(c));
5774            let (ins, maps, body) = match start {
5775                Some(start) => (vec![row_lim, start], vec![per_row_map(w1), fixed(0)], *sub(chunk, op(1))),
5776                None => (vec![row_lim], vec![per_row_map(w1)], *chunk),
5777            };
5778            t.stage("block_floor", generic(head_space(tr, 1, None), maps, ident_map(w1), body), ins, vec![floor_b]);
5779            Some(floor_b)
5780        }
5781        (Some(Window::Sliding(w)), Some(lim_b)) => {
5782            let floor_b = t.scratch_deep("block_floor", 1, 1, ix);
5783            t.stage(
5784                "block_floor",
5785                generic(head_space(tr, 1, None), vec![per_row_map(w1)], ident_map(w1), *sub(op(0), konst(w as f64))),
5786                vec![lim_b],
5787                vec![floor_b],
5788            );
5789            Some(floor_b)
5790        }
5791        _ => None,
5792    };
5793    let hd = g.hd as i64;
5794    // `S[r, j] = Σ_e Q[r, e]·K[j, e]/√hd` — the tile's rows against the block's positions, one contraction.
5795    let s_buf = t.scratch_deep("scores", 1, 1, acc);
5796    let s_scores = t.stage(
5797        "scores",
5798        generic(
5799            contract(tr, tr, g.hd),
5800            vec![
5801                m2(&[(LV_GEN_ROW, hd), (LV_GEN_K, 1)]),
5802                m2(&[(LV_GEN_COL, hd), (LV_GEN_K, 1)]),
5803            ],
5804            ident_map(w1),
5805            *mul(mul(op(0), op(1)), konst(scale)),
5806        ),
5807        vec![q_in, k_in],
5808        vec![s_buf],
5809    );
5810    // ★ A SOFT CAP takes every scaled score before the block's maximum and the mask.
5811    let s_buf = match g.cap {
5812        None => s_buf,
5813        Some(c) => {
5814            let capped_buf = t.scratch_deep("capped_scores", 1, 1, acc);
5815            t.stage(
5816                "score_cap",
5817                generic(head_space(tr, tr, None), vec![ident_map(w1)], ident_map(w1), *capped(op(0), c)),
5818                vec![s_buf],
5819                vec![capped_buf],
5820            );
5821            capped_buf
5822        }
5823    };
5824    // A position is live below its row's limit in the block and, windowed, at or past its floor; with no mask, every
5825    // position is.
5826    let masking = lim.is_some();
5827    let windowed = floor.is_some();
5828    let live = |then: Box<TensorBody>, otherwise: Box<TensorBody>, at: usize| {
5829        if masking {
5830            let kept = if windowed {
5831                Box::new(TensorBody::Select {
5832                    cond: Pred::AxisGe(LV_GEN_COL, PredRhs::Operand(at + 1)),
5833                    then,
5834                    otherwise: otherwise.clone(),
5835                })
5836            } else {
5837                then
5838            };
5839            TensorBody::Select {
5840                cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(at)),
5841                then: kept,
5842                otherwise,
5843            }
5844        } else {
5845            *then
5846        }
5847    };
5848    let with_limit = |mut ins: Vec<BufferId>, mut maps: Vec<ElemAccess>, map: ElemAccess| -> (Vec<BufferId>, Vec<ElemAccess>) {
5849        for b in lim.into_iter().chain(floor) {
5850            ins.push(b);
5851            maps.push(map.clone());
5852        }
5853        (ins, maps)
5854    };
5855    let m_blk = t.scratch_deep("block_max", 1, 1, acc);
5856    let (max_ins, max_maps) = with_limit(vec![s_buf, sc_in], vec![ident_map(w1), const_map()], per_row_map(w1));
5857    let s_mblk = t.stage(
5858        "block_max",
5859        generic(
5860            head_space(tr, tr, Some(ReducePool::Max)),
5861            max_maps,
5862            per_row_map(w1),
5863            live(op(0), konst(MASKED), 2),
5864        ),
5865        max_ins,
5866        vec![m_blk],
5867    );
5868    let m_new = t.scratch_deep("new_max", 1, 1, acc);
5869    let pair = vec![
5870        IterSpec {
5871            var: LV_GEN_ROW,
5872            extent: tr,
5873            kind: IterKind::Parallel,
5874        },
5875        IterSpec {
5876            var: LV_GEN_COL,
5877            extent: 2,
5878            kind: IterKind::Reduce {
5879                pool: ReducePool::Max,
5880            },
5881        },
5882    ];
5883    let pick = TensorBody::Select {
5884        cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Const(1)),
5885        then: op(0),
5886        otherwise: op(1),
5887    };
5888    t.stage(
5889        "new_max",
5890        generic(
5891            pair,
5892            vec![per_row_map(w1), per_row_map(w1), const_map()],
5893            per_row_map(w1),
5894            pick,
5895        ),
5896        vec![m_run, m_blk, sc_in],
5897        vec![m_new],
5898    );
5899    let alpha = t.scratch_deep("rescale", 1, 1, acc);
5900    t.stage(
5901        "rescale",
5902        generic(
5903            head_space(tr, 1, None),
5904            vec![per_row_map(w1), per_row_map(w1)],
5905            ident_map(w1),
5906            *unary(SfpuKind::Exp, sub(op(0), op(1))),
5907        ),
5908        vec![m_run, m_new],
5909        vec![alpha],
5910    );
5911    let p_buf = t.scratch_deep("probs", 1, 1, acc);
5912    let (p_ins, p_maps) = with_limit(vec![s_buf, m_new], vec![ident_map(w1), per_row_map(w1)], per_row_map(w1));
5913    t.stage(
5914        "probs",
5915        generic(
5916            head_space(tr, tr, None),
5917            p_maps,
5918            ident_map(w1),
5919            live(unary(SfpuKind::Exp, sub(op(0), op(1))), konst(0.0), 2),
5920        ),
5921        p_ins,
5922        vec![p_buf],
5923    );
5924    let l_blk = t.scratch_deep("block_sum", 1, 1, acc);
5925    t.stage(
5926        "block_sum",
5927        generic(
5928            head_space(tr, tr, Some(ReducePool::Sum)),
5929            vec![ident_map(w1), const_map()],
5930            per_row_map(w1),
5931            TensorBody::Operand(0),
5932        ),
5933        vec![p_buf, sc_in],
5934        vec![l_blk],
5935    );
5936    // `O_b[r, e] = Σ_j P[r, j]·V[j, e]` — one contraction over the block's positions, a value page at a time.
5937    let vw = (vt * tw) as i64;
5938    let o_blks: Vec<(StageId, BufferId)> = (0..nv)
5939        .map(|j| {
5940            let o_blk = t.scratch_deep(&part("block_out", j), 1, vt, acc);
5941            let s_oblk = t.stage(
5942                &part("block_out", j),
5943                generic(
5944                    contract(tr, vt * tw, tr),
5945                    vec![
5946                        m2(&[(LV_GEN_ROW, w1 as i64), (LV_GEN_K, 1)]),
5947                        m2(&[(LV_GEN_K, vw), (LV_GEN_COL, 1)]),
5948                    ],
5949                    ident_map(wd),
5950                    *mul(op(0), op(1)),
5951                ),
5952                vec![p_buf, v_ins[j]],
5953                vec![o_blk],
5954            );
5955            (s_oblk, o_blk)
5956        })
5957        .collect();
5958    // The updates, each reading the carry it writes: the block folds into the running state.
5959    t.stage(
5960        "fold_sum",
5961        generic(
5962            head_space(tr, 1, None),
5963            vec![per_row_map(w1), per_row_map(w1), per_row_map(w1)],
5964            ident_map(w1),
5965            *add(mul(op(0), op(1)), op(2)),
5966        ),
5967        vec![s_run, alpha, l_blk],
5968        vec![s_run],
5969    );
5970    for (j, &o_run) in o_runs.iter().enumerate() {
5971        t.stage(
5972            &part("fold_out", j),
5973            generic(
5974                head_space(tr, vt * tw, None),
5975                vec![ident_map(wd), per_row_map(w1), ident_map(wd)],
5976                ident_map(wd),
5977                *add(mul(op(0), op(1)), op(2)),
5978            ),
5979            vec![o_run, alpha, o_blks[j].1],
5980            vec![o_run],
5981        );
5982    }
5983    t.stage(
5984        "fold_max",
5985        generic(
5986            head_space(tr, 1, None),
5987            vec![per_row_map(w1), per_row_map(w1)],
5988            ident_map(w1),
5989            TensorBody::Operand(1),
5990        ),
5991        vec![m_run, m_new],
5992        vec![m_run],
5993    );
5994    if let Some(start) = start {
5995        t.stage(
5996            "next_start",
5997            generic(
5998                flat_space(1),
5999                vec![fixed(0)],
6000                ident_map(w1),
6001                *add(op(0), konst(tr as f64)),
6002            ),
6003            vec![start],
6004            vec![start],
6005        );
6006    }
6007
6008    // ── The epilogue, once per item: the attention, into the item's tile of the dense result.
6009    let mut epilogue = Vec::new();
6010    for (j, &o_run) in o_runs.iter().enumerate() {
6011        let out = t.buffer_deep(&part("out", j), pages, vt, precision);
6012        let s_norm = t.stage(
6013            &part("normalize", j),
6014            generic(
6015                head_space(tr, vt * tw, None),
6016                vec![ident_map(wd), per_row_map(w1)],
6017                ident_map(wd),
6018                *mul(op(0), unary(SfpuKind::Recip, op(1))),
6019            ),
6020            vec![o_run, s_run],
6021            vec![out],
6022        );
6023        let st = t.store_block(
6024            &part("store_ctx", j),
6025            c,
6026            j * vt,
6027            vec![(LV_ITEM_Y, part_t)],
6028            1,
6029            vt,
6030            out_width_t,
6031            out,
6032        );
6033        epilogue.push((s_norm, out, st));
6034    }
6035    for (from, buf, to) in [(l_sc, sc_in, s_mblk), (l_q, q_in, s_scores), (l_k, k_in, s_scores)] {
6036        t.edge(from, buf, to);
6037    }
6038    for j in 0..nv {
6039        t.edge(l_vs[j], v_ins[j], o_blks[j].0);
6040    }
6041    for (s_norm, out, st) in epilogue {
6042        t.edge(s_norm, out, st);
6043    }
6044    let work: Vec<WorkItem> = (0..items)
6045        .map(|y| {
6046            let (qt, h) = (y / per_row, y % per_row);
6047            let kv_head = if g.grouped { y } else { g.q.kv_of(h) };
6048            WorkItem {
6049                id: WorkItemId::new(y),
6050                coord: (y, kv_head),
6051                out: TileRect {
6052                    row0_t: qt,
6053                    col0_t: h * part_t,
6054                    h_t: 1,
6055                    w_t: part_t,
6056                },
6057            }
6058        })
6059        .collect();
6060    // ★ THE BLOCKS RUN TO THE POSITION. Query row `r` sees the positions below `P + r`, so the last of the `rows`
6061    // query rows reaches `P + rows − 1` and no row reads a block past `⌈(P + rows − 1) / tile⌉`: every later block is
6062    // masked whole, its update the identity. That count is the fold's trip, read from `P` before the program runs —
6063    // a decode step at position `P` streams `⌈P / tile⌉` blocks of its cache, not the `⌈l / tile⌉` it is declared with.
6064    // It admits exactly the `P` the guard does, `1 ≤ P ≤ rows − (T − 1)`, so an out-of-cache `P` is refused against the
6065    // cache before any op runs; its greatest count is the cache's `⌈l / tile⌉` tile rows.
6066    let k_trip = position.filter(|_| counted).map(|pt| {
6067        let reach = g.q.rows - 1;
6068        t.trip(Trip {
6069            tensor: pt,
6070            elem: 0,
6071            admit: (1, bound_max(g.rows, g.q.rows) as i64),
6072            bias: reach as i64,
6073            per: tr,
6074        })
6075    });
6076    let values = match g.v_part {
6077        Some((first, count)) => format!("_v{}_pages{first}+{count}", g.vd),
6078        None if g.vd != g.hd => format!("_v{}", g.vd),
6079        None => String::new(),
6080    };
6081    let mut prog = t.finish(
6082        format!("{tag}_flash_blocks_{items}x{}{values}_l{}_b{tr}", g.hd, g.l),
6083        work,
6084        AccPolicy::F32,
6085        approx("attention", precision),
6086        target,
6087    )?;
6088    prog.work.k_trip = k_trip;
6089    // ★ THE BLOCKS ARE THE CONTRACTION'S K BLOCKS. An item's result is one tile row of the dense result; the cache's
6090    // `⌈l / tile⌉` blocks are what it folds over.
6091    prog.work.plan = Some(BlockPlan {
6092        per_site_m: 1,
6093        per_site_n: part_t,
6094        grid: (tiles_q, per_row),
6095        inner: InnerLoop {
6096            num_blocks: nb,
6097            block_k_tiles: 1,
6098        },
6099        block: BlockShape {
6100            osh: 1,
6101            osw: part_t,
6102            ibw: 1,
6103        },
6104        acc: AccPolicy::F32,
6105        pack_accumulate: false,
6106        engines: EnginePacking::LoadsTogether,
6107        granularity: Granularity::PerTile,
6108    });
6109    prog.work.k_tiles = nb;
6110    Ok(prog)
6111}
6112
6113/// Attention in the flash form: one fold section ([`flash_fold`]), refused by name when a site cannot hold it.
6114fn flash_attention(
6115    names: &StageNames,
6116    g: FlashGeometry,
6117    tag: &str,
6118    target: &SpatialTarget,
6119) -> Result<Vec<SpatialProgram>, SpatialReject> {
6120    if target.tile.rows != target.tile.cols {
6121        return Err(SpatialReject::NotLowered("a flash attention fold on a tile that is not square — a block's positions are its scores' columns"));
6122    }
6123    let p = flash_fold(names, g, tag, target)?;
6124    if site_footprint(&p) > target.max_local_mem_bytes() {
6125        return value_pages(names, g, tag, target);
6126    }
6127    // ★ A SINK IS THE RUNNING SOFTMAX'S FIRST TERM. Each fold row starts from its head's sink logit as its maximum, with a
6128    // sum of one and no value, so every block's update carries the sink's term to the end. The sinks are first spread
6129    // one tile per item, a row's sink in its own column ([`flash_sinks_name`]).
6130    let Some(si) = g.sinks else { return Ok(vec![p]) };
6131    let (tw, (_, per_row, _, items)) = (target.tile.cols, fold_items(&g, target));
6132    let head_of = |e: usize| {
6133        let (y, row) = (e / tw, e % tw);
6134        if g.grouped {
6135            (row < g.q.group).then_some(y * g.q.group + row)
6136        } else {
6137            Some(y % per_row)
6138        }
6139    };
6140    let spread = spread_sinks(names.operand(si), g.q.per_row, &flash_sinks_name(names), items * tw, &head_of, target)?;
6141    Ok(vec![spread, p])
6142}
6143
6144/// A decode attention whose fold does not fit a site, as one fold per value page: fold `j` streams the keys and its
6145/// page of every head's values into `{out}__values{j}`, and tile joins lay the pages side by side per head into the
6146/// result. Each fold reads the keys again — the price of a site that cannot hold a head's every value page.
6147fn value_pages(
6148    names: &StageNames,
6149    g: FlashGeometry,
6150    tag: &str,
6151    target: &SpatialTarget,
6152) -> Result<Vec<SpatialProgram>, SpatialReject> {
6153    let too_large = SpatialReject::NotLowered("a flash attention fold — one tile of queries against one tile row of its K/V head — does not fit a site");
6154    let tw = target.tile.cols;
6155    let v_t = g.vd / tw;
6156    let vt = row_block(v_t, target);
6157    let pages = v_t / vt;
6158    if pages < 2 || g.grouped || g.q.rows != 1 || g.v_part.is_some() {
6159        return Err(too_large);
6160    }
6161    let page_name = |j: usize| format!("{}__values{j}", names.out);
6162    let mut sections = Vec::new();
6163    for j in 0..pages {
6164        let mut fold_names = names.clone();
6165        fold_names.out = page_name(j);
6166        let fold = flash_attention(&fold_names, FlashGeometry { v_part: Some((j, 1)), ..g }, tag, target)?;
6167        sections.extend(fold);
6168    }
6169    let page = vt * tw;
6170    let mut joined = page_name(0);
6171    for j in 1..pages {
6172        let out = if j + 1 == pages { names.out.clone() } else { format!("{}__values0to{j}", names.out) };
6173        let join_names = StageNames { out: out.clone(), ins: vec![joined, page_name(j)], fills: vec![Fill::Unknown, Fill::Unknown] };
6174        sections.push(tile_join(&join_names, (1, g.q.per_row), (j * page, page), target)?);
6175        joined = out;
6176    }
6177    Ok(sections)
6178}
6179
6180/// A one-row attention whose query heads share K/V heads, folded ONE ITEM PER K/V HEAD: three sections.
6181///   1. `{out}__gqa_q` — the row's heads regrouped, a `(kvw · tile) × hd` matrix: K/V head `x`'s `group` query heads
6182///      are rows `0 … group − 1` of tile row `x`, the rows past them zero. Item `x` loads its heads' `group · hd`
6183///      elements of the row as one page and reads row `j`, column `e` at `j·hd + e` of it.
6184///   2. The flash fold over it ([`flash_fold`], `grouped`): item `x` streams K/V head `x`'s cache once for all
6185///      `group` rows, every row limited by the row's `P`.
6186///   3. `out` — the fold's `{out}__gqa_ctx` back as the row: each query head's row lifted out of its tile row into its
6187///      `hd` elements ([`relayout_dense_picks_to_rows`]).
6188///
6189/// Refused when a regrouping page does not fit a site; the caller then folds an item per query head.
6190fn grouped_decode_attention(
6191    names: &StageNames,
6192    g: FlashGeometry,
6193    tag: &str,
6194    target: &SpatialTarget,
6195) -> Result<Vec<SpatialProgram>, SpatialReject> {
6196    if target.tile.rows != target.tile.cols {
6197        return Err(SpatialReject::NotLowered(
6198            "a grouped decode attention on a tile that is not square",
6199        ));
6200    }
6201    let (tr, tw) = (target.tile.rows, target.tile.cols);
6202    let (kv, group, hd) = (g.kvw, g.q.group, g.hd);
6203    let head_t = hd / tw;
6204    let (heads_t, row_t) = (group * head_t, g.q.per_row * head_t);
6205    let precision = elem_precision(target)?;
6206    let (q_name, q_fill) = names.operand(0);
6207    let grouped_q = format!("{}__gqa_q", names.out);
6208    let grouped_ctx = format!("{}__gqa_ctx", names.out);
6209    let flat = |terms: &[(LoopVar, i64)]| ElemAccess {
6210        base: 0,
6211        terms: terms.to_vec(),
6212        indirect: None,
6213    };
6214    let below = |bound: usize| Pred::AxisLt(LV_GEN_ROW, PredRhs::Const(bound));
6215    let budget = target.max_local_mem_bytes();
6216    let fits = |p: SpatialProgram| {
6217        if site_footprint(&p) <= budget {
6218            Ok(p)
6219        } else {
6220            Err(SpatialReject::NotLowered(
6221                "a regrouping of one K/V head's query heads does not fit a site",
6222            ))
6223        }
6224    };
6225
6226    let regroup = paged(
6227        |pages| {
6228            let mut t = Template::new(target)?;
6229            let q = t.tensor(
6230                q_name,
6231                TensorRole::Input,
6232                1,
6233                g.q.per_row * hd,
6234                precision,
6235                super::from_tensor_op::init_of(q_fill, g.q.per_row * hd),
6236            );
6237            let gq = t.tensor(
6238                &grouped_q,
6239                TensorRole::Output,
6240                kv * tr,
6241                hd,
6242                precision,
6243                TensorInit::Zero,
6244            );
6245            let inp = t.buffer_deep("in_heads", pages, heads_t, precision);
6246            let out = t.buffer_deep("out", pages, head_t, precision);
6247            let load = t.load_block_at(
6248                "load_heads",
6249                q,
6250                0,
6251                vec![(LV_ITEM_Y, heads_t)],
6252                1,
6253                heads_t,
6254                row_t,
6255                inp,
6256            );
6257            let s = t.stage(
6258                "regroup",
6259                StageKind::Compute(ComputeOp::Generic {
6260                    iters: head_space(tr, hd, None),
6261                    ins: vec![flat(&[(LV_GEN_ROW, hd as i64), (LV_GEN_COL, 1)])],
6262                    out: ident_map(wide(head_t, target)),
6263                    body: TensorBody::Select {
6264                        cond: below(group),
6265                        then: op(0),
6266                        otherwise: konst(0.0),
6267                    },
6268                }),
6269                vec![inp],
6270                vec![out],
6271            );
6272            let st = t.store_block(
6273                "store_groups",
6274                gq,
6275                0,
6276                vec![(LV_ITEM_Y, head_t)],
6277                1,
6278                head_t,
6279                head_t,
6280                out,
6281            );
6282            t.edge(load, inp, s);
6283            t.edge(s, out, st);
6284            let items = (0..kv)
6285                .map(|x| WorkItem {
6286                    id: WorkItemId::new(x),
6287                    coord: (x, 0),
6288                    out: TileRect {
6289                        row0_t: x,
6290                        col0_t: 0,
6291                        h_t: 1,
6292                        w_t: head_t,
6293                    },
6294                })
6295                .collect();
6296            let mut regroup = t.finish(
6297                format!("{tag}_gqa_groups_{kv}x{group}"),
6298                items,
6299                AccPolicy::for_result(precision),
6300                NumericsClaim::Exact,
6301                target,
6302            )?;
6303            regroup.work.plan = Some(BlockPlan {
6304                grid: (kv, 1),
6305                ..tile_row_plan(head_t, head_t, precision)
6306            });
6307            Ok(regroup)
6308        },
6309        target,
6310    )?;
6311
6312    let mut fold_names = names.clone();
6313    fold_names.out = grouped_ctx.clone();
6314    fold_names.ins[0] = grouped_q.clone();
6315    fold_names.fills[0] = Fill::Unknown;
6316    let fold = flash_attention(
6317        &fold_names,
6318        FlashGeometry { grouped: true, ..g },
6319        tag,
6320        target,
6321    )?;
6322
6323    // Query head `v` is row `v mod group` of tile row `v / group`.
6324    let picks: Vec<usize> = (0..g.q.per_row)
6325        .map(|v| (v / group) * tr + v % group)
6326        .collect();
6327    let ungroup = relayout_dense_picks_to_rows(
6328        (&grouped_ctx, &Fill::Unknown),
6329        &names.out,
6330        kv * tr,
6331        &picks,
6332        g.vd,
6333        precision,
6334        target,
6335    )?;
6336
6337    let mut sections = vec![fits(regroup)?];
6338    sections.extend(fold);
6339    sections.push(fits(ungroup)?);
6340    Ok(sections)
6341}
6342
6343/// ★ AN INDEX RIDES THE TARGET'S ADDRESS INTEGER, SIZED BY ITS INDEX SPACE. Never the float datapath — bf16
6344/// holds integers exactly only up to 256, and a Qwen3.8 token id reaches 248,319 — and never a width the chip
6345/// does not have: on an RDU every index is INT32, the widest integer the PCU computes in. `space` is the largest
6346/// magnitude the index arithmetic of the lowering can reach; a space past every integer the target has is
6347/// refused as the precision it would need (`MissingPrecision`), never wrapped.
6348pub(crate) fn index_type(target: &SpatialTarget, space: u64) -> Result<Precision, SpatialReject> {
6349    target
6350        .index_precision_for(space)
6351        .map_err(|needed| SpatialReject::MissingPrecision {
6352            needed,
6353            target: target.name,
6354        })
6355}
6356
6357/// The largest magnitude a recognized INTEGER fill declares — `None` for a fill the compiler cannot see.
6358fn fill_abs_bound(fill: &Fill) -> Option<u64> {
6359    match fill {
6360        Fill::Const(c) => Some(c.unsigned_abs()),
6361        Fill::Pattern { modulus, shift } => Some(
6362            shift
6363                .unsigned_abs()
6364                .max((modulus - 1 + shift).unsigned_abs()),
6365        ),
6366        Fill::Overrides { base, sets } => Some(
6367            sets.iter()
6368                .map(|(_, v)| v.unsigned_abs())
6369                .fold(fill_abs_bound(base)?, u64::max),
6370        ),
6371        Fill::ConstF(_) | Fill::FloatOverrides { .. } | Fill::RealPattern { .. } | Fill::Unknown => None,
6372    }
6373}
6374
6375/// The index space of an index OPERAND: the `span` the lowering's own arithmetic reaches for an in-range index,
6376/// widened by whatever value the operand's fill declares — a declared index past its table must still be HELD
6377/// exactly to be refused.
6378fn index_space(span: u64, per_unit: u64, operand_fill: &Fill) -> u64 {
6379    span.max(
6380        fill_abs_bound(operand_fill)
6381            .map_or(0, |v| v.saturating_add(1).saturating_mul(per_unit.max(1))),
6382    )
6383}
6384
6385/// A table's element precision in a section: the float datapath for a Real table; for an Int one, the
6386/// address-wide integer that holds every value its fill declares (an undeclared fill is its producer's bytes,
6387/// read at the datapath's integer).
6388fn table_precision(
6389    b: &FlatBuild,
6390    float: bool,
6391    fill: &Fill,
6392    target: &SpatialTarget,
6393) -> Result<Precision, SpatialReject> {
6394    if float {
6395        Ok(b.precision)
6396    } else {
6397        index_type(target, fill_abs_bound(fill).unwrap_or(0))
6398    }
6399}
6400
6401/// The element precision of a row WRITE: the table's, widened to hold the source row's declared values too — the
6402/// written table holds both.
6403fn row_write_precision(
6404    b: &FlatBuild,
6405    float: bool,
6406    table: &Fill,
6407    source: &Fill,
6408    target: &SpatialTarget,
6409) -> Result<Precision, SpatialReject> {
6410    if float {
6411        return Ok(b.precision);
6412    }
6413    index_type(
6414        target,
6415        fill_abs_bound(table)
6416            .unwrap_or(0)
6417            .max(fill_abs_bound(source).unwrap_or(0)),
6418    )
6419}
6420
6421/// `lo = (index − 1)·cols` — the flat offset of a 1-based row, computed on the INTEGER datapath (at `ix`, the
6422/// index type) from a one-element index operand. Returns the stages and the buffer holding `lo`.
6423///
6424/// ★ THE MULTIPLY MAY NOT WRAP. On a 32-bit datapath `(index − 1)·cols` wraps once the index passes
6425/// `INT_MAX / cols`, and a wrapped offset can land inside the table — a silently wrong row. So the row is GUARDED
6426/// first, by a gather of stride zero: `index − 1 + B ∈ [0, 2B)` with `B = ⌊INT_MAX / cols⌋`. That range is exactly
6427/// `|index − 1| ≤ B` even where the addition itself wraps, and inside it the multiply cannot overflow. A row just
6428/// outside the table passes the guard and faults at the offset's own range check, which names the offset.
6429fn row_offset_stages(
6430    b: &mut FlatBuild,
6431    index: BufferId,
6432    cols: usize,
6433    ix: Precision,
6434    target: &SpatialTarget,
6435) -> (Vec<StageId>, BufferId) {
6436    let one = wide(1, target);
6437    let bound = crate::spatial::target::int_precision_max(ix).expect("an index type is an integer")
6438        / cols as i64;
6439    let biased = b.t.scratch("row_biased", 1, ix);
6440    let s0 = b.compute(
6441        "row_biased",
6442        flat_space(1),
6443        vec![Slot {
6444            buf: index,
6445            map: fixed(0),
6446        }],
6447        ident_map(one),
6448        *add(sub(op(0), konst(1.0)), konst(bound as f64)),
6449        biased,
6450    );
6451    let guard = ElemAccess {
6452        base: 0,
6453        terms: Vec::new(),
6454        indirect: Some(IndirectTerm {
6455            operand: 0,
6456            stride: 0,
6457            modulus: 2 * bound as usize,
6458        }),
6459    };
6460    let row = b.t.scratch("row_guarded", 1, ix);
6461    let s1 = b.compute(
6462        "row_guarded",
6463        flat_space(1),
6464        vec![Slot {
6465            buf: biased,
6466            map: guard,
6467        }],
6468        ident_map(one),
6469        *sub(op(0), konst(bound as f64)),
6470        row,
6471    );
6472    let lo = b.t.scratch("row_offset", 1, ix);
6473    let s2 = b.compute(
6474        "row_offset",
6475        flat_space(1),
6476        vec![Slot {
6477            buf: row,
6478            map: fixed(0),
6479        }],
6480        ident_map(one),
6481        *mul(op(0), konst(cols as f64)),
6482        lo,
6483    );
6484    (vec![s0, s1, s2], lo)
6485}
6486
6487/// A literal row index, checked at compile time: the flat offset of row `index` of a `rows`-row table.
6488fn literal_row_offset(index: i64, rows: usize, cols: usize) -> Result<i64, SpatialReject> {
6489    if index < 1 || index as usize > rows {
6490        return Err(SpatialReject::NotLowered("a literal row index outside its table — the program fails at run time on every tier, so there is no section to lower"));
6491    }
6492    Ok((index - 1) * cols as i64)
6493}
6494
6495/// `x[j] = T[lo + j]` — the embedding lookup: row `index` (1-based) of a flat `rows·cols` table.
6496///
6497/// ★ THE ROW IS DATA. A named index is an operand: its flat offset `lo = (index − 1)·cols` is computed on the
6498/// integer datapath and the gather reads `T` through an [`IndirectTerm`] of extent `rows·cols` — so a token
6499/// id past the table (or below 1: `lo < 0`) FAULTS in the evaluator, exactly where every tier's run-time
6500/// check fails, and is never wrapped into somebody else's row.
6501///
6502/// `dense`: the table is laid out as a matrix, a tile row holding `tile` rows — see `dense_row_gather`.
6503pub fn from_row_gather(
6504    spec: &RowGatherSpec,
6505    names: &StageNames,
6506    float: bool,
6507    dense: bool,
6508    target: &SpatialTarget,
6509) -> Result<SpatialProgram, SpatialReject> {
6510    let (rows, cols) = (spec.rows, spec.cols);
6511    if rows == 0 || cols == 0 {
6512        return Err(SpatialReject::NotLowered(
6513            "a row gather from an empty table",
6514        ));
6515    }
6516    if dense {
6517        return dense_row_gather(spec, names, float, target);
6518    }
6519    if cols % target.tile.cols == 0 {
6520        return tile_row_gather(spec, names, float, target);
6521    }
6522    let n = rows * cols;
6523    let mut b = FlatBuild::new(target)?;
6524    let (t_name, t_fill) = names.operand(0);
6525    let dt = table_precision(&b, float, t_fill, target)?;
6526    let table = b.input_as(t_name, n, t_fill, dt, target);
6527    let out_buf = b.t.buffer("row", tiles_of(cols, target), dt);
6528    let ow = wide(tiles_of(cols, target), target);
6529    let chain = match spec.index {
6530        RowIndex::Literal(i) => {
6531            let lo = literal_row_offset(i, rows, cols)?;
6532            vec![b.compute(
6533                "gather_row",
6534                flat_space(cols),
6535                vec![Slot {
6536                    buf: table,
6537                    map: at(lo),
6538                }],
6539                ident_map(ow),
6540                TensorBody::Operand(0),
6541                out_buf,
6542            )]
6543        }
6544        RowIndex::Named(_) => {
6545            let (i_name, i_fill) = names.operand(1);
6546            let ix = index_type(target, index_space(n as u64, cols as u64, i_fill))?;
6547            let index = b.input_as(i_name, 1, i_fill, ix, target);
6548            let (mut chain, lo) = row_offset_stages(&mut b, index, cols, ix, target);
6549            let gather = ElemAccess {
6550                base: 0,
6551                terms: vec![(LV_GEN_COL, 1)],
6552                indirect: Some(IndirectTerm {
6553                    operand: 1,
6554                    stride: 1,
6555                    modulus: n,
6556                }),
6557            };
6558            chain.push(b.compute(
6559                "gather_row",
6560                flat_space(cols),
6561                vec![
6562                    Slot {
6563                        buf: table,
6564                        map: gather,
6565                    },
6566                    Slot {
6567                        buf: lo,
6568                        map: fixed(0),
6569                    },
6570                ],
6571                ident_map(ow),
6572                TensorBody::Operand(0),
6573                out_buf,
6574            ));
6575            chain
6576        }
6577    };
6578    let claim = if float {
6579        approx("row gather", dt)
6580    } else {
6581        NumericsClaim::Exact
6582    };
6583    b.finish_as(
6584        &names.out,
6585        cols,
6586        chain,
6587        out_buf,
6588        format!("row_gather_{rows}x{cols}"),
6589        claim,
6590        dt,
6591        target,
6592    )
6593}
6594
6595/// `X = rows I of T` — the block gather: row `k` of `X` is row `I[k]` (1-based) of a flat `rows·cols` table, for every
6596/// index of `I` at once, in ONE section; `X` is written as flat rows (the layout pass bridges it once where the program
6597/// lays `X` out as a matrix).
6598///
6599/// ★ EVERY INDEX IS DATA, CHECKED WHERE IT IS READ: taken less one and read through a guard of stride zero against the
6600/// table's rows, so an index outside the table faults where every tier's check fails instead of reading another row.
6601///   · Rows of WHOLE TILES (a real embedding table): item `(k, b)` moves row `k`'s `b`-th `blk`-tile block. It reads its
6602///     own index — element `k` of `I`, `k` taken from a compile-time table holding it in tile `k` — and loads the row's
6603///     block through an [`IndirectTile`] over the guarded `I[k] − 1`, so exactly the gathered rows' tiles move.
6604///   · Rows narrower than a tile: one item gathers the block from the table held whole — one band of the output per
6605///     row, each reading the table at its row's offset `(I[k] − 1)·cols` — refused by name when the table or the block
6606///     is wider than a site's page.
6607pub fn from_rows_gather(
6608    spec: &RowsGatherSpec,
6609    names: &StageNames,
6610    float: bool,
6611    dense: bool,
6612    target: &SpatialTarget,
6613) -> Result<SpatialProgram, SpatialReject> {
6614    let (rows, cols, n) = (spec.rows, spec.cols, spec.count);
6615    if rows == 0 || cols == 0 || n == 0 {
6616        return Err(SpatialReject::NotLowered(
6617            "a rows gather from an empty table, or of no indices",
6618        ));
6619    }
6620    if dense {
6621        return Err(SpatialReject::NotLowered(
6622            "a rows gather from a table laid out as a matrix",
6623        ));
6624    }
6625    if cols % target.tile.cols == 0 {
6626        return tile_rows_gather(spec, names, float, target);
6627    }
6628    let tn = rows * cols;
6629    if tiles_of(tn, target) > row_block_tiles(target)
6630        || tiles_of(n * cols, target) > row_block_tiles(target)
6631    {
6632        return Err(SpatialReject::NotLowered("a rows gather of rows narrower than a tile holds its table and its block whole, and this one is wider than a site's page"));
6633    }
6634    let mut b = FlatBuild::new(target)?;
6635    let (t_name, t_fill) = names.operand(0);
6636    let dt = table_precision(&b, float, t_fill, target)?;
6637    let table = b.input_as(t_name, tn, t_fill, dt, target);
6638    let (i_name, i_fill) = names.operand(1);
6639    let ix = index_type(target, index_space(tn as u64, cols as u64, i_fill))?;
6640    let idx = b.input_as(i_name, n, i_fill, ix, target);
6641    let (i_t, w_n) = (tiles_of(n, target), wide(tiles_of(n, target), target));
6642    let biased = b.t.scratch("rows_biased", i_t, ix);
6643    let s0 = b.compute(
6644        "rows_biased",
6645        flat_space(n),
6646        vec![Slot {
6647            buf: idx,
6648            map: at(0),
6649        }],
6650        ident_map(w_n),
6651        *sub(op(0), konst(1.0)),
6652        biased,
6653    );
6654    let guard = ElemAccess {
6655        base: 0,
6656        terms: vec![(LV_GEN_COL, 1)],
6657        indirect: Some(IndirectTerm {
6658            operand: 0,
6659            stride: 0,
6660            modulus: rows,
6661        }),
6662    };
6663    let guarded = b.t.scratch("rows_guarded", i_t, ix);
6664    let s1 = b.compute(
6665        "rows_guarded",
6666        flat_space(n),
6667        vec![Slot {
6668            buf: biased,
6669            map: guard,
6670        }],
6671        ident_map(w_n),
6672        TensorBody::Operand(0),
6673        guarded,
6674    );
6675    let lo = b.t.scratch("rows_offset", i_t, ix);
6676    let s2 = b.compute(
6677        "rows_offset",
6678        flat_space(n),
6679        vec![Slot {
6680            buf: guarded,
6681            map: at(0),
6682        }],
6683        ident_map(w_n),
6684        *mul(op(0), konst(cols as f64)),
6685        lo,
6686    );
6687    // Band `k` — the output's columns `[k·cols, (k + 1)·cols)` — reads the table at `c − k·cols + (I[k] − 1)·cols`.
6688    let (mut slots, mut bands) = (Vec::new(), Vec::new());
6689    for k in 0..n {
6690        slots.push(Slot {
6691            buf: lo,
6692            map: fixed(k as i64),
6693        });
6694        let offset = slots.len() - 1;
6695        slots.push(Slot {
6696            buf: table,
6697            map: ElemAccess {
6698                indirect: Some(IndirectTerm {
6699                    operand: offset,
6700                    stride: 1,
6701                    modulus: tn,
6702                }),
6703                ..at(-((k * cols) as i64))
6704            },
6705        });
6706        bands.push(((k + 1) * cols, TensorBody::Operand(slots.len() - 1)));
6707    }
6708    let out_t = tiles_of(n * cols, target);
6709    let out_buf = b.t.buffer("rows", out_t, dt);
6710    let s3 = b.compute(
6711        "gather_rows",
6712        flat_space(n * cols),
6713        slots,
6714        ident_map(wide(out_t, target)),
6715        banded(bands),
6716        out_buf,
6717    );
6718    let claim = if float {
6719        approx("rows gather", dt)
6720    } else {
6721        NumericsClaim::Exact
6722    };
6723    b.finish_as(
6724        &names.out,
6725        n * cols,
6726        vec![s0, s1, s2, s3],
6727        out_buf,
6728        format!("rows_gather_{n}x{rows}x{cols}"),
6729        claim,
6730        dt,
6731        target,
6732    )
6733}
6734
6735/// [`from_rows_gather`] for a table whose rows are WHOLE TILES: item `(k, b)` moves row `I[k]`'s `b`-th `blk`-tile block
6736/// into row `k` of `X`.
6737fn tile_rows_gather(
6738    spec: &RowsGatherSpec,
6739    names: &StageNames,
6740    float: bool,
6741    target: &SpatialTarget,
6742) -> Result<SpatialProgram, SpatialReject> {
6743    let (rows, cols, n) = (spec.rows, spec.cols, spec.count);
6744    let tw = target.tile.cols;
6745    let row_t = cols / tw;
6746    let blk = row_block(row_t, target);
6747    let tn = rows * cols;
6748    let (t_name, t_fill) = names.operand(0);
6749    let dt = if float {
6750        elem_precision(target)?
6751    } else {
6752        index_type(target, fill_abs_bound(t_fill).unwrap_or(0))?
6753    };
6754    let (i_name, i_fill) = names.operand(1);
6755    let ix = index_type(target, index_space(rows.max(n) as u64, 1, i_fill))?;
6756    let mut t = Template::new(target)?;
6757    let table = t.tensor(
6758        t_name,
6759        TensorRole::Input,
6760        1,
6761        tn,
6762        dt,
6763        super::from_tensor_op::init_of(t_fill, tn),
6764    );
6765    let indices = t.tensor(
6766        i_name,
6767        TensorRole::Input,
6768        1,
6769        n,
6770        ix,
6771        super::from_tensor_op::init_of(i_fill, n),
6772    );
6773    // Item `k`'s own number, in tile `k` of a compile-time table: what makes element `k` of `I` an item's own read.
6774    let item_of = t.tensor(
6775        &format!("{}__rows_item", names.out),
6776        TensorRole::Param,
6777        1,
6778        n * tw,
6779        ix,
6780        TensorInit::Values(
6781            (0..n as i64)
6782                .flat_map(|k| std::iter::once(k).chain(std::iter::repeat_n(0, tw - 1)))
6783                .collect(),
6784        ),
6785    );
6786    let w1 = wide(1, target);
6787    let scalar = |t: &mut Template,
6788                  name: &str,
6789                  ins: Vec<BufferId>,
6790                  maps: Vec<ElemAccess>,
6791                  body: TensorBody,
6792                  out: BufferId| {
6793        t.stage(
6794            name,
6795            StageKind::Compute(ComputeOp::Generic {
6796                iters: flat_space(1),
6797                ins: maps,
6798                out: ident_map(w1),
6799                body,
6800            }),
6801            ins,
6802            vec![out],
6803        )
6804    };
6805    let i_t = tiles_of(n, target);
6806    let iin = t.buffer("in_indices", i_t, ix);
6807    let li = t.load("load_indices", indices, vec![], i_t, iin);
6808    let kin = t.buffer("in_item", 1, ix);
6809    let lk = t.load_block_at("load_item", item_of, 0, vec![(LV_ITEM_Y, 1)], 1, 1, 1, kin);
6810    // `I[k] − 1`: element `k` of the indices, then checked against the table's rows.
6811    let pick = ElemAccess {
6812        base: 0,
6813        terms: Vec::new(),
6814        indirect: Some(IndirectTerm {
6815            operand: 1,
6816            stride: 1,
6817            modulus: n,
6818        }),
6819    };
6820    let biased = t.scratch_deep("index_biased", 1, 1, ix);
6821    let s_biased = scalar(
6822        &mut t,
6823        "index_biased",
6824        vec![iin, kin],
6825        vec![pick, fixed(0)],
6826        *sub(op(0), konst(1.0)),
6827        biased,
6828    );
6829    t.edge(li, iin, s_biased);
6830    t.edge(lk, kin, s_biased);
6831    let row0 = t.buffer("row_index", 1, ix);
6832    let guard = ElemAccess {
6833        base: 0,
6834        terms: Vec::new(),
6835        indirect: Some(IndirectTerm {
6836            operand: 0,
6837            stride: 0,
6838            modulus: rows,
6839        }),
6840    };
6841    let s_row = scalar(
6842        &mut t,
6843        "row_index",
6844        vec![biased],
6845        vec![guard],
6846        TensorBody::Operand(0),
6847        row0,
6848    );
6849    t.edge(s_biased, biased, s_row);
6850    let origin = TileAddr {
6851        base: 0,
6852        steps: vec![(LV_ITEM_X, blk)],
6853        indirect: Some(IndirectTile {
6854            index: row0,
6855            tile_stride: row_t,
6856            elems_per_tile: 1,
6857            bound: (0, rows - 1),
6858        }),
6859    };
6860    let row = t.buffer("in_row", blk, dt);
6861    let load = t.stage(
6862        "gather_row",
6863        StageKind::Load {
6864            tensor: table,
6865            pattern: AccessPattern {
6866                origin,
6867                block_h: 1,
6868                block_w: blk,
6869                stride_w: 1,
6870                stride_h: row_t,
6871            },
6872            num_blocks: 1,
6873            multicast: None,
6874        },
6875        vec![row0],
6876        vec![row],
6877    );
6878    t.edge(s_row, row0, load);
6879    let (copy, out) = copy_page(&mut t, row, blk, dt, target);
6880    t.edge(load, row, copy);
6881    let x = t.tensor(
6882        &names.out,
6883        TensorRole::Output,
6884        1,
6885        n * cols,
6886        dt,
6887        TensorInit::Zero,
6888    );
6889    let st = t.store(
6890        &format!("store_{}", names.out),
6891        x,
6892        vec![(LV_ITEM_Y, row_t), (LV_ITEM_X, blk)],
6893        blk,
6894        out,
6895    );
6896    t.edge(copy, out, st);
6897    let items: Vec<WorkItem> = (0..n)
6898        .flat_map(|k| (0..row_t / blk).map(move |b| (k, b)))
6899        .enumerate()
6900        .map(|(id, (k, b))| WorkItem {
6901            id: WorkItemId::new(id),
6902            coord: (k, b),
6903            out: TileRect {
6904                row0_t: 0,
6905                col0_t: k * row_t + b * blk,
6906                h_t: 1,
6907                w_t: blk,
6908            },
6909        })
6910        .collect();
6911    let claim = if float {
6912        approx("rows gather", dt)
6913    } else {
6914        NumericsClaim::Exact
6915    };
6916    t.finish(
6917        format!("rows_gather_{n}x{rows}x{cols}"),
6918        items,
6919        AccPolicy::F32,
6920        claim,
6921        target,
6922    )
6923}
6924
6925/// How many tiles of a whole-tile row ONE work item moves: the largest divisor of the row's tiles the
6926/// accumulator holds, since the row passes through the tile datapath one accumulator page at a time. A row that
6927/// fits is one item; Qwen3.8's `d = 5120` (160 tiles) is twenty items of eight.
6928fn row_block(row_t: usize, target: &SpatialTarget) -> usize {
6929    (1..=row_t.min(row_block_tiles(target)))
6930        .rev()
6931        .find(|d| row_t % d == 0)
6932        .unwrap_or(1)
6933}
6934
6935/// The address of item `i`'s `blk`-tile block of a whole-tile row of a flat `rows × (row_t · tile)` table. A
6936/// literal row is affine. A named row adds an [`IndirectTile`] over the 0-based row `index − 1`, computed per item
6937/// on the integer datapath (at the target's index type) into a FIFO the movement engine waits on and retires once
6938/// the address is resolved. Returns the address and the index path's `(stage, buffer)`.
6939fn row_block_address(
6940    t: &mut Template,
6941    index: RowIndex,
6942    operand: (&str, &Fill),
6943    rows: usize,
6944    row_t: usize,
6945    blk: usize,
6946    target: &SpatialTarget,
6947) -> Result<(TileAddr, Option<(StageId, BufferId)>), SpatialReject> {
6948    let steps = vec![(LV_ITEM_X, blk)];
6949    match index {
6950        RowIndex::Literal(i) => {
6951            if i < 1 || i as usize > rows {
6952                return Err(SpatialReject::NotLowered("a literal row index outside its table — the program fails at run time on every tier, so there is no section to lower"));
6953            }
6954            Ok((
6955                TileAddr {
6956                    base: (i as usize - 1) * row_t,
6957                    steps,
6958                    indirect: None,
6959                },
6960                None,
6961            ))
6962        }
6963        RowIndex::Named(_) => {
6964            let ix = index_type(target, index_space(rows as u64, 1, operand.1))?;
6965            let tok = t.tensor(
6966                operand.0,
6967                TensorRole::Input,
6968                1,
6969                1,
6970                ix,
6971                super::from_tensor_op::init_of(operand.1, 1),
6972            );
6973            let tin = t.buffer(&format!("in_{}", operand.0), 1, ix);
6974            let load = t.load(&format!("load_{}", operand.0), tok, vec![], 1, tin);
6975            let row0 = t.buffer("row_index", 1, ix);
6976            let w1 = wide(1, target);
6977            let s = t.stage(
6978                "row_index",
6979                StageKind::Compute(ComputeOp::Generic {
6980                    iters: flat_space(1),
6981                    ins: vec![fixed(0)],
6982                    out: ident_map(w1),
6983                    body: *sub(op(0), konst(1.0)),
6984                }),
6985                vec![tin],
6986                vec![row0],
6987            );
6988            t.edge(load, tin, s);
6989            let indirect = Some(IndirectTile {
6990                index: row0,
6991                tile_stride: row_t,
6992                elems_per_tile: 1,
6993                bound: (0, rows - 1),
6994            });
6995            Ok((
6996                TileAddr {
6997                    base: 0,
6998                    steps,
6999                    indirect,
7000                },
7001                Some((s, row0)),
7002            ))
7003        }
7004    }
7005}
7006
7007/// The items of a whole-tile row moved `blk` tiles at a time: item `i` owns tiles `[i·blk, (i+1)·blk)` of the row.
7008fn row_items(row_t: usize, blk: usize) -> Vec<WorkItem> {
7009    (0..row_t / blk)
7010        .map(|i| WorkItem {
7011            id: WorkItemId::new(i),
7012            coord: (0, i),
7013            out: TileRect {
7014                row0_t: 0,
7015                col0_t: i * blk,
7016                h_t: 1,
7017                w_t: blk,
7018            },
7019        })
7020        .collect()
7021}
7022
7023/// A tile copy of one `blk`-tile page through the tile datapath.
7024fn copy_page(
7025    t: &mut Template,
7026    from: BufferId,
7027    blk: usize,
7028    dt: Precision,
7029    target: &SpatialTarget,
7030) -> (StageId, BufferId) {
7031    let out = t.buffer("row", blk, dt);
7032    let w = wide(blk, target);
7033    let s = t.stage(
7034        "copy_row",
7035        StageKind::Compute(ComputeOp::Generic {
7036            iters: flat_space(w),
7037            ins: vec![ident_map(w)],
7038            out: ident_map(w),
7039            body: TensorBody::Operand(0),
7040        }),
7041        vec![from],
7042        vec![out],
7043    );
7044    (s, out)
7045}
7046
7047/// `x = row i of T` for a table whose rows are WHOLE TILES — the real-width embedding lookup. Each work item
7048/// Loads its `blk` tiles of the row (an [`IndirectTile`] for a named row) and passes them out, so exactly the
7049/// row's tiles move and a table far wider than a site is never brought on chip.
7050fn tile_row_gather(
7051    spec: &RowGatherSpec,
7052    names: &StageNames,
7053    float: bool,
7054    target: &SpatialTarget,
7055) -> Result<SpatialProgram, SpatialReject> {
7056    let (rows, cols) = (spec.rows, spec.cols);
7057    let row_t = cols / target.tile.cols;
7058    let blk = row_block(row_t, target);
7059    let n = rows * cols;
7060    let (t_name, t_fill) = names.operand(0);
7061    let dt = if float {
7062        elem_precision(target)?
7063    } else {
7064        index_type(target, fill_abs_bound(t_fill).unwrap_or(0))?
7065    };
7066    let mut t = Template::new(target)?;
7067    let table = t.tensor(
7068        t_name,
7069        TensorRole::Input,
7070        1,
7071        n,
7072        dt,
7073        super::from_tensor_op::init_of(t_fill, n),
7074    );
7075    let index_operand = if matches!(spec.index, RowIndex::Named(_)) {
7076        names.operand(1)
7077    } else {
7078        (t_name, t_fill)
7079    };
7080    let (addr, index) =
7081        row_block_address(&mut t, spec.index, index_operand, rows, row_t, blk, target)?;
7082    let x = t.tensor(
7083        &names.out,
7084        TensorRole::Output,
7085        1,
7086        cols,
7087        dt,
7088        TensorInit::Zero,
7089    );
7090    let row = t.buffer("in_row", blk, dt);
7091    let load = t.stage(
7092        "gather_row",
7093        StageKind::Load {
7094            tensor: table,
7095            pattern: AccessPattern {
7096                origin: addr,
7097                block_h: 1,
7098                block_w: blk,
7099                stride_w: 1,
7100                stride_h: row_t,
7101            },
7102            num_blocks: 1,
7103            multicast: None,
7104        },
7105        index.map(|(_, b)| vec![b]).unwrap_or_default(),
7106        vec![row],
7107    );
7108    if let Some((s, b)) = index {
7109        t.edge(s, b, load);
7110    }
7111    let (copy, out) = copy_page(&mut t, row, blk, dt, target);
7112    t.edge(load, row, copy);
7113    let st = t.store(
7114        &format!("store_{}", names.out),
7115        x,
7116        vec![(LV_ITEM_X, blk)],
7117        blk,
7118        out,
7119    );
7120    t.edge(copy, out, st);
7121    let claim = if float {
7122        approx("row gather", dt)
7123    } else {
7124        NumericsClaim::Exact
7125    };
7126    t.finish(
7127        format!("row_gather_{rows}x{cols}"),
7128        row_items(row_t, blk),
7129        AccPolicy::F32,
7130        claim,
7131        target,
7132    )
7133}
7134
7135/// `Set row i of T to S` for a table whose rows are WHOLE TILES — the real-width KV-cache append. Each work item
7136/// Stores its `blk` tiles of the source into the row (an [`IndirectTile`] scatter for a named row), so exactly
7137/// the row's tiles move.
7138///
7139/// ★ THE TABLE IS `State`, NOT AN `Output`: the step updates it IN PLACE. An output must be written whole; this
7140/// section writes one row and every other tile carries the value it came in with — which is why the table is
7141/// declared with its SOURCE fill, the value the untouched rows hold.
7142fn tile_set_row(
7143    spec: &SetRowSpec,
7144    names: &StageNames,
7145    float: bool,
7146    target: &SpatialTarget,
7147) -> Result<SpatialProgram, SpatialReject> {
7148    let (rows, cols, count) = (spec.rows, spec.cols, spec.count);
7149    let row_t = cols / target.tile.cols;
7150    // A block of `count` rows is `count·row_t` contiguous tiles from its first row's: the one-row scatter, wider, its
7151    // first row bounded so the last is the table's.
7152    let span_t = count * row_t;
7153    let blk = row_block(span_t, target);
7154    let n = rows * cols;
7155    let (t_name, t_fill) = names.operand(0);
7156    let (s_name, s_fill) = names.operand(1);
7157    let dt = if float {
7158        elem_precision(target)?
7159    } else {
7160        index_type(
7161            target,
7162            fill_abs_bound(t_fill)
7163                .unwrap_or(0)
7164                .max(fill_abs_bound(s_fill).unwrap_or(0)),
7165        )?
7166    };
7167    let mut t = Template::new(target)?;
7168    let src = t.tensor(
7169        s_name,
7170        TensorRole::Input,
7171        1,
7172        count * cols,
7173        dt,
7174        super::from_tensor_op::init_of(s_fill, count * cols),
7175    );
7176    let index_operand = if matches!(spec.index, RowIndex::Named(_)) {
7177        names.operand(2)
7178    } else {
7179        (s_name, s_fill)
7180    };
7181    let (addr, index) = row_block_address(
7182        &mut t,
7183        spec.index,
7184        index_operand,
7185        rows + 1 - count,
7186        row_t,
7187        blk,
7188        target,
7189    )?;
7190    let table = t.tensor(
7191        t_name,
7192        TensorRole::State,
7193        1,
7194        n,
7195        dt,
7196        super::from_tensor_op::init_of(t_fill, n),
7197    );
7198    let sin = t.buffer(&format!("in_{s_name}"), blk, dt);
7199    let ls = t.load(
7200        &format!("load_{s_name}"),
7201        src,
7202        vec![(LV_ITEM_X, blk)],
7203        blk,
7204        sin,
7205    );
7206    let (copy, out) = copy_page(&mut t, sin, blk, dt, target);
7207    t.edge(ls, sin, copy);
7208    let st = t.stage(
7209        &format!("scatter_{t_name}"),
7210        StageKind::Store {
7211            tensor: table,
7212            pattern: AccessPattern {
7213                origin: addr,
7214                block_h: 1,
7215                block_w: blk,
7216                stride_w: 1,
7217                stride_h: row_t,
7218            },
7219            sub_blocks: (1, 1),
7220        },
7221        std::iter::once(out).chain(index.map(|(_, b)| b)).collect(),
7222        vec![],
7223    );
7224    t.edge(copy, out, st);
7225    if let Some((s, b)) = index {
7226        t.edge(s, b, st);
7227    }
7228    let claim = if float {
7229        approx("row write", dt)
7230    } else {
7231        NumericsClaim::Exact
7232    };
7233    let mut prog = t.finish(
7234        set_row_name(spec, ""),
7235        row_items(span_t, blk),
7236        AccPolicy::F32,
7237        claim,
7238        target,
7239    )?;
7240    // ★ THE PLAN IS THE ROW'S, NOT THE TABLE'S. `trivial_plan` sizes its grid from the result tensor, and the
7241    // result here is the whole carried table — a grid of every row, where the work is one.
7242    prog.work.plan = Some(BlockPlan {
7243        per_site_m: 1,
7244        per_site_n: blk,
7245        grid: (1, span_t / blk),
7246        inner: InnerLoop {
7247            num_blocks: 1,
7248            block_k_tiles: 1,
7249        },
7250        block: BlockShape {
7251            osh: 1,
7252            osw: blk,
7253            ibw: 1,
7254        },
7255        acc: AccPolicy::for_result(dt),
7256        pack_accumulate: false,
7257        engines: EnginePacking::LoadsTogether,
7258        granularity: Granularity::PerTile,
7259    });
7260    Ok(prog)
7261}
7262
7263/// Where row `r = i − 1` of a DENSE `rows × cols` table lives: tile row `⌊r / tile⌋`, at in-tile row `r mod tile`.
7264///
7265/// ★ THE DIVISION STAYS IN THE ADDRESS. `⌊r / tile⌋` is what an [`IndirectTile`] with `elems_per_tile = tile`
7266/// computes, so the tile row is addressed by `r` itself. The in-tile row `r mod tile` is `r − tile·⌊r / tile⌋`:
7267/// the second term is read the same way from a compile-time table whose tile `q` holds `tile·q` — one tile moved,
7268/// no division on the datapath.
7269struct DenseRow {
7270    /// One tile-row address per movement that touches the row, each with the index copy it waits on.
7271    tile_rows: Vec<(TileAddr, Option<(StageId, BufferId)>)>,
7272    in_tile: InTileRow,
7273}
7274
7275/// The in-tile row of a dense row: known at compile time, or computed onto the integer datapath.
7276#[derive(Clone, Copy)]
7277enum InTileRow {
7278    Fixed(usize),
7279    Data {
7280        row: BufferId,
7281        stage: StageId,
7282        ix: Precision,
7283    },
7284}
7285
7286/// The addresses of row `index` of a dense `rows`-row table whose rows are `row_t` tiles, for `movements` tile
7287/// moves of `steps` — every consumer of the index waits on its own copy of `r = i − 1`.
7288fn dense_row(
7289    t: &mut Template,
7290    index: RowIndex,
7291    operand: (&str, &Fill),
7292    table: &str,
7293    rows: usize,
7294    row_t: usize,
7295    steps: &[(LoopVar, usize)],
7296    movements: usize,
7297    target: &SpatialTarget,
7298) -> Result<DenseRow, SpatialReject> {
7299    let tr = target.tile.rows;
7300    match index {
7301        RowIndex::Literal(i) => {
7302            if i < 1 || i as usize > rows {
7303                return Err(SpatialReject::NotLowered("a literal row index outside its table — the program fails at run time on every tier, so there is no section to lower"));
7304            }
7305            let r = i as usize - 1;
7306            let at = TileAddr {
7307                base: (r / tr) * row_t,
7308                steps: steps.to_vec(),
7309                indirect: None,
7310            };
7311            Ok(DenseRow {
7312                tile_rows: vec![(at, None); movements],
7313                in_tile: InTileRow::Fixed(r % tr),
7314            })
7315        }
7316        RowIndex::Named(_) => {
7317            let (i_name, i_fill) = operand;
7318            let ix = index_type(target, index_space(rows as u64, 1, i_fill))?;
7319            let pos = t.tensor(
7320                i_name,
7321                TensorRole::Input,
7322                1,
7323                1,
7324                ix,
7325                super::from_tensor_op::init_of(i_fill, 1),
7326            );
7327            let w1 = wide(1, target);
7328            let row_index = |t: &mut Template, what: &str| {
7329                let tin = t.buffer(&format!("in_{i_name}_{what}"), 1, ix);
7330                let load = t.load(&format!("load_{i_name}_{what}"), pos, vec![], 1, tin);
7331                let r = t.buffer(&format!("row_{what}"), 1, ix);
7332                let s = t.stage(
7333                    &format!("row_{what}"),
7334                    StageKind::Compute(ComputeOp::Generic {
7335                        iters: flat_space(1),
7336                        ins: vec![fixed(0)],
7337                        out: ident_map(w1),
7338                        body: *sub(op(0), konst(1.0)),
7339                    }),
7340                    vec![tin],
7341                    vec![r],
7342                );
7343                t.edge(load, tin, s);
7344                (s, r)
7345            };
7346            let names = ["gather", "scatter"];
7347            let feeds: Vec<(StageId, BufferId)> = (0..movements)
7348                .map(|m| row_index(t, names.get(m).copied().unwrap_or("move")))
7349                .collect();
7350            let (s_base, r_base) = row_index(t, "base");
7351            let (s_in, r_in) = row_index(t, "in_tile");
7352            let tile_rows = feeds
7353                .into_iter()
7354                .map(|(s, index)| {
7355                    (
7356                        TileAddr {
7357                            base: 0,
7358                            steps: steps.to_vec(),
7359                            indirect: Some(IndirectTile {
7360                                index,
7361                                tile_stride: row_t,
7362                                elems_per_tile: tr,
7363                                bound: (0, rows - 1),
7364                            }),
7365                        },
7366                        Some((s, index)),
7367                    )
7368                })
7369                .collect();
7370            let rows_t = rows.div_ceil(tr);
7371            let bases: Vec<i64> = (0..rows_t * target.tile.cols)
7372                .map(|e| ((e / target.tile.cols) * tr) as i64)
7373                .collect();
7374            let base_table = t.tensor(
7375                &format!("{table}__row_bases"),
7376                TensorRole::Param,
7377                1,
7378                rows_t * target.tile.cols,
7379                ix,
7380                TensorInit::Values(bases),
7381            );
7382            let base_in = t.buffer("in_row_base", 1, ix);
7383            let base_addr = TileAddr {
7384                base: 0,
7385                steps: Vec::new(),
7386                indirect: Some(IndirectTile {
7387                    index: r_base,
7388                    tile_stride: 1,
7389                    elems_per_tile: tr,
7390                    bound: (0, rows - 1),
7391                }),
7392            };
7393            let l_base = t.stage(
7394                "load_row_base",
7395                StageKind::Load {
7396                    tensor: base_table,
7397                    pattern: AccessPattern {
7398                        origin: base_addr,
7399                        block_h: 1,
7400                        block_w: 1,
7401                        stride_w: 1,
7402                        stride_h: 1,
7403                    },
7404                    num_blocks: 1,
7405                    multicast: None,
7406                },
7407                vec![r_base],
7408                vec![base_in],
7409            );
7410            t.edge(s_base, r_base, l_base);
7411            let row = t.scratch_deep("in_tile_row_index", 1, 1, ix);
7412            let stage = t.stage(
7413                "in_tile_row_index",
7414                StageKind::Compute(ComputeOp::Generic {
7415                    iters: flat_space(1),
7416                    ins: vec![fixed(0), fixed(0)],
7417                    out: ident_map(w1),
7418                    body: *sub(op(0), op(1)),
7419                }),
7420                vec![r_in, base_in],
7421                vec![row],
7422            );
7423            t.edge(s_in, r_in, stage);
7424            t.edge(l_base, base_in, stage);
7425            Ok(DenseRow {
7426                tile_rows,
7427                in_tile: InTileRow::Data { row, stage, ix },
7428            })
7429        }
7430    }
7431}
7432
7433/// The plan of a section that moves one tile row, `row_t / blk` items across it.
7434fn tile_row_plan(row_t: usize, blk: usize, dt: Precision) -> BlockPlan {
7435    BlockPlan {
7436        per_site_m: 1,
7437        per_site_n: blk,
7438        grid: (1, row_t / blk),
7439        inner: InnerLoop {
7440            num_blocks: 1,
7441            block_k_tiles: 1,
7442        },
7443        block: BlockShape {
7444            osh: 1,
7445            osw: blk,
7446            ibw: 1,
7447        },
7448        acc: AccPolicy::for_result(dt),
7449        pack_accumulate: false,
7450        engines: EnginePacking::LoadsTogether,
7451        granularity: Granularity::PerTile,
7452    }
7453}
7454
7455/// `Set row i of T to S` for a DENSE table — `rows × cols`, a tile row holding `tile` rows: the KV cache a flash
7456/// attention reads, a stacked batch a matmul reads. Row `r = i − 1` lives where [`dense_row`] says, so each work
7457/// item READS its `blk` tiles of that tile row, replaces in-tile row `r mod tile` with its `blk` tiles of the
7458/// source, and WRITES them back: `T'(y, c) = (y = r mod tile) ? S(c) : T(y, c)` over the tile row. A row that is
7459/// not whole tiles is the same write over its last, partial tile.
7460///
7461/// ★ THE TABLE IS `State`, read and written in one section: the write is downstream of the read on every site, so
7462/// the other rows of the tile row carry the value they came in with and nothing races the rewrite.
7463fn dense_set_row(
7464    spec: &SetRowSpec,
7465    names: &StageNames,
7466    float: bool,
7467    target: &SpatialTarget,
7468) -> Result<SpatialProgram, SpatialReject> {
7469    let (rows, cols) = (spec.rows, spec.cols);
7470    let row_t = cols.div_ceil(target.tile.cols);
7471    let blk = row_block(row_t, target);
7472    let n = rows * cols;
7473    let (t_name, t_fill) = names.operand(0);
7474    let (s_name, s_fill) = names.operand(1);
7475    let dt = if float {
7476        elem_precision(target)?
7477    } else {
7478        index_type(
7479            target,
7480            fill_abs_bound(t_fill)
7481                .unwrap_or(0)
7482                .max(fill_abs_bound(s_fill).unwrap_or(0)),
7483        )?
7484    };
7485    let mut t = Template::new(target)?;
7486    let src = t.tensor(
7487        s_name,
7488        TensorRole::Input,
7489        1,
7490        cols,
7491        dt,
7492        super::from_tensor_op::init_of(s_fill, cols),
7493    );
7494    let table = t.tensor(
7495        t_name,
7496        TensorRole::State,
7497        rows,
7498        cols,
7499        dt,
7500        super::from_tensor_op::init_of(t_fill, n),
7501    );
7502    let old = t.buffer("in_tile_row", blk, dt);
7503    let sin = t.buffer(&format!("in_{s_name}"), blk, dt);
7504    let out = t.buffer("tile_row", blk, dt);
7505    let steps = vec![(LV_ITEM_X, blk)];
7506    let pattern = |origin: TileAddr| AccessPattern {
7507        origin,
7508        block_h: 1,
7509        block_w: blk,
7510        stride_w: 1,
7511        stride_h: row_t,
7512    };
7513    let index_operand = if matches!(spec.index, RowIndex::Named(_)) {
7514        names.operand(2)
7515    } else {
7516        (s_name, s_fill)
7517    };
7518    let at = dense_row(
7519        &mut t,
7520        spec.index,
7521        index_operand,
7522        t_name,
7523        rows,
7524        row_t,
7525        &steps,
7526        2,
7527        target,
7528    )?;
7529    let (load_addr, load_feed) = at.tile_rows[0].clone();
7530    let (store_addr, store_feed) = at.tile_rows[1].clone();
7531    // The in-tile row the select replaces: `[lo, lo + 1)`.
7532    let w1 = wide(1, target);
7533    let in_tile = match at.in_tile {
7534        InTileRow::Fixed(y) => RowSelect::Fixed(y),
7535        InTileRow::Data {
7536            row: lo,
7537            stage: s_lo,
7538            ix,
7539        } => {
7540            let hi = t.scratch_deep("in_tile_row_end", 1, 1, ix);
7541            let s_hi = t.stage(
7542                "in_tile_row_end",
7543                StageKind::Compute(ComputeOp::Generic {
7544                    iters: flat_space(1),
7545                    ins: vec![fixed(0)],
7546                    out: ident_map(w1),
7547                    body: *add(op(0), konst(1.0)),
7548                }),
7549                vec![lo],
7550                vec![hi],
7551            );
7552            t.edge(s_lo, lo, s_hi);
7553            RowSelect::Data {
7554                lo,
7555                hi,
7556                stages: (s_lo, s_hi),
7557            }
7558        }
7559    };
7560    let l_old = t.stage(
7561        &format!("gather_{t_name}"),
7562        StageKind::Load {
7563            tensor: table,
7564            pattern: pattern(load_addr),
7565            num_blocks: 1,
7566            multicast: None,
7567        },
7568        load_feed.map(|(_, b)| vec![b]).unwrap_or_default(),
7569        vec![old],
7570    );
7571    if let Some((s, b)) = load_feed {
7572        t.edge(s, b, l_old);
7573    }
7574    let l_src = t.load(&format!("load_{s_name}"), src, steps.clone(), blk, sin);
7575    // Operand 0 the tile row, operand 1 the source row (its page's row 0, broadcast down the tile row).
7576    let w = wide(blk, target);
7577    let source_row = ElemAccess {
7578        base: 0,
7579        terms: vec![(LV_GEN_COL, 1)],
7580        indirect: None,
7581    };
7582    let (ins, maps, lo_rhs, hi_rhs) = match in_tile {
7583        RowSelect::Fixed(y) => (
7584            vec![old, sin],
7585            vec![ident_map(w), source_row],
7586            PredRhs::Const(y),
7587            PredRhs::Const(y + 1),
7588        ),
7589        RowSelect::Data { lo, hi, .. } => (
7590            vec![old, sin, lo, hi],
7591            vec![ident_map(w), source_row, fixed(0), fixed(0)],
7592            PredRhs::Operand(2),
7593            PredRhs::Operand(3),
7594        ),
7595    };
7596    let body = TensorBody::Select {
7597        cond: Pred::AxisGe(LV_GEN_ROW, lo_rhs),
7598        then: Box::new(TensorBody::Select {
7599            cond: Pred::AxisLt(LV_GEN_ROW, hi_rhs),
7600            then: op(1),
7601            otherwise: op(0),
7602        }),
7603        otherwise: op(0),
7604    };
7605    let s_rw = t.stage(
7606        "write_row",
7607        StageKind::Compute(ComputeOp::Generic {
7608            iters: super::from_tensor_op::space(target, w, None),
7609            ins: maps,
7610            out: ident_map(w),
7611            body,
7612        }),
7613        ins,
7614        vec![out],
7615    );
7616    t.edge(l_old, old, s_rw);
7617    t.edge(l_src, sin, s_rw);
7618    if let RowSelect::Data {
7619        hi,
7620        stages: (_, s_hi),
7621        ..
7622    } = in_tile
7623    {
7624        t.edge(s_hi, hi, s_rw);
7625    }
7626    let store_ins: Vec<BufferId> = std::iter::once(out)
7627        .chain(store_feed.map(|(_, b)| b))
7628        .collect();
7629    let st = t.stage(
7630        &format!("scatter_{t_name}"),
7631        StageKind::Store {
7632            tensor: table,
7633            pattern: pattern(store_addr),
7634            sub_blocks: (1, 1),
7635        },
7636        store_ins,
7637        vec![],
7638    );
7639    t.edge(s_rw, out, st);
7640    if let Some((s, b)) = store_feed {
7641        t.edge(s, b, st);
7642    }
7643    let claim = if float {
7644        approx("row write", dt)
7645    } else {
7646        NumericsClaim::Exact
7647    };
7648    let mut prog = t.finish(
7649        format!("set_row_{rows}x{cols}_dense"),
7650        row_items(row_t, blk),
7651        AccPolicy::F32,
7652        claim,
7653        target,
7654    )?;
7655    prog.work.plan = Some(tile_row_plan(row_t, blk, dt));
7656    Ok(prog)
7657}
7658
7659/// `Set rows from P of T to X` for a DENSE table — the KV cache a prefill chunk appends its `count` keys or values to,
7660/// `X` a dense `count × cols` block. Rows `r … r + count − 1` (`r = P − 1`) are replaced; the block starts at in-tile
7661/// row `o = r mod tile` of tile row `K0 = ⌊r / tile⌋`, both DATA, and touches at most `m = min(⌈count/tile⌉ + 1, n_t)`
7662/// tile rows of the table's `n_t`.
7663///
7664/// ★ THE ITEMS' TILE ROWS ARE DISTINCT BY THEIR COORDINATE, NOT BY THE DATA. Item `(i, b)` rewrites column block `b` of
7665/// tile row `K0' + i`, `K0' = min(K0, n_t − m)`: affine in `i` over a data start, so every item's tile is its own for
7666/// every `P` — which the certificate's dry run, taking each index at its low bound, sees — and `[K0', K0' + m)` always
7667/// lies in the table and covers the block's tile rows. With `c = r − tile·K0'` (data), row `y` of item `i` is the
7668/// block's row `s = y + i·tile − c`: membership is `c − i·tile ≤ y < count + c − i·tile`, and `s` lies in source tile
7669/// rows `i − δ − 1` and `i − δ` (`δ = K0 − K0'`), loaded as pages `A` and `B`, at in-tile rows `y + tile − o` (for
7670/// `y < o`) and `y − o`. A clamped source tile row only ever holds non-members, and a tile row the block misses is
7671/// written back unchanged.
7672///
7673/// Every number an address needs is a compile-time table read through the data: `K0` and `tile·K0` at tile `⌊r/tile⌋`;
7674/// `K0'`, `tile·K0'` and `m − 1 − δ` at tile `K0`; `i·tile`; the clamped source tile rows at tile `i + m − 1 − δ`. The
7675/// scalars are computed once into scratch and copied into one FIFO per movement that waits on them.
7676fn dense_set_rows(
7677    spec: &SetRowSpec,
7678    names: &StageNames,
7679    float: bool,
7680    target: &SpatialTarget,
7681) -> Result<SpatialProgram, SpatialReject> {
7682    let (rows, cols, count) = (spec.rows, spec.cols, spec.count);
7683    let (tr, tw) = (target.tile.rows, target.tile.cols);
7684    if cols % tw != 0 {
7685        return Err(SpatialReject::NotLowered(
7686            "a block of rows into a dense table whose rows are not whole tiles",
7687        ));
7688    }
7689    let row_t = cols / tw;
7690    let blk = row_block(row_t, target);
7691    let w = wide(blk, target);
7692    let n_t = rows.div_ceil(tr);
7693    let src_t = count.div_ceil(tr);
7694    let m = (src_t + 1).min(n_t);
7695    let (t_name, t_fill) = names.operand(0);
7696    let (s_name, s_fill) = names.operand(1);
7697    let dt = if float {
7698        elem_precision(target)?
7699    } else {
7700        index_type(
7701            target,
7702            fill_abs_bound(t_fill)
7703                .unwrap_or(0)
7704                .max(fill_abs_bound(s_fill).unwrap_or(0)),
7705        )?
7706    };
7707    let ix = index_type(target, (rows + tr * m) as u64)?;
7708    let mut t = Template::new(target)?;
7709    let table = t.tensor(
7710        t_name,
7711        TensorRole::State,
7712        rows,
7713        cols,
7714        dt,
7715        super::from_tensor_op::init_of(t_fill, rows * cols),
7716    );
7717    let src = t.tensor(
7718        s_name,
7719        TensorRole::Input,
7720        count,
7721        cols,
7722        dt,
7723        super::from_tensor_op::init_of(s_fill, count * cols),
7724    );
7725    // `P`: the operand, or the literal as a one-element table — either way checked below.
7726    let first = match spec.index {
7727        RowIndex::Named(_) => {
7728            let (p_name, p_fill) = names.operand(2);
7729            t.tensor(
7730                p_name,
7731                TensorRole::Input,
7732                1,
7733                1,
7734                ix,
7735                super::from_tensor_op::init_of(p_fill, 1),
7736            )
7737        }
7738        RowIndex::Literal(i) => {
7739            if i < 1 || i as usize + count - 1 > rows {
7740                return Err(SpatialReject::NotLowered("a literal block of rows outside its table — the program fails at run time on every tier, so there is no section to lower"));
7741            }
7742            t.tensor(
7743                &format!("{t_name}__block_first"),
7744                TensorRole::Param,
7745                1,
7746                1,
7747                ix,
7748                TensorInit::Values(vec![i]),
7749            )
7750        }
7751    };
7752    let param = |t: &mut Template, what: &str, values: Vec<i64>| {
7753        let n = values.len();
7754        let tiles = values
7755            .into_iter()
7756            .flat_map(|v| std::iter::once(v).chain(std::iter::repeat_n(0, tw - 1)))
7757            .collect();
7758        t.tensor(
7759            &format!("{t_name}__block_{what}"),
7760            TensorRole::Param,
7761            1,
7762            n * tw,
7763            ix,
7764            TensorInit::Values(tiles),
7765        )
7766    };
7767    let start = |q: usize| q.min(n_t - m);
7768    let tile_rows = param(&mut t, "tile_row", (0..n_t as i64).collect());
7769    let tile_bases = param(
7770        &mut t,
7771        "tile_base",
7772        (0..n_t).map(|q| (q * tr) as i64).collect(),
7773    );
7774    let starts = param(&mut t, "start", (0..n_t).map(|q| start(q) as i64).collect());
7775    let start_bases = param(
7776        &mut t,
7777        "start_base",
7778        (0..n_t).map(|q| (start(q) * tr) as i64).collect(),
7779    );
7780    let source_shift = param(
7781        &mut t,
7782        "source_shift",
7783        (0..n_t).map(|q| (m - 1 - (q - start(q))) as i64).collect(),
7784    );
7785    let item_rows = param(
7786        &mut t,
7787        "item_row",
7788        (0..m).map(|i| (i * tr) as i64).collect(),
7789    );
7790    let clamp = |k: i64| k.clamp(0, src_t as i64 - 1);
7791    let sources_a = param(
7792        &mut t,
7793        "source_a",
7794        (0..2 * m as i64 - 1).map(|k| clamp(k - m as i64)).collect(),
7795    );
7796    let sources_b = param(
7797        &mut t,
7798        "source_b",
7799        (0..2 * m as i64 - 1)
7800            .map(|k| clamp(k - m as i64 + 1))
7801            .collect(),
7802    );
7803    let w1 = wide(1, target);
7804    let scalar = |t: &mut Template,
7805                  name: &str,
7806                  ins: Vec<BufferId>,
7807                  maps: Vec<ElemAccess>,
7808                  body: TensorBody,
7809                  out: BufferId| {
7810        t.stage(
7811            name,
7812            StageKind::Compute(ComputeOp::Generic {
7813                iters: flat_space(1),
7814                ins: maps,
7815                out: ident_map(w1),
7816                body,
7817            }),
7818            ins,
7819            vec![out],
7820        )
7821    };
7822    // A scalar's FIFO copy for one movement to wait on.
7823    let fifo = |t: &mut Template, name: &str, from: (StageId, BufferId)| {
7824        let out = t.buffer(name, 1, ix);
7825        let s = scalar(
7826            t,
7827            name,
7828            vec![from.1],
7829            vec![fixed(0)],
7830            TensorBody::Operand(0),
7831            out,
7832        );
7833        t.edge(from.0, from.1, s);
7834        (s, out)
7835    };
7836    // One element of a table, at `steps` plus — when given — the tile a data index selects.
7837    let lookup = |t: &mut Template,
7838                  name: &str,
7839                  from: TensorId,
7840                  steps: Vec<(LoopVar, usize)>,
7841                  index: Option<((StageId, BufferId), usize, (usize, usize))>| {
7842        let out = t.buffer(name, 1, ix);
7843        let indirect = index.map(|((_, b), per, bound)| IndirectTile {
7844            index: b,
7845            tile_stride: 1,
7846            elems_per_tile: per,
7847            bound,
7848        });
7849        let load = t.stage(
7850            &format!("load_{name}"),
7851            StageKind::Load {
7852                tensor: from,
7853                pattern: AccessPattern {
7854                    origin: TileAddr {
7855                        base: 0,
7856                        steps,
7857                        indirect,
7858                    },
7859                    block_h: 1,
7860                    block_w: 1,
7861                    stride_w: 1,
7862                    stride_h: 1,
7863                },
7864                num_blocks: 1,
7865                multicast: None,
7866            },
7867            index.map(|((_, b), _, _)| vec![b]).unwrap_or_default(),
7868            vec![out],
7869        );
7870        if let Some(((s, b), _, _)) = index {
7871            t.edge(s, b, load);
7872        }
7873        (load, out)
7874    };
7875    // `r = P − 1`, checked so every row of the block is one of the table's: a guard of stride zero, as a row index is.
7876    let p_in = t.buffer("in_first", 1, ix);
7877    let l_p = t.load_block_at("load_first", first, 0, Vec::new(), 1, 1, 1, p_in);
7878    let biased = t.scratch_deep("first_biased", 1, 1, ix);
7879    let s_biased = scalar(
7880        &mut t,
7881        "first_biased",
7882        vec![p_in],
7883        vec![fixed(0)],
7884        *sub(op(0), konst(1.0)),
7885        biased,
7886    );
7887    t.edge(l_p, p_in, s_biased);
7888    let r = t.scratch_deep("first_row", 1, 1, ix);
7889    let guard = ElemAccess {
7890        base: 0,
7891        terms: Vec::new(),
7892        indirect: Some(IndirectTerm {
7893            operand: 0,
7894            stride: 0,
7895            modulus: rows - count + 1,
7896        }),
7897    };
7898    let s_r = scalar(
7899        &mut t,
7900        "first_row",
7901        vec![biased],
7902        vec![guard],
7903        TensorBody::Operand(0),
7904        r,
7905    );
7906    t.edge(s_biased, biased, s_r);
7907    let r_at = |t: &mut Template, what: &str| fifo(t, &format!("first_row_{what}"), (s_r, r));
7908    let span = (0, rows - count);
7909    // `K0`, into scratch for its own lookups.
7910    let r_k0 = r_at(&mut t, "tile_row");
7911    let (l_k0, k0_in) = lookup(
7912        &mut t,
7913        "tile_row",
7914        tile_rows,
7915        Vec::new(),
7916        Some((r_k0, tr, span)),
7917    );
7918    let k0 = t.scratch_deep("tile_row", 1, 1, ix);
7919    let s_k0 = scalar(
7920        &mut t,
7921        "tile_row",
7922        vec![k0_in],
7923        vec![fixed(0)],
7924        TensorBody::Operand(0),
7925        k0,
7926    );
7927    t.edge(l_k0, k0_in, s_k0);
7928    let by_k0 = |t: &mut Template,
7929                 what: &str,
7930                 from: TensorId,
7931                 steps: Vec<(LoopVar, usize)>,
7932                 bound: (usize, usize)| {
7933        let k = fifo(t, &format!("tile_row_{what}"), (s_k0, k0));
7934        lookup(t, what, from, steps, Some((k, 1, bound)))
7935    };
7936    // `o = r − tile·K0`.
7937    let r_base = r_at(&mut t, "tile_base");
7938    let (l_base, base) = lookup(
7939        &mut t,
7940        "tile_base",
7941        tile_bases,
7942        Vec::new(),
7943        Some((r_base, tr, span)),
7944    );
7945    let o = t.scratch_deep("block_offset", 1, 1, ix);
7946    let s_o = scalar(
7947        &mut t,
7948        "block_offset",
7949        vec![r, base],
7950        vec![fixed(0), fixed(0)],
7951        *sub(op(0), op(1)),
7952        o,
7953    );
7954    t.edge(l_base, base, s_o);
7955    // The membership thresholds: `lo = r − tile·K0' − i·tile`, `hi = lo + count`.
7956    let (l_sb, start_base) = by_k0(&mut t, "start_base", start_bases, Vec::new(), (0, n_t - 1));
7957    let (l_it, it) = lookup(&mut t, "item_row", item_rows, vec![(LV_ITEM_Y, 1)], None);
7958    let lo = t.scratch_deep("member_lo", 1, 1, ix);
7959    let s_lo = scalar(
7960        &mut t,
7961        "member_lo",
7962        vec![r, start_base, it],
7963        vec![fixed(0), fixed(0), fixed(0)],
7964        *sub(sub(op(0), op(1)), op(2)),
7965        lo,
7966    );
7967    t.edge(l_sb, start_base, s_lo);
7968    t.edge(l_it, it, s_lo);
7969    let hi = t.scratch_deep("member_hi", 1, 1, ix);
7970    let s_hi = scalar(
7971        &mut t,
7972        "member_hi",
7973        vec![lo],
7974        vec![fixed(0)],
7975        *add(op(0), konst(count as f64)),
7976        hi,
7977    );
7978    t.edge(s_lo, lo, s_hi);
7979    // The two source tile rows the item's rows come from: table tile `i + m − 1 − δ`.
7980    let source_page = |t: &mut Template, what: &str, rows_of: TensorId| {
7981        let shift = by_k0(
7982            t,
7983            &format!("{what}_shift"),
7984            source_shift,
7985            Vec::new(),
7986            (0, n_t - 1),
7987        );
7988        let (l_at, at) = lookup(
7989            t,
7990            &format!("{what}_row"),
7991            rows_of,
7992            vec![(LV_ITEM_Y, 1)],
7993            Some((shift, 1, (0, m - 1))),
7994        );
7995        let page = t.buffer(&format!("in_{what}"), blk, dt);
7996        let origin = TileAddr {
7997            base: 0,
7998            steps: vec![(LV_ITEM_X, blk)],
7999            indirect: Some(IndirectTile {
8000                index: at,
8001                tile_stride: row_t,
8002                elems_per_tile: 1,
8003                bound: (0, src_t - 1),
8004            }),
8005        };
8006        let load = t.stage(
8007            &format!("load_{what}"),
8008            StageKind::Load {
8009                tensor: src,
8010                pattern: AccessPattern {
8011                    origin,
8012                    block_h: 1,
8013                    block_w: blk,
8014                    stride_w: 1,
8015                    stride_h: row_t,
8016                },
8017                num_blocks: 1,
8018                multicast: None,
8019            },
8020            vec![at],
8021            vec![page],
8022        );
8023        t.edge(l_at, at, load);
8024        (load, page)
8025    };
8026    let (l_a, page_a) = source_page(&mut t, "source_a", sources_a);
8027    let (l_b, page_b) = source_page(&mut t, "source_b", sources_b);
8028    // The item's tile row `K0' + i`, once per movement that addresses it.
8029    let tile_row_at = |t: &mut Template, what: &str| {
8030        let (l_s, s) = by_k0(
8031            t,
8032            &format!("start_{what}"),
8033            starts,
8034            Vec::new(),
8035            (0, n_t - 1),
8036        );
8037        let addr = TileAddr {
8038            base: 0,
8039            steps: vec![(LV_ITEM_Y, row_t), (LV_ITEM_X, blk)],
8040            indirect: Some(IndirectTile {
8041                index: s,
8042                tile_stride: row_t,
8043                elems_per_tile: 1,
8044                bound: (0, n_t - m),
8045            }),
8046        };
8047        (addr, (l_s, s))
8048    };
8049    let (old_at, (l_s_old, s_old)) = tile_row_at(&mut t, "gather");
8050    let (new_at, (l_s_new, s_new)) = tile_row_at(&mut t, "scatter");
8051    let old = t.buffer("in_tile_row", blk, dt);
8052    let l_old = t.stage(
8053        &format!("gather_{t_name}"),
8054        StageKind::Load {
8055            tensor: table,
8056            pattern: AccessPattern {
8057                origin: old_at,
8058                block_h: 1,
8059                block_w: blk,
8060                stride_w: 1,
8061                stride_h: row_t,
8062            },
8063            num_blocks: 1,
8064            multicast: None,
8065        },
8066        vec![s_old],
8067        vec![old],
8068    );
8069    t.edge(l_s_old, s_old, l_old);
8070    // Operands: 0 the tile row, 1 and 2 the source halves, 3 `o`, 4 `lo`, 5 `hi`.
8071    let shifted = |base: i64| ElemAccess {
8072        base,
8073        terms: vec![(LV_GEN_ROW, w as i64), (LV_GEN_COL, 1)],
8074        indirect: Some(IndirectTerm {
8075            operand: 3,
8076            stride: -(w as i64),
8077            modulus: tr,
8078        }),
8079    };
8080    let from_block = TensorBody::Select {
8081        cond: Pred::AxisLt(LV_GEN_ROW, PredRhs::Operand(3)),
8082        then: op(1),
8083        otherwise: op(2),
8084    };
8085    let body = TensorBody::Select {
8086        cond: Pred::AxisGe(LV_GEN_ROW, PredRhs::Operand(4)),
8087        then: Box::new(TensorBody::Select {
8088            cond: Pred::AxisLt(LV_GEN_ROW, PredRhs::Operand(5)),
8089            then: Box::new(from_block),
8090            otherwise: op(0),
8091        }),
8092        otherwise: op(0),
8093    };
8094    let out = t.buffer("tile_row", blk, dt);
8095    let s_rw = t.stage(
8096        "write_rows",
8097        StageKind::Compute(ComputeOp::Generic {
8098            iters: super::from_tensor_op::space(target, w, None),
8099            ins: vec![
8100                ident_map(w),
8101                shifted((tr * w) as i64),
8102                shifted(0),
8103                fixed(0),
8104                fixed(0),
8105                fixed(0),
8106            ],
8107            out: ident_map(w),
8108            body,
8109        }),
8110        vec![old, page_a, page_b, o, lo, hi],
8111        vec![out],
8112    );
8113    for (from, buf) in [(l_old, old), (l_a, page_a), (l_b, page_b)] {
8114        t.edge(from, buf, s_rw);
8115    }
8116    t.edge(s_o, o, s_rw);
8117    t.edge(s_hi, hi, s_rw);
8118    let st = t.stage(
8119        &format!("scatter_{t_name}"),
8120        StageKind::Store {
8121            tensor: table,
8122            pattern: AccessPattern {
8123                origin: new_at,
8124                block_h: 1,
8125                block_w: blk,
8126                stride_w: 1,
8127                stride_h: row_t,
8128            },
8129            sub_blocks: (1, 1),
8130        },
8131        vec![out, s_new],
8132        vec![],
8133    );
8134    t.edge(s_rw, out, st);
8135    t.edge(l_s_new, s_new, st);
8136    let items: Vec<WorkItem> = (0..m)
8137        .flat_map(|i| (0..row_t / blk).map(move |b| (i, b)))
8138        .enumerate()
8139        .map(|(id, (i, b))| WorkItem {
8140            id: WorkItemId::new(id),
8141            coord: (i, b),
8142            out: TileRect {
8143                row0_t: i,
8144                col0_t: b * blk,
8145                h_t: 1,
8146                w_t: blk,
8147            },
8148        })
8149        .collect();
8150    let claim = if float {
8151        approx("row block write", dt)
8152    } else {
8153        NumericsClaim::Exact
8154    };
8155    let mut prog = t.finish(
8156        set_row_name(spec, "_dense"),
8157        items,
8158        AccPolicy::F32,
8159        claim,
8160        target,
8161    )?;
8162    // An in-place update of carried state: its grid is its items' blocks, one per item.
8163    prog.work.plan = Some(BlockPlan {
8164        grid: (1, m * (row_t / blk)),
8165        ..tile_row_plan(row_t, blk, dt)
8166    });
8167    Ok(prog)
8168}
8169
8170/// `x = row i of T` for a DENSE table — a stacked batch a matmul wrote, `rows × cols` with a tile row holding
8171/// `tile` rows. Each work item Loads its `blk` tiles of tile row `⌊r / tile⌋` ([`dense_row`]) and lifts in-tile row
8172/// `r mod tile` into row 0 of the result's page: `x(c) = T(r mod tile, c)`, the in-tile row a constant offset for
8173/// a literal row and an [`IndirectTerm`] one page row wide for a named one.
8174fn dense_row_gather(
8175    spec: &RowGatherSpec,
8176    names: &StageNames,
8177    float: bool,
8178    target: &SpatialTarget,
8179) -> Result<SpatialProgram, SpatialReject> {
8180    let (rows, cols) = (spec.rows, spec.cols);
8181    let tr = target.tile.rows;
8182    let row_t = cols.div_ceil(target.tile.cols);
8183    let blk = row_block(row_t, target);
8184    let n = rows * cols;
8185    let (t_name, t_fill) = names.operand(0);
8186    let dt = if float {
8187        elem_precision(target)?
8188    } else {
8189        index_type(target, fill_abs_bound(t_fill).unwrap_or(0))?
8190    };
8191    let mut t = Template::new(target)?;
8192    let table = t.tensor(
8193        t_name,
8194        TensorRole::Input,
8195        rows,
8196        cols,
8197        dt,
8198        super::from_tensor_op::init_of(t_fill, n),
8199    );
8200    let steps = vec![(LV_ITEM_X, blk)];
8201    let index_operand = if matches!(spec.index, RowIndex::Named(_)) {
8202        names.operand(1)
8203    } else {
8204        (t_name, t_fill)
8205    };
8206    let at = dense_row(
8207        &mut t,
8208        spec.index,
8209        index_operand,
8210        t_name,
8211        rows,
8212        row_t,
8213        &steps,
8214        1,
8215        target,
8216    )?;
8217    let (addr, feed) = at.tile_rows[0].clone();
8218    let x = t.tensor(
8219        &names.out,
8220        TensorRole::Output,
8221        1,
8222        cols,
8223        dt,
8224        TensorInit::Zero,
8225    );
8226    let tile_row = t.buffer("in_tile_row", blk, dt);
8227    let load = t.stage(
8228        "gather_tile_row",
8229        StageKind::Load {
8230            tensor: table,
8231            pattern: AccessPattern {
8232                origin: addr,
8233                block_h: 1,
8234                block_w: blk,
8235                stride_w: 1,
8236                stride_h: row_t,
8237            },
8238            num_blocks: 1,
8239            multicast: None,
8240        },
8241        feed.map(|(_, b)| vec![b]).unwrap_or_default(),
8242        vec![tile_row],
8243    );
8244    if let Some((s, b)) = feed {
8245        t.edge(s, b, load);
8246    }
8247    let w = wide(blk, target);
8248    let (ins, maps) = match at.in_tile {
8249        InTileRow::Fixed(y) => (vec![tile_row], vec![at_row(y, w)]),
8250        InTileRow::Data { row, .. } => (
8251            vec![tile_row, row],
8252            vec![
8253                ElemAccess {
8254                    base: 0,
8255                    terms: vec![(LV_GEN_COL, 1)],
8256                    indirect: Some(IndirectTerm {
8257                        operand: 1,
8258                        stride: w as i64,
8259                        modulus: tr * w,
8260                    }),
8261                },
8262                fixed(0),
8263            ],
8264        ),
8265    };
8266    let out = t.buffer("row", blk, dt);
8267    let s_row = t.stage(
8268        "lift_row",
8269        StageKind::Compute(ComputeOp::Generic {
8270            iters: flat_space(w),
8271            ins: maps,
8272            out: ident_map(w),
8273            body: TensorBody::Operand(0),
8274        }),
8275        ins,
8276        vec![out],
8277    );
8278    t.edge(load, tile_row, s_row);
8279    if let InTileRow::Data { row, stage, .. } = at.in_tile {
8280        t.edge(stage, row, s_row);
8281    }
8282    let st = t.store(&format!("store_{}", names.out), x, steps, blk, out);
8283    t.edge(s_row, out, st);
8284    let claim = if float {
8285        approx("row gather", dt)
8286    } else {
8287        NumericsClaim::Exact
8288    };
8289    let mut prog = t.finish(
8290        format!("row_gather_{rows}x{cols}_dense"),
8291        row_items(row_t, blk),
8292        AccPolicy::F32,
8293        claim,
8294        target,
8295    )?;
8296    prog.work.plan = Some(tile_row_plan(row_t, blk, dt));
8297    Ok(prog)
8298}
8299
8300/// Row `y` of a page `w` elements wide, walked along its columns.
8301fn at_row(y: usize, w: usize) -> ElemAccess {
8302    at((y * w) as i64)
8303}
8304
8305/// Which in-tile row a dense row write replaces: known at compile time, or computed into `[lo, hi)`.
8306#[derive(Clone, Copy)]
8307enum RowSelect {
8308    Fixed(usize),
8309    Data {
8310        lo: BufferId,
8311        hi: BufferId,
8312        stages: (StageId, StageId),
8313    },
8314}
8315
8316/// `t = argmax(X) + 1` — an `ArgMax` fold (first maximum, ties to the lowest index: the VM's `>` scan) onto
8317/// the integer datapath, then LOGOS's 1-based count.
8318pub fn from_argmax(
8319    spec: &ArgMaxSpec,
8320    names: &StageNames,
8321    float: bool,
8322    target: &SpatialTarget,
8323) -> Result<Vec<SpatialProgram>, SpatialReject> {
8324    let n = spec.n;
8325    if n == 0 {
8326        return Err(SpatialReject::NotLowered("an argmax of an empty tensor"));
8327    }
8328    if tiles_of(n, target) > row_block_tiles(target) {
8329        return blocked_argmax(spec, names, float, target);
8330    }
8331    let mut b = FlatBuild::on(Datapath::of(float), target)?;
8332    let (x_name, x_fill) = names.operand(0);
8333    let dt = table_precision(&b, float, x_fill, target)?;
8334    let x = b.input_as(x_name, n, x_fill, dt, target);
8335    let one_w = wide(1, target);
8336    // The result is a 1-based position in `[1, n]`.
8337    let ix = index_type(target, n as u64)?;
8338    let at0 = b.t.scratch("argmax", 1, ix);
8339    let s1 = b.compute(
8340        "argmax",
8341        head_space(1, n, Some(ReducePool::ArgMax)),
8342        vec![Slot {
8343            buf: x,
8344            map: map2(0, 0, 1),
8345        }],
8346        per_row_map(one_w),
8347        TensorBody::Operand(0),
8348        at0,
8349    );
8350    let out_buf = b.t.buffer("index", 1, ix);
8351    let s2 = b.compute(
8352        "one_based",
8353        flat_space(1),
8354        vec![Slot {
8355            buf: at0,
8356            map: fixed(0),
8357        }],
8358        ident_map(one_w),
8359        *add(op(0), konst(1.0)),
8360        out_buf,
8361    );
8362    b.finish_as(
8363        &names.out,
8364        1,
8365        vec![s1, s2],
8366        out_buf,
8367        format!("argmax_{n}"),
8368        NumericsClaim::Exact,
8369        ix,
8370        target,
8371    )
8372    .map(|p| vec![p])
8373}
8374
8375/// `t = argmax(X) + 1` over a row wider than a site holds — the decode step's `argmax of logits` over a real
8376/// vocabulary. Sections: the BLOCKS, then a merge tree.
8377///
8378/// ★ EXACTLY THE VM'S SCAN, NaN INCLUDED. The VM seeds its best with item 1 and replaces it only on a strictly
8379/// greater item, and the [`ReducePool::ArgMax`] fold is that scan. Cut into blocks, the row's scan is: block 0's
8380/// scan (seeded by item 1 — a NaN there is never replaced, which makes it the answer), then every later block's
8381/// scan seeded by `−∞` (a NaN at a later block's start must not capture the block, since the VM never lets a NaN
8382/// replace anything), then a scan over the blocks' results in position order. So every block folds `blk` tiles
8383/// behind ONE seed element — item 1 for block 0 (a duplicate the scan cannot be moved by: nothing is `>` itself),
8384/// `−∞` for the rest — and reports the VALUE it holds (read back at its index, so block 0's NaN travels) and that
8385/// value's position in the row. Each merge level scans `fan` results the same way; only the first result of the
8386/// first group can be a NaN, and it is exactly the one the VM keeps.
8387///
8388/// The row's last block starts `blk` tiles before its end, overlapping the one before it — harmless, the overlap
8389/// reports a position twice and the earlier report is kept. Its padding past `n` is masked, and a pad block of the
8390/// tree (past the row) is masked whole. ★ A MASKED POSITION FOLDS AS THE SCAN'S IDENTITY — `−∞` for a float row,
8391/// the type's minimum for an integer one — never a large finite negative: a row of logits below `−1e30` is a row
8392/// like any other. A masked position never replaces the seed, and a live value equal to the mask is kept because
8393/// it comes first.
8394///
8395/// A float row's partials are fp32: its values are the datapath's (exact in fp32) and a position below `2^24` is
8396/// exact too. An integer row's partials are its own integer type.
8397fn blocked_argmax(
8398    spec: &ArgMaxSpec,
8399    names: &StageNames,
8400    float: bool,
8401    target: &SpatialTarget,
8402) -> Result<Vec<SpatialProgram>, SpatialReject> {
8403    let n = spec.n;
8404    let tw = target.tile.cols;
8405    let blk = row_block_tiles(target);
8406    let block = blk * tw;
8407    let x_tiles = tiles_of(n, target);
8408    let nb = x_tiles.div_ceil(blk);
8409    let (x_name, x_fill) = names.operand(0);
8410    let dt = if float {
8411        elem_precision(target)?
8412    } else {
8413        index_type(target, fill_abs_bound(x_fill).unwrap_or(0))?
8414    };
8415    let ix = index_type(target, (n + block + 1) as u64)?;
8416    let (pd, mask) = if float {
8417        if n >= 1 << 24 {
8418            return Err(SpatialReject::MissingPrecision {
8419                needed: Precision::Int64,
8420                target: target.name,
8421            });
8422        }
8423        (AccPolicy::F32.interm_precision(), f64::NEG_INFINITY)
8424    } else {
8425        let wide = index_type(
8426            target,
8427            (fill_abs_bound(x_fill).unwrap_or(0) as u64).max((n + block + 1) as u64),
8428        )?;
8429        (
8430            wide,
8431            -(crate::spatial::target::int_precision_max(wide).expect("an integer type") as f64)
8432                - 1.0,
8433        )
8434    };
8435    let fans = merge_fans(nb, blk);
8436    let slots: usize = fans.iter().product::<usize>().max(1);
8437    let last_start = (x_tiles - blk) * tw;
8438    // Per block, four index tiles: the start twice (the gather's index, and the position base), the live bound
8439    // `limit + 1` over the scan's `1 + block` elements (element 0 is the seed), and whether the block is the
8440    // row's first (its seed is item 1; every other block's is the mask).
8441    const FIELDS: usize = 4;
8442    let mut table = vec![0i64; FIELDS * slots * tw];
8443    for b in 0..slots {
8444        let (start, live, first) = if b < nb {
8445            let start = (b * block).min(last_start);
8446            (start, (n - start).min(block) + 1, usize::from(b == 0))
8447        } else {
8448            (0, 1, 0)
8449        };
8450        for (field, v) in [start, start, live, first].into_iter().enumerate() {
8451            table[(FIELDS * b + field) * tw] = v as i64;
8452        }
8453    }
8454    let one = wide(1, target);
8455    let partial_name = |level: usize| format!("{}__argmax_part{level}", names.out);
8456    // Pack `(value, position)` into one partial tile.
8457    let pack = |t: &mut Template, m: BufferId, at: BufferId, into: BufferId| {
8458        let body = TensorBody::Select {
8459            cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Const(1)),
8460            then: op(0),
8461            otherwise: op(1),
8462        };
8463        t.stage(
8464            "pack_partial",
8465            StageKind::Compute(ComputeOp::Generic {
8466                iters: flat_space(2),
8467                ins: vec![fixed(0), fixed(0)],
8468                out: ident_map(one),
8469                body,
8470            }),
8471            vec![m, at],
8472            vec![into],
8473        )
8474    };
8475    // `0 < operand` at the single point of a one-element stage: whether a scan moved off its seed.
8476    let past_seed = |operand: usize| Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(operand));
8477
8478    let blocks = {
8479        let mut t = Template::new(target)?;
8480        let x = t.tensor(
8481            x_name,
8482            TensorRole::Input,
8483            1,
8484            n,
8485            dt,
8486            super::from_tensor_op::init_of(x_fill, n),
8487        );
8488        let tab = t.tensor(
8489            &format!("{}__argmax_blocks", names.out),
8490            TensorRole::Param,
8491            1,
8492            FIELDS * slots * tw,
8493            ix,
8494            TensorInit::Values(table),
8495        );
8496        let part = t.tensor(
8497            &partial_name(0),
8498            TensorRole::Output,
8499            1,
8500            slots * tw,
8501            pd,
8502            TensorInit::Zero,
8503        );
8504        let fields: Vec<BufferId> = ["in_start", "in_base", "in_live", "in_first"]
8505            .iter()
8506            .map(|name| t.buffer(name, 1, ix))
8507            .collect();
8508        let loads: Vec<StageId> = ["load_start", "load_base", "load_live", "load_first"]
8509            .iter()
8510            .enumerate()
8511            .map(|(field, name)| {
8512                t.load_block_at(
8513                    name,
8514                    tab,
8515                    field,
8516                    vec![(LV_ITEM_X, FIELDS)],
8517                    1,
8518                    1,
8519                    1,
8520                    fields[field],
8521                )
8522            })
8523            .collect();
8524        let (gather_in, base_in, live_in, first_in) = (fields[0], fields[1], fields[2], fields[3]);
8525        let x_in = t.buffer("in_x", blk, dt);
8526        let origin = TileAddr {
8527            base: 0,
8528            steps: Vec::new(),
8529            indirect: Some(IndirectTile {
8530                index: gather_in,
8531                tile_stride: 1,
8532                elems_per_tile: tw,
8533                bound: (0, last_start),
8534            }),
8535        };
8536        let l_x = t.stage(
8537            "load_x",
8538            StageKind::Load {
8539                tensor: x,
8540                pattern: AccessPattern {
8541                    origin,
8542                    block_h: 1,
8543                    block_w: blk,
8544                    stride_w: 1,
8545                    stride_h: x_tiles,
8546                },
8547                num_blocks: 1,
8548                multicast: None,
8549            },
8550            vec![gather_in],
8551            vec![x_in],
8552        );
8553        t.edge(loads[0], gather_in, l_x);
8554        // The seed: item 1 for the row's first block, the mask for every other.
8555        let seed = t.scratch_deep("block_seed", 1, 1, pd);
8556        let s_seed = t.stage(
8557            "block_seed",
8558            StageKind::Compute(ComputeOp::Generic {
8559                iters: flat_space(1),
8560                ins: vec![fixed(0), fixed(0)],
8561                out: ident_map(one),
8562                body: TensorBody::Select {
8563                    cond: past_seed(1),
8564                    then: op(0),
8565                    otherwise: konst(mask),
8566                },
8567            }),
8568            vec![x_in, first_in],
8569            vec![seed],
8570        );
8571        // The scan over `[seed, x_0, …, x_{block−1}]`: element `c ≥ 1` is position `c − 1`, live below the bound.
8572        let scanned = TensorBody::Select {
8573            cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Const(1)),
8574            then: op(2),
8575            otherwise: Box::new(TensorBody::Select {
8576                cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(1)),
8577                then: op(0),
8578                otherwise: konst(mask),
8579            }),
8580        };
8581        let at = t.scratch_deep("block_argmax", 1, 1, ix);
8582        let s_at = t.stage(
8583            "block_argmax",
8584            StageKind::Compute(ComputeOp::Generic {
8585                iters: head_space(1, block + 1, Some(ReducePool::ArgMax)),
8586                ins: vec![map2(-1, 0, 1), fixed(0), fixed(0)],
8587                out: per_row_map(one),
8588                body: scanned,
8589            }),
8590            vec![x_in, live_in, seed],
8591            vec![at],
8592        );
8593        // The value the scan holds: element `at − 1` of the block, or the seed.
8594        let held = ElemAccess {
8595            base: -1,
8596            terms: Vec::new(),
8597            indirect: Some(IndirectTerm {
8598                operand: 1,
8599                stride: 1,
8600                modulus: block + 1,
8601            }),
8602        };
8603        let v = t.scratch_deep("block_value", 1, 1, pd);
8604        let s_v = t.stage(
8605            "block_value",
8606            StageKind::Compute(ComputeOp::Generic {
8607                iters: flat_space(1),
8608                ins: vec![held, fixed(0), fixed(0)],
8609                out: ident_map(one),
8610                body: TensorBody::Select {
8611                    cond: past_seed(1),
8612                    then: op(0),
8613                    otherwise: op(2),
8614                },
8615            }),
8616            vec![x_in, at, seed],
8617            vec![v],
8618        );
8619        // Its position in the row: `base + at − 1`, or `base` for the seed (item 1 itself in the first block).
8620        let pos = t.scratch_deep("block_position", 1, 1, ix);
8621        let s_pos = t.stage(
8622            "block_position",
8623            StageKind::Compute(ComputeOp::Generic {
8624                iters: flat_space(1),
8625                ins: vec![fixed(0), fixed(0)],
8626                out: ident_map(one),
8627                body: TensorBody::Select {
8628                    cond: past_seed(0),
8629                    then: add(sub(op(0), konst(1.0)), op(1)),
8630                    otherwise: op(1),
8631                },
8632            }),
8633            vec![at, base_in],
8634            vec![pos],
8635        );
8636        let out = t.buffer("out", 1, pd);
8637        let s_pack = pack(&mut t, v, pos, out);
8638        // Each hand-off's edge goes to its FIRST reader; a later stage naming the same page reads it before the
8639        // last reader pops it.
8640        for (from, buf, to) in [
8641            (l_x, x_in, s_seed),
8642            (loads[3], first_in, s_seed),
8643            (s_seed, seed, s_at),
8644            (loads[2], live_in, s_at),
8645            (s_at, at, s_v),
8646            (loads[1], base_in, s_pos),
8647            (s_v, v, s_pack),
8648            (s_pos, pos, s_pack),
8649        ] {
8650            t.edge(from, buf, to);
8651        }
8652        let st = t.store("store_partial", part, vec![(LV_ITEM_X, 1)], 1, out);
8653        t.edge(s_pack, out, st);
8654        t.finish(
8655            format!("argmax_blocks_{n}_b{blk}"),
8656            super::from_tensor_op::items_per_tile(slots),
8657            AccPolicy::F32,
8658            NumericsClaim::Exact,
8659            target,
8660        )?
8661    };
8662    let mut out = vec![blocks];
8663
8664    let mut n_in = slots;
8665    for (level, fan) in fans.iter().enumerate() {
8666        let root = level + 1 == fans.len();
8667        let n_out = n_in / fan;
8668        let mut t = Template::new(target)?;
8669        let part_in = t.tensor(
8670            &partial_name(level),
8671            TensorRole::Input,
8672            1,
8673            n_in * tw,
8674            pd,
8675            TensorInit::External,
8676        );
8677        let p_in = t.buffer("in_parts", *fan, pd);
8678        let l_p = t.load("load_parts", part_in, vec![(LV_ITEM_X, *fan)], *fan, p_in);
8679        // Partial `f` sits at column `f·tile` of the page: its value at `+0`, its position at `+1`.
8680        let values = ElemAccess {
8681            base: 0,
8682            terms: vec![(LV_GEN_COL, tw as i64)],
8683            indirect: None,
8684        };
8685        let f_at = t.scratch_deep("merge_argmax", 1, 1, ix);
8686        let s_f = t.stage(
8687            "merge_argmax",
8688            StageKind::Compute(ComputeOp::Generic {
8689                iters: head_space(1, *fan, Some(ReducePool::ArgMax)),
8690                ins: vec![values],
8691                out: per_row_map(one),
8692                body: TensorBody::Operand(0),
8693            }),
8694            vec![p_in],
8695            vec![f_at],
8696        );
8697        let field = |base: i64| ElemAccess {
8698            base,
8699            terms: Vec::new(),
8700            indirect: Some(IndirectTerm {
8701                operand: 1,
8702                stride: tw as i64,
8703                modulus: *fan,
8704            }),
8705        };
8706        let pos = t.scratch_deep("merge_position", 1, 1, pd);
8707        let s_pos = t.stage(
8708            "merge_position",
8709            StageKind::Compute(ComputeOp::Generic {
8710                iters: flat_space(1),
8711                ins: vec![field(1), fixed(0)],
8712                out: ident_map(one),
8713                body: TensorBody::Operand(0),
8714            }),
8715            vec![p_in, f_at],
8716            vec![pos],
8717        );
8718        let program = if root {
8719            t.edge(l_p, p_in, s_f);
8720            t.edge(s_f, f_at, s_pos);
8721            // The row's first maximum, 1-based.
8722            let res = t.tensor(&names.out, TensorRole::Output, 1, 1, ix, TensorInit::Zero);
8723            let o = t.buffer("index", 1, ix);
8724            let s_o = t.stage(
8725                "one_based",
8726                StageKind::Compute(ComputeOp::Generic {
8727                    iters: flat_space(1),
8728                    ins: vec![fixed(0)],
8729                    out: ident_map(one),
8730                    body: *add(op(0), konst(1.0)),
8731                }),
8732                vec![pos],
8733                vec![o],
8734            );
8735            t.edge(s_pos, pos, s_o);
8736            let st = t.store(&format!("store_{}", names.out), res, vec![], 1, o);
8737            t.edge(s_o, o, st);
8738            t.finish(
8739                format!("argmax_merge{}_f{fan}", level + 1),
8740                super::from_tensor_op::items_per_tile(1),
8741                AccPolicy::for_result(ix),
8742                NumericsClaim::Exact,
8743                target,
8744            )?
8745        } else {
8746            let part_out = t.tensor(
8747                &partial_name(level + 1),
8748                TensorRole::Output,
8749                1,
8750                n_out * tw,
8751                pd,
8752                TensorInit::Zero,
8753            );
8754            // The value the scan holds, read back at its index — a NaN the first partial carries travels on.
8755            let v = t.scratch_deep("merge_value", 1, 1, pd);
8756            let s_v = t.stage(
8757                "merge_value",
8758                StageKind::Compute(ComputeOp::Generic {
8759                    iters: flat_space(1),
8760                    ins: vec![field(0), fixed(0)],
8761                    out: ident_map(one),
8762                    body: TensorBody::Operand(0),
8763                }),
8764                vec![p_in, f_at],
8765                vec![v],
8766            );
8767            t.edge(l_p, p_in, s_f);
8768            t.edge(s_f, f_at, s_pos);
8769            let o = t.buffer("out", 1, pd);
8770            let s_pack = pack(&mut t, v, pos, o);
8771            t.edge(s_v, v, s_pack);
8772            t.edge(s_pos, pos, s_pack);
8773            let st = t.store("store_partial", part_out, vec![(LV_ITEM_X, 1)], 1, o);
8774            t.edge(s_pack, o, st);
8775            t.finish(
8776                format!("argmax_merge{}_f{fan}", level + 1),
8777                super::from_tensor_op::items_per_tile(n_out),
8778                AccPolicy::F32,
8779                NumericsClaim::Exact,
8780                target,
8781            )?
8782        };
8783        out.push(program);
8784        n_in = n_out;
8785    }
8786    Ok(out)
8787}
8788
8789/// `Set row index of T to S` — the KV-cache append, as ONE rewrite of the carried table:
8790/// `T'[c] = (lo ≤ c < lo + cols) ? S[c − lo] : T[c]`.
8791///
8792/// ★ AN OUT-OF-RANGE ROW MUST STILL FAULT, AND HERE NOTHING WOULD. The source is read through an
8793/// [`IndirectTerm`] only where the row is selected; with `lo` out of range no position is, so the gather
8794/// never runs and the write silently vanishes — where every tier fails. So the offset is CHECKED first by a
8795/// gather of STRIDE ZERO (it reads `lo` itself, and its extent `rows·cols` is the range check): a bounds
8796/// check is a gather that moves nothing.
8797pub fn from_set_row(
8798    spec: &SetRowSpec,
8799    names: &StageNames,
8800    float: bool,
8801    dense: bool,
8802    target: &SpatialTarget,
8803) -> Result<SpatialProgram, SpatialReject> {
8804    let (rows, cols, count) = (spec.rows, spec.cols, spec.count);
8805    if rows == 0 || cols == 0 || count == 0 || count > rows {
8806        return Err(SpatialReject::NotLowered(
8807            "a row write into an empty table, or a block of rows taller than its table",
8808        ));
8809    }
8810    if dense {
8811        return if count == 1 {
8812            dense_set_row(spec, names, float, target)
8813        } else {
8814            dense_set_rows(spec, names, float, target)
8815        };
8816    }
8817    if cols % target.tile.cols == 0 {
8818        return tile_set_row(spec, names, float, target);
8819    }
8820    // ★ A BLOCK OF ROWS IS ONE CONTIGUOUS RUN of `count·cols` elements from row `P − 1`: the one-row write, wider.
8821    let n = rows * cols;
8822    let span = count * cols;
8823    let mut b = FlatBuild::columns(n, target)?;
8824    let (t_name, t_fill) = names.operand(0);
8825    let (s_name, s_fill) = names.operand(1);
8826    let dt = row_write_precision(&b, float, t_fill, s_fill, target)?;
8827    let table = b.input_as(t_name, n, t_fill, dt, target);
8828    let source = b.input_as(s_name, span, s_fill, dt, target);
8829    let out_buf = b.result_page("table", n, dt, target);
8830    let one_w = wide(1, target);
8831    // `(c ≥ lo) ? ((c < hi) ? S : T) : T` — operand 0 is the table, operand 1 the source.
8832    let in_row = |lo_rhs: PredRhs, hi_rhs: PredRhs| TensorBody::Select {
8833        cond: Pred::AxisGe(LV_GEN_COL, lo_rhs),
8834        then: Box::new(TensorBody::Select {
8835            cond: Pred::AxisLt(LV_GEN_COL, hi_rhs),
8836            then: op(1),
8837            otherwise: op(0),
8838        }),
8839        otherwise: op(0),
8840    };
8841    let chain = match spec.index {
8842        RowIndex::Literal(i) => {
8843            let lo = literal_row_offset(i, rows + 1 - count, cols)?;
8844            let body = in_row(
8845                PredRhs::Const(lo as usize),
8846                PredRhs::Const(lo as usize + span),
8847            );
8848            b.pointwise(
8849                "write_row",
8850                n,
8851                vec![
8852                    Slot {
8853                        buf: table,
8854                        map: at(0),
8855                    },
8856                    Slot {
8857                        buf: source,
8858                        map: at(-lo),
8859                    },
8860                ],
8861                body,
8862                out_buf,
8863                target,
8864            )?
8865        }
8866        RowIndex::Named(_) => {
8867            let (i_name, i_fill) = names.operand(2);
8868            // `row_end = lo + span` reaches one row past the table.
8869            let ix = index_type(target, index_space((n + cols) as u64, cols as u64, i_fill))?;
8870            let index = b.input_as(i_name, 1, i_fill, ix, target);
8871            let (mut chain, lo_raw) = row_offset_stages(&mut b, index, cols, ix, target);
8872            let lo = b.t.scratch("row_checked", 1, ix);
8873            // The block's first offset `lo` is checked against the offsets a block of `count` rows may start at.
8874            let check = ElemAccess {
8875                base: 0,
8876                terms: Vec::new(),
8877                indirect: Some(IndirectTerm {
8878                    operand: 0,
8879                    stride: 0,
8880                    modulus: n - (count - 1) * cols,
8881                }),
8882            };
8883            let s2 = b.compute(
8884                "row_checked",
8885                flat_space(1),
8886                vec![Slot {
8887                    buf: lo_raw,
8888                    map: check,
8889                }],
8890                ident_map(one_w),
8891                TensorBody::Operand(0),
8892                lo,
8893            );
8894            let hi = b.t.scratch("row_end", 1, ix);
8895            let s3 = b.compute(
8896                "row_end",
8897                flat_space(1),
8898                vec![Slot {
8899                    buf: lo,
8900                    map: fixed(0),
8901                }],
8902                ident_map(one_w),
8903                *add(op(0), konst(span as f64)),
8904                hi,
8905            );
8906            let src = ElemAccess {
8907                base: 0,
8908                terms: vec![(LV_GEN_COL, 1)],
8909                indirect: Some(IndirectTerm {
8910                    operand: 2,
8911                    stride: -1,
8912                    modulus: n,
8913                }),
8914            };
8915            let body = in_row(PredRhs::Operand(2), PredRhs::Operand(3));
8916            let write = b.pointwise(
8917                "write_row",
8918                n,
8919                vec![
8920                    Slot {
8921                        buf: table,
8922                        map: at(0),
8923                    },
8924                    Slot {
8925                        buf: source,
8926                        map: src,
8927                    },
8928                    Slot {
8929                        buf: lo,
8930                        map: fixed(0),
8931                    },
8932                    Slot {
8933                        buf: hi,
8934                        map: fixed(0),
8935                    },
8936                ],
8937                body,
8938                out_buf,
8939                target,
8940            )?;
8941            chain.extend([s2, s3]);
8942            chain.extend(write);
8943            chain
8944        }
8945    };
8946    let claim = if float {
8947        approx("row write", dt)
8948    } else {
8949        NumericsClaim::Exact
8950    };
8951    b.finish_as(
8952        t_name,
8953        n,
8954        chain,
8955        out_buf,
8956        set_row_name(spec, ""),
8957        claim,
8958        dt,
8959        target,
8960    )
8961}
8962
8963/// A row write's program name: the table, a block's row count, and its layout's suffix.
8964fn set_row_name(spec: &SetRowSpec, layout: &str) -> String {
8965    let block = if spec.count > 1 {
8966        format!("_block{}", spec.count)
8967    } else {
8968        String::new()
8969    };
8970    format!("set_row_{}x{}{block}{layout}", spec.rows, spec.cols)
8971}
8972
8973/// Where a mixture's optional operands sit past `[x, wr, wg, wu, wd]`: the router's bias, then the gate, up and down
8974/// biases that are present, in that order.
8975struct MoeOperands {
8976    router_bias: Option<usize>,
8977    /// The gate, up and down biases' operands, each a row per expert.
8978    biases: [Option<usize>; 3],
8979    /// The selection bias's operand (`biased by B`).
8980    selection_bias: Option<usize>,
8981}
8982
8983impl MoeOperands {
8984    fn of(spec: &MoeSpec) -> MoeOperands {
8985        let mut next = 5;
8986        let mut take = |present: bool| {
8987            present.then(|| {
8988                next += 1;
8989                next - 1
8990            })
8991        };
8992        let router_bias = take(spec.router_bias.is_some());
8993        let biases = [take(spec.biases[0].is_some()), take(spec.biases[1].is_some()), take(spec.biases[2].is_some())];
8994        let selection_bias = take(spec.selection_bias.is_some());
8995        MoeOperands { router_bias, biases, selection_bias }
8996    }
8997}
8998
8999/// A mixture's sigmoid scores, `P[e] = 1/(1 + e^(0 − L[e]))`, at `f`'s precision — the desugar's arrangement.
9000fn moe_sigmoid(
9001    logits: &str,
9002    probs: &str,
9003    ne: usize,
9004    mut f: FlatBuild,
9005    target: &SpatialTarget,
9006) -> Result<SpatialProgram, SpatialReject> {
9007    let l = f.input(logits, ne, &Fill::Unknown, target);
9008    let out = f.t.buffer("probs", tiles_of(ne, target), f.precision);
9009    let s = f.compute(
9010        "sigmoid",
9011        flat_space(ne),
9012        vec![Slot { buf: l, map: at(0) }],
9013        ident_map(wide(tiles_of(ne, target), target)),
9014        *unary(SfpuKind::Recip, add(konst(1.0), unary(SfpuKind::Exp, sub(konst(0.0), op(0))))),
9015        out,
9016    );
9017    let claim = approx("moe router sigmoid", f.precision);
9018    f.finish(probs, ne, vec![s], out, format!("moe_sigmoid_{ne}"), claim, target)
9019}
9020
9021/// The scores the routing's arg-max reads: `max(P[e], −1e30)`, a NaN score as `−1e30`. The arg-max fold is seeded by
9022/// the first score where the desugar's scan seeds `−1e30` and never picks a NaN; on a softmax's scores the two agree
9023/// (one NaN logit makes every score NaN), but a sigmoid scores each logit alone, so a NaN first score would seed the
9024/// fold and win. The weights are still read from `P`. The copy keeps the scores' own precision: it decides the
9025/// experts, so it is never rounded to the section's datapath.
9026fn routed_scores(f: &mut FlatBuild, p: BufferId, ne: usize, target: &SpatialTarget) -> (StageId, BufferId) {
9027    let e_t = tiles_of(ne, target);
9028    let routed = f.t.scratch_deep("routed", 1, e_t, f.t.dtype(p));
9029    let stage = f.compute(
9030        "routed",
9031        flat_space(ne),
9032        vec![Slot { buf: p, map: at(0) }],
9033        ident_map(wide(e_t, target)),
9034        TensorBody::Max(op(0), konst(-1.0e30)),
9035        routed,
9036    );
9037    (stage, routed)
9038}
9039
9040/// The scores a mixture's routing rounds choose by, from its scores `p` and at their precision: `p` itself for a softmax;
9041/// `max(P[e] + B[e], −1e30)` with a selection bias and `max(P[e], −1e30)` for a sigmoid without one (a NaN is never
9042/// routed ahead of a number); then, under a group limit, [`moe_groups`]'s. The weights are still read from `P`. Returns
9043/// the stages, in order, and the buffer the rounds scan.
9044fn moe_choosing(
9045    f: &mut FlatBuild,
9046    spec: &MoeSpec,
9047    names: &StageNames,
9048    operands: &MoeOperands,
9049    p: BufferId,
9050    target: &SpatialTarget,
9051) -> Result<(Vec<StageId>, BufferId), SpatialReject> {
9052    let ne = spec.ne;
9053    let e_t = tiles_of(ne, target);
9054    let dt = f.t.dtype(p);
9055    let mut stages = Vec::new();
9056    let sel = match (operands.selection_bias, spec.score) {
9057        (Some(bi), _) => {
9058            let (b_name, b_fill) = names.operand(bi);
9059            // Read by the one stage of the section's one item: one page.
9060            let b = f.input_paged(b_name, ne, b_fill, dt, 1, target);
9061            let sel = f.t.scratch_deep("selection", 1, e_t, dt);
9062            stages.push(f.compute(
9063                "selection",
9064                flat_space(ne),
9065                vec![Slot { buf: p, map: at(0) }, Slot { buf: b, map: at(0) }],
9066                ident_map(wide(e_t, target)),
9067                TensorBody::Max(add(op(0), op(1)), konst(MASKED)),
9068                sel,
9069            ));
9070            sel
9071        }
9072        (None, MoeScore::Sigmoid) => {
9073            let (stage, routed) = routed_scores(f, p, ne, target);
9074            stages.push(stage);
9075            routed
9076        }
9077        (None, MoeScore::Softmax) => p,
9078    };
9079    match spec.groups {
9080        Some(groups) => {
9081            let (more, chosen) = moe_groups(f, spec, groups, sel, target)?;
9082            stages.extend(more);
9083            Ok((stages, chosen))
9084        }
9085        None => Ok((stages, sel)),
9086    }
9087}
9088
9089/// The penalty that masks the earlier picks of a routing: `Σ_j −1e30·[p_j ≤ c < e_j]` over the column `c`, the picks'
9090/// bounds at operands `first`, `first + 1`, … two by two — zero on every other column. Added to a score, a picked column
9091/// falls to `−1e30` (an f32 sum absorbs the score), below the arg-max's seed, so no column is picked twice; the body
9092/// grows by one term a pick, never by a page.
9093fn picked_penalty(picks: usize, first: usize) -> Option<Box<TensorBody>> {
9094    (0..picks)
9095        .map(|j| {
9096            Box::new(TensorBody::Select {
9097                cond: Pred::AxisGe(LV_GEN_COL, PredRhs::Operand(first + 2 * j)),
9098                then: Box::new(TensorBody::Select {
9099                    cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(first + 2 * j + 1)),
9100                    then: konst(MASKED),
9101                    otherwise: konst(0.0),
9102                }),
9103                otherwise: konst(0.0),
9104            })
9105        })
9106        .reduce(|a, b| add(a, b))
9107}
9108
9109/// Round `r` of a routing: the arg-max over `n` columns of `scores` (read through `read`) less the earlier picks
9110/// ([`picked_penalty`]) into the scalar `{name}{r}`, then — where `more` rounds follow — its end `{name}_end{r}`, one past
9111/// it. The fold is seeded by the first score where the desugar's scan seeds `−1e30`; the two agree on every vector a
9112/// softmax produces — entries in `[0, 1]` (the picked ones at `−1e30`), or all NaN, since one NaN logit makes the sum
9113/// NaN — and on an all-NaN vector both pick the first. Returns the pick and its end.
9114#[allow(clippy::too_many_arguments)]
9115fn moe_pick(
9116    f: &mut FlatBuild,
9117    chain: &mut Vec<StageId>,
9118    (name, stage): (&str, &str),
9119    r: usize,
9120    (scores, read, n): (BufferId, ElemAccess, usize),
9121    earlier: &[(BufferId, BufferId)],
9122    more: bool,
9123    ix: Precision,
9124    target: &SpatialTarget,
9125) -> (BufferId, Option<BufferId>) {
9126    let one_w = wide(1, target);
9127    let mut slots = vec![Slot { buf: scores, map: read }];
9128    for (pick, end) in earlier {
9129        slots.push(Slot { buf: *pick, map: fixed(0) });
9130        slots.push(Slot { buf: *end, map: fixed(0) });
9131    }
9132    let body = match picked_penalty(earlier.len(), 1) {
9133        Some(penalty) => *add(op(0), penalty),
9134        None => TensorBody::Operand(0),
9135    };
9136    let pick = f.t.scratch_deep(&format!("{name}{r}"), 1, 1, ix);
9137    chain.push(f.compute(&format!("{stage}{r}"), head_space(1, n, Some(ReducePool::ArgMax)), slots, per_row_map(one_w), body, pick));
9138    if !more {
9139        return (pick, None);
9140    }
9141    let end = f.t.scratch_deep(&format!("{name}_end{r}"), 1, 1, ix);
9142    chain.push(f.compute(&format!("{name}_end{r}"), flat_space(1), vec![Slot { buf: pick, map: fixed(0) }], ident_map(one_w), *add(op(0), konst(1.0)), end));
9143    (pick, Some(end))
9144}
9145
9146/// The group limit over the selection scores `sel`, as llama.cpp's `build_moe_ffn` takes it: each of `count` groups of
9147/// `per = ne / count` experts scored by its two best selection scores — the best, then the best of the rest with the
9148/// best's slot masked (a group of one by its one) — then `kept` rounds each choosing the best group not yet chosen
9149/// ([`moe_pick`]), and the selection scores of the chosen groups' experts kept in a row that is `−1e30` elsewhere: an
9150/// expert of a group not chosen is never picked. Returns the stages and that row.
9151///
9152/// ★ A GROUP'S VALUE IS ITS ROW. A fold over each group's `per` selection scores writes group `g`'s value on row `g` of
9153/// its page — a per-row result lands on its own row — so every stage over the groups walks rows, the arg-max over them
9154/// reads row `c` as its column `c`, and the groups are at most a tile's rows.
9155fn moe_groups(
9156    f: &mut FlatBuild,
9157    spec: &MoeSpec,
9158    groups: MoeGroups,
9159    sel: BufferId,
9160    target: &SpatialTarget,
9161) -> Result<(Vec<StageId>, BufferId), SpatialReject> {
9162    let MoeGroups { count, kept } = groups;
9163    if count > target.tile.rows {
9164        return Err(SpatialReject::NotLowered("a group limit over more groups than a page has rows"));
9165    }
9166    let (ne, per) = (spec.ne, spec.ne / count);
9167    let e_t = tiles_of(ne, target);
9168    let (ew, one_w) = (wide(e_t, target), wide(1, target));
9169    let dt = f.t.dtype(sel);
9170    let ix = index_type(target, ne as u64 + 1)?;
9171    let rows = || map2(0, per as i64, 1);
9172    let own = || per_row_map(one_w);
9173    let each = |reduce: Option<ReducePool>| head_space(count, if reduce.is_some() { per } else { 1 }, reduce);
9174    let mut stages = Vec::new();
9175    let best = f.t.scratch_deep("group_best", 1, 1, dt);
9176    stages.push(f.compute("group_best", each(Some(ReducePool::Max)), vec![Slot { buf: sel, map: rows() }], own(), TensorBody::Operand(0), best));
9177    let score = if per == 1 {
9178        best
9179    } else {
9180        let at_best = f.t.scratch_deep("group_at_best", 1, 1, ix);
9181        stages.push(f.compute("group_at_best", each(Some(ReducePool::ArgMax)), vec![Slot { buf: sel, map: rows() }], own(), TensorBody::Operand(0), at_best));
9182        let past_best = f.t.scratch_deep("group_past_best", 1, 1, ix);
9183        stages.push(f.compute("group_past_best", each(None), vec![Slot { buf: at_best, map: own() }], own(), *add(op(0), konst(1.0)), past_best));
9184        let second = f.t.scratch_deep("group_second", 1, 1, dt);
9185        let rest = TensorBody::Select {
9186            cond: Pred::AxisGe(LV_GEN_COL, PredRhs::Operand(1)),
9187            then: Box::new(TensorBody::Select { cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(2)), then: konst(MASKED), otherwise: op(0) }),
9188            otherwise: op(0),
9189        };
9190        stages.push(f.compute(
9191            "group_second",
9192            each(Some(ReducePool::Max)),
9193            vec![Slot { buf: sel, map: rows() }, Slot { buf: at_best, map: own() }, Slot { buf: past_best, map: own() }],
9194            own(),
9195            rest,
9196            second,
9197        ));
9198        let score = f.t.scratch_deep("group_score", 1, 1, dt);
9199        stages.push(f.compute("group_score", each(None), vec![Slot { buf: best, map: own() }, Slot { buf: second, map: own() }], own(), *add(op(0), op(1)), score));
9200        score
9201    };
9202    // Group `c`'s score, read as column `c` of the arg-max: row `c` of the page.
9203    let by_row = ElemAccess { base: 0, terms: vec![(LV_GEN_COL, one_w as i64)], indirect: None };
9204    let mut picked = Vec::new();
9205    let mut spans = Vec::new();
9206    for r in 0..kept {
9207        let (g, end) = moe_pick(f, &mut stages, ("group", "group"), r, (score, by_row.clone(), count), &picked, r + 1 < kept, ix, target);
9208        picked.extend(end.map(|end| (g, end)));
9209        // Its experts, `[g·per, g·per + per)`.
9210        let from = f.t.scratch_deep(&format!("group{r}_from"), 1, 1, ix);
9211        stages.push(f.compute(&format!("group{r}_from"), flat_space(1), vec![Slot { buf: g, map: fixed(0) }], ident_map(one_w), *mul(op(0), konst(per as f64)), from));
9212        let to = f.t.scratch_deep(&format!("group{r}_to"), 1, 1, ix);
9213        stages.push(f.compute(&format!("group{r}_to"), flat_space(1), vec![Slot { buf: from, map: fixed(0) }], ident_map(one_w), *add(op(0), konst(per as f64)), to));
9214        spans.push((from, to));
9215    }
9216    // `sel[c] + (1 − Σ_j [from_j ≤ c < to_j])·(−1e30)`: the chosen groups are disjoint, so an expert of one keeps its
9217    // selection score exactly (`+ 0·(−1e30)`), and every other falls to `−1e30`.
9218    let inside = (0..kept)
9219        .map(|j| {
9220            Box::new(TensorBody::Select {
9221                cond: Pred::AxisGe(LV_GEN_COL, PredRhs::Operand(1 + 2 * j)),
9222                then: Box::new(TensorBody::Select { cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(2 + 2 * j)), then: konst(1.0), otherwise: konst(0.0) }),
9223                otherwise: konst(0.0),
9224            })
9225        })
9226        .reduce(|a, b| add(a, b))
9227        .expect("a group limit keeps at least one group");
9228    let mut slots = vec![Slot { buf: sel, map: at(0) }];
9229    for (from, to) in spans {
9230        slots.push(Slot { buf: from, map: fixed(0) });
9231        slots.push(Slot { buf: to, map: fixed(0) });
9232    }
9233    let chosen = f.t.scratch_deep("chosen", 1, e_t, dt);
9234    stages.push(f.compute("chosen", flat_space(ne), slots, ident_map(ew), *add(op(0), mul(sub(konst(1.0), inside), konst(MASKED))), chosen));
9235    Ok((stages, chosen))
9236}
9237
9238/// The divisor a renormalization takes: the kept weights' sum, clamped at 2⁻¹⁴ as llama.cpp clamps it where the sum
9239/// can fall below that ([`MoeSpec::sum_can_clamp`]).
9240fn renormalizing_sum(spec: &MoeSpec, sum: Box<TensorBody>) -> Box<TensorBody> {
9241    if spec.sum_can_clamp() {
9242        Box::new(TensorBody::Max(sum, konst(6.103515625e-5)))
9243    } else {
9244        sum
9245    }
9246}
9247
9248/// A round's weight read out of the scores: the routed scale multiplies it here when nothing renormalizes it after.
9249fn round_weight(spec: &MoeSpec) -> TensorBody {
9250    match spec.scale {
9251        Some(_) if !spec.renormalize => *mul(op(0), konst(spec.scale())),
9252        _ => TensorBody::Operand(0),
9253    }
9254}
9255
9256/// A mixture's gated activation of `gate` and `up`: SwiGLU's `silu(gate)·up`, or gpt-oss's
9257/// `min(gate, L)·σ(α·min(gate, L))·(clamp(up, −L, L) + 1)` for its slope `α` — the sigmoid term as `silu(α·g)/α`, one
9258/// SFPU op.
9259fn moe_gated(act: crate::tensor_recognize::qwen::MoeAct, gate: Box<TensorBody>, up: Box<TensorBody>) -> Box<TensorBody> {
9260    match act {
9261        crate::tensor_recognize::qwen::MoeAct::Swiglu => mul(unary(SfpuKind::Silu, gate), up),
9262        crate::tensor_recognize::qwen::MoeAct::ClampedSwiglu { limit, slope } => {
9263            let (limit, slope) = (f64::from_bits(limit), f64::from_bits(slope));
9264            let g = Box::new(TensorBody::Min(gate, konst(limit)));
9265            let u = Box::new(TensorBody::Max(Box::new(TensorBody::Min(up, konst(limit))), konst(-limit)));
9266            mul(mul(unary(SfpuKind::Silu, mul(konst(slope), g)), konst(1.0 / slope)), add(u, konst(1.0)))
9267        }
9268    }
9269}
9270
9271/// The Mixture-of-Experts FFN — route `x` to its top-`k` of `ne` experts and combine their SwiGLUs by the
9272/// renormalized router weight:
9273///
9274/// ```text
9275///   L = x·Wr                  router logits (1 × ne)
9276///   P = softmax(L)
9277///   b_1..b_k                  the k largest of P, first maximum on ties (k rounds of arg-max over a mask)
9278///   out = Σ_i P[b_i]·down_{b_i}(silu(gate_{b_i}·x) ⊙ up_{b_i}·x)  /  Σ_i P[b_i]
9279/// ```
9280///
9281/// ★ ONLY THE CHOSEN EXPERTS' RUNS MOVE — THAT IS WHAT THE GATHER IS FOR. The desugar computes every expert
9282/// and weights the unchosen ones by zero (a mask the VM can run); a chip that did the same would stream all
9283/// `ne` experts' weights to use `k` of them. Expert `b`'s gate row `o` lives at `(b·dff + o)·d + t` — the
9284/// affine `o·d + t` past its run's first element `b·(dff·d)` — so each round loads the whole tiles around its
9285/// expert's run of each table and reads the run at its offset in the page (`ExpertTable`): `k` runs move,
9286/// and an expert index past the table faults. A table no more tiles than its `k` windows loads once, and the
9287/// rounds read their runs out of it through an [`IndirectTerm`] of stride the run and extent `ne`.
9288///
9289/// Three sections: the router (a reduction), the row softmax (`flat_softmax_rows`), and the routing plus
9290/// the experts. The weights sum is divided out once at the end: `Σ P[b]·y_b / Σ P[b]` is the desugar's
9291/// `Σ (P[b]/Σ P)·y_b`.
9292pub fn from_moe(
9293    spec: &MoeSpec,
9294    names: &StageNames,
9295    target: &SpatialTarget,
9296) -> Result<Vec<SpatialProgram>, SpatialReject> {
9297    let (d, dff, ne, k) = (spec.d, spec.dff, spec.ne, spec.k);
9298    if d == 0 || dff == 0 || ne == 0 || k == 0 || k > ne {
9299        return Err(SpatialReject::NotLowered("a mixture-of-experts block with no width, no experts, or more chosen experts than there are"));
9300    }
9301    if spec.weigh_inputs && spec.renormalize {
9302        return Err(SpatialReject::NotLowered(
9303            "a mixture whose renormalized weights scale its experts' inputs — each round is weighed as it is picked",
9304        ));
9305    }
9306    let operands = MoeOperands::of(spec);
9307    if moe_blocks(spec, target) {
9308        return blocked_moe(spec, names, target);
9309    }
9310    let (x_name, x_fill) = names.operand(0);
9311    let (wr_name, wr_fill) = names.operand(1);
9312    let (wg_name, wg_fill) = names.operand(2);
9313    let (wu_name, wu_fill) = names.operand(3);
9314    let (wd_name, wd_fill) = names.operand(4);
9315    let logits_name = format!("{}_logits", names.out);
9316    let probs_name = format!("{}_probs", names.out);
9317    let row = |var: LoopVar, stride: i64| ElemAccess {
9318        base: 0,
9319        terms: vec![(var, stride)],
9320        indirect: None,
9321    };
9322
9323    // ★ THE ROUTER DECIDES THE EXPERTS, SO ITS LOGITS AND SCORES STAY AT THE ACCUMULATOR'S WIDTH, as the blocked form
9324    // keeps them: two experts whose scores differ by less than the datapath's rounding would otherwise be picked as the
9325    // datapath rounds them, not as the model scores them. X and Wr are read as the datapath stores them — X is one
9326    // tensor, so one datatype, in every section of the block.
9327    let route = AccPolicy::for_result(Datapath::Float.precision(target)?).interm_precision();
9328    // ── Section 1: the router logits, `L[e] = Σ_t x[t]·Wr[t·ne + e]`.
9329    let router = {
9330        let mut f = FlatBuild::new(target)?;
9331        let x = f.input(x_name, d, x_fill, target);
9332        let wr = f.input(wr_name, d * ne, wr_fill, target);
9333        let out_buf = f.t.buffer("logits", tiles_of(ne, target), route);
9334        // A router bias adds after the product: the product lands in its own page first.
9335        let product = match operands.router_bias {
9336            Some(_) => f.t.scratch_deep("router_product", 1, tiles_of(ne, target), route),
9337            None => out_buf,
9338        };
9339        let iters = vec![
9340            IterSpec {
9341                var: LV_GEN_ROW,
9342                extent: d,
9343                kind: IterKind::Reduce {
9344                    pool: ReducePool::Sum,
9345                },
9346            },
9347            IterSpec {
9348                var: LV_GEN_COL,
9349                extent: ne,
9350                kind: IterKind::Parallel,
9351            },
9352        ];
9353        let s = f.compute(
9354            "router",
9355            iters,
9356            vec![
9357                Slot {
9358                    buf: x,
9359                    map: row(LV_GEN_ROW, 1),
9360                },
9361                Slot {
9362                    buf: wr,
9363                    map: map2(0, ne as i64, 1),
9364                },
9365            ],
9366            per_col_map(),
9367            *mul(op(0), op(1)),
9368            product,
9369        );
9370        let mut chain = vec![s];
9371        if let Some(bi) = operands.router_bias {
9372            let (b_name, b_fill) = names.operand(bi);
9373            let b = f.input_as(b_name, ne, b_fill, route, target);
9374            chain.push(f.compute(
9375                "router_bias",
9376                flat_space(ne),
9377                vec![Slot { buf: product, map: at(0) }, Slot { buf: b, map: at(0) }],
9378                ident_map(wide(tiles_of(ne, target), target)),
9379                *add(op(0), op(1)),
9380                out_buf,
9381            ));
9382        }
9383        let claim = approx("moe router", route);
9384        f.finish_as(&logits_name, ne, chain, out_buf, format!("moe_router_{d}x{ne}"), claim, route, target)?
9385    };
9386
9387    // ── Section 2: the router's scores, a softmax or each logit's sigmoid.
9388    let softmax = match spec.score {
9389        MoeScore::Softmax => {
9390            flat_softmax_rows_at((&logits_name, &Fill::Unknown), &probs_name, 1, ne, format!("moe_softmax_{ne}"), route, None, target)?
9391        }
9392        MoeScore::Sigmoid => moe_sigmoid(&logits_name, &probs_name, ne, FlatBuild::at(route, target)?, target)?,
9393    };
9394
9395    // ── Section 3: `k` routing rounds, each gathering one expert's weights, then the renormalized combine.
9396    let experts = {
9397        let mut f = FlatBuild::new(target)?;
9398        // ★ ONE PAGE OF EACH INTERMEDIATE. The section is one work item, and every round's intermediates are written
9399        // once and read within the item's own pass, so none needs a second page — `K` rounds of double-buffered
9400        // pages are what overflow a site once a mixture's biases add three more a round.
9401        let acc_p = AccPolicy::for_result(f.precision).interm_precision();
9402        let page = |f: &mut FlatBuild, name: &str, tiles: usize| f.t.scratch_deep(name, 1, tiles, acc_p);
9403        let p = f.input_as(&probs_name, ne, &Fill::Unknown, route, target);
9404        let x = f.input(x_name, d, x_fill, target);
9405        let wg = ExpertTable::new(&mut f, (wg_name, wg_fill), ne, dff * d, k, target);
9406        let wu = ExpertTable::new(&mut f, (wu_name, wu_fill), ne, dff * d, k, target);
9407        let wd = ExpertTable::new(&mut f, (wd_name, wd_fill), ne, d * dff, k, target);
9408        let one_w = wide(1, target);
9409        let (dff_t, d_t) = (tiles_of(dff, target), tiles_of(d, target));
9410        let (dffw, dw) = (wide(dff_t, target), wide(d_t, target));
9411        // `into[o] = Σ_t x[t]·W[o·fold + t]` over the chosen expert's run of `W`.
9412        let expert_gemv = |f: &mut FlatBuild, name: &str, run: ExpertRun, x: BufferId, fold: usize, width: usize, into: BufferId| {
9413            let iters = vec![
9414                IterSpec {
9415                    var: LV_GEN_ROW,
9416                    extent: fold,
9417                    kind: IterKind::Reduce {
9418                        pool: ReducePool::Sum,
9419                    },
9420                },
9421                IterSpec {
9422                    var: LV_GEN_COL,
9423                    extent: width,
9424                    kind: IterKind::Parallel,
9425                },
9426            ];
9427            let mut slots = vec![
9428                Slot {
9429                    buf: x,
9430                    map: row(LV_GEN_ROW, 1),
9431                },
9432                Slot {
9433                    buf: run.page,
9434                    map: run.map(vec![(LV_GEN_COL, fold as i64), (LV_GEN_ROW, 1)], 2),
9435                },
9436            ];
9437            slots.extend(run.slot());
9438            f.compute(name, iters, slots, per_col_map(), *mul(op(0), op(1)), into)
9439        };
9440        // Each expert bias table, `ne` runs of its projection's width.
9441        let bias_tables: Vec<Option<ExpertTable>> = (0..3)
9442            .map(|i| {
9443                operands.biases[i].map(|bi| {
9444                    let width = if i == 2 { d } else { dff };
9445                    ExpertTable::new(&mut f, names.operand(bi), ne, width, k, target)
9446                })
9447            })
9448            .collect();
9449        // `out[c] = y[c] + B[c]` over the chosen expert's run of `B`: a projection's output and its expert's bias row.
9450        let with_bias = |f: &mut FlatBuild, name: &str, y: BufferId, run: ExpertRun, width: usize, tiles: usize| {
9451            let into = page(f, name, tiles);
9452            let mut slots = vec![Slot { buf: y, map: at(0) }, Slot { buf: run.page, map: run.map(vec![(LV_GEN_COL, 1)], 2) }];
9453            slots.extend(run.slot());
9454            let stage = f.compute(name, flat_space(width), slots, ident_map(wide(tiles, target)), *add(op(0), op(1)), into);
9455            (stage, into)
9456        };
9457        // An expert pick is in `[0, ne)`; its end is one past.
9458        let ix = index_type(target, ne as u64 + 1)?;
9459        let (mut chain, choosing) = moe_choosing(&mut f, spec, names, &operands, p, target)?;
9460        let mut acc: Option<BufferId> = None;
9461        let mut wsum: Option<BufferId> = None;
9462        let mut picked = Vec::new();
9463        for r in 0..k {
9464            // The arg-max of what is left: the chosen expert, 0-based, on the integer datapath.
9465            let (pick, end) = moe_pick(&mut f, &mut chain, ("pick", "route"), r, (choosing, map2(0, 0, 1), ne), &picked, r + 1 < k, ix, target);
9466            picked.extend(end.map(|end| (pick, end)));
9467            // Its router weight, gathered out of P.
9468            let weight = page(&mut f, &format!("weight{r}"), 1);
9469            let from_p = ElemAccess {
9470                base: 0,
9471                terms: Vec::new(),
9472                indirect: Some(IndirectTerm {
9473                    operand: 1,
9474                    stride: 1,
9475                    modulus: ne,
9476                }),
9477            };
9478            chain.push(f.compute(
9479                &format!("weight{r}"),
9480                flat_space(1),
9481                vec![
9482                    Slot {
9483                        buf: p,
9484                        map: from_p,
9485                    },
9486                    Slot {
9487                        buf: pick,
9488                        map: fixed(0),
9489                    },
9490                ],
9491                ident_map(one_w),
9492                round_weight(spec),
9493                weight,
9494            ));
9495            // The input the expert reads: `x`, or `x` scaled by its weight when the weights scale the inputs.
9496            let x_in = if spec.weigh_inputs {
9497                let xs = page(&mut f, &format!("scaled_x{r}"), d_t);
9498                chain.push(f.compute(
9499                    &format!("scaled_x{r}"),
9500                    flat_space(d),
9501                    vec![Slot { buf: weight, map: fixed(0) }, Slot { buf: x, map: at(0) }],
9502                    ident_map(dw),
9503                    *mul(op(0), op(1)),
9504                    xs,
9505                ));
9506                xs
9507            } else {
9508                x
9509            };
9510            // The chosen expert's SwiGLU: gate and up gathered GEMVs, the gated product, the gathered down GEMV.
9511            let gate = page(&mut f, &format!("gate{r}"), dff_t);
9512            let run = wg.run(&mut f, &mut chain, &format!("gate{r}"), pick, ne, target)?;
9513            chain.push(expert_gemv(&mut f, &format!("gate{r}"), run, x_in, d, dff, gate));
9514            let up = page(&mut f, &format!("up{r}"), dff_t);
9515            let run = wu.run(&mut f, &mut chain, &format!("up{r}"), pick, ne, target)?;
9516            chain.push(expert_gemv(&mut f, &format!("up{r}"), run, x_in, d, dff, up));
9517            let (gate, up) = {
9518                let mut biased = [gate, up];
9519                for (i, part) in ["gate", "up"].into_iter().enumerate() {
9520                    if let Some(table) = &bias_tables[i] {
9521                        let run = table.run(&mut f, &mut chain, &format!("{part}_bias{r}"), pick, ne, target)?;
9522                        let (stage, into) = with_bias(&mut f, &format!("{part}_bias{r}"), biased[i], run, dff, dff_t);
9523                        chain.push(stage);
9524                        biased[i] = into;
9525                    }
9526                }
9527                (biased[0], biased[1])
9528            };
9529            let hidden = page(&mut f, &format!("hidden{r}"), dff_t);
9530            chain.push(f.compute(
9531                &format!("hidden{r}"),
9532                flat_space(dff),
9533                vec![
9534                    Slot {
9535                        buf: gate,
9536                        map: at(0),
9537                    },
9538                    Slot {
9539                        buf: up,
9540                        map: at(0),
9541                    },
9542                ],
9543                ident_map(dffw),
9544                *moe_gated(spec.act, op(0), op(1)),
9545                hidden,
9546            ));
9547            let y = page(&mut f, &format!("down{r}"), d_t);
9548            let run = wd.run(&mut f, &mut chain, &format!("down{r}"), pick, ne, target)?;
9549            chain.push(expert_gemv(&mut f, &format!("down{r}"), run, hidden, dff, d, y));
9550            let y = match &bias_tables[2] {
9551                Some(table) => {
9552                    let run = table.run(&mut f, &mut chain, &format!("down_bias{r}"), pick, ne, target)?;
9553                    let (stage, into) = with_bias(&mut f, &format!("down_bias{r}"), y, run, d, d_t);
9554                    chain.push(stage);
9555                    into
9556                }
9557                None => y,
9558            };
9559            // Accumulate the weighted expert output and the weight; an expert whose input took its weight adds as it is.
9560            let next_acc = page(&mut f, &format!("acc{r}"), d_t);
9561            chain.push(match acc {
9562                None if spec.weigh_inputs => f.compute(
9563                    &format!("acc{r}"),
9564                    flat_space(d),
9565                    vec![Slot { buf: y, map: at(0) }],
9566                    ident_map(dw),
9567                    TensorBody::Operand(0),
9568                    next_acc,
9569                ),
9570                Some(prev) if spec.weigh_inputs => f.compute(
9571                    &format!("acc{r}"),
9572                    flat_space(d),
9573                    vec![Slot { buf: prev, map: at(0) }, Slot { buf: y, map: at(0) }],
9574                    ident_map(dw),
9575                    *add(op(0), op(1)),
9576                    next_acc,
9577                ),
9578                None => f.compute(
9579                    &format!("acc{r}"),
9580                    flat_space(d),
9581                    vec![
9582                        Slot {
9583                            buf: weight,
9584                            map: fixed(0),
9585                        },
9586                        Slot { buf: y, map: at(0) },
9587                    ],
9588                    ident_map(dw),
9589                    *mul(op(0), op(1)),
9590                    next_acc,
9591                ),
9592                Some(prev) => f.compute(
9593                    &format!("acc{r}"),
9594                    flat_space(d),
9595                    vec![
9596                        Slot {
9597                            buf: prev,
9598                            map: at(0),
9599                        },
9600                        Slot {
9601                            buf: weight,
9602                            map: fixed(0),
9603                        },
9604                        Slot { buf: y, map: at(0) },
9605                    ],
9606                    ident_map(dw),
9607                    *add(op(0), mul(op(1), op(2))),
9608                    next_acc,
9609                ),
9610            });
9611            acc = Some(next_acc);
9612            wsum = Some(match wsum {
9613                None => weight,
9614                Some(prev) => {
9615                    let next = page(&mut f, &format!("wsum{r}"), 1);
9616                    chain.push(f.compute(
9617                        &format!("wsum{r}"),
9618                        flat_space(1),
9619                        vec![
9620                            Slot {
9621                                buf: prev,
9622                                map: fixed(0),
9623                            },
9624                            Slot {
9625                                buf: weight,
9626                                map: fixed(0),
9627                            },
9628                        ],
9629                        ident_map(one_w),
9630                        *add(op(0), op(1)),
9631                        next,
9632                    ));
9633                    next
9634                }
9635            });
9636        }
9637        let (acc, wsum) = (acc.expect("k ≥ 1 rounds"), wsum.expect("k ≥ 1 rounds"));
9638        let out_buf = f.t.buffer("out", d_t, f.precision);
9639        chain.push(f.compute(
9640            "combine",
9641            flat_space(d),
9642            vec![
9643                Slot {
9644                    buf: acc,
9645                    map: at(0),
9646                },
9647                Slot {
9648                    buf: wsum,
9649                    map: fixed(0),
9650                },
9651            ],
9652            ident_map(dw),
9653            match (spec.renormalize, spec.scale) {
9654                (true, None) => *mul(op(0), unary(SfpuKind::Recip, renormalizing_sum(spec, op(1)))),
9655                (true, Some(_)) => *mul(mul(op(0), unary(SfpuKind::Recip, renormalizing_sum(spec, op(1)))), konst(spec.scale())),
9656                (false, _) => TensorBody::Operand(0),
9657            },
9658            out_buf,
9659        ));
9660        let claim = approx("mixture of experts", f.precision);
9661        f.finish(
9662            &names.out,
9663            d,
9664            chain,
9665            out_buf,
9666            format!("moe_{ne}x{k}_{d}x{dff}"),
9667            claim,
9668            target,
9669        )?
9670    };
9671    Ok(vec![router, softmax, experts])
9672}
9673
9674// ── Mixture of experts at real width ──────────────────────────────────────────────────────────────────
9675//
9676// ★ ONLY THE CHOSEN EXPERTS' WEIGHTS MOVE, AND THEY MOVE BY TILE. Expert `e`'s gate rows are tile rows
9677// `[e·dff/32, (e+1)·dff/32)` of the stacked `(E·dff) × d` weight — one contiguous run of `(dff/32)·(d/32)` tiles —
9678// so each chosen expert's GEMV reads its weights through an [`IndirectTile`] whose index is the expert's pick,
9679// and the other `E − K` experts never leave memory. A GEMV row is a DOT PRODUCT across a tile row of the weight,
9680// and a row of Qwen3-30B-A3B's `d = 2048` is 64 tiles; so the reduction streams in CHUNKS of the row, each item a
9681// block of 32 outputs, each chunk a partial sum, and a second section adds the partials. Seven sections:
9682//
9683//   router (the matmul it is) → softmax → routing (top-K picks and weights, at the index datatype)
9684//   → gate/up partials (tile-gathered) → hidden = silu(Σ gate)·(Σ up) → down partials (tile-gathered)
9685//   → combine: Σ_r w_r·(Σ down_r) / Σ_r w_r
9686//
9687// The partials and the hidden activation stay at the accumulator's width; only the block's output leaves at the
9688// datapath's precision.
9689
9690/// A flat table of `ne` expert runs of `span` elements each, read by every routing round at its chosen expert's run.
9691///
9692/// ★ A ROUND MOVES ITS EXPERT'S RUN, OR THE TABLE MOVES ONCE, WHICHEVER IS FEWER TILES. Run `e` starts at element
9693/// `e·span`, mid-tile unless `span` is whole tiles. A round loads the window of whole tiles around its run — `W` tiles,
9694/// the most any run reaches from its first element's offset in its tile, started at `min(⌊e·span/32⌋, T − W)` so it ends
9695/// inside the table — and reads the run at its offset `e·span − 32·start` in the page; a run of whole tiles is its own
9696/// window, addressed by the pick alone. `k` rounds move `k·W` tiles of a `T`-tile table; where that is not fewer, the
9697/// table loads once and each round reads its run out of it by element. Either way the section's one item fills each
9698/// page once, so each is one page deep.
9699enum ExpertTable {
9700    Whole { page: BufferId, span: usize },
9701    Windows { tensor: TensorId, span: usize, window_t: usize, table_t: usize },
9702}
9703
9704/// How a round reads its chosen expert's run: the page holding it, and the element term that moves the read to the
9705/// run — `(index buffer, stride, extent)`: the pick by the run's length in a whole table, the run's offset in a window
9706/// that does not start at it.
9707struct ExpertRun {
9708    page: BufferId,
9709    by: Option<(BufferId, i64, usize)>,
9710}
9711
9712impl ExpertRun {
9713    /// The map reading element `Σ terms` of the run, its moving term reading the stage's operand `operand`.
9714    fn map(&self, terms: Vec<(LoopVar, i64)>, operand: usize) -> ElemAccess {
9715        ElemAccess {
9716            base: 0,
9717            terms,
9718            indirect: self.by.map(|(_, stride, modulus)| IndirectTerm { operand, stride, modulus }),
9719        }
9720    }
9721
9722    /// The operand the moving term reads, when the run has one.
9723    fn slot(&self) -> Option<Slot> {
9724        self.by.map(|(buf, _, _)| Slot { buf, map: fixed(0) })
9725    }
9726}
9727
9728impl ExpertTable {
9729    fn new(f: &mut FlatBuild, (name, fill): (&str, &Fill), ne: usize, span: usize, k: usize, target: &SpatialTarget) -> ExpertTable {
9730        let tw = target.tile.cols;
9731        let table_t = tiles_of(ne * span, target);
9732        let window_t = (0..ne).map(|e| (e * span % tw + span).div_ceil(tw)).max().unwrap_or(1).min(table_t);
9733        if k * window_t >= table_t {
9734            let precision = f.precision;
9735            return ExpertTable::Whole { page: f.input_paged(name, ne * span, fill, precision, 1, target), span };
9736        }
9737        let tensor = f.t.tensor(name, TensorRole::Input, 1, ne * span, f.precision, super::from_tensor_op::init_of(fill, ne * span));
9738        ExpertTable::Windows { tensor, span, window_t, table_t }
9739    }
9740
9741    /// Round `part`'s run, chosen by `pick`: the window's index and offset stages pushed onto `chain`, its load wired to
9742    /// the stage that first reads the page.
9743    fn run(
9744        &self,
9745        f: &mut FlatBuild,
9746        chain: &mut Vec<StageId>,
9747        part: &str,
9748        pick: BufferId,
9749        ne: usize,
9750        target: &SpatialTarget,
9751    ) -> Result<ExpertRun, SpatialReject> {
9752        let (tensor, span, window_t, table_t) = match *self {
9753            ExpertTable::Whole { page, span } => return Ok(ExpertRun { page, by: Some((pick, span as i64, ne)) }),
9754            ExpertTable::Windows { tensor, span, window_t, table_t } => (tensor, span, window_t, table_t),
9755        };
9756        let tw = target.tile.cols;
9757        let one = wide(1, target);
9758        let ix = index_type(target, (ne * span + tw) as u64)?;
9759        let whole_tiles = span % tw == 0;
9760        let first = || TensorBody::Min(Box::new(TensorBody::FloorDiv(mul(op(0), konst(span as f64)), konst(tw as f64))), konst((table_t - window_t) as f64));
9761        // Each gather retires its index page, so the load reads a page of its own: the pick, or the window's first tile.
9762        let index = f.t.scratch_deep(&format!("{part}_window"), 1, 1, ix);
9763        let body = if whole_tiles { TensorBody::Operand(0) } else { first() };
9764        let s = f.compute(&format!("{part}_window"), flat_space(1), vec![Slot { buf: pick, map: fixed(0) }], ident_map(one), body, index);
9765        chain.push(s);
9766        let indirect = if whole_tiles {
9767            IndirectTile { index, tile_stride: window_t, elems_per_tile: 1, bound: (0, ne - 1) }
9768        } else {
9769            IndirectTile { index, tile_stride: 1, elems_per_tile: 1, bound: (0, table_t - window_t) }
9770        };
9771        let page = f.t.buffer_deep(&format!("in_{part}"), 1, window_t, f.precision);
9772        let l = f.t.stage(
9773            &format!("load_{part}"),
9774            StageKind::Load {
9775                tensor,
9776                pattern: AccessPattern {
9777                    origin: TileAddr { base: 0, steps: Vec::new(), indirect: Some(indirect) },
9778                    block_h: 1,
9779                    block_w: window_t,
9780                    stride_w: 1,
9781                    stride_h: window_t,
9782                },
9783                num_blocks: 1,
9784                multicast: None,
9785            },
9786            vec![index],
9787            vec![page],
9788        );
9789        f.t.edge(s, index, l);
9790        f.loads.push((l, page));
9791        if whole_tiles {
9792            return Ok(ExpertRun { page, by: None });
9793        }
9794        let offset = f.t.scratch_deep(&format!("{part}_offset"), 1, 1, ix);
9795        chain.push(f.compute(
9796            &format!("{part}_offset"),
9797            flat_space(1),
9798            vec![Slot { buf: pick, map: fixed(0) }],
9799            ident_map(one),
9800            *sub(mul(op(0), konst(span as f64)), mul(Box::new(first()), konst(tw as f64))),
9801            offset,
9802        ));
9803        Ok(ExpertRun { page, by: Some((offset, 1, window_t * tw - span + 1)) })
9804    }
9805}
9806
9807/// Whether a mixture-of-experts block takes the blocked form: widths in whole tiles, and a router softmax and a
9808/// pick table that one accumulator page holds.
9809fn moe_blocks(spec: &MoeSpec, target: &SpatialTarget) -> bool {
9810    let tw = target.tile.cols;
9811    spec.d % tw == 0
9812        && spec.dff % tw == 0
9813        && (!row_needs_blocking(spec.ne, target) || spec.groups.is_none())
9814}
9815
9816/// The names of a blocked MoE's internal tensors, all derived from the block's binder.
9817struct MoeNames {
9818    /// The router's product before its bias, when it has one.
9819    router_product: String,
9820    logits: String,
9821    probs: String,
9822    picks: String,
9823    weights: String,
9824    gate_part: String,
9825    up_part: String,
9826    hidden: String,
9827    down_part: String,
9828    scaler: String,
9829}
9830
9831impl MoeNames {
9832    fn of(names: &StageNames) -> MoeNames {
9833        let n = |s: &str| format!("{}__moe_{s}", names.out);
9834        MoeNames {
9835            router_product: n("router_product"),
9836            logits: n("logits"),
9837            probs: n("probs"),
9838            picks: n("picks"),
9839            weights: n("weights"),
9840            gate_part: n("gate_part"),
9841            up_part: n("up_part"),
9842            hidden: n("hidden"),
9843            down_part: n("down_part"),
9844            scaler: n("scaler"),
9845        }
9846    }
9847}
9848
9849/// The largest divisor of `n` (a chunk of a row's tiles) whose section fits a site — each candidate is BUILT and
9850/// measured, so the sizing can never disagree with the section it sizes. When none fits, a chunk of one tile that
9851/// cannot be BUILT is refused with its own reason, and one that is built but too large as `too_large`.
9852fn fitting_chunk(
9853    n: usize,
9854    build: &dyn Fn(usize) -> Result<SpatialProgram, SpatialReject>,
9855    too_large: &'static str,
9856    target: &SpatialTarget,
9857) -> Result<usize, SpatialReject> {
9858    let budget = target.max_local_mem_bytes();
9859    if let Some(c) = (1..=n)
9860        .rev()
9861        .filter(|c| n % c == 0)
9862        .find(|c| build(*c).is_ok_and(|p| site_footprint(&p) <= budget))
9863    {
9864        return Ok(c);
9865    }
9866    build(1)?;
9867    Err(SpatialReject::NotLowered(too_large))
9868}
9869
9870fn blocked_moe(
9871    spec: &MoeSpec,
9872    names: &StageNames,
9873    target: &SpatialTarget,
9874) -> Result<Vec<SpatialProgram>, SpatialReject> {
9875    let (d, dff, ne) = (spec.d, spec.dff, spec.ne);
9876    let tw = target.tile.cols;
9877    let (d_t, f_t) = (d / tw, dff / tw);
9878    let mn = MoeNames::of(names);
9879    let (x_name, x_fill) = names.operand(0);
9880    let (wr_name, wr_fill) = names.operand(1);
9881    // The router is `x · Wr` — a `1 × d` by `d × E` product, streamed by the matmul lowering.
9882    let router_spec = crate::tensor_recognize::MatmulSpec {
9883        m: 1,
9884        n: ne,
9885        k: d,
9886        a: Default::default(),
9887        b: Default::default(),
9888        c: Default::default(),
9889    };
9890    let operands = MoeOperands::of(spec);
9891    let router_names = MatmulNames {
9892        a: x_name.to_string(),
9893        b: wr_name.to_string(),
9894        c: if operands.router_bias.is_some() { mn.router_product.clone() } else { mn.logits.clone() },
9895    };
9896    let router = super::from_tensor_op::from_matmul(
9897        &router_spec,
9898        &router_names,
9899        (x_fill, wr_fill),
9900        (None, None),
9901        NumericsPolicy::Exact,
9902        target,
9903    )?;
9904    // ★ X IS ONE TENSOR, SO ONE DATATYPE. The router reads it as a matmul operand; the expert chunks read the same
9905    // bytes, so every section declares X at the operand precision the router's matmul chose.
9906    let x_dtype = router
9907        .tensors
9908        .iter()
9909        .find(|t| t.name == x_name)
9910        .map(|t| t.dtype)
9911        .ok_or(SpatialReject::NotLowered(
9912            "the router matmul does not read the block's input",
9913        ))?;
9914    // ★ A ROUTER ROW WIDER THAN ONE ACCUMULATOR PAGE (more than 256 experts: Kimi K2's 384, the 512 of Qwen3-Next, the
9915    // Qwen 3.5 mixtures and qwen4_exp) is scored and routed in blocks of the row.
9916    let wide_row = row_needs_blocking(ne, target);
9917    // A router bias adds to the product, in a section of its own.
9918    let router_bias = match operands.router_bias {
9919        None => None,
9920        Some(bi) if wide_row => Some(wide_router_bias(&mn, names.operand(bi), ne, target)?),
9921        Some(bi) => {
9922            let mut f = FlatBuild::at(AccPolicy::F32.interm_precision(), target)?;
9923            let product = f.input(&mn.router_product, ne, &Fill::Unknown, target);
9924            let (b_name, b_fill) = names.operand(bi);
9925            let b = f.input(b_name, ne, b_fill, target);
9926            let out = f.t.buffer("logits", tiles_of(ne, target), f.precision);
9927            let s = f.compute(
9928                "router_bias",
9929                flat_space(ne),
9930                vec![Slot { buf: product, map: at(0) }, Slot { buf: b, map: at(0) }],
9931                ident_map(wide(tiles_of(ne, target), target)),
9932                *add(op(0), op(1)),
9933                out,
9934            );
9935            let claim = approx("moe router bias", f.precision);
9936            Some(f.finish(&mn.logits, ne, vec![s], out, format!("moe_router_bias_{ne}"), claim, target)?)
9937        }
9938    };
9939    // The router's logits and scores stay at the accumulator's width: they decide the experts.
9940    let scores = match (spec.score, wide_row) {
9941        (MoeScore::Softmax, true) => wide_softmax(&mn, ne, target)?,
9942        (MoeScore::Sigmoid, true) => vec![wide_sigmoid(&mn, ne, target)?],
9943        (MoeScore::Softmax, false) => vec![flat_softmax_rows_at(
9944            (&mn.logits, &Fill::Unknown),
9945            &mn.probs,
9946            1,
9947            ne,
9948            format!("moe_softmax_{ne}"),
9949            AccPolicy::F32.interm_precision(),
9950            None,
9951            target,
9952        )?],
9953        (MoeScore::Sigmoid, false) => vec![moe_sigmoid(&mn.logits, &mn.probs, ne, FlatBuild::at(AccPolicy::F32.interm_precision(), target)?, target)?],
9954    };
9955    let routing = if wide_row { wide_moe_routing(spec, names, &operands, &mn, target)? } else { vec![moe_routing(spec, names, &operands, &mn, target)?] };
9956    let cd = fitting_chunk(
9957        d_t,
9958        &|c| moe_up_partials(spec, names, &mn, c, x_dtype, target),
9959        "a mixture-of-experts gate/up chunk of one tile does not fit a site",
9960        target,
9961    )?;
9962    let up = moe_up_partials(spec, names, &mn, cd, x_dtype, target)?;
9963    let hidden = moe_hidden(spec, names, &operands, &mn, d_t / cd, target)?;
9964    let cdd = fitting_chunk(
9965        f_t,
9966        &|c| moe_down_partials(spec, names, &mn, c, target),
9967        "a mixture-of-experts down chunk of one tile does not fit a site",
9968        target,
9969    )?;
9970    let down = moe_down_partials(spec, names, &mn, cdd, target)?;
9971    let combine = moe_combine(spec, names, &operands, &mn, f_t / cdd, target)?;
9972    Ok([router].into_iter().chain(router_bias).chain(scores).chain(routing).chain([up, hidden, down, combine]).collect())
9973}
9974
9975/// The routing section: `K` rounds of arg-max over what is left of the router probabilities, each pick (at the
9976/// index datatype) and its probability written to tile `r` of the pick and weight tables — a tile per pick, so
9977/// each is element 0 of the page an expert's tile gather waits on.
9978fn moe_routing(
9979    spec: &MoeSpec,
9980    names: &StageNames,
9981    operands: &MoeOperands,
9982    mn: &MoeNames,
9983    target: &SpatialTarget,
9984) -> Result<SpatialProgram, SpatialReject> {
9985    let (ne, k) = (spec.ne, spec.k);
9986    let tw = target.tile.cols;
9987    let acc = AccPolicy::F32.interm_precision();
9988    let ix = index_type(target, ne as u64 + 1)?;
9989    let mut f = FlatBuild::new(target)?;
9990    f.precision = acc;
9991    // ★ ONE PAGE OF EVERYTHING. The section is one item: each page is filled once and read within the item's pass, so a
9992    // second page buys nothing, and at f32 every scalar is a whole tile — two pages of each is what overflowed a site at
9993    // DeepSeek-V3's 256 experts, top 8 of the top 4 of 8 groups.
9994    let p = f.input_paged(&mn.probs, ne, &Fill::Unknown, acc, 1, target);
9995    let one_w = wide(1, target);
9996    let (mut chain, choosing) = moe_choosing(&mut f, spec, names, operands, p, target)?;
9997    let mut picks = Vec::new();
9998    let mut weights = Vec::new();
9999    // ★ NO PAGE A ROUND. Each round's arg-max reads the one choosing row less a penalty on the earlier picks, scalars all;
10000    // a row of masks a round is what overflowed a site at 256 experts, top 8, at f32.
10001    let mut picked = Vec::new();
10002    for r in 0..k {
10003        let (pick, end) = moe_pick(&mut f, &mut chain, ("pick", "route"), r, (choosing, map2(0, 0, 1), ne), &picked, r + 1 < k, ix, target);
10004        picked.extend(end.map(|end| (pick, end)));
10005        let weight = f.t.scratch_deep(&format!("weight{r}"), 1, 1, acc);
10006        let from_p = ElemAccess {
10007            base: 0,
10008            terms: Vec::new(),
10009            indirect: Some(IndirectTerm {
10010                operand: 1,
10011                stride: 1,
10012                modulus: ne,
10013            }),
10014        };
10015        chain.push(f.compute(
10016            &format!("weight{r}"),
10017            flat_space(1),
10018            vec![
10019                Slot {
10020                    buf: p,
10021                    map: from_p,
10022                },
10023                Slot {
10024                    buf: pick,
10025                    map: fixed(0),
10026                },
10027            ],
10028            ident_map(one_w),
10029            round_weight(spec),
10030            weight,
10031        ));
10032        picks.push(pick);
10033        weights.push(weight);
10034    }
10035    let kw = wide(k, target);
10036    let pack = |f: &mut FlatBuild, name: &str, parts: &[BufferId], dt: Precision| {
10037        let slots: Vec<Slot> = parts
10038            .iter()
10039            .map(|b| Slot {
10040                buf: *b,
10041                map: fixed(0),
10042            })
10043            .collect();
10044        let bands: Vec<(usize, TensorBody)> = (0..parts.len())
10045            .map(|r| ((r + 1) * tw, TensorBody::Operand(r)))
10046            .collect();
10047        let out = f.t.buffer_deep(name, 1, k, dt);
10048        (
10049            f.compute(
10050                &format!("pack_{name}"),
10051                flat_space(k * tw),
10052                slots,
10053                ident_map(kw),
10054                banded(bands),
10055                out,
10056            ),
10057            out,
10058        )
10059    };
10060    let (sp, pick_out) = pack(&mut f, "picks", &picks, ix);
10061    let (sw, weight_out) = pack(&mut f, "weights", &weights, acc);
10062    chain.push(sp);
10063    chain.push(sw);
10064    let pick_t = f.t.tensor(
10065        &mn.picks,
10066        TensorRole::Output,
10067        1,
10068        k * tw,
10069        ix,
10070        TensorInit::Zero,
10071    );
10072    let weight_t = f.t.tensor(
10073        &mn.weights,
10074        TensorRole::Output,
10075        1,
10076        k * tw,
10077        acc,
10078        TensorInit::Zero,
10079    );
10080    let st_p = f.t.store("store_picks", pick_t, Vec::new(), k, pick_out);
10081    let st_w =
10082        f.t.store("store_weights", weight_t, Vec::new(), k, weight_out);
10083    let claim = approx("mixture-of-experts routing", f.precision);
10084    f.wire_and_finish_many(
10085        chain,
10086        &[(pick_out, st_p), (weight_out, st_w)],
10087        k,
10088        format!("moe_routing_{ne}x{k}"),
10089        claim,
10090        target,
10091    )
10092}
10093
10094/// The blocks of a router row of `ne` experts, `(first tile, experts)` each: an accumulator page of experts a block, the
10095/// last the rest.
10096fn router_blocks(ne: usize, target: &SpatialTarget) -> Vec<(usize, usize)> {
10097    let blk = row_block_tiles(target);
10098    let span = blk * target.tile.cols;
10099    (0..ne.div_ceil(span)).map(|b| (b * blk, span.min(ne - b * span))).collect()
10100}
10101
10102/// A wide router's bias, `L[e] = (x·Wr)[e] + B[e]`, a block of the row an item.
10103fn wide_router_bias(mn: &MoeNames, (b_name, b_fill): (&str, &Fill), ne: usize, target: &SpatialTarget) -> Result<SpatialProgram, SpatialReject> {
10104    let mut f = FlatBuild::columns(ne, target)?;
10105    f.precision = AccPolicy::F32.interm_precision();
10106    let product = f.input(&mn.router_product, ne, &Fill::Unknown, target);
10107    let b = f.input(b_name, ne, b_fill, target);
10108    let out = f.result_page("logits", ne, f.precision, target);
10109    let chain = f.pointwise("router_bias", ne, vec![Slot { buf: product, map: at(0) }, Slot { buf: b, map: at(0) }], *add(op(0), op(1)), out, target)?;
10110    let claim = approx("moe router bias", f.precision);
10111    f.finish(&mn.logits, ne, chain, out, format!("moe_router_bias_{ne}"), claim, target)
10112}
10113
10114/// A wide router's sigmoid scores, `P[e] = 1/(1 + e^(0 − L[e]))` as [`moe_sigmoid`] writes them, a block of the row an
10115/// item.
10116fn wide_sigmoid(mn: &MoeNames, ne: usize, target: &SpatialTarget) -> Result<SpatialProgram, SpatialReject> {
10117    let mut f = FlatBuild::columns(ne, target)?;
10118    f.precision = AccPolicy::F32.interm_precision();
10119    let l = f.input(&mn.logits, ne, &Fill::Unknown, target);
10120    let out = f.result_page("probs", ne, f.precision, target);
10121    let body = *unary(SfpuKind::Recip, add(konst(1.0), unary(SfpuKind::Exp, sub(konst(0.0), op(0)))));
10122    let chain = f.pointwise("sigmoid", ne, vec![Slot { buf: l, map: at(0) }], body, out, target)?;
10123    let claim = approx("moe router sigmoid", f.precision);
10124    f.finish(&mn.probs, ne, chain, out, format!("moe_sigmoid_{ne}"), claim, target)
10125}
10126
10127/// A wide router's softmax, `P = e^(L − m)/Z` over its row, at the accumulator's width: each block's maximum `m_b` and
10128/// sum `s_b = Σ e^(L − m_b)` (a section a block), the row's `m = max m_b` and `Z = Σ_b s_b·e^(m_b − m)` (one section),
10129/// then `e^(L − m)·(1/Z)`, a block of the row an item.
10130fn wide_softmax(mn: &MoeNames, ne: usize, target: &SpatialTarget) -> Result<Vec<SpatialProgram>, SpatialReject> {
10131    let acc = AccPolicy::F32.interm_precision();
10132    let one = wide(1, target);
10133    let blocks = router_blocks(ne, target);
10134    let stats = |b: usize| format!("{}__stats{b}", mn.probs);
10135    let norm = format!("{}__norm", mn.probs);
10136    let mut out = Vec::new();
10137    for (b, &(first, n)) in blocks.iter().enumerate() {
10138        let mut f = FlatBuild::at(acc, target)?;
10139        let l = f.input_slice(&mn.logits, ne, first, n, &Fill::Unknown, acc, target);
10140        let m = f.t.scratch_deep("block_max", 1, 1, acc);
10141        let s1 = f.compute("block_max", head_space(1, n, Some(ReducePool::Max)), vec![Slot { buf: l, map: map2(0, 0, 1) }], per_row_map(one), TensorBody::Operand(0), m);
10142        let sum = f.t.scratch_deep("block_sum", 1, 1, acc);
10143        let s2 = f.compute(
10144            "block_sum",
10145            head_space(1, n, Some(ReducePool::Sum)),
10146            vec![Slot { buf: l, map: map2(0, 0, 1) }, Slot { buf: m, map: fixed(0) }],
10147            per_row_map(one),
10148            *unary(SfpuKind::Exp, sub(op(0), op(1))),
10149            sum,
10150        );
10151        let page = f.t.buffer("stats", 1, acc);
10152        let s3 = f.compute(
10153            "pack_stats",
10154            flat_space(2),
10155            vec![Slot { buf: m, map: fixed(0) }, Slot { buf: sum, map: fixed(0) }],
10156            ident_map(one),
10157            banded(vec![(1, TensorBody::Operand(0)), (2, TensorBody::Operand(1))]),
10158            page,
10159        );
10160        let claim = approx("moe router softmax block", acc);
10161        out.push(f.finish(&stats(b), 2, vec![s1, s2, s3], page, format!("moe_softmax_block{b}_{ne}"), claim, target)?);
10162    }
10163    // The row's maximum and sum, from the blocks' (m_b, s_b) at elements 0 and 1 of their pages.
10164    let mut f = FlatBuild::at(acc, target)?;
10165    let pages: Vec<BufferId> = (0..blocks.len()).map(|b| f.input_paged(&stats(b), 2, &Fill::Unknown, acc, 1, target)).collect();
10166    let m = f.t.scratch_deep("row_max", 1, 1, acc);
10167    let maxes = (1..pages.len()).fold(TensorBody::Operand(0), |acc_body, b| TensorBody::Max(Box::new(acc_body), op(b)));
10168    let s1 = f.compute("row_max", flat_space(1), pages.iter().map(|&buf| Slot { buf, map: fixed(0) }).collect(), ident_map(one), maxes, m);
10169    let mut slots: Vec<Slot> = pages.iter().flat_map(|&buf| [Slot { buf, map: fixed(0) }, Slot { buf, map: fixed(1) }]).collect();
10170    slots.push(Slot { buf: m, map: fixed(0) });
10171    let row_m = 2 * pages.len();
10172    let z_body = (0..pages.len())
10173        .map(|b| mul(op(2 * b + 1), unary(SfpuKind::Exp, sub(op(2 * b), op(row_m)))))
10174        .reduce(add)
10175        .expect("a wide row has blocks");
10176    let z = f.t.scratch_deep("row_sum", 1, 1, acc);
10177    let s2 = f.compute("row_sum", flat_space(1), slots, ident_map(one), *z_body, z);
10178    let page = f.t.buffer("norm", 1, acc);
10179    let s3 = f.compute(
10180        "pack_norm",
10181        flat_space(2),
10182        vec![Slot { buf: m, map: fixed(0) }, Slot { buf: z, map: fixed(0) }],
10183        ident_map(one),
10184        banded(vec![(1, TensorBody::Operand(0)), (2, TensorBody::Operand(1))]),
10185        page,
10186    );
10187    let claim = approx("moe router softmax row", acc);
10188    out.push(f.finish(&norm, 2, vec![s1, s2, s3], page, format!("moe_softmax_row_{ne}"), claim, target)?);
10189    // `e^(L − m)·(1/Z)`, a block of the row an item.
10190    let mut f = FlatBuild::columns(ne, target)?;
10191    f.precision = acc;
10192    let l = f.input(&mn.logits, ne, &Fill::Unknown, target);
10193    let nrm = f.input_paged(&norm, 2, &Fill::Unknown, acc, 1, target);
10194    let probs = f.result_page("probs", ne, acc, target);
10195    let body = *mul(unary(SfpuKind::Exp, sub(op(0), op(1))), unary(SfpuKind::Recip, op(2)));
10196    let chain = f.pointwise(
10197        "softmax",
10198        ne,
10199        vec![Slot { buf: l, map: at(0) }, Slot { buf: nrm, map: fixed(0) }, Slot { buf: nrm, map: fixed(1) }],
10200        body,
10201        probs,
10202        target,
10203    )?;
10204    let claim = approx("moe router softmax", acc);
10205    out.push(f.finish(&mn.probs, ne, chain, probs, format!("moe_softmax_{ne}"), claim, target)?);
10206    Ok(out)
10207}
10208
10209/// A wide router's routing. Each block of the row (a section a block) computes its choosing scores as [`moe_choosing`]
10210/// does and takes its top `K` by [`moe_pick`], packing each pick's expert, choosing score and weight `P[e]` into a tile
10211/// of candidates, an unused slot's score masked. One section then takes `K` rounds of [`moe_pick`] over the blocks'
10212/// candidates in row order, each round's pick the candidate's expert and its weight the candidate's, and writes the
10213/// pick and weight tables as [`moe_routing`] does. The row's top `K` are each among their block's, and among equal
10214/// scores a round's lowest candidate is its lowest expert, so the picks are the row scan's. A group limit over a row
10215/// this wide is refused.
10216fn wide_moe_routing(
10217    spec: &MoeSpec,
10218    names: &StageNames,
10219    operands: &MoeOperands,
10220    mn: &MoeNames,
10221    target: &SpatialTarget,
10222) -> Result<Vec<SpatialProgram>, SpatialReject> {
10223    if spec.groups.is_some() {
10224        return Err(SpatialReject::NotLowered("a group limit over a router row wider than an accumulator page"));
10225    }
10226    let (ne, k) = (spec.ne, spec.k);
10227    let tw = target.tile.cols;
10228    let acc = AccPolicy::F32.interm_precision();
10229    let ix = index_type(target, ne as u64 + 1)?;
10230    let one_w = wide(1, target);
10231    let blocks = router_blocks(ne, target);
10232    let cand = |kind: &str, b: usize| format!("{}__cand_{kind}{b}", mn.picks);
10233    // A tile of `parts`' scalars, part `r` at element `r`, every later element `rest`.
10234    let pack = |f: &mut FlatBuild, chain: &mut Vec<StageId>, name: &str, parts: &[BufferId], rest: f64, dt: Precision| {
10235        let slots = parts.iter().map(|&buf| Slot { buf, map: fixed(0) }).collect();
10236        let mut bands: Vec<(usize, TensorBody)> = (0..parts.len()).map(|r| (r + 1, TensorBody::Operand(r))).collect();
10237        bands.push((tw, TensorBody::Const(rest.to_bits())));
10238        let out = f.t.buffer_deep(name, 1, 1, dt);
10239        chain.push(f.compute(&format!("pack_{name}"), flat_space(tw), slots, ident_map(one_w), banded(bands), out));
10240        out
10241    };
10242    let mut out = Vec::new();
10243    for (b, &(first, n)) in blocks.iter().enumerate() {
10244        let mut f = FlatBuild::at(acc, target)?;
10245        let p = f.input_slice(&mn.probs, ne, first, n, &Fill::Unknown, acc, target);
10246        let n_t = tiles_of(n, target);
10247        let mut chain = Vec::new();
10248        let choosing = match (operands.selection_bias, spec.score) {
10249            (Some(bi), _) => {
10250                let (b_name, b_fill) = names.operand(bi);
10251                let bias = f.input_slice(b_name, ne, first, n, b_fill, acc, target);
10252                let sel = f.t.scratch_deep("selection", 1, n_t, acc);
10253                chain.push(f.compute(
10254                    "selection",
10255                    flat_space(n),
10256                    vec![Slot { buf: p, map: at(0) }, Slot { buf: bias, map: at(0) }],
10257                    ident_map(wide(n_t, target)),
10258                    TensorBody::Max(add(op(0), op(1)), konst(MASKED)),
10259                    sel,
10260                ));
10261                sel
10262            }
10263            (None, MoeScore::Sigmoid) => {
10264                let (stage, routed) = routed_scores(&mut f, p, n, target);
10265                chain.push(stage);
10266                routed
10267            }
10268            (None, MoeScore::Softmax) => p,
10269        };
10270        let kb = k.min(n);
10271        let at_pick = ElemAccess { base: 0, terms: Vec::new(), indirect: Some(IndirectTerm { operand: 1, stride: 1, modulus: n }) };
10272        let (mut picked, mut ids, mut scores, mut weights) = (Vec::new(), Vec::new(), Vec::new(), Vec::new());
10273        for r in 0..kb {
10274            let (pick, end) = moe_pick(&mut f, &mut chain, ("pick", "route"), r, (choosing, map2(0, 0, 1), n), &picked, r + 1 < kb, ix, target);
10275            picked.extend(end.map(|end| (pick, end)));
10276            let id = f.t.scratch_deep(&format!("expert{r}"), 1, 1, ix);
10277            chain.push(f.compute(&format!("expert{r}"), flat_space(1), vec![Slot { buf: pick, map: fixed(0) }], ident_map(one_w), *add(op(0), konst((first * tw) as f64)), id));
10278            let score = f.t.scratch_deep(&format!("score{r}"), 1, 1, acc);
10279            chain.push(f.compute(
10280                &format!("score{r}"),
10281                flat_space(1),
10282                vec![Slot { buf: choosing, map: at_pick.clone() }, Slot { buf: pick, map: fixed(0) }],
10283                ident_map(one_w),
10284                TensorBody::Operand(0),
10285                score,
10286            ));
10287            let weight = f.t.scratch_deep(&format!("weight{r}"), 1, 1, acc);
10288            chain.push(f.compute(
10289                &format!("weight{r}"),
10290                flat_space(1),
10291                vec![Slot { buf: p, map: at_pick.clone() }, Slot { buf: pick, map: fixed(0) }],
10292                ident_map(one_w),
10293                TensorBody::Operand(0),
10294                weight,
10295            ));
10296            ids.push(id);
10297            scores.push(score);
10298            weights.push(weight);
10299        }
10300        let ids_out = pack(&mut f, &mut chain, "ids", &ids, 0.0, ix);
10301        let scores_out = pack(&mut f, &mut chain, "scores", &scores, MASKED, acc);
10302        let weights_out = pack(&mut f, &mut chain, "weights", &weights, 0.0, acc);
10303        let mut stores = Vec::new();
10304        for (kind, buf, dt) in [("ids", ids_out, ix), ("scores", scores_out, acc), ("weights", weights_out, acc)] {
10305            let tensor = f.t.tensor(&cand(kind, b), TensorRole::Output, 1, tw, dt, TensorInit::Zero);
10306            stores.push((buf, f.t.store(&format!("store_{kind}"), tensor, Vec::new(), 1, buf)));
10307        }
10308        let claim = approx("mixture-of-experts routing block", acc);
10309        out.push(f.wire_and_finish_many(chain, &stores, 1, format!("moe_candidates{b}_{ne}x{k}"), claim, target)?);
10310    }
10311    // The candidates of every block in row order, `K` rounds over them.
10312    let nb = blocks.len();
10313    let mut f = FlatBuild::at(acc, target)?;
10314    let mut chain = Vec::new();
10315    let mut row = |f: &mut FlatBuild, kind: &str, dt: Precision| {
10316        let pages: Vec<BufferId> = (0..nb).map(|b| f.input_paged(&cand(kind, b), tw, &Fill::Unknown, dt, 1, target)).collect();
10317        let slots = pages.iter().enumerate().map(|(b, &buf)| Slot { buf, map: at(-((b * tw) as i64)) }).collect();
10318        let bands = (0..nb).map(|b| ((b + 1) * tw, TensorBody::Operand(b))).collect();
10319        let all = f.t.scratch_deep(&format!("{kind}_row"), 1, nb, dt);
10320        chain.push(f.compute(&format!("{kind}_row"), flat_space(nb * tw), slots, ident_map(wide(nb, target)), banded(bands), all));
10321        all
10322    };
10323    let (ids, scores, weights) = (row(&mut f, "ids", ix), row(&mut f, "scores", acc), row(&mut f, "weights", acc));
10324    let at_pick = ElemAccess { base: 0, terms: Vec::new(), indirect: Some(IndirectTerm { operand: 1, stride: 1, modulus: nb * tw }) };
10325    let (mut picked, mut picks, mut chosen) = (Vec::new(), Vec::new(), Vec::new());
10326    for r in 0..k {
10327        let (at, end) = moe_pick(&mut f, &mut chain, ("candidate", "choose"), r, (scores, map2(0, 0, 1), nb * tw), &picked, r + 1 < k, ix, target);
10328        picked.extend(end.map(|end| (at, end)));
10329        let pick = f.t.scratch_deep(&format!("pick{r}"), 1, 1, ix);
10330        chain.push(f.compute(&format!("pick{r}"), flat_space(1), vec![Slot { buf: ids, map: at_pick.clone() }, Slot { buf: at, map: fixed(0) }], ident_map(one_w), TensorBody::Operand(0), pick));
10331        let weight = f.t.scratch_deep(&format!("weight{r}"), 1, 1, acc);
10332        chain.push(f.compute(&format!("weight{r}"), flat_space(1), vec![Slot { buf: weights, map: at_pick.clone() }, Slot { buf: at, map: fixed(0) }], ident_map(one_w), round_weight(spec), weight));
10333        picks.push(pick);
10334        chosen.push(weight);
10335    }
10336    // A tile of the pick and weight tables a pick, as [`moe_routing`] writes them, packed and stored in parts of at most
10337    // an accumulator page: qwen4_exp's top 10 are more picks than a page holds.
10338    let blk = row_block_tiles(target);
10339    let mut stores = Vec::new();
10340    for (name, parts, dt, tensor) in [("picks", &picks, ix, &mn.picks), ("weights", &chosen, acc, &mn.weights)] {
10341        let t = f.t.tensor(tensor, TensorRole::Output, 1, k * tw, dt, TensorInit::Zero);
10342        for (c, first) in (0..k).step_by(blk).enumerate() {
10343            let part = &parts[first..(first + blk).min(k)];
10344            let slots = part.iter().map(|&buf| Slot { buf, map: fixed(0) }).collect();
10345            let bands = (0..part.len()).map(|r| ((r + 1) * tw, TensorBody::Operand(r))).collect();
10346            let buf = f.t.buffer_deep(&format!("{name}{c}"), 1, part.len(), dt);
10347            chain.push(f.compute(&format!("pack_{name}{c}"), flat_space(part.len() * tw), slots, ident_map(wide(part.len(), target)), banded(bands), buf));
10348            stores.push((buf, f.t.store_block_over(&format!("store_{name}{c}"), t, first, Vec::new(), 1, part.len(), k, (1, 1), buf)));
10349        }
10350    }
10351    let claim = approx("mixture-of-experts routing", acc);
10352    out.push(f.wire_and_finish_many(chain, &stores, k, format!("moe_routing_{ne}x{k}"), claim, target)?);
10353    Ok(out)
10354}
10355
10356/// The plan of a section whose item walks `chunks` sub-blocks, one output tile per pass.
10357fn chunked_plan(chunks: usize, items: usize) -> BlockPlan {
10358    BlockPlan {
10359        per_site_m: chunks,
10360        per_site_n: 1,
10361        grid: (items, 1),
10362        inner: InnerLoop {
10363            num_blocks: 1,
10364            block_k_tiles: 1,
10365        },
10366        block: BlockShape {
10367            osh: 1,
10368            osw: 1,
10369            ibw: 1,
10370        },
10371        acc: AccPolicy::F32,
10372        pack_accumulate: false,
10373        engines: EnginePacking::LoadsTogether,
10374        granularity: Granularity::PerTile,
10375    }
10376}
10377
10378/// A partial GEMV of one expert's rows: 32 outputs of a tile row of the weight against one `chunk`-tile piece
10379/// of the input — `Σ_t W[o, t]·v[t]` over the chunk, the weight page read as `o·(chunk·32) + t`.
10380fn row_dot(chunk_elems: usize) -> (Vec<IterSpec>, Vec<ElemAccess>) {
10381    (
10382        vec![
10383            IterSpec {
10384                var: LV_GEN_ROW,
10385                extent: chunk_elems,
10386                kind: IterKind::Reduce {
10387                    pool: ReducePool::Sum,
10388                },
10389            },
10390            IterSpec {
10391                var: LV_GEN_COL,
10392                extent: 32,
10393                kind: IterKind::Parallel,
10394            },
10395        ],
10396        vec![
10397            ElemAccess {
10398                base: 0,
10399                terms: vec![(LV_GEN_COL, chunk_elems as i64), (LV_GEN_ROW, 1)],
10400                indirect: None,
10401            },
10402            ElemAccess {
10403                base: 0,
10404                terms: vec![(LV_GEN_ROW, 1)],
10405                indirect: None,
10406            },
10407            const_map(),
10408        ],
10409    )
10410}
10411
10412/// A load of `chunk` tiles of a chosen expert's weight: tile row `ITEM_X` of the expert's `rows_t × cols_t` slab,
10413/// chunk `LV_SUB_H` of the row, the slab picked by the index in `pick`.
10414#[allow(clippy::too_many_arguments)]
10415fn expert_chunk(
10416    t: &mut Template,
10417    name: &str,
10418    weight: TensorId,
10419    pick: BufferId,
10420    experts: usize,
10421    rows_t: usize,
10422    cols_t: usize,
10423    chunk: usize,
10424    into: BufferId,
10425) -> StageId {
10426    let origin = TileAddr {
10427        base: 0,
10428        steps: vec![(LV_ITEM_X, cols_t), (super::LV_SUB_H, chunk)],
10429        indirect: Some(IndirectTile {
10430            index: pick,
10431            tile_stride: rows_t * cols_t,
10432            elems_per_tile: 1,
10433            bound: (0, experts - 1),
10434        }),
10435    };
10436    t.stage(
10437        name,
10438        StageKind::Load {
10439            tensor: weight,
10440            pattern: AccessPattern {
10441                origin,
10442                block_h: 1,
10443                block_w: chunk,
10444                stride_w: 1,
10445                stride_h: cols_t,
10446            },
10447            num_blocks: 1,
10448            multicast: None,
10449        },
10450        vec![pick],
10451        vec![into],
10452    )
10453}
10454
10455/// Gate and up partial sums: item `(r, ot)` is round `r`'s tile row `ot` of outputs; its sub-blocks walk the `d`
10456/// input in chunks of `cd` tiles.
10457fn moe_up_partials(
10458    spec: &MoeSpec,
10459    names: &StageNames,
10460    mn: &MoeNames,
10461    cd: usize,
10462    x_dtype: Precision,
10463    target: &SpatialTarget,
10464) -> Result<SpatialProgram, SpatialReject> {
10465    let (d, dff, ne, k) = (spec.d, spec.dff, spec.ne, spec.k);
10466    let tw = target.tile.cols;
10467    let (d_t, f_t) = (d / tw, dff / tw);
10468    let nc = d_t / cd;
10469    let precision = elem_precision(target)?;
10470    let acc = AccPolicy::F32.interm_precision();
10471    let ix = index_type(target, ne as u64 + 1)?;
10472    let (x_name, x_fill) = names.operand(0);
10473    let (wg_name, wg_fill) = names.operand(2);
10474    let (wu_name, wu_fill) = names.operand(3);
10475    let mut t = Template::new(target)?;
10476    let scaler = t.tensor(
10477        &mn.scaler,
10478        TensorRole::Param,
10479        target.tile.rows,
10480        tw,
10481        precision,
10482        TensorInit::ConstF(1.0f64.to_bits()),
10483    );
10484    let picks = t.tensor(
10485        &mn.picks,
10486        TensorRole::Input,
10487        1,
10488        k * tw,
10489        ix,
10490        TensorInit::External,
10491    );
10492    let x = t.tensor(
10493        x_name,
10494        TensorRole::Input,
10495        1,
10496        d,
10497        x_dtype,
10498        super::from_tensor_op::init_of(x_fill, d),
10499    );
10500    let wg = t.tensor(
10501        wg_name,
10502        TensorRole::Param,
10503        ne * dff,
10504        d,
10505        precision,
10506        super::from_tensor_op::init_of(wg_fill, ne * dff * d),
10507    );
10508    let wu = t.tensor(
10509        wu_name,
10510        TensorRole::Param,
10511        ne * dff,
10512        d,
10513        precision,
10514        super::from_tensor_op::init_of(wu_fill, ne * dff * d),
10515    );
10516    let gp = t.tensor(
10517        &mn.gate_part,
10518        TensorRole::Output,
10519        k * f_t * nc * target.tile.rows,
10520        tw,
10521        acc,
10522        TensorInit::Zero,
10523    );
10524    let up = t.tensor(
10525        &mn.up_part,
10526        TensorRole::Output,
10527        k * f_t * nc * target.tile.rows,
10528        tw,
10529        acc,
10530        TensorInit::Zero,
10531    );
10532    let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
10533    let pg_in = t.buffer_deep("in_pick_gate", 1, 1, ix);
10534    let pu_in = t.buffer_deep("in_pick_up", 1, 1, ix);
10535    let x_in = t.buffer("in_x", cd, x_dtype);
10536    let wg_in = t.buffer("in_gate", cd, precision);
10537    let wu_in = t.buffer("in_up", cd, precision);
10538    let g_out = t.buffer("gate_out", 1, acc);
10539    let u_out = t.buffer("up_out", 1, acc);
10540    let hold = |mut v: Vec<(LoopVar, usize)>| {
10541        v.push((super::LV_SUB_H, 0));
10542        v
10543    };
10544    let l_sc = t.load_block_at("load_scaler", scaler, 0, hold(Vec::new()), 1, 1, 1, sc_in);
10545    // Every gather retires its index page, so each expert load waits on a pick of its own.
10546    let l_pg = t.load_block_at(
10547        "load_pick_gate",
10548        picks,
10549        0,
10550        hold(vec![(LV_ITEM_Y, 1)]),
10551        1,
10552        1,
10553        1,
10554        pg_in,
10555    );
10556    let l_pu = t.load_block_at(
10557        "load_pick_up",
10558        picks,
10559        0,
10560        hold(vec![(LV_ITEM_Y, 1)]),
10561        1,
10562        1,
10563        1,
10564        pu_in,
10565    );
10566    let l_x = t.load_block_at("load_x", x, 0, vec![(super::LV_SUB_H, cd)], 1, cd, cd, x_in);
10567    let l_wg = expert_chunk(&mut t, "load_gate", wg, pg_in, ne, f_t, d_t, cd, wg_in);
10568    let l_wu = expert_chunk(&mut t, "load_up", wu, pu_in, ne, f_t, d_t, cd, wu_in);
10569    let (iters, maps) = row_dot(cd * tw);
10570    let s_g = t.stage(
10571        "gate_dot",
10572        StageKind::Compute(ComputeOp::Generic {
10573            iters: iters.clone(),
10574            ins: maps.clone(),
10575            out: per_col_map(),
10576            body: *mul(op(0), op(1)),
10577        }),
10578        vec![wg_in, x_in, sc_in],
10579        vec![g_out],
10580    );
10581    let s_u = t.stage(
10582        "up_dot",
10583        StageKind::Compute(ComputeOp::Generic {
10584            iters,
10585            ins: maps,
10586            out: per_col_map(),
10587            body: *mul(op(0), op(1)),
10588        }),
10589        vec![wu_in, x_in, sc_in],
10590        vec![u_out],
10591    );
10592    let steps = vec![(LV_ITEM_Y, f_t * nc), (LV_ITEM_X, nc), (super::LV_SUB_H, 1)];
10593    let st_g = t.store_block_over("store_gate", gp, 0, steps.clone(), 1, 1, 1, (nc, 1), g_out);
10594    let st_u = t.store_block_over("store_up", up, 0, steps, 1, 1, 1, (nc, 1), u_out);
10595    for (from, buf, to) in [
10596        (l_pg, pg_in, l_wg),
10597        (l_pu, pu_in, l_wu),
10598        (l_wg, wg_in, s_g),
10599        (l_x, x_in, s_g),
10600        (l_sc, sc_in, s_g),
10601        (l_wu, wu_in, s_u),
10602        (s_g, g_out, st_g),
10603        (s_u, u_out, st_u),
10604    ] {
10605        t.edge(from, buf, to);
10606    }
10607    let items: Vec<WorkItem> = (0..k)
10608        .flat_map(|r| (0..f_t).map(move |ot| (r, ot)))
10609        .enumerate()
10610        .map(|(id, (r, ot))| WorkItem {
10611            id: WorkItemId::new(id),
10612            coord: (r, ot),
10613            out: TileRect {
10614                row0_t: (r * f_t + ot) * nc,
10615                col0_t: 0,
10616                h_t: nc,
10617                w_t: 1,
10618            },
10619        })
10620        .collect();
10621    let mut prog = t.finish(
10622        format!("moe_up_{ne}x{k}_{d}x{dff}_c{cd}"),
10623        items,
10624        AccPolicy::F32,
10625        approx("mixture-of-experts gate/up partials", precision),
10626        target,
10627    )?;
10628    prog.work.plan = Some(chunked_plan(nc, k * f_t));
10629    prog.work.k_tiles = 1;
10630    Ok(prog)
10631}
10632
10633/// The tile `LV_ITEM_X` of a chosen expert's row of `table` (`ne` rows of `row_t` tiles): the round's pick loaded
10634/// from the pick table — tile `round.1`, plus `LV_ITEM_Y` when `round.0` is it — then the bias tile gathered by it.
10635/// Each gather retires its index page, so every bias load waits on a pick page of its own.
10636#[allow(clippy::too_many_arguments)]
10637fn expert_bias_tile(
10638    t: &mut Template,
10639    part: &str,
10640    mn: &MoeNames,
10641    table: TensorId,
10642    round: (LoopVar, usize),
10643    row_t: usize,
10644    spec: &MoeSpec,
10645    target: &SpatialTarget,
10646) -> Result<(StageId, BufferId), SpatialReject> {
10647    let ix = index_type(target, spec.ne as u64 + 1)?;
10648    let acc = AccPolicy::F32.interm_precision();
10649    let picks = t.tensor(&mn.picks, TensorRole::Input, 1, spec.k * target.tile.cols, ix, TensorInit::External);
10650    let pick_in = t.buffer_deep(&format!("in_pick_{part}_bias"), 1, 1, ix);
10651    let bias_in = t.buffer_deep(&format!("in_{part}_bias"), 1, 1, acc);
10652    let steps = if round.0 == LV_ITEM_Y { vec![(LV_ITEM_Y, 1)] } else { Vec::new() };
10653    let l_p = t.load_block_at(&format!("load_pick_{part}_bias"), picks, round.1, steps, 1, 1, 1, pick_in);
10654    let l_b = t.stage(
10655        &format!("load_{part}_bias"),
10656        StageKind::Load {
10657            tensor: table,
10658            pattern: AccessPattern {
10659                origin: TileAddr {
10660                    base: 0,
10661                    steps: vec![(LV_ITEM_X, 1)],
10662                    indirect: Some(IndirectTile { index: pick_in, tile_stride: row_t, elems_per_tile: 1, bound: (0, spec.ne - 1) }),
10663                },
10664                block_h: 1,
10665                block_w: 1,
10666                stride_w: 1,
10667                stride_h: 1,
10668            },
10669            num_blocks: 1,
10670            multicast: None,
10671        },
10672        vec![pick_in],
10673        vec![bias_in],
10674    );
10675    t.edge(l_p, pick_in, l_b);
10676    Ok((l_b, bias_in))
10677}
10678
10679/// `hidden = act(Σ gate + b_gate, Σ up + b_up)` for each round's tile of outputs — the partials added, the chosen
10680/// expert's biases when it has them, then the mixture's gated activation.
10681fn moe_hidden(
10682    spec: &MoeSpec,
10683    names: &StageNames,
10684    operands: &MoeOperands,
10685    mn: &MoeNames,
10686    nc: usize,
10687    target: &SpatialTarget,
10688) -> Result<SpatialProgram, SpatialReject> {
10689    let (dff, k) = (spec.dff, spec.k);
10690    let tw = target.tile.cols;
10691    let f_t = dff / tw;
10692    let precision = elem_precision(target)?;
10693    let acc = AccPolicy::F32.interm_precision();
10694    let mut t = Template::new(target)?;
10695    let scaler = t.tensor(
10696        &mn.scaler,
10697        TensorRole::Param,
10698        target.tile.rows,
10699        tw,
10700        precision,
10701        TensorInit::ConstF(1.0f64.to_bits()),
10702    );
10703    let gp = t.tensor(
10704        &mn.gate_part,
10705        TensorRole::Input,
10706        k * f_t * nc * target.tile.rows,
10707        tw,
10708        acc,
10709        TensorInit::External,
10710    );
10711    let up = t.tensor(
10712        &mn.up_part,
10713        TensorRole::Input,
10714        k * f_t * nc * target.tile.rows,
10715        tw,
10716        acc,
10717        TensorInit::External,
10718    );
10719    let hidden = t.tensor(
10720        &mn.hidden,
10721        TensorRole::Output,
10722        1,
10723        k * dff,
10724        acc,
10725        TensorInit::Zero,
10726    );
10727    let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
10728    let g_in = t.buffer("in_gate", nc, acc);
10729    let u_in = t.buffer("in_up", nc, acc);
10730    let gs = t.scratch_deep("gate_sum", 1, 1, acc);
10731    let us = t.scratch_deep("up_sum", 1, 1, acc);
10732    let out = t.buffer("out", 1, acc);
10733    let l_sc = t.load_block_at("load_scaler", scaler, 0, Vec::new(), 1, 1, 1, sc_in);
10734    let parts = vec![(LV_ITEM_Y, f_t * nc), (LV_ITEM_X, nc)];
10735    let l_g = t.load_block_at("load_gate", gp, 0, parts.clone(), nc, 1, 1, g_in);
10736    let l_u = t.load_block_at("load_up", up, 0, parts, nc, 1, 1, u_in);
10737    // Partial `c`'s output `o` sits at `c·32 + o` of the page.
10738    let fold = |t: &mut Template, name: &str, from: BufferId, into: BufferId| {
10739        let iters = vec![
10740            IterSpec {
10741                var: LV_GEN_ROW,
10742                extent: nc,
10743                kind: IterKind::Reduce {
10744                    pool: ReducePool::Sum,
10745                },
10746            },
10747            IterSpec {
10748                var: LV_GEN_COL,
10749                extent: tw,
10750                kind: IterKind::Parallel,
10751            },
10752        ];
10753        let maps = vec![
10754            ElemAccess {
10755                base: 0,
10756                terms: vec![(LV_GEN_ROW, tw as i64), (LV_GEN_COL, 1)],
10757                indirect: None,
10758            },
10759            const_map(),
10760        ];
10761        t.stage(
10762            name,
10763            StageKind::Compute(ComputeOp::Generic {
10764                iters,
10765                ins: maps,
10766                out: per_col_map(),
10767                body: TensorBody::Operand(0),
10768            }),
10769            vec![from, sc_in],
10770            vec![into],
10771        )
10772    };
10773    let s_g = fold(&mut t, "gate_sum", g_in, gs);
10774    let s_u = fold(&mut t, "up_sum", u_in, us);
10775    let w1 = wide(1, target);
10776    let mut ins = vec![gs, us];
10777    let mut terms: Vec<Box<TensorBody>> = vec![op(0), op(1)];
10778    let mut bias_loads = Vec::new();
10779    // The round's weight, when it scales the input: `w·(x·W)` is `(w·x)·W`, so it scales the summed products.
10780    if spec.weigh_inputs {
10781        let weights = t.tensor(&mn.weights, TensorRole::Input, 1, k * tw, acc, TensorInit::External);
10782        let w_in = t.buffer_deep("in_weight", 1, 1, acc);
10783        let l_w = t.load_block_at("load_weight", weights, 0, vec![(LV_ITEM_Y, 1)], 1, 1, 1, w_in);
10784        terms = vec![mul(op(2), op(0)), mul(op(2), op(1))];
10785        ins.push(w_in);
10786        bias_loads.push((l_w, w_in));
10787    }
10788    // Each bias the chosen expert has: tile `ot` of its row, gathered by the round's pick.
10789    for (i, part) in ["gate", "up"].into_iter().enumerate() {
10790        let Some(bi) = operands.biases[i] else { continue };
10791        let (b_name, b_fill) = names.operand(bi);
10792        let table = t.tensor(b_name, TensorRole::Param, 1, spec.ne * dff, acc, super::from_tensor_op::init_of(b_fill, spec.ne * dff));
10793        let (l_b, bias_in) = expert_bias_tile(&mut t, part, mn, table, (LV_ITEM_Y, 0), f_t, spec, target)?;
10794        terms[i] = add(terms[i].clone(), op(ins.len()));
10795        ins.push(bias_in);
10796        bias_loads.push((l_b, bias_in));
10797    }
10798    let [gate, up]: [Box<TensorBody>; 2] = terms.try_into().expect("a gate and an up");
10799    let s_h = t.stage(
10800        "swiglu",
10801        StageKind::Compute(ComputeOp::Generic {
10802            iters: flat_space(tw),
10803            ins: vec![at(0); ins.len()],
10804            out: ident_map(w1),
10805            body: *moe_gated(spec.act, gate, up),
10806        }),
10807        ins,
10808        vec![out],
10809    );
10810    for (from, buf) in bias_loads {
10811        t.edge(from, buf, s_h);
10812    }
10813    let st = t.store_block(
10814        "store_hidden",
10815        hidden,
10816        0,
10817        vec![(LV_ITEM_Y, f_t), (LV_ITEM_X, 1)],
10818        1,
10819        1,
10820        1,
10821        out,
10822    );
10823    for (from, buf, to) in [
10824        (l_sc, sc_in, s_g),
10825        (l_g, g_in, s_g),
10826        (l_u, u_in, s_u),
10827        (s_g, gs, s_h),
10828        (s_u, us, s_h),
10829        (s_h, out, st),
10830    ] {
10831        t.edge(from, buf, to);
10832    }
10833    let items: Vec<WorkItem> = (0..k)
10834        .flat_map(|r| (0..f_t).map(move |ot| (r, ot)))
10835        .enumerate()
10836        .map(|(id, (r, ot))| WorkItem {
10837            id: WorkItemId::new(id),
10838            coord: (r, ot),
10839            out: TileRect {
10840                row0_t: 0,
10841                col0_t: r * f_t + ot,
10842                h_t: 1,
10843                w_t: 1,
10844            },
10845        })
10846        .collect();
10847    t.finish(
10848        format!("moe_hidden_{k}x{dff}"),
10849        items,
10850        AccPolicy::F32,
10851        approx("mixture-of-experts SwiGLU", precision),
10852        target,
10853    )
10854}
10855
10856/// Down partial sums: item `(r, ot)` is round `r`'s tile row `ot` of the block's `d` outputs; its sub-blocks walk
10857/// the hidden activation in chunks of `cd` tiles.
10858fn moe_down_partials(
10859    spec: &MoeSpec,
10860    names: &StageNames,
10861    mn: &MoeNames,
10862    cd: usize,
10863    target: &SpatialTarget,
10864) -> Result<SpatialProgram, SpatialReject> {
10865    let (d, dff, ne, k) = (spec.d, spec.dff, spec.ne, spec.k);
10866    let tw = target.tile.cols;
10867    let (d_t, f_t) = (d / tw, dff / tw);
10868    let nc = f_t / cd;
10869    let precision = elem_precision(target)?;
10870    let acc = AccPolicy::F32.interm_precision();
10871    let ix = index_type(target, ne as u64 + 1)?;
10872    let (wd_name, wd_fill) = names.operand(4);
10873    let mut t = Template::new(target)?;
10874    let scaler = t.tensor(
10875        &mn.scaler,
10876        TensorRole::Param,
10877        target.tile.rows,
10878        tw,
10879        precision,
10880        TensorInit::ConstF(1.0f64.to_bits()),
10881    );
10882    let picks = t.tensor(
10883        &mn.picks,
10884        TensorRole::Input,
10885        1,
10886        k * tw,
10887        ix,
10888        TensorInit::External,
10889    );
10890    let hidden = t.tensor(
10891        &mn.hidden,
10892        TensorRole::Input,
10893        1,
10894        k * dff,
10895        acc,
10896        TensorInit::External,
10897    );
10898    let wd = t.tensor(
10899        wd_name,
10900        TensorRole::Param,
10901        ne * d,
10902        dff,
10903        precision,
10904        super::from_tensor_op::init_of(wd_fill, ne * d * dff),
10905    );
10906    let dp = t.tensor(
10907        &mn.down_part,
10908        TensorRole::Output,
10909        k * d_t * nc * target.tile.rows,
10910        tw,
10911        acc,
10912        TensorInit::Zero,
10913    );
10914    let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
10915    let pk_in = t.buffer_deep("in_pick", 1, 1, ix);
10916    let h_in = t.buffer("in_hidden", cd, acc);
10917    let w_in = t.buffer("in_down", cd, precision);
10918    let d_out = t.buffer("down_out", 1, acc);
10919    let hold = |mut v: Vec<(LoopVar, usize)>| {
10920        v.push((super::LV_SUB_H, 0));
10921        v
10922    };
10923    let l_sc = t.load_block_at("load_scaler", scaler, 0, hold(Vec::new()), 1, 1, 1, sc_in);
10924    let l_pk = t.load_block_at(
10925        "load_pick",
10926        picks,
10927        0,
10928        hold(vec![(LV_ITEM_Y, 1)]),
10929        1,
10930        1,
10931        1,
10932        pk_in,
10933    );
10934    let l_h = t.load_block_at(
10935        "load_hidden",
10936        hidden,
10937        0,
10938        vec![(LV_ITEM_Y, f_t), (super::LV_SUB_H, cd)],
10939        1,
10940        cd,
10941        cd,
10942        h_in,
10943    );
10944    let l_w = expert_chunk(&mut t, "load_down", wd, pk_in, ne, d_t, f_t, cd, w_in);
10945    let (iters, maps) = row_dot(cd * tw);
10946    let s_d = t.stage(
10947        "down_dot",
10948        StageKind::Compute(ComputeOp::Generic {
10949            iters,
10950            ins: maps,
10951            out: per_col_map(),
10952            body: *mul(op(0), op(1)),
10953        }),
10954        vec![w_in, h_in, sc_in],
10955        vec![d_out],
10956    );
10957    let st = t.store_block_over(
10958        "store_down",
10959        dp,
10960        0,
10961        vec![(LV_ITEM_Y, d_t * nc), (LV_ITEM_X, nc), (super::LV_SUB_H, 1)],
10962        1,
10963        1,
10964        1,
10965        (nc, 1),
10966        d_out,
10967    );
10968    for (from, buf, to) in [
10969        (l_pk, pk_in, l_w),
10970        (l_w, w_in, s_d),
10971        (l_h, h_in, s_d),
10972        (l_sc, sc_in, s_d),
10973        (s_d, d_out, st),
10974    ] {
10975        t.edge(from, buf, to);
10976    }
10977    let items: Vec<WorkItem> = (0..k)
10978        .flat_map(|r| (0..d_t).map(move |ot| (r, ot)))
10979        .enumerate()
10980        .map(|(id, (r, ot))| WorkItem {
10981            id: WorkItemId::new(id),
10982            coord: (r, ot),
10983            out: TileRect {
10984                row0_t: (r * d_t + ot) * nc,
10985                col0_t: 0,
10986                h_t: nc,
10987                w_t: 1,
10988            },
10989        })
10990        .collect();
10991    let mut prog = t.finish(
10992        format!("moe_down_{ne}x{k}_{d}x{dff}_c{cd}"),
10993        items,
10994        AccPolicy::F32,
10995        approx("mixture-of-experts down partials", precision),
10996        target,
10997    )?;
10998    prog.work.plan = Some(chunked_plan(nc, k * d_t));
10999    prog.work.k_tiles = 1;
11000    Ok(prog)
11001}
11002
11003/// The combine: per tile of the `d` outputs, each round's down partials added, plus its expert's down bias when the
11004/// mixture has one, weighted by its router probability, and divided by the kept probabilities' sum when the router
11005/// renormalizes — the desugar's `Σ (P[b]/Σ P)·y_b`, or `Σ P[b]·y_b`.
11006fn moe_combine(
11007    spec: &MoeSpec,
11008    names: &StageNames,
11009    operands: &MoeOperands,
11010    mn: &MoeNames,
11011    nc: usize,
11012    target: &SpatialTarget,
11013) -> Result<SpatialProgram, SpatialReject> {
11014    let (d, k) = (spec.d, spec.k);
11015    let tw = target.tile.cols;
11016    let d_t = d / tw;
11017    let precision = elem_precision(target)?;
11018    let acc = AccPolicy::F32.interm_precision();
11019    let mut t = Template::new(target)?;
11020    let scaler = t.tensor(
11021        &mn.scaler,
11022        TensorRole::Param,
11023        target.tile.rows,
11024        tw,
11025        precision,
11026        TensorInit::ConstF(1.0f64.to_bits()),
11027    );
11028    // The kept experts' weights, unless each already scaled its expert's input.
11029    let weights = (!spec.weigh_inputs).then(|| t.tensor(&mn.weights, TensorRole::Input, 1, k * tw, acc, TensorInit::External));
11030    let dp = t.tensor(
11031        &mn.down_part,
11032        TensorRole::Input,
11033        k * d_t * nc * target.tile.rows,
11034        tw,
11035        acc,
11036        TensorInit::External,
11037    );
11038    let y = t.tensor(
11039        &names.out,
11040        TensorRole::Output,
11041        1,
11042        d,
11043        precision,
11044        TensorInit::Zero,
11045    );
11046    let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
11047    let w_in = weights.map(|_| t.buffer("in_weights", k, acc));
11048    let p_in = t.buffer("in_down", k * nc, acc);
11049    // ★ A ROUND'S SUM IS A TILE OF THE ACCUMULATOR: more rounds than a page holds (qwen4_exp's top 10) are summed and
11050    // weighed in groups of at most a page of rounds, and the groups' weighted sums added.
11051    let blk = row_block_tiles(target);
11052    let groups: Vec<(usize, usize)> = (0..k).step_by(blk).map(|g0| (g0, (g0 + blk).min(k))).collect();
11053    let named = |base: &str, g: usize| if groups.len() == 1 { base.to_string() } else { format!("{base}{g}") };
11054    let pages: Vec<(BufferId, BufferId)> = groups
11055        .iter()
11056        .enumerate()
11057        .map(|(g, &(g0, g1))| (t.scratch_deep(&named("down_rounds", g), 1, g1 - g0, acc), t.scratch_deep(&named("weighted", g), 1, 1, acc)))
11058        .collect();
11059    let ws = t.scratch_deep("weight_sum", 1, 1, acc);
11060    let out = t.buffer("out", 1, precision);
11061    let l_sc = t.load_block_at("load_scaler", scaler, 0, Vec::new(), 1, 1, 1, sc_in);
11062    let l_w = weights.zip(w_in).map(|(weights, w_in)| t.load_block_at("load_weights", weights, 0, Vec::new(), 1, k, k, w_in));
11063    // Round `r`'s partials for this output tile: `nc` tiles at `(r·d_t + ot)·nc` — `k` rows `d_t·nc` apart.
11064    let l_p = t.load_block_at(
11065        "load_down",
11066        dp,
11067        0,
11068        vec![(LV_ITEM_X, nc)],
11069        k,
11070        nc,
11071        d_t * nc,
11072        p_in,
11073    );
11074    let mut edges = Vec::new();
11075    let mut weighed = Vec::new();
11076    for (g, &(g0, g1)) in groups.iter().enumerate() {
11077        let len = g1 - g0;
11078        let (dr, ys) = pages[g];
11079        // Page tile `r·nc + c` holds round `r`'s partial `c`; output column `q = (r − g0)·32 + o` of the group's round
11080        // sums reads it at `q + (r·(nc − 1) + g0)·32 + c·32`, one band per round.
11081        let mut slots = Vec::new();
11082        let mut bands = Vec::new();
11083        for (j, r) in (g0..g1).enumerate() {
11084            slots.push(ElemAccess {
11085                base: ((r * (nc - 1) + g0) * tw) as i64,
11086                terms: vec![(LV_GEN_COL, 1), (LV_GEN_ROW, tw as i64)],
11087                indirect: None,
11088            });
11089            bands.push(((j + 1) * tw, TensorBody::Operand(j)));
11090        }
11091        slots.push(const_map());
11092        let mut ins = vec![p_in; len];
11093        ins.push(sc_in);
11094        let s_r = t.stage(
11095            &named("round_sums", g),
11096            StageKind::Compute(ComputeOp::Generic {
11097                iters: vec![
11098                    IterSpec {
11099                        var: LV_GEN_ROW,
11100                        extent: nc,
11101                        kind: IterKind::Reduce {
11102                            pool: ReducePool::Sum,
11103                        },
11104                    },
11105                    IterSpec {
11106                        var: LV_GEN_COL,
11107                        extent: len * tw,
11108                        kind: IterKind::Parallel,
11109                    },
11110                ],
11111                ins: slots,
11112                out: per_col_map(),
11113                body: banded(bands),
11114            }),
11115            ins,
11116            vec![dr],
11117        );
11118        // `Σ_r w_r·down_r`, or `Σ_r down_r` when each round's weight scaled its expert's input.
11119        let round_sum = ElemAccess {
11120            base: 0,
11121            terms: vec![(LV_GEN_ROW, tw as i64), (LV_GEN_COL, 1)],
11122            indirect: None,
11123        };
11124        let (sum_ins, sum_maps, sum_body) = match w_in {
11125            None => (vec![dr, sc_in], vec![round_sum, const_map()], TensorBody::Operand(0)),
11126            Some(w_in) => (
11127                vec![w_in, dr, sc_in],
11128                vec![
11129                    ElemAccess {
11130                        base: (g0 * tw) as i64,
11131                        terms: vec![(LV_GEN_ROW, tw as i64)],
11132                        indirect: None,
11133                    },
11134                    round_sum,
11135                    const_map(),
11136                ],
11137                *mul(op(0), op(1)),
11138            ),
11139        };
11140        let s_y = t.stage(
11141            &named("weighted", g),
11142            StageKind::Compute(ComputeOp::Generic {
11143                iters: vec![
11144                    IterSpec {
11145                        var: LV_GEN_ROW,
11146                        extent: len,
11147                        kind: IterKind::Reduce {
11148                            pool: ReducePool::Sum,
11149                        },
11150                    },
11151                    IterSpec {
11152                        var: LV_GEN_COL,
11153                        extent: tw,
11154                        kind: IterKind::Parallel,
11155                    },
11156                ],
11157                ins: sum_maps,
11158                out: per_col_map(),
11159                body: sum_body,
11160            }),
11161            sum_ins,
11162            vec![ys],
11163        );
11164        if g == 0 {
11165            edges.extend([(l_sc, sc_in, s_r), (l_p, p_in, s_r)]);
11166            if let (Some(l_w), Some(w_in)) = (l_w, w_in) {
11167                edges.push((l_w, w_in, s_y));
11168            }
11169        }
11170        edges.push((s_r, dr, s_y));
11171        weighed.push((s_y, ys));
11172    }
11173    let one_w = wide(1, target);
11174    // The groups' weighted sums added, when there is more than one.
11175    let (s_y, ys) = match weighed.as_slice() {
11176        [one] => *one,
11177        many => {
11178            let ys = t.scratch_deep("weighted", 1, 1, acc);
11179            let body = (1..many.len()).fold(TensorBody::Operand(0), |b, g| *add(Box::new(b), op(g)));
11180            let s = t.stage(
11181                "weighted",
11182                StageKind::Compute(ComputeOp::Generic {
11183                    iters: flat_space(tw),
11184                    ins: vec![at(0); many.len()],
11185                    out: ident_map(one_w),
11186                    body,
11187                }),
11188                many.iter().map(|&(_, b)| b).collect(),
11189                vec![ys],
11190            );
11191            edges.extend(many.iter().map(|&(from, b)| (from, b, s)));
11192            (s, ys)
11193        }
11194    };
11195    // The final stage's operands: the weighted sum, then the bias mix and the weights' sum when each is there.
11196    let mut ins = vec![ys];
11197    let mut maps = vec![at(0)];
11198    let mut total = op(0);
11199    let mut into_out = vec![(s_y, ys)];
11200    // `Σ_r P[b_r]·B[b_r]`: each round's chosen down-bias tile, weighted by its probability.
11201    if let Some(bi) = operands.biases[2] {
11202        let (b_name, b_fill) = names.operand(bi);
11203        let table = t.tensor(b_name, TensorRole::Param, 1, spec.ne * d, acc, super::from_tensor_op::init_of(b_fill, spec.ne * d));
11204        let mut mix_ins = Vec::new();
11205        let mut mix_maps = Vec::new();
11206        let mut mix: Option<Box<TensorBody>> = None;
11207        let mut loads = Vec::new();
11208        for r in 0..k {
11209            let (l_b, bias_in) = expert_bias_tile(&mut t, &format!("down{r}"), mn, table, (LV_ITEM_X, r), d_t, spec, target)?;
11210            let term = match w_in {
11211                Some(w_in) => {
11212                    let term = mul(op(mix_ins.len()), op(mix_ins.len() + 1));
11213                    mix_ins.extend([w_in, bias_in]);
11214                    mix_maps.extend([fixed((r * tw) as i64), at(0)]);
11215                    term
11216                }
11217                None => {
11218                    mix_ins.push(bias_in);
11219                    mix_maps.push(at(0));
11220                    op(mix_ins.len() - 1)
11221                }
11222            };
11223            mix = Some(match mix {
11224                None => term,
11225                Some(m) => add(m, term),
11226            });
11227            loads.push((l_b, bias_in));
11228        }
11229        let bm = t.scratch_deep("bias_mix", 1, 1, acc);
11230        let s_b = t.stage(
11231            "bias_mix",
11232            StageKind::Compute(ComputeOp::Generic {
11233                iters: flat_space(tw),
11234                ins: mix_maps,
11235                out: ident_map(one_w),
11236                body: *mix.expect("k ≥ 1 rounds"),
11237            }),
11238            mix_ins,
11239            vec![bm],
11240        );
11241        edges.extend(loads.into_iter().map(|(l, b)| (l, b, s_b)));
11242        total = add(total, op(ins.len()));
11243        ins.push(bm);
11244        maps.push(at(0));
11245        into_out.push((s_b, bm));
11246    }
11247    if spec.renormalize {
11248        let s_w = t.stage(
11249            "weight_sum",
11250            StageKind::Compute(ComputeOp::Generic {
11251                iters: head_space(1, k, Some(ReducePool::Sum)),
11252                ins: vec![
11253                    ElemAccess {
11254                        base: 0,
11255                        terms: vec![(LV_GEN_COL, tw as i64)],
11256                        indirect: None,
11257                    },
11258                    const_map(),
11259                ],
11260                out: per_row_map(one_w),
11261                body: TensorBody::Operand(0),
11262            }),
11263            vec![w_in.expect("a renormalized mixture weighs its outputs"), sc_in],
11264            vec![ws],
11265        );
11266        total = mul(total, unary(SfpuKind::Recip, renormalizing_sum(spec, op(ins.len()))));
11267        if spec.scale.is_some() {
11268            total = mul(total, konst(spec.scale()));
11269        }
11270        ins.push(ws);
11271        maps.push(fixed(0));
11272        into_out.push((s_w, ws));
11273    }
11274    let s_o = t.stage(
11275        if spec.renormalize { "renormalize" } else { "combine" },
11276        StageKind::Compute(ComputeOp::Generic {
11277            iters: flat_space(tw),
11278            ins: maps,
11279            out: ident_map(one_w),
11280            body: *total,
11281        }),
11282        ins,
11283        vec![out],
11284    );
11285    let st = t.store_block("store_out", y, 0, vec![(LV_ITEM_X, 1)], 1, 1, 1, out);
11286    edges.extend(into_out.into_iter().map(|(from, buf)| (from, buf, s_o)));
11287    edges.push((s_o, out, st));
11288    for (from, buf, to) in edges {
11289        t.edge(from, buf, to);
11290    }
11291    let items: Vec<WorkItem> = (0..d_t)
11292        .map(|ot| WorkItem {
11293            id: WorkItemId::new(ot),
11294            coord: (0, ot),
11295            out: TileRect {
11296                row0_t: 0,
11297                col0_t: ot,
11298                h_t: 1,
11299                w_t: 1,
11300            },
11301        })
11302        .collect();
11303    t.finish(
11304        format!("moe_combine_{k}x{d}"),
11305        items,
11306        AccPolicy::F32,
11307        approx("mixture-of-experts combine", precision),
11308        target,
11309    )
11310}
11311
11312/// The same nesting as [`banded`], predicated on a chosen axis — the reductions put the head on the ROW
11313/// axis (so the folded axis can be the column), and a band is a band whichever axis carries it.
11314fn banded_on(axis: LoopVar, bands: Vec<(usize, TensorBody)>) -> TensorBody {
11315    let mut it = bands.into_iter().rev();
11316    let (_, mut acc) = it.next().expect("a banded body has at least one band");
11317    for (hi, body) in it {
11318        acc = TensorBody::Select {
11319            cond: Pred::AxisLt(axis, PredRhs::Const(hi)),
11320            then: Box::new(body),
11321            otherwise: Box::new(acc),
11322        };
11323    }
11324    acc
11325}
11326
11327/// Every compute stage of a composition is one of the CLOSED algebra's two shapes — the property
11328/// `each_composition_is_a_chain_of_the_closed_algebra` asserts, exposed so a lock and a ledger read the
11329/// same predicate rather than two copies of it.
11330pub fn is_closed_algebra(p: &SpatialProgram) -> bool {
11331    p.stages.iter().all(|s| match &s.kind {
11332        StageKind::Compute(op) => matches!(
11333            op,
11334            ComputeOp::Generic { .. } | ComputeOp::MatmulTiles { .. }
11335        ),
11336        _ => true,
11337    })
11338}
11339
11340/// The compute shapes a program uses, named — for a ledger row or a refusal.
11341pub fn compute_shapes(p: &SpatialProgram) -> BTreeSet<&'static str> {
11342    p.stages
11343        .iter()
11344        .filter_map(|s| match &s.kind {
11345            StageKind::Compute(ComputeOp::Generic { .. }) => Some("Generic"),
11346            StageKind::Compute(ComputeOp::MatmulTiles { .. }) => Some("MatmulTiles"),
11347            StageKind::Compute(ComputeOp::Requant { .. }) => Some("Requant"),
11348            StageKind::Compute(ComputeOp::Bias { .. }) => Some("Bias"),
11349            StageKind::Compute(ComputeOp::Relu) => Some("Relu"),
11350            StageKind::Compute(ComputeOp::Scalar(_)) | StageKind::Scalar(_) => Some("Scalar"),
11351            _ => None,
11352        })
11353        .collect()
11354}
11355
11356// ── The blocked row reduction ─────────────────────────────────────────────────────────────────────
11357//
11358// ★ A ROW REDUCTION OVER A ROW THAT DOES NOT FIT IS NOT A HARDER REDUCTION — IT IS A STREAMED ONE.
11359//
11360// The whole-row rmsnorm keeps `cols_t` tiles resident and folds them in one pass. At Qwen3.8's real
11361// `d = 5120` that is 160 tiles against an eight-tile accumulator, so the form is refused outright
11362// (`AccTilesExceeded`) — and raising the budget is not available to us: it is a property of the silicon.
11363//
11364// The way out is not a new primitive. `sum(x²)` IS a matmul — `x² · 1` — and a matmul over a long
11365// contraction is precisely the thing this compiler already streams, block by block, with the
11366// accumulation living in the register file across blocks. So the reduction is expressed in the operation
11367// that already solved this problem, and the row op becomes four SECTIONS, each of which independently
11368// fits a site:
11369//
11370//   1. `xsq  = (x + residual)²`     pointwise, one item per tile
11371//   2. `ssq  = xsq · ones`          matmul, 1xN by Nx1 — streams over K blocks
11372//   3. `inv  = rsqrt(ssq / n)`      pointwise on a single tile
11373//   4. `y    = x · w · inv`         pointwise, one item per tile, `inv` broadcast
11374//
11375// ★ THE NUMERICS ARE THE ONES THE FUSED FORM ALREADY EARNED. The scale is applied ONCE, at step 4, and
11376// the division by `n` rides inside `Rsqrt { over }` rather than being smeared across the reduction's
11377// terms — the same arrangement `from_fused_row` documents, and for the same measured reason. The
11378// accumulation path stays at the accumulator's width; only the output rounds.
11379//
11380// ★ A STACK OF ROWS IS THE SAME FOUR SECTIONS OVER A MATRIX. `B` sequences normalized at once are `B × d`:
11381// the squares and the rescale map tile by tile over the 2-D grid, the sum is the `B × d` by `d × 1` matmul —
11382// every row's contraction in the one weight-free pass — and the scale is one per row, broadcast along it.
11383pub fn blocked_rms_norm(
11384    spec: &FusedRowSpec,
11385    names: &StageNames,
11386    target: &SpatialTarget,
11387) -> Result<Vec<SpatialProgram>, SpatialReject> {
11388    use crate::tensor_recognize::{MatmulSpec, RowReduceKind};
11389    if spec.reduce != RowReduceKind::RmsNorm {
11390        return Err(SpatialReject::NotLowered(
11391            "a blocked softmax or layernorm row op — only rmsnorm streams today",
11392        ));
11393    }
11394    if spec.bias.is_some() {
11395        return Err(SpatialReject::NotLowered(
11396            "a blocked rmsnorm with a bias term",
11397        ));
11398    }
11399    let (rows, n) = (spec.rows, spec.cols);
11400    if rows == 0 {
11401        return Err(SpatialReject::NotLowered("a blocked rmsnorm over no rows"));
11402    }
11403    let precision = elem_precision(target)?;
11404    let (xn, xf) = names.operand(0);
11405    // The spec's operand order is `[x, weight?, residual?, bias?]`.
11406    let mut rest = 1usize;
11407    let weight = spec.weight.map(|_| {
11408        let i = rest;
11409        rest += 1;
11410        names.operand(i)
11411    });
11412    let residual = spec.residual.map(|_| {
11413        let i = rest;
11414        rest += 1;
11415        names.operand(i)
11416    });
11417    let y = &names.out;
11418    let sq_name = format!("{y}__xsq");
11419    let ssq_name = format!("{y}__ssq");
11420    let inv_name = format!("{y}__inv");
11421    let ones_name = format!("{y}__ones");
11422    // The value that is normalized: `x`, or `x + residual` when the add is fused in. Both the reduction
11423    // and the rescale read it, so when a residual is present it is materialized once under its own name.
11424    // ★ AND WITH ITS REAL FILL. `x` is read by two sections (the squares and the rescale) and stored by
11425    // neither, so if both declare it as `External` the composition unifies a tensor nobody ever gave data
11426    // to. When the residual add is fused in, the base IS produced by a section and carries no fill of its
11427    // own; when it is not, the base is the source operand and must keep the fill the source gave it.
11428    let (base_name, base_fill): (String, Fill) = match residual {
11429        Some(_) => (format!("{y}__xr"), Fill::Unknown),
11430        None => (xn.to_string(), xf.clone()),
11431    };
11432
11433    let mut out = Vec::new();
11434    if let Some((rn, rf)) = residual {
11435        out.push(match rows {
11436            1 => blocked_binary_map(
11437                "xr",
11438                "add",
11439                |a, b| add(a, b),
11440                (xn, xf),
11441                (rn, rf),
11442                &base_name,
11443                n,
11444                target,
11445            )?,
11446            _ => matrix_binary_map(
11447                "xr",
11448                "add",
11449                |a, b| add(a, b),
11450                (xn, xf),
11451                (rn, rf),
11452                &base_name,
11453                (rows, n),
11454                target,
11455            )?,
11456        });
11457    }
11458    // 1. the squares
11459    let square = |x: Box<TensorBody>| mul(x.clone(), x);
11460    out.push(match rows {
11461        1 => blocked_unary_map(
11462            "xsq",
11463            "square",
11464            Datapath::Float,
11465            square,
11466            (&base_name, &base_fill),
11467            &sq_name,
11468            n,
11469            target,
11470        )?,
11471        _ => matrix_unary_map(
11472            "xsq",
11473            "square",
11474            Datapath::Float,
11475            square,
11476            (&base_name, &base_fill),
11477            &sq_name,
11478            (rows, n),
11479            target,
11480        )?,
11481    });
11482    // 2. the sum, as the contraction it is
11483    let ones = MatmulNames {
11484        a: sq_name.clone(),
11485        b: ones_name,
11486        c: ssq_name.clone(),
11487    };
11488    let mm = MatmulSpec {
11489        m: rows,
11490        n: 1,
11491        k: n,
11492        a: Default::default(),
11493        b: Default::default(),
11494        c: Default::default(),
11495    };
11496    out.push(super::from_tensor_op::from_matmul(
11497        &mm,
11498        &ones,
11499        (&Fill::Unknown, &Fill::ConstF(1.0f64.to_bits())),
11500        (None, None),
11501        NumericsPolicy::Exact,
11502        target,
11503    )?);
11504    // 3. the scale, computed once per row
11505    let eps = eps_nano(spec.eps)?;
11506    let rsqrt = |s| {
11507        Box::new(TensorBody::Unary {
11508            kind: SfpuKind::Rsqrt {
11509                eps_nano: eps,
11510                over: n,
11511            },
11512            arg: s,
11513        })
11514    };
11515    out.push(match rows {
11516        1 => blocked_unary_map(
11517            "rms_scale",
11518            "rsqrt",
11519            Datapath::Float,
11520            rsqrt,
11521            (&ssq_name, &Fill::Unknown),
11522            &inv_name,
11523            1,
11524            target,
11525        )?,
11526        _ => matrix_unary_map(
11527            "rms_scale",
11528            "rsqrt",
11529            Datapath::Float,
11530            rsqrt,
11531            (&ssq_name, &Fill::Unknown),
11532            &inv_name,
11533            (rows, 1),
11534            target,
11535        )?,
11536    });
11537    // 4. the rescale, with each row's scalar broadcast across every tile of its row
11538    out.push(match rows {
11539        1 => rescale_row(
11540            &base_name, &base_fill, weight, &inv_name, y, n, precision, target,
11541        )?,
11542        _ => rescale_rows(
11543            &base_name,
11544            &base_fill,
11545            weight,
11546            &inv_name,
11547            y,
11548            (rows, n),
11549            precision,
11550            target,
11551        )?,
11552    });
11553    Ok(out)
11554}
11555
11556/// [`blocked_binary_map`] over a MATRIX: one item per tile of the 2-D grid, as [`matrix_unary_map`].
11557fn matrix_binary_map(
11558    name: &str,
11559    kind: &str,
11560    body: impl Fn(Box<TensorBody>, Box<TensorBody>) -> Box<TensorBody>,
11561    a: (&str, &Fill),
11562    b: (&str, &Fill),
11563    out: &str,
11564    shape: (usize, usize),
11565    target: &SpatialTarget,
11566) -> Result<SpatialProgram, SpatialReject> {
11567    let (rows, cols) = shape;
11568    let precision = elem_precision(target)?;
11569    let mut t = Template::new(target)?;
11570    let at_ = t.tensor(
11571        a.0,
11572        TensorRole::Input,
11573        rows,
11574        cols,
11575        precision,
11576        super::from_tensor_op::init_of(a.1, rows * cols),
11577    );
11578    let bt = t.tensor(
11579        b.0,
11580        TensorRole::Input,
11581        rows,
11582        cols,
11583        precision,
11584        super::from_tensor_op::init_of(b.1, rows * cols),
11585    );
11586    let yt = t.tensor(
11587        out,
11588        TensorRole::Output,
11589        rows,
11590        cols,
11591        precision,
11592        TensorInit::Zero,
11593    );
11594    let (rows_t, cols_t) = (
11595        rows.div_ceil(target.tile.rows),
11596        cols.div_ceil(target.tile.cols),
11597    );
11598    let steps = super::from_tensor_op::tile_grid_steps(rows_t, cols_t);
11599    let ain = t.buffer("in_a", 1, precision);
11600    let bin = t.buffer("in_b", 1, precision);
11601    let ob = t.buffer("out", 1, precision);
11602    let la = t.load("load_a", at_, steps.clone(), 1, ain);
11603    let lb = t.load("load_b", bt, steps.clone(), 1, bin);
11604    let w = target.tile.cols;
11605    let s = t.stage(
11606        kind,
11607        StageKind::Compute(ComputeOp::Generic {
11608            iters: super::from_tensor_op::space(target, w, None),
11609            ins: vec![ident_map(w), ident_map(w)],
11610            out: ident_map(w),
11611            body: *body(op(0), op(1)),
11612        }),
11613        vec![ain, bin],
11614        vec![ob],
11615    );
11616    let st = t.store(&format!("store_{out}"), yt, steps, 1, ob);
11617    t.edge(la, ain, s);
11618    t.edge(lb, bin, s);
11619    t.edge(s, ob, st);
11620    t.finish(
11621        format!("{name}_{rows}x{cols}"),
11622        super::from_tensor_op::items_per_tile_grid(rows_t, cols_t),
11623        AccPolicy::F32,
11624        approx(kind, precision),
11625        target,
11626    )
11627}
11628
11629/// [`rescale_row`] over a `rows × n` MATRIX: `y[r][c] = x[r][c] · w[c] · inv[r]`, one item per tile of the 2-D
11630/// grid. The weight is the item's tile of the one weight row, broadcast down the tile's rows; the scale is the
11631/// item's tile of the `rows × 1` column, broadcast along them.
11632fn rescale_rows(
11633    x: &str,
11634    x_fill: &Fill,
11635    weight: Option<(&str, &Fill)>,
11636    inv: &str,
11637    y: &str,
11638    shape: (usize, usize),
11639    precision: Precision,
11640    target: &SpatialTarget,
11641) -> Result<SpatialProgram, SpatialReject> {
11642    let (rows, n) = shape;
11643    let (rows_t, cols_t) = (
11644        rows.div_ceil(target.tile.rows),
11645        n.div_ceil(target.tile.cols),
11646    );
11647    let steps = super::from_tensor_op::tile_grid_steps(rows_t, cols_t);
11648    let mut t = Template::new(target)?;
11649    let xt = t.tensor(
11650        x,
11651        TensorRole::Input,
11652        rows,
11653        n,
11654        precision,
11655        super::from_tensor_op::init_of(x_fill, rows * n),
11656    );
11657    let it = t.tensor(
11658        inv,
11659        TensorRole::Input,
11660        rows,
11661        1,
11662        precision,
11663        TensorInit::External,
11664    );
11665    let yt = t.tensor(y, TensorRole::Output, rows, n, precision, TensorInit::Zero);
11666    let xin = t.buffer("in_x", 1, precision);
11667    let iin = t.buffer("in_inv", 1, precision);
11668    let ob = t.buffer("out", 1, precision);
11669    let lx = t.load("load_x", xt, steps.clone(), 1, xin);
11670    let row_steps: Vec<(LoopVar, usize)> = if rows_t > 1 {
11671        vec![(LV_ITEM_Y, 1)]
11672    } else {
11673        Vec::new()
11674    };
11675    let li = t.load("load_inv", it, row_steps, 1, iin);
11676    let w = target.tile.cols;
11677    let (wbuf, lw) = match weight {
11678        Some((wn, wf)) => {
11679            let wt = t.tensor(
11680                wn,
11681                TensorRole::Param,
11682                1,
11683                n,
11684                precision,
11685                super::from_tensor_op::init_of(wf, n),
11686            );
11687            let b = t.buffer("in_w", 1, precision);
11688            let l = t.load("load_w", wt, vec![(LV_ITEM_X, 1)], 1, b);
11689            (Some(b), Some(l))
11690        }
11691        None => (None, None),
11692    };
11693    let mut ins = vec![xin];
11694    let mut maps = vec![ident_map(w)];
11695    let mut body = op(0);
11696    if let Some(b) = wbuf {
11697        ins.push(b);
11698        maps.push(super::from_tensor_op::per_col_map());
11699        body = mul(body, op(ins.len() - 1));
11700    }
11701    ins.push(iin);
11702    maps.push(super::from_tensor_op::per_row_map(w));
11703    body = mul(body, op(ins.len() - 1));
11704    let s = t.stage(
11705        "rescale",
11706        StageKind::Compute(ComputeOp::Generic {
11707            iters: super::from_tensor_op::space(target, w, None),
11708            ins: maps,
11709            out: ident_map(w),
11710            body: *body,
11711        }),
11712        ins,
11713        vec![ob],
11714    );
11715    let st = t.store(&format!("store_{y}"), yt, steps, 1, ob);
11716    t.edge(lx, xin, s);
11717    if let (Some(l), Some(b)) = (lw, wbuf) {
11718        t.edge(l, b, s);
11719    }
11720    t.edge(li, iin, s);
11721    t.edge(s, ob, st);
11722    t.finish(
11723        format!("rmsnorm_rescale_{rows}x{n}"),
11724        super::from_tensor_op::items_per_tile_grid(rows_t, cols_t),
11725        AccPolicy::F32,
11726        approx("rmsnorm", precision),
11727        target,
11728    )
11729}
11730
11731/// A blocked pointwise map of TWO equal-length rows — the fused residual add.
11732fn blocked_binary_map(
11733    name: &str,
11734    kind: &str,
11735    body: impl Fn(Box<TensorBody>, Box<TensorBody>) -> Box<TensorBody>,
11736    a: (&str, &Fill),
11737    b: (&str, &Fill),
11738    out: &str,
11739    n: usize,
11740    target: &SpatialTarget,
11741) -> Result<SpatialProgram, SpatialReject> {
11742    let precision = elem_precision(target)?;
11743    let cols_t = tiles_of(n, target);
11744    let mut t = Template::new(target)?;
11745    let at_ = t.tensor(
11746        a.0,
11747        TensorRole::Input,
11748        1,
11749        n,
11750        precision,
11751        super::from_tensor_op::init_of(a.1, n),
11752    );
11753    let bt = t.tensor(
11754        b.0,
11755        TensorRole::Input,
11756        1,
11757        n,
11758        precision,
11759        super::from_tensor_op::init_of(b.1, n),
11760    );
11761    let yt = t.tensor(out, TensorRole::Output, 1, n, precision, TensorInit::Zero);
11762    let ain = t.buffer("in_a", 1, precision);
11763    let bin = t.buffer("in_b", 1, precision);
11764    let ob = t.buffer("out", 1, precision);
11765    let la = t.load("load_a", at_, vec![(LV_ITEM_X, 1)], 1, ain);
11766    let lb = t.load("load_b", bt, vec![(LV_ITEM_X, 1)], 1, bin);
11767    let w = wide(1, target);
11768    let s = t.stage(
11769        kind,
11770        StageKind::Compute(ComputeOp::Generic {
11771            iters: flat_space(w),
11772            ins: vec![ident_map(w), ident_map(w)],
11773            out: ident_map(w),
11774            body: *body(op(0), op(1)),
11775        }),
11776        vec![ain, bin],
11777        vec![ob],
11778    );
11779    let st = t.store(&format!("store_{out}"), yt, vec![(LV_ITEM_X, 1)], 1, ob);
11780    t.edge(la, ain, s);
11781    t.edge(lb, bin, s);
11782    t.edge(s, ob, st);
11783    t.finish(
11784        format!("{name}_{n}"),
11785        items_per_tile(cols_t),
11786        AccPolicy::F32,
11787        approx(kind, precision),
11788        target,
11789    )
11790}
11791
11792/// `y[i] = x[i] · w[i] · inv` — one item per tile, with `inv` a single tile read through a fixed map so
11793/// the same scalar reaches every block of the row.
11794fn rescale_row(
11795    x: &str,
11796    x_fill: &Fill,
11797    weight: Option<(&str, &Fill)>,
11798    inv: &str,
11799    y: &str,
11800    n: usize,
11801    precision: Precision,
11802    target: &SpatialTarget,
11803) -> Result<SpatialProgram, SpatialReject> {
11804    let cols_t = tiles_of(n, target);
11805    let mut t = Template::new(target)?;
11806    let xt = t.tensor(
11807        x,
11808        TensorRole::Input,
11809        1,
11810        n,
11811        precision,
11812        super::from_tensor_op::init_of(x_fill, n),
11813    );
11814    let it = t.tensor(
11815        inv,
11816        TensorRole::Input,
11817        1,
11818        1,
11819        precision,
11820        TensorInit::External,
11821    );
11822    let yt = t.tensor(y, TensorRole::Output, 1, n, precision, TensorInit::Zero);
11823    let xin = t.buffer("in_x", 1, precision);
11824    let iin = t.buffer("in_inv", 1, precision);
11825    let ob = t.buffer("out", 1, precision);
11826    let lx = t.load("load_x", xt, vec![(LV_ITEM_X, 1)], 1, xin);
11827    let li = t.load("load_inv", it, vec![], 1, iin);
11828    let w = wide(1, target);
11829    let (wbuf, lw) = match weight {
11830        Some((wn, wf)) => {
11831            let wt = t.tensor(
11832                wn,
11833                TensorRole::Param,
11834                1,
11835                n,
11836                precision,
11837                super::from_tensor_op::init_of(wf, n),
11838            );
11839            let b = t.buffer("in_w", 1, precision);
11840            let l = t.load("load_w", wt, vec![(LV_ITEM_X, 1)], 1, b);
11841            (Some(b), Some(l))
11842        }
11843        None => (None, None),
11844    };
11845    let mut ins = vec![xin];
11846    let mut maps = vec![ident_map(w)];
11847    let mut body = op(0);
11848    if let Some(b) = wbuf {
11849        ins.push(b);
11850        maps.push(ident_map(w));
11851        body = mul(body, op(ins.len() - 1));
11852    }
11853    ins.push(iin);
11854    maps.push(fixed(0));
11855    body = mul(body, op(ins.len() - 1));
11856    let s = t.stage(
11857        "rescale",
11858        StageKind::Compute(ComputeOp::Generic {
11859            iters: flat_space(w),
11860            ins: maps,
11861            out: ident_map(w),
11862            body: *body,
11863        }),
11864        ins.clone(),
11865        vec![ob],
11866    );
11867    let st = t.store(&format!("store_{y}"), yt, vec![(LV_ITEM_X, 1)], 1, ob);
11868    t.edge(lx, xin, s);
11869    if let (Some(l), Some(b)) = (lw, wbuf) {
11870        t.edge(l, b, s);
11871    }
11872    t.edge(li, iin, s);
11873    t.edge(s, ob, st);
11874    t.finish(
11875        format!("rmsnorm_rescale_{n}"),
11876        items_per_tile(cols_t),
11877        AccPolicy::F32,
11878        approx("rmsnorm", precision),
11879        target,
11880    )
11881}
11882
11883/// ★ L2NORM REDUCES PER HEAD, SO IT BLOCKS PER HEAD. The flat form keeps all `rows·cols` elements
11884/// resident and bands the rescale with constant-offset maps — correct, and impossible once the flat row
11885/// is 64 tiles. But the REDUCTION was never over the flat row: each head folds only its own `cols`
11886/// elements, which at Qwen3.8's `gdn_head_dim = 128` is four tiles and fits a site with room to spare.
11887///
11888/// So the item becomes the HEAD. Each work item loads one band, folds it, takes the reciprocal root and
11889/// rescales it — the same chain the flat form runs, with the band as the page instead of the whole row.
11890/// No banded map is needed at all, because a band is no longer a region inside a bigger page.
11891fn blocked_l2norm(
11892    spec: &L2NormSpec,
11893    names: &StageNames,
11894    target: &SpatialTarget,
11895) -> Result<SpatialProgram, SpatialReject> {
11896    let (rows, cols) = (spec.rows, spec.cols);
11897    let n = rows * cols;
11898    let precision = elem_precision(target)?;
11899    let band_t = tiles_of(cols, target);
11900    let mut t = Template::new(target)?;
11901    let (xn, xf) = names.operand(0);
11902    let xt = t.tensor(
11903        xn,
11904        TensorRole::Input,
11905        1,
11906        n,
11907        precision,
11908        super::from_tensor_op::init_of(xf, n),
11909    );
11910    let yt = t.tensor(
11911        &names.out,
11912        TensorRole::Output,
11913        1,
11914        n,
11915        precision,
11916        TensorInit::Zero,
11917    );
11918    let acc = AccPolicy::F32.interm_precision();
11919    let xin = t.buffer("in_x", band_t, precision);
11920    let ones_t = t.tensor(
11921        &format!("{}__ones", names.out),
11922        TensorRole::Param,
11923        target.tile.rows,
11924        target.tile.cols,
11925        precision,
11926        TensorInit::FloatValues(vec![1.0f64.to_bits(); target.tile.rows * target.tile.cols]),
11927    );
11928    let ones_in = t.buffer("in_ones", 1, precision);
11929    let sq = t.scratch("band_sq", band_t, acc);
11930    let ssq = t.scratch("band_ssq", 1, acc);
11931    let inv = t.scratch("band_inv", 1, acc);
11932    let ob = t.buffer("out", band_t, precision);
11933    // The band this item owns: `band_t` tiles starting at item `h`.
11934    let step = vec![(LV_ITEM_X, band_t)];
11935    let lx = t.load("load_band", xt, step.clone(), band_t, xin);
11936    let lo = t.load("load_ones", ones_t, vec![], 1, ones_in);
11937    let w = wide(band_t, target);
11938    let one = wide(1, target);
11939    // ★ THE BAND IS ONE HEAD, SO THE 2-D VIEW HAS EXACTLY ONE ROW. `head_space(rows, cols)` is how the
11940    // flat form views `rows·cols` as a plane whose row is the head; an item that owns a single head sees
11941    // that plane with `rows = 1`. Reducing with the whole-row `space()` instead would fold 32 tile rows
11942    // where the band has one, summing 31 rows of padding into the norm.
11943    let s_sq = t.stage(
11944        "square",
11945        StageKind::Compute(ComputeOp::Generic {
11946            iters: head_space(1, cols, None),
11947            ins: vec![map2(0, cols as i64, 1), map2(0, cols as i64, 1)],
11948            out: ident_map(w),
11949            body: *mul(op(0), op(1)),
11950        }),
11951        vec![xin, xin],
11952        vec![sq],
11953    );
11954    let s_sum = t.stage(
11955        "band_sum",
11956        StageKind::Compute(ComputeOp::Generic {
11957            iters: head_space(1, cols, Some(ReducePool::Sum)),
11958            ins: vec![map2(0, w as i64, 1), const_map()],
11959            out: per_row_map(one),
11960            body: TensorBody::Operand(0),
11961        }),
11962        vec![sq, ones_in],
11963        vec![ssq],
11964    );
11965    let s_inv = t.stage(
11966        "band_rsqrt",
11967        StageKind::Compute(ComputeOp::Generic {
11968            iters: head_space(1, target.tile.cols, None),
11969            ins: vec![per_row_map(one)],
11970            out: per_row_map(one),
11971            body: TensorBody::Unary {
11972                kind: SfpuKind::Rsqrt {
11973                    eps_nano: eps_nano(spec.eps)?,
11974                    over: 1,
11975                },
11976                arg: op(0),
11977            },
11978        }),
11979        vec![ssq],
11980        vec![inv],
11981    );
11982    let s_out = t.stage(
11983        "rescale",
11984        StageKind::Compute(ComputeOp::Generic {
11985            iters: head_space(1, cols, None),
11986            ins: vec![map2(0, cols as i64, 1), fixed(0)],
11987            out: ident_map(w),
11988            body: *mul(op(0), op(1)),
11989        }),
11990        vec![xin, inv],
11991        vec![ob],
11992    );
11993    let st = t.store(&format!("store_{}", names.out), yt, step, band_t, ob);
11994    t.edge(lx, xin, s_sq);
11995    t.edge(lo, ones_in, s_sum);
11996    t.edge(s_sq, sq, s_sum);
11997    t.edge(s_sum, ssq, s_inv);
11998    t.edge(s_inv, inv, s_out);
11999    t.edge(s_out, ob, st);
12000    let items: Vec<WorkItem> = (0..rows)
12001        .map(|h| WorkItem {
12002            id: WorkItemId::new(h),
12003            coord: (0, h),
12004            out: TileRect {
12005                row0_t: 0,
12006                col0_t: h * band_t,
12007                h_t: 1,
12008                w_t: band_t,
12009            },
12010        })
12011        .collect();
12012    t.finish(
12013        format!("l2norm_{rows}x{cols}"),
12014        items,
12015        AccPolicy::F32,
12016        approx("per-head l2 normalization", precision),
12017        target,
12018    )
12019}
12020
12021/// ★ A CAUSAL CONV BLOCKS BY CHANNEL, AND NEEDS NO HALO. `y[c] = Σ_t W[t,c]·CS[t,c]` — every channel
12022/// reads only its OWN column of the ring, so channels are independent and a block of them is a complete
12023/// sub-problem. (The "causal" axis is the ring's TAP axis, which is already resident in full; the halo a
12024/// sliding-window convolution would need does not arise.)
12025///
12026/// The ring is `taps × chan` stored flat, so one channel tile's taps sit `chan/tile` TILES apart — a
12027/// strided `block_h = taps` load. The loaded page is a `32 × taps·32` plane whose tile index is the tap,
12028/// so inside the page the tap's stride is one TILE WIDTH, not `chan`. That last sentence is the whole
12029/// difference between this and the flat form, and getting it wrong reads another channel's history.
12030fn blocked_causal_conv(
12031    spec: &CausalConvSpec,
12032    names: &StageNames,
12033    target: &SpatialTarget,
12034) -> Result<Vec<SpatialProgram>, SpatialReject> {
12035    let (taps, chan) = (spec.taps, spec.chan);
12036    let tw = target.tile.cols;
12037    if chan % tw != 0 {
12038        return Err(SpatialReject::NotLowered(
12039            "a blocked causal conv whose channel count is not a whole number of tiles",
12040        ));
12041    }
12042    let chan_t = chan / tw;
12043    let ring_len = taps * chan;
12044    let precision = elem_precision(target)?;
12045    let (x_name, x_fill) = names.operand(0);
12046    let (w_name, w_fill) = names.operand(1);
12047    let (cs_name, cs_fill) = names.operand(2);
12048
12049    // ── Section 1: the ring shift, one channel tile per item.
12050    let shift = {
12051        let mut t = Template::new(target)?;
12052        let cs = t.tensor(
12053            cs_name,
12054            TensorRole::Input,
12055            taps,
12056            chan,
12057            precision,
12058            super::from_tensor_op::init_of(cs_fill, ring_len),
12059        );
12060        let x = t.tensor(
12061            x_name,
12062            TensorRole::Input,
12063            1,
12064            chan,
12065            precision,
12066            super::from_tensor_op::init_of(x_fill, chan),
12067        );
12068        let outt = t.tensor(
12069            cs_name,
12070            TensorRole::Output,
12071            taps,
12072            chan,
12073            precision,
12074            TensorInit::Zero,
12075        );
12076        let cin = t.buffer("in_cs", 1, precision);
12077        let xin = t.buffer("in_x", 1, precision);
12078        let ob = t.buffer("ring", 1, precision);
12079        let lc = t.load("load_ring", cs, vec![(LV_ITEM_X, 1)], 1, cin);
12080        let lx = t.load("load_x", x, vec![(LV_ITEM_X, 1)], 1, xin);
12081        // Inside the page the tap is the tile, so one tap up is one tile width along the flat plane.
12082        let body = TensorBody::Select {
12083            cond: Pred::AxisLt(LV_GEN_ROW, PredRhs::Const(taps - 1)),
12084            then: Box::new(TensorBody::Operand(0)),
12085            otherwise: Box::new(TensorBody::Operand(1)),
12086        };
12087        let s = t.stage(
12088            "ring_shift",
12089            StageKind::Compute(ComputeOp::Generic {
12090                iters: head_space(taps, tw, None),
12091                ins: vec![map2(tw as i64, tw as i64, 1), map2(0, 0, 1)],
12092                out: map2(0, tw as i64, 1),
12093                body,
12094            }),
12095            vec![cin, xin],
12096            vec![ob],
12097        );
12098        let st = t.store("store_ring", outt, vec![(LV_ITEM_X, 1)], 1, ob);
12099        t.edge(lc, cin, s);
12100        t.edge(lx, xin, s);
12101        t.edge(s, ob, st);
12102        let items: Vec<WorkItem> = (0..chan_t)
12103            .map(|c| WorkItem {
12104                id: WorkItemId::new(c),
12105                coord: (0, c),
12106                out: TileRect {
12107                    row0_t: 0,
12108                    col0_t: c,
12109                    h_t: 1,
12110                    w_t: 1,
12111                },
12112            })
12113            .collect();
12114        t.finish(
12115            format!("conv_ring_{taps}x{chan}"),
12116            items,
12117            AccPolicy::F32,
12118            NumericsClaim::Exact,
12119            target,
12120        )?
12121    };
12122
12123    // ── Section 2: the tap reduction against the updated ring, same blocking.
12124    let reduce = {
12125        let mut t = Template::new(target)?;
12126        let cs = t.tensor(
12127            cs_name,
12128            TensorRole::Input,
12129            taps,
12130            chan,
12131            precision,
12132            super::from_tensor_op::init_of(cs_fill, ring_len),
12133        );
12134        let wc = t.tensor(
12135            w_name,
12136            TensorRole::Param,
12137            taps,
12138            chan,
12139            precision,
12140            super::from_tensor_op::init_of(w_fill, ring_len),
12141        );
12142        let yt = t.tensor(
12143            &names.out,
12144            TensorRole::Output,
12145            1,
12146            chan,
12147            precision,
12148            TensorInit::Zero,
12149        );
12150        let tile_elems = target.tile.rows * target.tile.cols;
12151        let ones = t.tensor(
12152            &format!("{}__conv_scaler", names.out),
12153            TensorRole::Param,
12154            target.tile.rows,
12155            target.tile.cols,
12156            precision,
12157            TensorInit::FloatValues(vec![1.0f64.to_bits(); tile_elems]),
12158        );
12159        let cin = t.buffer("in_cs", 1, precision);
12160        let win = t.buffer("in_w", 1, precision);
12161        let oin = t.buffer("in_scaler", 1, precision);
12162        let ob = t.buffer("out", 1, precision);
12163        let lc = t.load("load_ring", cs, vec![(LV_ITEM_X, 1)], 1, cin);
12164        let lw = t.load("load_w", wc, vec![(LV_ITEM_X, 1)], 1, win);
12165        let lo = t.load("load_scaler", ones, vec![], 1, oin);
12166        let iters = vec![
12167            IterSpec {
12168                var: LV_GEN_ROW,
12169                extent: taps,
12170                kind: IterKind::Reduce {
12171                    pool: ReducePool::Sum,
12172                },
12173            },
12174            IterSpec {
12175                var: LV_GEN_COL,
12176                extent: tw,
12177                kind: IterKind::Parallel,
12178            },
12179        ];
12180        let s = t.stage(
12181            "tap_reduce",
12182            StageKind::Compute(ComputeOp::Generic {
12183                iters,
12184                ins: vec![map2(0, tw as i64, 1), map2(0, tw as i64, 1), const_map()],
12185                out: per_col_map(),
12186                body: *mul(op(0), op(1)),
12187            }),
12188            vec![cin, win, oin],
12189            vec![ob],
12190        );
12191        let st = t.store("store_y", yt, vec![(LV_ITEM_X, 1)], 1, ob);
12192        t.edge(lc, cin, s);
12193        t.edge(lw, win, s);
12194        t.edge(lo, oin, s);
12195        t.edge(s, ob, st);
12196        let items: Vec<WorkItem> = (0..chan_t)
12197            .map(|c| WorkItem {
12198                id: WorkItemId::new(c),
12199                coord: (0, c),
12200                out: TileRect {
12201                    row0_t: 0,
12202                    col0_t: c,
12203                    h_t: 1,
12204                    w_t: 1,
12205                },
12206            })
12207            .collect();
12208        t.finish(
12209            format!("conv_taps_{taps}x{chan}"),
12210            items,
12211            AccPolicy::F32,
12212            approx("causal convolution", precision),
12213            target,
12214        )?
12215    };
12216    Ok(vec![shift, reduce])
12217}
12218
12219/// ★ GATED DELTANET AT REAL WIDTH — the head becomes the WORK ITEM.
12220///
12221/// The flat form keeps the whole `hv·d²` memory resident and distinguishes heads with a `banded` body:
12222/// at Qwen3.8's real geometry that is a 786,432-element buffer and 6,144 bands. Neither survives contact
12223/// with a site.
12224///
12225/// ★ WHY THE HEAD HAS TO RIDE THE ITEM'S COORDINATE. A `Generic`'s maps and body cannot see WHERE the
12226/// work item sits — `offset_at` substitutes the two iteration axes and nothing else, so every other loop
12227/// variable reads as zero. That is precisely why the flat forms use one enormous item: the body can only
12228/// tell heads apart if every head is inside the page. A LOAD, by contrast, addresses tiles through
12229/// `TileAddr.steps`, which DO read the item's coordinate — so the head must be selected when the data is
12230/// fetched, not when it is used.
12231///
12232/// ★ AND WHY THE COORDINATE IS `(kh, g)` RATHER THAN `(h, ·)`. The memory is indexed by the v-head `h`
12233/// while `k` and `q` are indexed by `h mod hk` — 48 v-heads sharing 16 k-heads — and `h mod hk` is not
12234/// affine in `h`, so no step can express it. Splitting the head as `h = g·hk + kh` and carrying BOTH
12235/// halves in the item's coordinate makes each index affine in its own axis: `k` steps by `kh`, the memory
12236/// steps by both. The non-affine operation disappears because it was only ever a decomposition.
12237///
12238/// The per-head scalars `α` and `β` are the one thing left that a load cannot reach — element `h` of a
12239/// 48-long vector lives at tile `h/32`, and that is not affine either — so they are broadcast once, ahead
12240/// of the step, into one tile per head.
12241///
12242/// `head` is the width the step's `1/√d` and its norm read: `spec.d`, or a head's own width where
12243/// [`widened_gated_delta`] widened it to `spec.d`. `state(i)` names the memory after `i` rows (or chunks): the carried
12244/// memory's one name for every `i`, updated in place, or a name per version where the memory is the op's own.
12245fn blocked_gated_delta(
12246    spec: &GatedDeltaSpec,
12247    names: &StageNames,
12248    dense: &BTreeSet<String>,
12249    head: usize,
12250    state: &dyn Fn(usize) -> String,
12251    target: &SpatialTarget,
12252) -> Result<Vec<SpatialProgram>, SpatialReject> {
12253    let (hv, hk, d, rows) = (spec.hv, spec.hk, spec.d, spec.rows);
12254    let (tr, tw) = (target.tile.rows, target.tile.cols);
12255    if d % tw != 0 || hv % hk != 0 {
12256        return Err(SpatialReject::NotLowered("a blocked gated delta needs a head dimension in whole tiles and v-heads a multiple of k-heads"));
12257    }
12258    let nb = d / tw; // tile-rows per head's memory, and tiles per head's vectors
12259    if takes_chunks(rows, tr) {
12260        return chunked_gated_delta(spec, names, dense, head, state, target);
12261    }
12262    // ★ THIS ONCE DID NOT FINISH, AND THE REASON WAS NOT HERE. Blocking gated delta by head produces a
12263    // layer of ~14,800 stages, and certifying it ran for sixteen minutes at 13 GB. Every candidate cause
12264    // was wrong until it was measured: the sections certify in MILLISECONDS, `place` clones 864 stages in
12265    // 0.6ms, and `compose` joins all 45 sections in 5.1s. The cost was one line in `validate_schedule`
12266    // asking `per_stage` for every store stage — `stores × ops`, with an allocation per call. Built as a
12267    // set in one pass it is linear, and the whole layer now lowers in ~77s.
12268    // ★ THE A-BLOCK IS A STAGE, NOT AN ITEM. A head's memory is an `nb × nb` tile square — sixteen tiles against an
12269    // eight-tile accumulator — so no one pass writes it. Each a-block (a tile row of the square) is its own chain of
12270    // stages inside the head's item — its prediction, its look `S·q`, its update, its store — each page `nb` tiles
12271    // wide, so the accumulator holds one a-block at a time while the item stays the head: `hv` items, not `hv·nb`.
12272    let precision = elem_precision(target)?;
12273    let (q_name, q_fill) = names.operand(0);
12274    let (k_name, k_fill) = names.operand(1);
12275    let (v_name, v_fill) = names.operand(2);
12276    let (z_name, z_fill) = names.operand(3);
12277    let (b_name, b_fill) = names.operand(4);
12278    let (a_name, a_fill) = names.operand(5);
12279    let st_fill = names.operand(6).1;
12280    let (gn_name, gn_fill) = names.operand(7);
12281    let out_len = hv * d;
12282    let s_rows = hv * d;
12283    // Tiles one row spans: of `Q` and `K`, and of `V`, `Z` and the output.
12284    let (qk_t, v_t) = (hk * nb, hv * nb);
12285    let a_exp = format!("{}__alpha_t", names.out);
12286    let b_exp = format!("{}__beta_t", names.out);
12287    // A block's rows each take the step in turn, so each has its own section.
12288    let tag = |r: usize| {
12289        if rows == 1 {
12290            String::new()
12291        } else {
12292            format!("_row{r}")
12293        }
12294    };
12295
12296    // One item per v-head, coordinate `(kh, g)` — what rebuilds `h` for the loads. The item's RECT is its strip of
12297    // the output, the step's result; the memory's `nb × nb` square it also writes is covered by the tiles its stores
12298    // address.
12299    let head_items = |rect: &dyn Fn(usize) -> TileRect| -> Vec<WorkItem> {
12300        (0..hv)
12301            .map(|h| WorkItem {
12302                id: WorkItemId::new(h),
12303                coord: (h % hk, h / hk),
12304                out: rect(h),
12305            })
12306            .collect()
12307    };
12308    // Tile steps that rebuild `h = g·hk + kh` from the coordinate.
12309    let head_steps = |per_head: usize| vec![(LV_ITEM_Y, per_head), (LV_ITEM_X, hk * per_head)];
12310
12311    // ★ A BROADCAST IS A MATMUL AGAINST A SELECTION MATRIX.
12312    //
12313    // `α` and `β` are per-head SCALARS, and a scalar is the one thing a blocked section cannot reach:
12314    // element `h` of an `hv`-long vector lives at tile `h/32` offset `h%32`, a LOAD selects only tiles,
12315    // and a MAP cannot see which item it is running for (`offset_at` substitutes the two iteration axes
12316    // and nothing else). Every attempt to broadcast them with a `Generic` relocates that problem into the
12317    // broadcast's own body.
12318    //
12319    // But `alpha_t[h·tw + j] = α[h]` is exactly `α · SEL` for the constant `hv × (hv·tw)` matrix with
12320    // `SEL[h, c] = 1 iff c/tw == h` — and a matmul is the one operation this compiler already streams
12321    // over a long output, blocking it across items and sites without being asked. So the scalar becomes
12322    // ONE TILE PER HEAD, which a tile address CAN reach, and no gather is needed at all. The selection
12323    // matrix is a compile-time fill; the work is `hv·tw` MACs per row, and a block's rows are ONE matmul —
12324    // row `r`'s tiles are tile row `r / tile`, its values plane row `r mod tile` of each.
12325    let sel_fill = || -> Fill {
12326        Fill::Overrides {
12327            base: Box::new(Fill::Const(0)),
12328            sets: (0..hv)
12329                .flat_map(|h| (0..tw).map(move |j| (h * hv * tw + h * tw + j, 1)))
12330                .collect(),
12331        }
12332    };
12333    let broadcast =
12334        |src: &str, fill: &Fill, dst: &str, sel: &str| -> Result<SpatialProgram, SpatialReject> {
12335            let spec = crate::tensor_recognize::MatmulSpec {
12336                m: rows,
12337                n: hv * tw,
12338                k: hv,
12339                a: Default::default(),
12340                b: Default::default(),
12341                c: Default::default(),
12342            };
12343            let names = MatmulNames {
12344                a: src.to_string(),
12345                b: sel.to_string(),
12346                c: dst.to_string(),
12347            };
12348            super::from_tensor_op::from_matmul(
12349                &spec,
12350                &names,
12351                (fill, &sel_fill()),
12352                (None, None),
12353                NumericsPolicy::Exact,
12354                target,
12355            )
12356        };
12357    let alpha_t = broadcast(a_name, a_fill, &a_exp, &format!("{}__sel_a", names.out))?;
12358    let beta_t = broadcast(b_name, b_fill, &b_exp, &format!("{}__sel_b", names.out))?;
12359
12360    // ── Row `r`'s step, ONE PASS OVER EACH HEAD'S MEMORY.
12361    //
12362    // ★ EVERY TERM OF THE STEP IS LOCAL TO AN A-BLOCK. Row `a` of the prediction `u = S·k`, of the update
12363    // `S'[a, ·] = α·S[a, ·] + β·(v[a] − α·u[a])·k` and of the output `(S'·q)[a]` read only row `a` of the memory, and the
12364    // output needs no second read of it: `S'·q = α·(S·q) + β·(v − α·u)·(k·q)`, `k·q` one number per head. So the item —
12365    // the head — loads each a-block of its memory ONCE, forms `u` and `w = S·q` from it, writes the updated a-block
12366    // back, and assembles the output from `w`, `u` and `k·q`: the memory is read once and written once a step, where a
12367    // prediction, an update and a read section each moved it. The output is the section's result; the memory's a-blocks
12368    // are a second tensor, each store covering the tiles it addresses.
12369    let o_dense = rows > 1 && dense.contains(&names.out);
12370    let inv_sqrt_d = 1.0f64 / (head as f64).sqrt();
12371    // `pages`: how many pages each once-an-item buffer holds (each a-block's store, the norm weight, the gated output).
12372    let step = |r: usize, pages: usize| -> Result<SpatialProgram, SpatialReject> {
12373        let mut t = Template::new(target)?;
12374        let acc = AccPolicy::F32.interm_precision();
12375        // The output first: the step's result, whose strips are the items' rectangles.
12376        let y = match (rows, o_dense) {
12377            (1, _) => t.tensor(
12378                &names.out,
12379                TensorRole::Output,
12380                1,
12381                out_len,
12382                precision,
12383                TensorInit::Zero,
12384            ),
12385            (_, false) => t.tensor(
12386                &names.out,
12387                TensorRole::State,
12388                1,
12389                rows * out_len,
12390                precision,
12391                TensorInit::Zero,
12392            ),
12393            (_, true) => t.tensor(
12394                &names.out,
12395                TensorRole::State,
12396                rows,
12397                out_len,
12398                precision,
12399                TensorInit::Zero,
12400            ),
12401        };
12402        let st = t.tensor(
12403            &state(r),
12404            TensorRole::Input,
12405            s_rows,
12406            d,
12407            precision,
12408            super::from_tensor_op::init_of(st_fill, s_rows * d),
12409        );
12410        let stout = t.tensor(
12411            &state(r + 1),
12412            TensorRole::State,
12413            s_rows,
12414            d,
12415            precision,
12416            TensorInit::Zero,
12417        );
12418        let q = t.tensor(
12419            q_name,
12420            TensorRole::Input,
12421            1,
12422            rows * hk * d,
12423            precision,
12424            super::from_tensor_op::init_of(q_fill, rows * hk * d),
12425        );
12426        let k = t.tensor(
12427            k_name,
12428            TensorRole::Input,
12429            1,
12430            rows * hk * d,
12431            precision,
12432            super::from_tensor_op::init_of(k_fill, rows * hk * d),
12433        );
12434        let v = t.tensor(
12435            v_name,
12436            TensorRole::Input,
12437            1,
12438            rows * out_len,
12439            precision,
12440            super::from_tensor_op::init_of(v_fill, rows * out_len),
12441        );
12442        let z = t.tensor(
12443            z_name,
12444            TensorRole::Input,
12445            1,
12446            rows * out_len,
12447            precision,
12448            super::from_tensor_op::init_of(z_fill, rows * out_len),
12449        );
12450        let gn = t.tensor(
12451            gn_name,
12452            TensorRole::Param,
12453            1,
12454            d,
12455            precision,
12456            super::from_tensor_op::init_of(gn_fill, d),
12457        );
12458        let alpha = t.tensor(
12459            &a_exp,
12460            TensorRole::Input,
12461            rows,
12462            hv * tw,
12463            precision,
12464            TensorInit::External,
12465        );
12466        let beta = t.tensor(
12467            &b_exp,
12468            TensorRole::Input,
12469            rows,
12470            hv * tw,
12471            precision,
12472            TensorInit::External,
12473        );
12474        let ones = t.tensor(
12475            &format!("{}__gd_scaler", names.out),
12476            TensorRole::Param,
12477            target.tile.rows,
12478            tw,
12479            precision,
12480            TensorInit::FloatValues(vec![1.0f64.to_bits(); target.tile.rows * tw]),
12481        );
12482        // The head's operands, each loaded once: `k` and `q` at its k-head, `v` at its v-head, `α` and `β` at plane row
12483        // `r mod tile` of the head's tile.
12484        let kin = t.buffer_deep("in_k", 1, nb, precision);
12485        let qin = t.buffer_deep("in_q", 1, nb, precision);
12486        let vin = t.buffer_deep("in_v", 1, nb, precision);
12487        let ain = t.buffer_deep("in_a", 1, 1, precision);
12488        let bin = t.buffer_deep("in_b", 1, 1, precision);
12489        let sc_in = t.buffer_deep("in_scaler", 1, 1, precision);
12490        let lk = t.load_block_at("load_k", k, r * qk_t, vec![(LV_ITEM_Y, nb)], 1, nb, nb, kin);
12491        let lq = t.load_block_at("load_q", q, r * qk_t, vec![(LV_ITEM_Y, nb)], 1, nb, nb, qin);
12492        let lv = t.load_block_at("load_v", v, r * v_t, head_steps(nb), 1, nb, nb, vin);
12493        let la = t.load_block_at("load_a", alpha, (r / tr) * hv, head_steps(1), 1, 1, 1, ain);
12494        let lb = t.load_block_at("load_b", beta, (r / tr) * hv, head_steps(1), 1, 1, 1, bin);
12495        let l_sc = t.load("load_scaler", ones, vec![], 1, sc_in);
12496        let gin = t.buffer_deep("in_gn", pages, nb, precision);
12497        let lg = t.load("load_gn", gn, vec![], nb, gin);
12498        let scalar = fixed(((r % tr) * tw) as i64);
12499        let column = ElemAccess {
12500            base: 0,
12501            terms: vec![(LV_GEN_COL, 1)],
12502            indirect: None,
12503        };
12504        let down = ElemAccess {
12505            base: 0,
12506            terms: vec![(LV_GEN_ROW, 1)],
12507            indirect: None,
12508        };
12509        // Fold `b` (the ROW), keep `a` (the COLUMN): `Σ_b S[a, b]·x[b]` lands at column `a` of the one-tile page.
12510        let matvec = |t: &mut Template, name: &str, s: BufferId, x: BufferId, out: BufferId| {
12511            let iters = vec![
12512                IterSpec {
12513                    var: LV_GEN_ROW,
12514                    extent: d,
12515                    kind: IterKind::Reduce {
12516                        pool: ReducePool::Sum,
12517                    },
12518                },
12519                IterSpec {
12520                    var: LV_GEN_COL,
12521                    extent: tw,
12522                    kind: IterKind::Parallel,
12523                },
12524            ];
12525            let ins = vec![
12526                ElemAccess {
12527                    base: 0,
12528                    terms: vec![(LV_GEN_COL, d as i64), (LV_GEN_ROW, 1)],
12529                    indirect: None,
12530                },
12531                ElemAccess {
12532                    base: 0,
12533                    terms: vec![(LV_GEN_ROW, 1)],
12534                    indirect: None,
12535                },
12536                const_map(),
12537            ];
12538            t.stage(
12539                name,
12540                StageKind::Compute(ComputeOp::Generic {
12541                    iters,
12542                    ins,
12543                    out: per_col_map(),
12544                    body: *mul(op(0), op(1)),
12545                }),
12546                vec![s, x, sc_in],
12547                vec![out],
12548            )
12549        };
12550        let mut parts = Vec::with_capacity(nb);
12551        for j in 0..nb {
12552            let sin = t.buffer_deep(&format!("in_s{j}"), 1, nb, precision);
12553            let ls = t.load_block_at(
12554                &format!("load_s{j}"),
12555                st,
12556                j * nb,
12557                head_steps(nb * nb),
12558                1,
12559                nb,
12560                nb,
12561                sin,
12562            );
12563            let u = t.scratch_deep(&format!("u{j}"), 1, 1, acc);
12564            let w = t.scratch_deep(&format!("w{j}"), 1, 1, acc);
12565            let s_u = matvec(&mut t, &format!("gd_predict{j}"), sin, kin, u);
12566            matvec(&mut t, &format!("gd_look{j}"), sin, qin, w);
12567            // `S'[a, b] = α·S[a, b] + β·(v[a] − α·u[a])·k[b]` — 0:S 1:k 2:α 3:β 4:v 5:u.
12568            let sout = t.buffer_deep(&format!("out_s{j}"), pages, nb, precision);
12569            let coef = mul(op(3), sub(op(4), mul(op(2), op(5))));
12570            let s_upd = t.stage(
12571                &format!("gd_update{j}"),
12572                StageKind::Compute(ComputeOp::Generic {
12573                    iters: vec![
12574                        IterSpec {
12575                            var: LV_GEN_ROW,
12576                            extent: tr,
12577                            kind: IterKind::Parallel,
12578                        },
12579                        IterSpec {
12580                            var: LV_GEN_COL,
12581                            extent: d,
12582                            kind: IterKind::Parallel,
12583                        },
12584                    ],
12585                    ins: vec![
12586                        map2(0, d as i64, 1),
12587                        column.clone(),
12588                        scalar.clone(),
12589                        scalar.clone(),
12590                        ElemAccess {
12591                            base: (j * tw) as i64,
12592                            terms: vec![(LV_GEN_ROW, 1)],
12593                            indirect: None,
12594                        },
12595                        down.clone(),
12596                    ],
12597                    out: map2(0, d as i64, 1),
12598                    body: *add(mul(op(2), op(0)), mul(coef, op(1))),
12599                }),
12600                vec![sin, kin, ain, bin, vin, u],
12601                vec![sout],
12602            );
12603            let store = t.store_block(
12604                &format!("store_s{j}"),
12605                stout,
12606                j * nb,
12607                head_steps(nb * nb),
12608                1,
12609                nb,
12610                nb,
12611                sout,
12612            );
12613            // Each a-block's write depends on its read: the update reads the prediction formed from the loaded block.
12614            t.edge(ls, sin, s_u);
12615            t.edge(s_u, u, s_upd);
12616            t.edge(s_upd, sout, store);
12617            parts.push((u, w, j));
12618        }
12619        // `k·q`, one number per head, across row 0 of its page.
12620        let one_w = wide(1, target);
12621        let kq = t.scratch("kq", 1, acc);
12622        let s_kq = t.stage(
12623            "gd_kq",
12624            StageKind::Compute(ComputeOp::Generic {
12625                iters: head_space(1, d, Some(ReducePool::Sum)),
12626                ins: vec![column.clone(), column.clone(), const_map()],
12627                out: per_row_map(one_w),
12628                body: *mul(op(0), op(1)),
12629            }),
12630            vec![kin, qin, sc_in],
12631            vec![kq],
12632        );
12633        // `o[c] = (α·w[a] + β·(v[c] − α·u[a])·(k·q)) / √d`, `a = c − j·tile` in a-block `j`'s band.
12634        let ovec = t.scratch("gd_o", nb, acc);
12635        let (mut o_ins, mut o_maps, mut bands) = (
12636            vec![ain, bin, vin, kq],
12637            vec![scalar.clone(), scalar.clone(), column.clone(), fixed(0)],
12638            Vec::new(),
12639        );
12640        for (u, w, j) in &parts {
12641            let shift = ElemAccess {
12642                base: -((j * tw) as i64),
12643                terms: vec![(LV_GEN_COL, 1)],
12644                indirect: None,
12645            };
12646            o_ins.extend([*u, *w]);
12647            o_maps.extend([shift.clone(), shift]);
12648            let (us, ws) = (o_ins.len() - 2, o_ins.len() - 1);
12649            let value = mul(
12650                add(
12651                    mul(op(0), op(ws)),
12652                    mul(mul(op(1), sub(op(2), mul(op(0), op(us)))), op(3)),
12653                ),
12654                konst(inv_sqrt_d),
12655            );
12656            bands.push(((j + 1) * tw, *value));
12657        }
12658        let s_o = t.stage(
12659            "gd_o",
12660            StageKind::Compute(ComputeOp::Generic {
12661                iters: flat_space(d),
12662                ins: o_maps,
12663                out: ident_map(wide(nb, target)),
12664                body: banded(bands),
12665            }),
12666            o_ins,
12667            vec![ovec],
12668        );
12669        for (l, b, to) in [
12670            (lk, kin, s_kq),
12671            (lq, qin, s_kq),
12672            (lv, vin, s_o),
12673            (la, ain, s_o),
12674            (lb, bin, s_o),
12675            (l_sc, sc_in, s_kq),
12676        ] {
12677            t.edge(l, b, to);
12678        }
12679        // The head's RMSNorm, then the weight and the gate.
12680        let w = wide(nb, target);
12681        let sq = t.scratch("o_sq", nb, acc);
12682        let s_sq = t.stage(
12683            "square",
12684            StageKind::Compute(ComputeOp::Generic {
12685                iters: flat_space(w),
12686                ins: vec![ident_map(w), ident_map(w)],
12687                out: ident_map(w),
12688                body: *mul(op(0), op(1)),
12689            }),
12690            vec![ovec, ovec],
12691            vec![sq],
12692        );
12693        let ssum = t.scratch("o_sum", 1, acc);
12694        let s_sum = t.stage(
12695            "o_sum",
12696            StageKind::Compute(ComputeOp::Generic {
12697                iters: head_space(1, d, Some(ReducePool::Sum)),
12698                ins: vec![map2(0, w as i64, 1), const_map()],
12699                out: per_row_map(one_w),
12700                body: TensorBody::Operand(0),
12701            }),
12702            vec![sq, sc_in],
12703            vec![ssum],
12704        );
12705        let inv = t.scratch("o_inv", 1, acc);
12706        let s_inv = t.stage(
12707            "o_rsqrt",
12708            StageKind::Compute(ComputeOp::Generic {
12709                iters: head_space(1, tw, None),
12710                ins: vec![per_row_map(one_w)],
12711                out: per_row_map(one_w),
12712                body: TensorBody::Unary {
12713                    kind: SfpuKind::Rsqrt {
12714                        eps_nano: eps_nano(spec.eps)?,
12715                        over: head,
12716                    },
12717                    arg: op(0),
12718                },
12719            }),
12720            vec![ssum],
12721            vec![inv],
12722        );
12723        let ob = t.buffer_deep("out", pages, nb, precision);
12724        let mut zins = Vec::new();
12725        for j in 0..nb {
12726            let zb = t.buffer_deep(&format!("in_z{j}"), 1, 1, precision);
12727            let lz = t.load_block_at(
12728                &format!("load_z{j}"),
12729                z,
12730                r * v_t + j,
12731                head_steps(nb),
12732                1,
12733                1,
12734                1,
12735                zb,
12736            );
12737            zins.push((lz, zb, j));
12738        }
12739        // `y = o · inv · GN · silu(z)` — the gate read through per-a-block tiles, banded back together.
12740        let mut ins = vec![ovec, inv, gin];
12741        let mut maps = vec![ident_map(w), fixed(0), ident_map(w)];
12742        let mut bands = Vec::new();
12743        for (_, zb, j) in &zins {
12744            ins.push(*zb);
12745            maps.push(ElemAccess {
12746                base: -((j * tw) as i64),
12747                terms: vec![(LV_GEN_COL, 1)],
12748                indirect: None,
12749            });
12750            let zslot = ins.len() - 1;
12751            let gated = mul(mul(mul(op(0), op(1)), op(2)), delta_gate(spec.gate, op(zslot)));
12752            bands.push(((j + 1) * tw, *gated));
12753        }
12754        // Where the head's strip of row `r` sits in the output, and — for a dense output — the tile row it shares.
12755        let (y_row, y_col) = if o_dense { (r / tr, 0) } else { (0, r * v_t) };
12756        let y_base = y_row * v_t + y_col;
12757        let shared_row = if o_dense {
12758            let yin = t.buffer("in_y", nb, precision);
12759            let ly = t.load_block_at("load_y", y, y_base, head_steps(nb), 1, nb, nb, yin);
12760            Some((ly, yin))
12761        } else {
12762            None
12763        };
12764        let s_out = match shared_row {
12765            None => t.stage(
12766                "gd_gate",
12767                StageKind::Compute(ComputeOp::Generic {
12768                    iters: flat_space(w),
12769                    ins: maps,
12770                    out: ident_map(w),
12771                    body: banded(bands),
12772                }),
12773                ins,
12774                vec![ob],
12775            ),
12776            // Every plane row of the page, row `r mod tile` the new values and the rest the tile row's other rows. The
12777            // new values are read down the page (their maps have no row term), so plane row `k` holds row `r`'s.
12778            Some((_, yin)) => {
12779                let k = r % tr;
12780                let old = ins.len();
12781                ins.push(yin);
12782                maps.push(ident_map(w));
12783                let maps = maps
12784                    .into_iter()
12785                    .enumerate()
12786                    .map(|(i, m)| if i == 0 || i == 2 { at(0) } else { m })
12787                    .collect();
12788                let body = TensorBody::Select {
12789                    cond: Pred::AxisLt(LV_GEN_ROW, PredRhs::Const(k)),
12790                    then: op(old),
12791                    otherwise: Box::new(TensorBody::Select {
12792                        cond: Pred::AxisLt(LV_GEN_ROW, PredRhs::Const(k + 1)),
12793                        then: Box::new(banded(bands)),
12794                        otherwise: op(old),
12795                    }),
12796                };
12797                t.stage(
12798                    "gd_gate",
12799                    StageKind::Compute(ComputeOp::Generic {
12800                        iters: super::from_tensor_op::space(target, w, None),
12801                        ins: maps,
12802                        out: ident_map(w),
12803                        body,
12804                    }),
12805                    ins,
12806                    vec![ob],
12807                )
12808            }
12809        };
12810        let store = t.store_block("store_y", y, y_base, head_steps(nb), 1, nb, nb, ob);
12811        t.edge(s_o, ovec, s_sq);
12812        t.edge(s_sq, sq, s_sum);
12813        t.edge(s_sum, ssum, s_inv);
12814        t.edge(s_inv, inv, s_out);
12815        t.edge(lg, gin, s_out);
12816        for (lz, zb, _) in &zins {
12817            t.edge(*lz, *zb, s_out);
12818        }
12819        if let Some((ly, yin)) = shared_row {
12820            t.edge(ly, yin, s_out);
12821        }
12822        t.edge(s_out, ob, store);
12823        let items = head_items(&|h| TileRect {
12824            row0_t: y_row,
12825            col0_t: y_col + h * nb,
12826            h_t: 1,
12827            w_t: nb,
12828        });
12829        t.finish(
12830            format!("gd_step_{hv}x{d}{}", tag(r)),
12831            items,
12832            AccPolicy::F32,
12833            approx("gated-delta step", precision),
12834            target,
12835        )
12836    };
12837
12838    let mut sections = vec![alpha_t, beta_t];
12839    for r in 0..rows {
12840        sections.push(paged(|pages| step(r, pages), target)?);
12841    }
12842    Ok(sections)
12843}
12844
12845/// ★ A HEAD THAT IS NOT WHOLE TILES IS WIDENED TO THEM. The blocked form addresses a head by whole tiles (its memory
12846/// an `nb × nb` square of them, its vectors `nb` of them), so a head of 16 or 80 whose memory is past the flat form's one
12847/// tile has no form of its own. Zeros past a head's `d` are a fixed point of the step: zero entries of `k`, `q` and `v` keep
12848/// the memory zero past `d` in both axes and the prediction and the output zero past `d`, and leave the norm's sum of
12849/// squares as it was. So each head is widened to `dp`, the next whole tiles (`Q`, `K`, `V`, `Z` and the norm weight per
12850/// head, the memory per head in both axes), the blocked step runs at `dp` reading the head's own `d` for its `1/√d`
12851/// and its norm, and the output and the memory are narrowed back. Every move copies a head's run or writes zeros, so
12852/// the moves are exact and a non-finite value stays in its own head. Every tensor the program names keeps the program's
12853/// layout (an edge the program lays out as a matrix is bridged where every op's is, `link::Prelude::agree_layouts`); the
12854/// widened ones are the op's own.
12855fn widened_gated_delta(
12856    spec: &GatedDeltaSpec,
12857    names: &StageNames,
12858    target: &SpatialTarget,
12859) -> Result<Vec<SpatialProgram>, SpatialReject> {
12860    let (hv, hk, d, rows) = (spec.hv, spec.hk, spec.d, spec.rows);
12861    if hv % hk != 0 {
12862        return Err(SpatialReject::NotLowered("a gated delta whose v-heads are not a multiple of its k-heads"));
12863    }
12864    let (tr, tw) = (target.tile.rows, target.tile.cols);
12865    let dp = d.next_multiple_of(tw);
12866    let dt = elem_precision(target)?;
12867    let wide = |what: &str| format!("{}__wide_{what}", names.out);
12868    // `n` runs of `from` elements, each copied to the start of `to`.
12869    let runs = |n: usize, from: usize, to: usize| -> Vec<(usize, Option<usize>)> {
12870        joined_zero_runs((0..n).flat_map(|u| [(u * to + from.min(to), Some(u * from)), ((u + 1) * to, None)]).collect())
12871    };
12872    let chunked = takes_chunks(rows, tr);
12873    let mut sections = Vec::new();
12874    // `Q`, `K`, `V` and `Z`, each row's heads in turn; the chunked step reads them as matrices.
12875    let mut widen = |i: usize, heads: usize, what: &str| -> Result<String, SpatialReject> {
12876        let to = wide(what);
12877        let flat = if chunked { format!("{to}__rows") } else { to.clone() };
12878        sections.push(moved_runs(names.operand(i), rows * heads * d, &flat, &runs(rows * heads, d, dp), target)?);
12879        if chunked {
12880            sections.extend(relayout_rows_to_dense((&flat, &Fill::Unknown), &to, rows, heads * dp, dt, target)?);
12881        }
12882        Ok(to)
12883    };
12884    let (q, k, v, z) = (widen(0, hk, "q")?, widen(1, hk, "k")?, widen(2, hv, "v")?, widen(3, hv, "z")?);
12885    let gn = wide("gn");
12886    sections.push(moved_runs(names.operand(7), d, &gn, &runs(1, d, dp), target)?);
12887    // The memory: head `h`'s row `a` of `d` lands at row `h·dp + a` of `dp`, its rows past `d` zero; the step reads it as
12888    // the `hv·dp × dp` matrix.
12889    let (st_name, st_fill) = names.operand(6);
12890    let st = wide("state");
12891    let st_rows = format!("{st}__rows");
12892    let into: Vec<(usize, Option<usize>)> = (0..hv * dp)
12893        .flat_map(|y| {
12894            let (h, a) = (y / dp, y % dp);
12895            let copied = (a < d).then_some((y * dp + d, Some((h * d + a) * d)));
12896            copied.into_iter().chain(std::iter::once(((y + 1) * dp, None)))
12897        })
12898        .collect();
12899    sections.push(moved_runs((st_name, st_fill), hv * d * d, &st_rows, &joined_zero_runs(into), target)?);
12900    sections.extend(relayout_rows_to_dense((&st_rows, &Fill::Unknown), &st, hv * dp, dp, dt, target)?);
12901    let out = wide("out");
12902    let (b, a) = (names.operand(4), names.operand(5));
12903    let step = StageNames {
12904        out: out.clone(),
12905        ins: vec![q, k, v, z, b.0.to_string(), a.0.to_string(), st.clone(), gn],
12906        fills: vec![Fill::Unknown, Fill::Unknown, Fill::Unknown, Fill::Unknown, b.1.clone(), a.1.clone(), Fill::Unknown, Fill::Unknown],
12907    };
12908    // The widened memory is the op's own, so each row's (or chunk's) version has a name of its own and one writer.
12909    let version = |i: usize| if i == 0 { st.clone() } else { format!("{st}__{i}") };
12910    sections.extend(blocked_gated_delta(&GatedDeltaSpec { d: dp, ..spec.clone() }, &step, &BTreeSet::new(), d, &version, target)?);
12911    let last = if chunked { rows.div_ceil(tr) } else { rows };
12912    // Back to the program's layout: the output's heads, and the memory's rows.
12913    sections.push(moved_runs((&out, &Fill::Unknown), rows * hv * dp, &names.out, &runs(rows * hv, dp, d), target)?);
12914    let back = format!("{st}__back");
12915    sections.push(relayout_dense_to_rows((&version(last), &Fill::Unknown), &back, hv * dp, dp, dt, target)?);
12916    let from: Vec<(usize, Option<usize>)> = (0..hv * d).map(|y| (y * d + d, Some((y / d * dp + y % d) * dp))).collect();
12917    sections.push(moved_runs((&back, &Fill::Unknown), hv * dp * dp, st_name, &from, target)?);
12918    Ok(sections)
12919}
12920
12921/// `runs` with every run of zeros joined to the zeros before it.
12922fn joined_zero_runs(runs: Vec<(usize, Option<usize>)>) -> Vec<(usize, Option<usize>)> {
12923    let mut out: Vec<(usize, Option<usize>)> = Vec::with_capacity(runs.len());
12924    for (end, from) in runs {
12925        match (out.last_mut(), from) {
12926            (Some(last), None) if last.1.is_none() => last.0 = end,
12927            (Some(last), _) if last.0 == end => {}
12928            _ => out.push((end, from)),
12929        }
12930    }
12931    out
12932}
12933
12934/// A section writing `out`, flat, from runs of `src`, flat: each `(end, from)` ends at element `end` of `out` and copies
12935/// `src` from element `from` (the run's first element reads `src[from]`), or writes zeros. Movement only: exact, and a
12936/// non-finite value lands only where its run puts it.
12937///
12938/// ★ AN ITEM IS AN OUTPUT TILE, AND IT LOADS ONLY THE SOURCE IT MOVES. A flat row spends a tile per tile-width of
12939/// elements, so a section that loads its whole source into every item outgrows a site at a widened memory's length.
12940/// Output tile `i` loads a window of `W` source tiles (`W` the most any tile's runs reach) from the tile a compile-time
12941/// table names for it, through an [`IndirectTile`], and reads each of its runs at the offset the same table gives: a run
12942/// that starts mid-tile, in either array, costs nothing more. Tile `i`'s table page holds the window's first tile, then
12943/// for each of its runs (`R` of them, the most any tile holds, the unused ones empty) the column the zeros before it end
12944/// at, the column it ends at, and its offset in the window. The page is loaded twice, once as the window's address and
12945/// once for the body, a loaded page having one reader.
12946fn moved_runs(
12947    src: (&str, &Fill),
12948    src_len: usize,
12949    out: &str,
12950    runs: &[(usize, Option<usize>)],
12951    target: &SpatialTarget,
12952) -> Result<SpatialProgram, SpatialReject> {
12953    let n = runs.last().map_or(0, |r| r.0);
12954    if n == 0 || src_len == 0 {
12955        return Err(SpatialReject::NotLowered("a move of nothing"));
12956    }
12957    let tw = target.tile.cols;
12958    let (n_t, src_t) = (tiles_of(n, target), tiles_of(src_len, target));
12959    let dt = elem_precision(target)?;
12960    // Each output tile's runs: its first column in the tile, the column it ends at, the source element it starts at.
12961    let mut pieces: Vec<Vec<(usize, usize, usize)>> = vec![Vec::new(); n_t];
12962    let mut start = 0;
12963    for &(end, from) in runs {
12964        if let Some(f) = from {
12965            let mut e = start;
12966            while e < end {
12967                let i = e / tw;
12968                let stop = end.min((i + 1) * tw);
12969                pieces[i].push((e - i * tw, stop - i * tw, f + (e - start)));
12970                e = stop;
12971            }
12972        }
12973        start = end;
12974    }
12975    let reach = |ps: &[(usize, usize, usize)]| -> Option<(usize, usize)> {
12976        let lo = ps.iter().map(|p| p.2).min()?;
12977        let hi = ps.iter().map(|p| p.2 + (p.1 - p.0)).max()?;
12978        Some((lo / tw, hi.div_ceil(tw)))
12979    };
12980    let window = pieces.iter().filter_map(|ps| reach(ps)).map(|(a, b)| b - a).max().unwrap_or(1).clamp(1, src_t);
12981    let r_max = pieces.iter().map(Vec::len).max().unwrap_or(0).max(1);
12982    if 1 + 3 * r_max > tw {
12983        return Err(SpatialReject::NotLowered("a move whose output tiles each take more runs than a table page names"));
12984    }
12985    // A run's offset in the window, lifted by a tile so it is never negative: column `c` reads `c + off − tile`.
12986    let lift = tw as i64;
12987    let mut values: Vec<i64> = Vec::with_capacity(n_t * tw);
12988    for ps in &pieces {
12989        let first = reach(ps).map_or(0, |(a, _)| a.min(src_t - window));
12990        let mut page = vec![first as i64];
12991        let mut last = 0usize;
12992        for k in 0..r_max {
12993            match ps.get(k) {
12994                Some(&(a, b, f)) => {
12995                    page.extend([a as i64, b as i64, f as i64 - (first * tw) as i64 - a as i64 + lift]);
12996                    last = b;
12997                }
12998                None => page.extend([last as i64, last as i64, lift]),
12999            }
13000        }
13001        page.resize(tw, 0);
13002        values.extend(page);
13003    }
13004    let reach_max = values.iter().copied().max().unwrap_or(0).max(0) as usize;
13005    let ix = index_type(target, reach_max.max(src_t) as u64)?;
13006    let mut t = Template::new(target)?;
13007    let source = t.tensor(src.0, TensorRole::Input, 1, src_len, dt, super::from_tensor_op::init_of(src.1, src_len));
13008    let y = t.tensor(out, TensorRole::Output, 1, n, dt, TensorInit::Zero);
13009    let table = t.tensor(&format!("{out}__runs"), TensorRole::Param, 1, n_t * tw, ix, TensorInit::Values(values));
13010    let addr = t.buffer("in_runs_window", 1, ix);
13011    let l_addr = t.load_block_at("load_runs_window", table, 0, vec![(LV_ITEM_X, 1)], 1, 1, n_t, addr);
13012    let at = t.buffer("in_runs", 1, ix);
13013    let l_at = t.load_block_at("load_runs", table, 0, vec![(LV_ITEM_X, 1)], 1, 1, n_t, at);
13014    let page = t.buffer("in_window", window, dt);
13015    let l_page = t.stage(
13016        "load_window",
13017        StageKind::Load {
13018            tensor: source,
13019            pattern: AccessPattern {
13020                origin: TileAddr {
13021                    base: 0,
13022                    steps: vec![],
13023                    indirect: Some(IndirectTile { index: addr, tile_stride: 1, elems_per_tile: 1, bound: (0, src_t - window) }),
13024                },
13025                block_h: 1,
13026                block_w: window,
13027                stride_w: 1,
13028                stride_h: src_t,
13029            },
13030            num_blocks: 1,
13031            multicast: None,
13032        },
13033        vec![addr],
13034        vec![page],
13035    );
13036    t.edge(l_addr, addr, l_page);
13037    // Operands: the table's entries for each run, then the window read at each run's offset.
13038    let (mut bufs, mut maps) = (Vec::new(), Vec::new());
13039    let entry = |bufs: &mut Vec<BufferId>, maps: &mut Vec<ElemAccess>, k: usize| {
13040        bufs.push(at);
13041        maps.push(fixed(k as i64));
13042        bufs.len() - 1
13043    };
13044    let mut bands = Vec::with_capacity(r_max);
13045    for k in 0..r_max {
13046        let (a, b, off) = (entry(&mut bufs, &mut maps, 1 + 3 * k), entry(&mut bufs, &mut maps, 2 + 3 * k), entry(&mut bufs, &mut maps, 3 + 3 * k));
13047        bufs.push(page);
13048        maps.push(ElemAccess {
13049            base: -lift,
13050            terms: vec![(LV_GEN_COL, 1)],
13051            indirect: Some(IndirectTerm { operand: off, stride: 1, modulus: reach_max + 1 }),
13052        });
13053        bands.push((a, b, bufs.len() - 1));
13054    }
13055    let below = |bound: usize, then: Box<TensorBody>, otherwise: Box<TensorBody>| {
13056        Box::new(TensorBody::Select { cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(bound)), then, otherwise })
13057    };
13058    let body = bands.iter().rev().fold(konst(0.0), |past, &(a, b, read)| below(a, konst(0.0), below(b, op(read), past)));
13059    let ob = t.buffer("out", 1, dt);
13060    let s_move = t.stage(
13061        "move_runs",
13062        StageKind::Compute(ComputeOp::Generic { iters: flat_space(tw), ins: maps, out: ident_map(wide(1, target)), body: *body }),
13063        bufs,
13064        vec![ob],
13065    );
13066    t.edge(l_at, at, s_move);
13067    t.edge(l_page, page, s_move);
13068    let st = t.store_block("store_runs", y, 0, vec![(LV_ITEM_X, 1)], 1, 1, n_t, ob);
13069    t.edge(s_move, ob, st);
13070    t.finish(format!("move_runs_{out}"), row_items(n_t, 1), AccPolicy::for_result(dt), NumericsClaim::Exact, target)
13071}
13072
13073/// `⌈log₂ c⌉`, at least 1: how many factors `(I + N)(I + N²)(I + N⁴)…` invert `I + N` for a strictly lower `c × c` `N`
13074/// — `N^{2^m}` vanishes once `2^m ≥ c`.
13075fn solve_factors(c: usize) -> usize {
13076    (usize::BITS - c.saturating_sub(1).leading_zeros()).max(1) as usize
13077}
13078
13079/// Whether a block of `rows` rows is lowered a chunk of `tile` rows at a time: its chunks take fewer sections than its
13080/// rows stepped one at a time.
13081fn takes_chunks(rows: usize, tile: usize) -> bool {
13082    rows > 1 && chunked_sections(rows, tile) < stepped_sections(rows)
13083}
13084
13085/// Sections the step takes over a block of `rows` rows, a row at a time: per row one pass over the memory; per block
13086/// the two gate broadcasts.
13087fn stepped_sections(rows: usize) -> usize {
13088    2 + rows
13089}
13090
13091/// Sections a chunked block of `rows` rows takes on `tile`-row chunks: per chunk the gates, `A` and `QKΓ`, the solve's
13092/// factors, `U`/`W`, `Δ`, the output and the memory update; per block the two gate broadcasts.
13093fn chunked_sections(rows: usize, tile: usize) -> usize {
13094    2 + (0..rows.div_ceil(tile))
13095        .map(|c| 6 + solve_factors((rows - c * tile).min(tile)))
13096        .sum::<usize>()
13097}
13098
13099/// ★ A BLOCK OF ROWS IS A CHUNK AT A TIME, NOT A ROW AT A TIME. Unrolling the step over a chunk of `C` rows (the
13100/// WY/UT form of the delta rule — the Gated DeltaNet paper's chunked algorithm) turns its recurrence into per-head
13101/// products: with `g_t = Σ_{s≤t} ln α_s`, `Γ_{ti} = e^{g_t − g_i}` (`i ≤ t`), `A_{ti} = β_t Γ_{ti} k_t·k_i` (`i < t`),
13102/// `T = (I + A)⁻¹ diag(β)`, `U = T V`, `W = T diag(e^g) K`, `Δ = U − W S₀ᵀ`:
13103///
13104/// ```text
13105///   O   = (diag(e^g) Q S₀ᵀ + (QKᵀ ⊙ Γ) Δ) / √d        every row's output, before its norm and gate
13106///   S_C = e^{g_C} S₀ + Δᵀ diag(e^{g_C − g}) K          the memory the chunk leaves
13107/// ```
13108///
13109/// Every product is ONE contraction (`LV_GEN_K`), and the unit lower-triangular solve is the product
13110/// `(I + N)(I + N²)(I + N⁴)…` with `N = −A` — `⌈log₂ C⌉` doublings, not `C` steps. The memory is read and written once
13111/// per chunk instead of three times per row.
13112///
13113/// A chunk is ONE TILE ROW of the block's rows, so a head's chunk of `Q`, `K`, `V`, `Z` is a `1 × nb` page of the DENSE
13114/// table — the tables are read dense, and the layout pass bridges a flat one once. A computed page has one consumer, so
13115/// what the algebra uses twice (`G`, `A`, the solve's factors, `U`, `W`, `Δ`) crosses a section as a per-chunk tensor,
13116/// loaded once per reader. A partial last chunk masks its rows on the contraction variable; its padding rows compute to
13117/// zero, never to a NaN (a masked exponential is SELECTED away, never multiplied by zero).
13118fn chunked_gated_delta(
13119    spec: &GatedDeltaSpec,
13120    names: &StageNames,
13121    dense: &BTreeSet<String>,
13122    head: usize,
13123    state: &dyn Fn(usize) -> String,
13124    target: &SpatialTarget,
13125) -> Result<Vec<SpatialProgram>, SpatialReject> {
13126    let (hv, hk, d, rows) = (spec.hv, spec.hk, spec.d, spec.rows);
13127    let (tr, tw) = (target.tile.rows, target.tile.cols);
13128    let nb = d / tw;
13129    let precision = elem_precision(target)?;
13130    let (q_name, q_fill) = names.operand(0);
13131    let (k_name, k_fill) = names.operand(1);
13132    let (v_name, v_fill) = names.operand(2);
13133    let (z_name, z_fill) = names.operand(3);
13134    let (b_name, b_fill) = names.operand(4);
13135    let (a_name, a_fill) = names.operand(5);
13136    let st_fill = names.operand(6).1;
13137    let (gn_name, gn_fill) = names.operand(7);
13138    let (s_rows, qk_w, v_w) = (hv * d, hk * d, hv * d);
13139    let (qk_t, v_t) = (hk * nb, hv * nb);
13140    let w1 = wide(1, target);
13141    let wd = wide(nb, target);
13142    let a_exp = format!("{}__alpha_t", names.out);
13143    let b_exp = format!("{}__beta_t", names.out);
13144    // The output: the program's matrix itself, or a dense twin a relayout turns into the program's flat row.
13145    let o_dense = dense.contains(&names.out);
13146    let o_name = if o_dense {
13147        names.out.clone()
13148    } else {
13149        format!("{}__dense", names.out)
13150    };
13151
13152    let head_items = |rect: &dyn Fn(usize) -> TileRect| -> Vec<WorkItem> {
13153        (0..hv)
13154            .map(|h| WorkItem {
13155                id: WorkItemId::new(h),
13156                coord: (h % hk, h / hk),
13157                out: rect(h),
13158            })
13159            .collect()
13160    };
13161    let tile_items = || {
13162        head_items(&|h| TileRect {
13163            row0_t: 0,
13164            col0_t: h,
13165            h_t: 1,
13166            w_t: 1,
13167        })
13168    };
13169    let strip_items = || {
13170        head_items(&|h| TileRect {
13171            row0_t: 0,
13172            col0_t: h * nb,
13173            h_t: 1,
13174            w_t: nb,
13175        })
13176    };
13177    let head_steps = |per_head: usize| vec![(LV_ITEM_Y, per_head), (LV_ITEM_X, hk * per_head)];
13178    let k_steps = || vec![(LV_ITEM_Y, nb)];
13179
13180    // ── Per block: the gates as one tile per head — `α_t[r, h·tw + j] = α[r, h]`, row `r` of tile row `r / tile`.
13181    let sel_fill = || -> Fill {
13182        Fill::Overrides {
13183            base: Box::new(Fill::Const(0)),
13184            sets: (0..hv)
13185                .flat_map(|h| (0..tw).map(move |j| (h * hv * tw + h * tw + j, 1)))
13186                .collect(),
13187        }
13188    };
13189    let broadcast =
13190        |src: &str, fill: &Fill, dst: &str, sel: &str| -> Result<SpatialProgram, SpatialReject> {
13191            let spec = crate::tensor_recognize::MatmulSpec {
13192                m: rows,
13193                n: hv * tw,
13194                k: hv,
13195                a: Default::default(),
13196                b: Default::default(),
13197                c: Default::default(),
13198            };
13199            let names = MatmulNames {
13200                a: src.to_string(),
13201                b: sel.to_string(),
13202                c: dst.to_string(),
13203            };
13204            super::from_tensor_op::from_matmul(
13205                &spec,
13206                &names,
13207                (fill, &sel_fill()),
13208                (None, None),
13209                NumericsPolicy::Exact,
13210                target,
13211            )
13212        };
13213    let mut sections = vec![
13214        broadcast(a_name, a_fill, &a_exp, &format!("{}__sel_a", names.out))?,
13215        broadcast(b_name, b_fill, &b_exp, &format!("{}__sel_b", names.out))?,
13216    ];
13217
13218    // Index tiles for the masks, and the identity: `R0[t, i] = t`, `R1[t, i] = t + 1`, `L[t, s] = [s ≤ t]`.
13219    let index_tile = |f: &dyn Fn(usize, usize) -> f64| -> TensorInit {
13220        TensorInit::FloatValues((0..tr * tw).map(|e| f(e / tw, e % tw).to_bits()).collect())
13221    };
13222    let param = |t: &mut Template, what: &str, init: TensorInit| {
13223        t.tensor(
13224            &format!("{}__gd_{what}", names.out),
13225            TensorRole::Param,
13226            tr,
13227            tw,
13228            precision,
13229            init,
13230        )
13231    };
13232    let row_index = |plus: usize| index_tile(&move |t, _| (t + plus) as f64);
13233    let tril = || index_tile(&|t, s| if s <= t { 1.0 } else { 0.0 });
13234    let eye = || index_tile(&|t, i| if t == i { 1.0 } else { 0.0 });
13235
13236    // Map helpers over a one-tile page (`w1` wide) and over a head's `1 × nb` page (`wd` wide).
13237    let m = |terms: &[(LoopVar, i64)]| ElemAccess {
13238        base: 0,
13239        terms: terms.to_vec(),
13240        indirect: None,
13241    };
13242    let contract = |rows_x: usize, cols_x: usize, k: usize| {
13243        vec![
13244            IterSpec {
13245                var: LV_GEN_ROW,
13246                extent: rows_x,
13247                kind: IterKind::Parallel,
13248            },
13249            IterSpec {
13250                var: LV_GEN_COL,
13251                extent: cols_x,
13252                kind: IterKind::Parallel,
13253            },
13254            IterSpec {
13255                var: LV_GEN_K,
13256                extent: k,
13257                kind: IterKind::Reduce {
13258                    pool: ReducePool::Sum,
13259                },
13260            },
13261        ]
13262    };
13263    let pointwise = |cols_x: usize| {
13264        vec![
13265            IterSpec {
13266                var: LV_GEN_ROW,
13267                extent: tr,
13268                kind: IterKind::Parallel,
13269            },
13270            IterSpec {
13271                var: LV_GEN_COL,
13272                extent: cols_x,
13273                kind: IterKind::Parallel,
13274            },
13275        ]
13276    };
13277    let generic =
13278        |iters: Vec<IterSpec>, ins: Vec<ElemAccess>, out: ElemAccess, body: TensorBody| {
13279            StageKind::Compute(ComputeOp::Generic {
13280                iters,
13281                ins,
13282                out,
13283                body,
13284            })
13285        };
13286    let chunk_tensor = |t: &mut Template, c: usize, what: &str, role: TensorRole, cols: usize| {
13287        let init = if role == TensorRole::Output {
13288            TensorInit::Zero
13289        } else {
13290            TensorInit::External
13291        };
13292        t.tensor(
13293            &format!("{}__c{c}_{what}", names.out),
13294            role,
13295            tr,
13296            cols,
13297            precision,
13298            init,
13299        )
13300    };
13301    let gate_tensor = |t: &mut Template, name: &str| {
13302        t.tensor(
13303            name,
13304            TensorRole::Input,
13305            rows,
13306            hv * tw,
13307            precision,
13308            TensorInit::External,
13309        )
13310    };
13311    let q_tensor = |t: &mut Template| {
13312        t.tensor(
13313            q_name,
13314            TensorRole::Input,
13315            rows,
13316            qk_w,
13317            precision,
13318            super::from_tensor_op::init_of(q_fill, rows * qk_w),
13319        )
13320    };
13321    let k_tensor = |t: &mut Template| {
13322        t.tensor(
13323            k_name,
13324            TensorRole::Input,
13325            rows,
13326            qk_w,
13327            precision,
13328            super::from_tensor_op::init_of(k_fill, rows * qk_w),
13329        )
13330    };
13331    let finish = |t: Template, what: &str, c: usize, items: Vec<WorkItem>| {
13332        t.finish(
13333            format!("gd_chunk{c}_{what}_{hv}x{d}"),
13334            items,
13335            AccPolicy::F32,
13336            approx("chunked gated delta", precision),
13337            target,
13338        )
13339    };
13340
13341    for c in 0..rows.div_ceil(tr) {
13342        let cc = (rows - c * tr).min(tr);
13343        let factors = solve_factors(cc);
13344        // Row `r` of the block, as a tile of a `rows`-row table: tile row `c`.
13345        let gate_base = c * hv;
13346
13347        // ── G and LG: `G[t, ·] = Σ_{s ≤ t} ln α_s`, the cumulative log decay down the rows of a row-broadcast tile, and
13348        // `LG[t, i] = Σ_{i < s ≤ t} ln α_s`, the log of every ratio `γ_t / γ_i` — each over the chunk's rows `s < C`.
13349        //
13350        // ★ A RATIO IS SUMMED, NEVER SUBTRACTED. `g_t` grows with the chunk (`−22` after 32 rows of `α ≈ ½`), where a
13351        // bf16 page keeps it only to `1/8`; `e^{g_t − g_i}` from two stored sums is then off by up to `e^{±1/8}` exactly
13352        // where it matters — near the diagonal, where the ratio is close to one. Summed over `(i, t]` alone, a
13353        // near-diagonal log ratio is a few small terms and keeps its precision.
13354        sections.push({
13355            let mut t = Template::new(target)?;
13356            let at = gate_tensor(&mut t, &a_exp);
13357            let l = param(&mut t, "tril", tril());
13358            let (col1, row1) = (
13359                param(&mut t, "col1", index_tile(&|_, i| (i + 1) as f64)),
13360                param(&mut t, "row1", row_index(1)),
13361            );
13362            let g = chunk_tensor(&mut t, c, "g", TensorRole::Output, hv * tw);
13363            let lg = chunk_tensor(&mut t, c, "lg", TensorRole::Output, hv * tw);
13364            let (ain, lin, gout) = (
13365                t.buffer("in_alpha", 1, precision),
13366                t.buffer("in_tril", 1, precision),
13367                t.buffer("g", 1, precision),
13368            );
13369            let la = t.load_block_at("load_alpha", at, gate_base, head_steps(1), 1, 1, hv, ain);
13370            let ll = t.load("load_tril", l, vec![], 1, lin);
13371            let body = TensorBody::Select {
13372                cond: Pred::AxisLt(LV_GEN_K, PredRhs::Const(cc)),
13373                then: mul(op(0), unary(SfpuKind::Ln, op(1))),
13374                otherwise: konst(0.0),
13375            };
13376            let s = t.stage(
13377                "gd_log_decay",
13378                generic(
13379                    contract(tr, tw, tr),
13380                    vec![
13381                        m(&[(LV_GEN_ROW, w1 as i64), (LV_GEN_K, 1)]),
13382                        m(&[(LV_GEN_K, w1 as i64), (LV_GEN_COL, 1)]),
13383                    ],
13384                    ident_map(w1),
13385                    body,
13386                ),
13387                vec![lin, ain],
13388                vec![gout],
13389            );
13390            let st = t.store_block("store_g", g, 0, head_steps(1), 1, 1, hv, gout);
13391            t.edge(ll, lin, s);
13392            t.edge(la, ain, s);
13393            t.edge(s, gout, st);
13394            let (ain2, cin, rin, lgout) = (
13395                t.buffer("in_alpha_r", 1, precision),
13396                t.buffer("in_col1", 1, precision),
13397                t.buffer("in_row1", 1, precision),
13398                t.buffer("lg", 1, precision),
13399            );
13400            let la2 = t.load_block_at("load_alpha_r", at, gate_base, head_steps(1), 1, 1, hv, ain2);
13401            let lc = t.load("load_col1", col1, vec![], 1, cin);
13402            let lr = t.load("load_row1", row1, vec![], 1, rin);
13403            // 0:α_s  1:i + 1  2:t + 1 — `s` counted when `s < C`, `s ≥ i + 1` and `s < t + 1`.
13404            let within = |then: Box<TensorBody>| {
13405                let upper = TensorBody::Select {
13406                    cond: Pred::AxisLt(LV_GEN_K, PredRhs::Operand(2)),
13407                    then,
13408                    otherwise: konst(0.0),
13409                };
13410                let lower = TensorBody::Select {
13411                    cond: Pred::AxisGe(LV_GEN_K, PredRhs::Operand(1)),
13412                    then: Box::new(upper),
13413                    otherwise: konst(0.0),
13414                };
13415                TensorBody::Select {
13416                    cond: Pred::AxisLt(LV_GEN_K, PredRhs::Const(cc)),
13417                    then: Box::new(lower),
13418                    otherwise: konst(0.0),
13419                }
13420            };
13421            let s2 = t.stage(
13422                "gd_log_ratio",
13423                generic(
13424                    contract(tr, tw, tr),
13425                    vec![m(&[(LV_GEN_K, w1 as i64)]), ident_map(w1), ident_map(w1)],
13426                    ident_map(w1),
13427                    within(unary(SfpuKind::Ln, op(0))),
13428                ),
13429                vec![ain2, cin, rin],
13430                vec![lgout],
13431            );
13432            let st2 = t.store_block("store_lg", lg, 0, head_steps(1), 1, 1, hv, lgout);
13433            t.edge(la2, ain2, s2);
13434            t.edge(lc, cin, s2);
13435            t.edge(lr, rin, s2);
13436            t.edge(s2, lgout, st2);
13437            finish(t, "g", c, tile_items())?
13438        });
13439
13440        // ── A and QKΓ: `A[t, i] = β_t Γ[t, i] Σ_b k_t[b] k_i[b]` below the diagonal; `QKΓ[t, i] = Γ[t, i] Σ_b q_t[b] k_i[b]`
13441        // on and below it.
13442        sections.push({
13443            let mut t = Template::new(target)?;
13444            let (q, k) = (q_tensor(&mut t), k_tensor(&mut t));
13445            let bt = gate_tensor(&mut t, &b_exp);
13446            let lg = chunk_tensor(&mut t, c, "lg", TensorRole::Input, hv * tw);
13447            let (r0, r1) = (
13448                param(&mut t, "row", row_index(0)),
13449                param(&mut t, "row1", row_index(1)),
13450            );
13451            let a = chunk_tensor(&mut t, c, "a", TensorRole::Output, hv * tw);
13452            let qk = chunk_tensor(&mut t, c, "qk", TensorRole::Output, hv * tw);
13453            let part = |t: &mut Template,
13454                        what: &str,
13455                        left: TensorId,
13456                        beta: bool,
13457                        strict: bool,
13458                        out_t: TensorId| {
13459                let lin = t.buffer_deep(&format!("in_{what}_l"), 1, nb, precision);
13460                let kin = t.buffer_deep(&format!("in_{what}_k"), 1, nb, precision);
13461                let gin = t.buffer(&format!("in_{what}_lg"), 1, precision);
13462                let rin = t.buffer(&format!("in_{what}_r"), 1, precision);
13463                let ll = t.load_block_at(
13464                    &format!("load_{what}_l"),
13465                    left,
13466                    c * qk_t,
13467                    k_steps(),
13468                    1,
13469                    nb,
13470                    qk_t,
13471                    lin,
13472                );
13473                let lk = t.load_block_at(
13474                    &format!("load_{what}_k"),
13475                    k,
13476                    c * qk_t,
13477                    k_steps(),
13478                    1,
13479                    nb,
13480                    qk_t,
13481                    kin,
13482                );
13483                let lg = t.load_block_at(
13484                    &format!("load_{what}_lg"),
13485                    lg,
13486                    0,
13487                    head_steps(1),
13488                    1,
13489                    1,
13490                    hv,
13491                    gin,
13492                );
13493                let lr = t.load(
13494                    &format!("load_{what}_r"),
13495                    if strict { r0 } else { r1 },
13496                    vec![],
13497                    1,
13498                    rin,
13499                );
13500                // 0:left_t[b]  1:k_i[b]  2:LG[t, i]  3:the mask's row bound  (4:β_t)
13501                let mut ins = vec![lin, kin, gin, rin];
13502                let mut maps = vec![
13503                    m(&[(LV_GEN_ROW, wd as i64), (LV_GEN_K, 1)]),
13504                    m(&[(LV_GEN_COL, wd as i64), (LV_GEN_K, 1)]),
13505                    ident_map(w1),
13506                    ident_map(w1),
13507                ];
13508                let mut value = mul(unary(SfpuKind::Exp, op(2)), mul(op(0), op(1)));
13509                let mut loads = vec![(ll, lin), (lk, kin), (lg, gin), (lr, rin)];
13510                if beta {
13511                    let bin = t.buffer(&format!("in_{what}_b"), 1, precision);
13512                    loads.push((
13513                        t.load_block_at(
13514                            &format!("load_{what}_b"),
13515                            bt,
13516                            gate_base,
13517                            head_steps(1),
13518                            1,
13519                            1,
13520                            hv,
13521                            bin,
13522                        ),
13523                        bin,
13524                    ));
13525                    ins.push(bin);
13526                    maps.push(m(&[(LV_GEN_ROW, w1 as i64)]));
13527                    value = mul(op(4), value);
13528                }
13529                let body = TensorBody::Select {
13530                    cond: Pred::AxisLt(LV_GEN_COL, PredRhs::Operand(3)),
13531                    then: value,
13532                    otherwise: konst(0.0),
13533                };
13534                let out = t.buffer(what, 1, precision);
13535                let s = t.stage(
13536                    &format!("gd_{what}"),
13537                    generic(contract(tr, tw, d), maps, ident_map(w1), body),
13538                    ins,
13539                    vec![out],
13540                );
13541                let st = t.store_block(
13542                    &format!("store_{what}"),
13543                    out_t,
13544                    0,
13545                    head_steps(1),
13546                    1,
13547                    1,
13548                    hv,
13549                    out,
13550                );
13551                for (l, b) in loads {
13552                    t.edge(l, b, s);
13553                }
13554                t.edge(s, out, st);
13555            };
13556            part(&mut t, "a", k, true, true, a);
13557            part(&mut t, "qk", q, false, false, qk);
13558            finish(t, "a", c, tile_items())?
13559        });
13560
13561        // ── The solve: `X = (I + N)(I + N²)(I + N⁴)…` with `N = −A`, factor by factor. Section `j` forms `X_{j−1}` from
13562        // `X_{j−2}` and `P_{j−1} = N^{2^{j−1}}`, and squares `P_{j−1}` into `P_j` for the section after it.
13563        for j in 0..factors {
13564            sections.push({
13565                let mut t = Template::new(target)?;
13566                // `P_j`: `A²` at the first factor (`N² = A²`), then the previous power squared.
13567                if j + 1 < factors {
13568                    let src = chunk_tensor(
13569                        &mut t,
13570                        c,
13571                        &if j == 0 {
13572                            "a".to_string()
13573                        } else {
13574                            format!("p{j}")
13575                        },
13576                        TensorRole::Input,
13577                        hv * tw,
13578                    );
13579                    let dst = chunk_tensor(
13580                        &mut t,
13581                        c,
13582                        &format!("p{}", j + 1),
13583                        TensorRole::Output,
13584                        hv * tw,
13585                    );
13586                    let (pin, pin2) = (
13587                        t.buffer("in_p", 1, precision),
13588                        t.buffer("in_p2", 1, precision),
13589                    );
13590                    let lp = t.load_block_at("load_p", src, 0, head_steps(1), 1, 1, hv, pin);
13591                    let lp2 = t.load_block_at("load_p2", src, 0, head_steps(1), 1, 1, hv, pin2);
13592                    let out = t.buffer("p_sq", 1, precision);
13593                    let s = t.stage(
13594                        "gd_square",
13595                        generic(
13596                            contract(tr, tw, tr),
13597                            vec![
13598                                m(&[(LV_GEN_ROW, w1 as i64), (LV_GEN_K, 1)]),
13599                                m(&[(LV_GEN_K, w1 as i64), (LV_GEN_COL, 1)]),
13600                            ],
13601                            ident_map(w1),
13602                            *mul(op(0), op(1)),
13603                        ),
13604                        vec![pin, pin2],
13605                        vec![out],
13606                    );
13607                    let st = t.store_block("store_p", dst, 0, head_steps(1), 1, 1, hv, out);
13608                    t.edge(lp, pin, s);
13609                    t.edge(lp2, pin2, s);
13610                    t.edge(s, out, st);
13611                }
13612                let x_out = chunk_tensor(&mut t, c, &format!("x{j}"), TensorRole::Output, hv * tw);
13613                let out = t.buffer("x", 1, precision);
13614                if j == 0 {
13615                    // `X₀ = I − A`.
13616                    let a = chunk_tensor(&mut t, c, "a", TensorRole::Input, hv * tw);
13617                    let id = param(&mut t, "eye", eye());
13618                    let (ain, iin) = (
13619                        t.buffer("in_a", 1, precision),
13620                        t.buffer("in_eye", 1, precision),
13621                    );
13622                    let la = t.load_block_at("load_a", a, 0, head_steps(1), 1, 1, hv, ain);
13623                    let li = t.load("load_eye", id, vec![], 1, iin);
13624                    let s = t.stage(
13625                        "gd_x0",
13626                        generic(
13627                            pointwise(tw),
13628                            vec![ident_map(w1), ident_map(w1)],
13629                            ident_map(w1),
13630                            *sub(op(1), op(0)),
13631                        ),
13632                        vec![ain, iin],
13633                        vec![out],
13634                    );
13635                    t.edge(la, ain, s);
13636                    t.edge(li, iin, s);
13637                    let st = t.store_block("store_x", x_out, 0, head_steps(1), 1, 1, hv, out);
13638                    t.edge(s, out, st);
13639                } else {
13640                    // `X_j = X_{j−1} + X_{j−1}·P_j`.
13641                    let x_prev = chunk_tensor(
13642                        &mut t,
13643                        c,
13644                        &format!("x{}", j - 1),
13645                        TensorRole::Input,
13646                        hv * tw,
13647                    );
13648                    let p = chunk_tensor(&mut t, c, &format!("p{j}"), TensorRole::Input, hv * tw);
13649                    let (xin, pin, xin2) = (
13650                        t.buffer("in_x", 1, precision),
13651                        t.buffer("in_xp_p", 1, precision),
13652                        t.buffer("in_x_add", 1, precision),
13653                    );
13654                    let lx = t.load_block_at("load_x", x_prev, 0, head_steps(1), 1, 1, hv, xin);
13655                    let lp = t.load_block_at("load_xp_p", p, 0, head_steps(1), 1, 1, hv, pin);
13656                    let lx2 =
13657                        t.load_block_at("load_x_add", x_prev, 0, head_steps(1), 1, 1, hv, xin2);
13658                    let xp = t.buffer("xp", 1, precision);
13659                    let s1 = t.stage(
13660                        "gd_xp",
13661                        generic(
13662                            contract(tr, tw, tr),
13663                            vec![
13664                                m(&[(LV_GEN_ROW, w1 as i64), (LV_GEN_K, 1)]),
13665                                m(&[(LV_GEN_K, w1 as i64), (LV_GEN_COL, 1)]),
13666                            ],
13667                            ident_map(w1),
13668                            *mul(op(0), op(1)),
13669                        ),
13670                        vec![xin, pin],
13671                        vec![xp],
13672                    );
13673                    let s2 = t.stage(
13674                        "gd_x",
13675                        generic(
13676                            pointwise(tw),
13677                            vec![ident_map(w1), ident_map(w1)],
13678                            ident_map(w1),
13679                            *add(op(0), op(1)),
13680                        ),
13681                        vec![xp, xin2],
13682                        vec![out],
13683                    );
13684                    t.edge(lx, xin, s1);
13685                    t.edge(lp, pin, s1);
13686                    t.edge(s1, xp, s2);
13687                    t.edge(lx2, xin2, s2);
13688                    let st = t.store_block("store_x", x_out, 0, head_steps(1), 1, 1, hv, out);
13689                    t.edge(s2, out, st);
13690                }
13691                finish(t, &format!("solve{j}"), c, tile_items())?
13692            });
13693        }
13694        let x_last = format!("x{}", factors - 1);
13695
13696        // ── U and W: `U[t, e] = Σ_i X[t, i] β_i v_i[e]`, `W[t, b] = Σ_i X[t, i] β_i e^{g_i} k_i[b]`.
13697        sections.push({
13698            let mut t = Template::new(target)?;
13699            let x = chunk_tensor(&mut t, c, &x_last, TensorRole::Input, hv * tw);
13700            let bt = gate_tensor(&mut t, &b_exp);
13701            let g = chunk_tensor(&mut t, c, "g", TensorRole::Input, hv * tw);
13702            let v = t.tensor(
13703                v_name,
13704                TensorRole::Input,
13705                rows,
13706                v_w,
13707                precision,
13708                super::from_tensor_op::init_of(v_fill, rows * v_w),
13709            );
13710            let k = k_tensor(&mut t);
13711            let u = chunk_tensor(&mut t, c, "u", TensorRole::Output, hv * d);
13712            let w = chunk_tensor(&mut t, c, "w", TensorRole::Output, hv * d);
13713            let part = |t: &mut Template,
13714                        what: &str,
13715                        right: TensorId,
13716                        right_base: usize,
13717                        right_steps: Vec<(LoopVar, usize)>,
13718                        right_stride: usize,
13719                        gated: bool,
13720                        out_t: TensorId| {
13721                let (xin, bin, rin) = (
13722                    t.buffer(&format!("in_{what}_x"), 1, precision),
13723                    t.buffer(&format!("in_{what}_b"), 1, precision),
13724                    t.buffer_deep(&format!("in_{what}_r"), 1, nb, precision),
13725                );
13726                let lx = t.load_block_at(
13727                    &format!("load_{what}_x"),
13728                    x,
13729                    0,
13730                    head_steps(1),
13731                    1,
13732                    1,
13733                    hv,
13734                    xin,
13735                );
13736                let lb = t.load_block_at(
13737                    &format!("load_{what}_b"),
13738                    bt,
13739                    gate_base,
13740                    head_steps(1),
13741                    1,
13742                    1,
13743                    hv,
13744                    bin,
13745                );
13746                let lr = t.load_block_at(
13747                    &format!("load_{what}_r"),
13748                    right,
13749                    right_base,
13750                    right_steps,
13751                    1,
13752                    nb,
13753                    right_stride,
13754                    rin,
13755                );
13756                let mut ins = vec![xin, bin, rin];
13757                let mut maps = vec![
13758                    m(&[(LV_GEN_ROW, w1 as i64), (LV_GEN_K, 1)]),
13759                    m(&[(LV_GEN_K, w1 as i64)]),
13760                    m(&[(LV_GEN_K, wd as i64), (LV_GEN_COL, 1)]),
13761                ];
13762                let mut value = mul(mul(op(0), op(1)), op(2));
13763                let mut loads = vec![(lx, xin), (lb, bin), (lr, rin)];
13764                if gated {
13765                    let gin = t.buffer(&format!("in_{what}_g"), 1, precision);
13766                    loads.push((
13767                        t.load_block_at(
13768                            &format!("load_{what}_g"),
13769                            g,
13770                            0,
13771                            head_steps(1),
13772                            1,
13773                            1,
13774                            hv,
13775                            gin,
13776                        ),
13777                        gin,
13778                    ));
13779                    ins.push(gin);
13780                    maps.push(m(&[(LV_GEN_K, w1 as i64)]));
13781                    value = mul(value, unary(SfpuKind::Exp, op(3)));
13782                }
13783                let out = t.buffer_deep(what, 1, nb, precision);
13784                let s = t.stage(
13785                    &format!("gd_{what}"),
13786                    generic(contract(tr, wd, tr), maps, ident_map(wd), *value),
13787                    ins,
13788                    vec![out],
13789                );
13790                let st = t.store_block(
13791                    &format!("store_{what}"),
13792                    out_t,
13793                    0,
13794                    head_steps(nb),
13795                    1,
13796                    nb,
13797                    v_t,
13798                    out,
13799                );
13800                for (l, b) in loads {
13801                    t.edge(l, b, s);
13802                }
13803                t.edge(s, out, st);
13804            };
13805            part(&mut t, "u", v, c * v_t, head_steps(nb), v_t, false, u);
13806            part(&mut t, "w", k, c * qk_t, k_steps(), qk_t, true, w);
13807            finish(t, "uw", c, strip_items())?
13808        });
13809
13810        // ── Δ: `Δ[t, a] = U[t, a] − Σ_b W[t, b] S₀[a, b]`, one a-block (tile) at a time, joined into the head's strip.
13811        sections.push({
13812            let mut t = Template::new(target)?;
13813            let u = chunk_tensor(&mut t, c, "u", TensorRole::Input, hv * d);
13814            let w = chunk_tensor(&mut t, c, "w", TensorRole::Input, hv * d);
13815            let st = t.tensor(
13816                &state(c),
13817                TensorRole::Input,
13818                s_rows,
13819                d,
13820                precision,
13821                super::from_tensor_op::init_of(st_fill, s_rows * d),
13822            );
13823            let dt = chunk_tensor(&mut t, c, "d", TensorRole::Output, hv * d);
13824            let mut parts = Vec::new();
13825            for j in 0..nb {
13826                let (win, sin, uin) = (
13827                    t.buffer_deep(&format!("in_w{j}"), 1, nb, precision),
13828                    t.buffer_deep(&format!("in_s{j}"), 1, nb, precision),
13829                    t.buffer(&format!("in_u{j}"), 1, precision),
13830                );
13831                let lw =
13832                    t.load_block_at(&format!("load_w{j}"), w, 0, head_steps(nb), 1, nb, v_t, win);
13833                let ls = t.load_block_at(
13834                    &format!("load_s{j}"),
13835                    st,
13836                    j * nb,
13837                    head_steps(nb * nb),
13838                    1,
13839                    nb,
13840                    nb,
13841                    sin,
13842                );
13843                let lu =
13844                    t.load_block_at(&format!("load_u{j}"), u, j, head_steps(nb), 1, 1, v_t, uin);
13845                let ws = t.buffer(&format!("ws{j}"), 1, precision);
13846                let s1 = t.stage(
13847                    &format!("gd_ws{j}"),
13848                    generic(
13849                        contract(tr, tw, d),
13850                        vec![
13851                            m(&[(LV_GEN_ROW, wd as i64), (LV_GEN_K, 1)]),
13852                            m(&[(LV_GEN_COL, wd as i64), (LV_GEN_K, 1)]),
13853                        ],
13854                        ident_map(w1),
13855                        *mul(op(0), op(1)),
13856                    ),
13857                    vec![win, sin],
13858                    vec![ws],
13859                );
13860                let dj = t.scratch_deep(&format!("d{j}"), 1, 1, precision);
13861                let s2 = t.stage(
13862                    &format!("gd_d{j}"),
13863                    generic(
13864                        pointwise(tw),
13865                        vec![ident_map(w1), ident_map(w1)],
13866                        ident_map(w1),
13867                        *sub(op(0), op(1)),
13868                    ),
13869                    vec![uin, ws],
13870                    vec![dj],
13871                );
13872                t.edge(lw, win, s1);
13873                t.edge(ls, sin, s1);
13874                t.edge(s1, ws, s2);
13875                t.edge(lu, uin, s2);
13876                parts.push((s2, dj));
13877            }
13878            let strip = t.buffer_deep("d_strip", 1, nb, precision);
13879            let jmaps: Vec<ElemAccess> = (0..nb)
13880                .map(|j| ElemAccess {
13881                    base: -((j * tw) as i64),
13882                    terms: vec![(LV_GEN_ROW, w1 as i64), (LV_GEN_COL, 1)],
13883                    indirect: None,
13884                })
13885                .collect();
13886            let jbands: Vec<(usize, TensorBody)> = (0..nb)
13887                .map(|j| ((j + 1) * tw, TensorBody::Operand(j)))
13888                .collect();
13889            let join = t.stage(
13890                "d_join",
13891                generic(pointwise(wd), jmaps, ident_map(wd), banded(jbands)),
13892                parts.iter().map(|(_, b)| *b).collect(),
13893                vec![strip],
13894            );
13895            for (s, b) in &parts {
13896                t.edge(*s, *b, join);
13897            }
13898            let sto = t.store_block("store_d", dt, 0, head_steps(nb), 1, nb, v_t, strip);
13899            t.edge(join, strip, sto);
13900            finish(t, "delta", c, strip_items())?
13901        });
13902
13903        // ── The output rows: `o = (e^{g_t} q_t·S₀ᵀ + Σ_i QKΓ[t, i] Δ[i]) / √d`, the head's RMSNorm per row, `GN`, `silu(z)`.
13904        sections.push({
13905            let mut t = Template::new(target)?;
13906            let q = q_tensor(&mut t);
13907            let st = t.tensor(
13908                &state(c),
13909                TensorRole::Input,
13910                s_rows,
13911                d,
13912                precision,
13913                super::from_tensor_op::init_of(st_fill, s_rows * d),
13914            );
13915            let qk = chunk_tensor(&mut t, c, "qk", TensorRole::Input, hv * tw);
13916            let dt = chunk_tensor(&mut t, c, "d", TensorRole::Input, hv * d);
13917            let g = chunk_tensor(&mut t, c, "g", TensorRole::Input, hv * tw);
13918            let z = t.tensor(
13919                z_name,
13920                TensorRole::Input,
13921                rows,
13922                v_w,
13923                precision,
13924                super::from_tensor_op::init_of(z_fill, rows * v_w),
13925            );
13926            let gn = t.tensor(
13927                gn_name,
13928                TensorRole::Param,
13929                1,
13930                d,
13931                precision,
13932                super::from_tensor_op::init_of(gn_fill, d),
13933            );
13934            let y = t.tensor(
13935                &o_name,
13936                TensorRole::State,
13937                rows,
13938                v_w,
13939                precision,
13940                TensorInit::Zero,
13941            );
13942            let inv_sqrt_d = 1.0f64 / (head as f64).sqrt();
13943            let mut parts = Vec::new();
13944            for j in 0..nb {
13945                let (qin, sin) = (
13946                    t.buffer_deep(&format!("in_q{j}"), 1, nb, precision),
13947                    t.buffer_deep(&format!("in_s{j}"), 1, nb, precision),
13948                );
13949                let (qkin, din, gin) = (
13950                    t.buffer(&format!("in_qk{j}"), 1, precision),
13951                    t.buffer(&format!("in_d{j}"), 1, precision),
13952                    t.buffer(&format!("in_g{j}"), 1, precision),
13953                );
13954                let lq = t.load_block_at(
13955                    &format!("load_q{j}"),
13956                    q,
13957                    c * qk_t,
13958                    k_steps(),
13959                    1,
13960                    nb,
13961                    qk_t,
13962                    qin,
13963                );
13964                let ls = t.load_block_at(
13965                    &format!("load_s{j}"),
13966                    st,
13967                    j * nb,
13968                    head_steps(nb * nb),
13969                    1,
13970                    nb,
13971                    nb,
13972                    sin,
13973                );
13974                let lqk =
13975                    t.load_block_at(&format!("load_qk{j}"), qk, 0, head_steps(1), 1, 1, hv, qkin);
13976                let ld =
13977                    t.load_block_at(&format!("load_d{j}"), dt, j, head_steps(nb), 1, 1, v_t, din);
13978                let lg = t.load_block_at(&format!("load_g{j}"), g, 0, head_steps(1), 1, 1, hv, gin);
13979                let qs = t.buffer(&format!("qs{j}"), 1, precision);
13980                let s1 = t.stage(
13981                    &format!("gd_qs{j}"),
13982                    generic(
13983                        contract(tr, tw, d),
13984                        vec![
13985                            m(&[(LV_GEN_ROW, wd as i64), (LV_GEN_K, 1)]),
13986                            m(&[(LV_GEN_COL, wd as i64), (LV_GEN_K, 1)]),
13987                        ],
13988                        ident_map(w1),
13989                        *mul(op(0), op(1)),
13990                    ),
13991                    vec![qin, sin],
13992                    vec![qs],
13993                );
13994                let qd = t.buffer(&format!("qd{j}"), 1, precision);
13995                let s2 = t.stage(
13996                    &format!("gd_qd{j}"),
13997                    generic(
13998                        contract(tr, tw, tr),
13999                        vec![
14000                            m(&[(LV_GEN_ROW, w1 as i64), (LV_GEN_K, 1)]),
14001                            m(&[(LV_GEN_K, w1 as i64), (LV_GEN_COL, 1)]),
14002                        ],
14003                        ident_map(w1),
14004                        *mul(op(0), op(1)),
14005                    ),
14006                    vec![qkin, din],
14007                    vec![qd],
14008                );
14009                let oj = t.scratch_deep(&format!("o{j}"), 1, 1, precision);
14010                let s3 = t.stage(
14011                    &format!("gd_o{j}"),
14012                    generic(
14013                        pointwise(tw),
14014                        vec![ident_map(w1), ident_map(w1), m(&[(LV_GEN_ROW, w1 as i64)])],
14015                        ident_map(w1),
14016                        *mul(
14017                            add(mul(unary(SfpuKind::Exp, op(2)), op(0)), op(1)),
14018                            konst(inv_sqrt_d),
14019                        ),
14020                    ),
14021                    vec![qs, qd, gin],
14022                    vec![oj],
14023                );
14024                for (l, b, s) in [
14025                    (lq, qin, s1),
14026                    (ls, sin, s1),
14027                    (lqk, qkin, s2),
14028                    (ld, din, s2),
14029                    (lg, gin, s3),
14030                ] {
14031                    t.edge(l, b, s);
14032                }
14033                t.edge(s1, qs, s3);
14034                t.edge(s2, qd, s3);
14035                parts.push((s3, oj));
14036            }
14037            let acc = AccPolicy::F32.interm_precision();
14038            let ovec = t.scratch_deep("o_strip", 1, nb, acc);
14039            let jmaps: Vec<ElemAccess> = (0..nb)
14040                .map(|j| ElemAccess {
14041                    base: -((j * tw) as i64),
14042                    terms: vec![(LV_GEN_ROW, w1 as i64), (LV_GEN_COL, 1)],
14043                    indirect: None,
14044                })
14045                .collect();
14046            let jbands: Vec<(usize, TensorBody)> = (0..nb)
14047                .map(|j| ((j + 1) * tw, TensorBody::Operand(j)))
14048                .collect();
14049            let join = t.stage(
14050                "o_join",
14051                generic(pointwise(wd), jmaps, ident_map(wd), banded(jbands)),
14052                parts.iter().map(|(_, b)| *b).collect(),
14053                vec![ovec],
14054            );
14055            for (s, b) in &parts {
14056                t.edge(*s, *b, join);
14057            }
14058            // Each row's RMSNorm over the head: the row's sum of squares (a column fold, one value per row), its
14059            // reciprocal root, then the weight and the gate. The strip is a scratch the sum and the gate both read.
14060            let ssum = t.scratch("o_sum", 1, acc);
14061            let s_sum = t.stage(
14062                "o_sum",
14063                generic(
14064                    vec![
14065                        IterSpec {
14066                            var: LV_GEN_ROW,
14067                            extent: tr,
14068                            kind: IterKind::Parallel,
14069                        },
14070                        IterSpec {
14071                            var: LV_GEN_COL,
14072                            extent: d,
14073                            kind: IterKind::Reduce {
14074                                pool: ReducePool::Sum,
14075                            },
14076                        },
14077                    ],
14078                    vec![ident_map(wd), ident_map(wd)],
14079                    per_row_map(w1),
14080                    *mul(op(0), op(1)),
14081                ),
14082                vec![ovec, ovec],
14083                vec![ssum],
14084            );
14085            t.edge(join, ovec, s_sum);
14086            let inv = t.scratch("o_inv", 1, acc);
14087            let s_inv = t.stage(
14088                "o_rsqrt",
14089                generic(
14090                    pointwise(tw),
14091                    vec![per_row_map(w1)],
14092                    per_row_map(w1),
14093                    TensorBody::Unary {
14094                        kind: SfpuKind::Rsqrt {
14095                            eps_nano: eps_nano(spec.eps)?,
14096                            over: head,
14097                        },
14098                        arg: op(0),
14099                    },
14100                ),
14101                vec![ssum],
14102                vec![inv],
14103            );
14104            t.edge(s_sum, ssum, s_inv);
14105            let (gin, zin) = (
14106                t.buffer("in_gn", nb, precision),
14107                t.buffer_deep("in_z", 1, nb, precision),
14108            );
14109            let lgn = t.load("load_gn", gn, vec![], nb, gin);
14110            let lz = t.load_block_at("load_z", z, c * v_t, head_steps(nb), 1, nb, v_t, zin);
14111            let ob = t.buffer_deep("out", 1, nb, precision);
14112            let gate = t.stage(
14113                "gd_gate",
14114                generic(
14115                    pointwise(wd),
14116                    vec![
14117                        ident_map(wd),
14118                        per_row_map(w1),
14119                        m(&[(LV_GEN_COL, 1)]),
14120                        ident_map(wd),
14121                    ],
14122                    ident_map(wd),
14123                    *mul(mul(mul(op(0), op(1)), op(2)), delta_gate(spec.gate, op(3))),
14124                ),
14125                vec![ovec, inv, gin, zin],
14126                vec![ob],
14127            );
14128            t.edge(s_inv, inv, gate);
14129            t.edge(lgn, gin, gate);
14130            t.edge(lz, zin, gate);
14131            let sto = t.store_block("store_o", y, c * v_t, head_steps(nb), 1, nb, v_t, ob);
14132            t.edge(gate, ob, sto);
14133            let items = head_items(&|h| TileRect {
14134                row0_t: c,
14135                col0_t: h * nb,
14136                h_t: 1,
14137                w_t: nb,
14138            });
14139            finish(t, "out", c, items)?
14140        });
14141
14142        // ── The memory the chunk leaves: `S_C[a, b] = e^{g_C} S₀[a, b] + Σ_t Δ[t, a] e^{g_C − g_t} k_t[b]`, the head's
14143        // a-blocks as sub-blocks of its item.
14144        sections.push({
14145            let mut t = Template::new(target)?;
14146            let st = t.tensor(
14147                &state(c),
14148                TensorRole::Input,
14149                s_rows,
14150                d,
14151                precision,
14152                super::from_tensor_op::init_of(st_fill, s_rows * d),
14153            );
14154            let stout = t.tensor(
14155                &state(c + 1),
14156                TensorRole::Output,
14157                s_rows,
14158                d,
14159                precision,
14160                TensorInit::Zero,
14161            );
14162            let dt = chunk_tensor(&mut t, c, "d", TensorRole::Input, hv * d);
14163            let k = k_tensor(&mut t);
14164            let g = chunk_tensor(&mut t, c, "g", TensorRole::Input, hv * tw);
14165            let lg = chunk_tensor(&mut t, c, "lg", TensorRole::Input, hv * tw);
14166            let hold = |mut v: Vec<(LoopVar, usize)>| {
14167                v.push((super::LV_SUB_H, 0));
14168                v
14169            };
14170            let row_steps = vec![
14171                (LV_ITEM_Y, nb * nb),
14172                (LV_ITEM_X, hk * nb * nb),
14173                (super::LV_SUB_H, nb),
14174            ];
14175            let d_steps = vec![(LV_ITEM_Y, nb), (LV_ITEM_X, hk * nb), (super::LV_SUB_H, 1)];
14176            let (din, kin, lgin, sin, gin2) = (
14177                t.buffer_deep("in_d", 1, 1, precision),
14178                t.buffer_deep("in_k", 1, nb, precision),
14179                t.buffer_deep("in_lg", 1, 1, precision),
14180                t.buffer_deep("in_s", 1, nb, precision),
14181                t.buffer_deep("in_g_c", 1, 1, precision),
14182            );
14183            let ld = t.load_block_at("load_d", dt, 0, d_steps, 1, 1, v_t, din);
14184            let lk = t.load_block_at("load_k", k, c * qk_t, hold(k_steps()), 1, nb, qk_t, kin);
14185            let llg = t.load_block_at("load_lg", lg, 0, hold(head_steps(1)), 1, 1, hv, lgin);
14186            let ls = t.load_block_at("load_s", st, 0, row_steps.clone(), 1, nb, nb, sin);
14187            let lg2 = t.load_block_at("load_g_c", g, 0, hold(head_steps(1)), 1, 1, hv, gin2);
14188            let g_c = fixed(((cc - 1) * w1) as i64);
14189            // 0:Δ[t, a]  1:k_t[b]  2:LG[C − 1, t] = ln(γ_C / γ_t) — folded over the chunk's rows `t`.
14190            let body = TensorBody::Select {
14191                cond: Pred::AxisLt(LV_GEN_K, PredRhs::Const(cc)),
14192                then: mul(mul(op(0), op(1)), unary(SfpuKind::Exp, op(2))),
14193                otherwise: konst(0.0),
14194            };
14195            let upd = t.buffer_deep("upd", 1, nb, precision);
14196            let s1 = t.stage(
14197                "gd_s_update",
14198                generic(
14199                    contract(tr, d, tr),
14200                    vec![
14201                        m(&[(LV_GEN_K, w1 as i64), (LV_GEN_ROW, 1)]),
14202                        m(&[(LV_GEN_K, wd as i64), (LV_GEN_COL, 1)]),
14203                        ElemAccess {
14204                            base: ((cc - 1) * w1) as i64,
14205                            terms: vec![(LV_GEN_K, 1)],
14206                            indirect: None,
14207                        },
14208                    ],
14209                    ident_map(wd),
14210                    body,
14211                ),
14212                vec![din, kin, lgin],
14213                vec![upd],
14214            );
14215            let ob = t.buffer_deep("out", 1, nb, precision);
14216            let s2 = t.stage(
14217                "gd_s_decay",
14218                generic(
14219                    pointwise(d),
14220                    vec![ident_map(wd), ident_map(wd), g_c],
14221                    ident_map(wd),
14222                    *add(op(0), mul(unary(SfpuKind::Exp, op(2)), op(1))),
14223                ),
14224                vec![upd, sin, gin2],
14225                vec![ob],
14226            );
14227            for (l, b, s) in [
14228                (ld, din, s1),
14229                (lk, kin, s1),
14230                (llg, lgin, s1),
14231                (ls, sin, s2),
14232                (lg2, gin2, s2),
14233            ] {
14234                t.edge(l, b, s);
14235            }
14236            t.edge(s1, upd, s2);
14237            let sto = t.store_block("store_s", stout, 0, row_steps, 1, nb, nb, ob);
14238            t.edge(s2, ob, sto);
14239            let items = head_items(&|h| TileRect {
14240                row0_t: h * nb,
14241                col0_t: 0,
14242                h_t: nb,
14243                w_t: nb,
14244            });
14245            let mut prog = finish(t, "memory", c, items)?;
14246            prog.work.plan = Some(BlockPlan {
14247                per_site_m: nb,
14248                per_site_n: nb,
14249                grid: (hv, 1),
14250                inner: InnerLoop {
14251                    num_blocks: 1,
14252                    block_k_tiles: 1,
14253                },
14254                block: BlockShape {
14255                    osh: 1,
14256                    osw: nb,
14257                    ibw: 1,
14258                },
14259                acc: AccPolicy::F32,
14260                pack_accumulate: false,
14261                engines: EnginePacking::LoadsTogether,
14262                granularity: Granularity::PerTile,
14263            });
14264            prog.work.k_tiles = 1;
14265            prog
14266        });
14267    }
14268    if !o_dense {
14269        sections.push(relayout_dense_to_rows(
14270            (&o_name, &Fill::Unknown),
14271            &names.out,
14272            rows,
14273            v_w,
14274            precision,
14275            target,
14276        )?);
14277    }
14278    Ok(sections)
14279}
14280
14281/// A `Generic` compute stage's kind.
14282fn generic_op(iters: Vec<IterSpec>, ins: Vec<ElemAccess>, out: ElemAccess, body: TensorBody) -> StageKind {
14283    StageKind::Compute(ComputeOp::Generic { iters, ins, out, body })
14284}
14285
14286/// `C = A ⊕ B` over a flat block of `n / w` rows of `w`, `B` one row: a band per row, each reading `B` from its row's
14287/// start — the one pointwise stage, blocked by the result's columns past an accumulator page.
14288#[allow(clippy::too_many_arguments)]
14289pub(crate) fn flat_row_broadcast(
14290    kind: crate::tensor_recognize::ElementwiseOp,
14291    names: &super::from_tensor_op::ElementwiseNames,
14292    fills: (&Fill, &Fill),
14293    n: usize,
14294    w: usize,
14295    precision: Precision,
14296    claim: NumericsClaim,
14297    target: &SpatialTarget,
14298) -> Result<SpatialProgram, SpatialReject> {
14299    let mut b = FlatBuild::columns(n, target)?;
14300    b.precision = precision;
14301    let a_in = b.input(&names.a, n, fills.0, target);
14302    let row = b.input(&names.b, w, fills.1, target);
14303    let out = b.result_page("out", n, precision, target);
14304    let mut slots = vec![Slot { buf: a_in, map: at(0) }];
14305    let mut bands = Vec::new();
14306    for r in 0..n / w {
14307        slots.push(Slot { buf: row, map: at(-((r * w) as i64)) });
14308        let rhs = op(slots.len() - 1);
14309        let combined = match kind {
14310            crate::tensor_recognize::ElementwiseOp::Add => *add(op(0), rhs),
14311            crate::tensor_recognize::ElementwiseOp::Mul => *mul(op(0), rhs),
14312        };
14313        bands.push(((r + 1) * w, combined));
14314    }
14315    let chain = b.pointwise("elementwise_row", n, slots, banded(bands), out, target)?;
14316    b.finish_as(&names.c, n, chain, out, format!("elementwise_{}_row_{}x{w}", kind.c_op(), n / w), claim, precision, target)
14317}
14318
14319/// [`from_rope`] with the HEAD as the work item — see the note there for why that is the natural unit.
14320fn blocked_rope(
14321    spec: &RopeSpec,
14322    names: &StageNames,
14323    target: &SpatialTarget,
14324) -> Result<SpatialProgram, SpatialReject> {
14325    let (heads, hd, half, rows) = (spec.h, spec.hd, spec.half, spec.rows);
14326    let n = heads * hd;
14327    let head_t = tiles_of(hd, target);
14328    let half_t = tiles_of(half, target);
14329    // ★ A HEAD TURNS BY ITS ROW'S TABLE ROW. Item `(y, x)` is head `x` of query row `y`, so the query steps by a row's
14330    // heads per `y` and the tables by one table row per `y` — which is a whole number of tiles only when `half` is.
14331    let per_row = heads / rows.max(1);
14332    // ★ A TABLE ROW THAT IS NOT WHOLE TILES IS READ THROUGH A WINDOW. Row `ρ` starts at element `ρ·half`, mid-tile; the
14333    // item loads the whole tiles around it and reads the row at its offset in the page — both from the row's first
14334    // element, computed per item.
14335    let windowed = (rows > 1 || spec.from.is_some()) && half % target.tile.cols != 0;
14336    let table_rows = spec.from.map_or(rows, |(_, cache)| cache);
14337    let row_steps = |per_y: usize, mut v: Vec<(LoopVar, usize)>| {
14338        if rows > 1 {
14339            v.insert(0, (LV_ITEM_Y, per_y));
14340        }
14341        v
14342    };
14343    let precision = elem_precision(target)?;
14344    let mut t = Template::new(target)?;
14345    let (qn, qf) = names.operand(0);
14346    let (cn, cf) = names.operand(1);
14347    let (sn, sf) = names.operand(2);
14348    let qt = t.tensor(
14349        qn,
14350        TensorRole::Input,
14351        1,
14352        n,
14353        precision,
14354        super::from_tensor_op::init_of(qf, n),
14355    );
14356    let ct = t.tensor(
14357        cn,
14358        TensorRole::Input,
14359        1,
14360        table_rows * half,
14361        precision,
14362        super::from_tensor_op::init_of(cf, table_rows * half),
14363    );
14364    let st = t.tensor(
14365        sn,
14366        TensorRole::Input,
14367        1,
14368        table_rows * half,
14369        precision,
14370        super::from_tensor_op::init_of(sf, table_rows * half),
14371    );
14372    let yt = t.tensor(
14373        &names.out,
14374        TensorRole::Output,
14375        1,
14376        n,
14377        precision,
14378        TensorInit::Zero,
14379    );
14380    let tw = target.tile.cols;
14381    let table_t = tiles_of(table_rows * half, target);
14382    let window_t = if windowed { (tw - 1 + half).div_ceil(tw).min(table_t) } else { half_t };
14383    let qin = t.buffer_deep("in_q", 1, head_t, precision);
14384    let cin = t.buffer_deep("in_rc", 1, window_t, precision);
14385    let sin = t.buffer_deep("in_rs", 1, window_t, precision);
14386    let ob = t.buffer_deep("out", 1, head_t, precision);
14387    // One table row is the SAME for every head of its row, so the tables' tile address moves only with the row — and,
14388    // reading a cache, starts at the position's row `P − 1`: a literal in the base; a named one an [`IndirectTile`] over
14389    // `P − 1`, computed per item on the integer datapath into a FIFO its load retires (one per table: a FIFO has one
14390    // reader) — behind a guard of stride zero that keeps the block's last row among the cache's, as a row index is.
14391    let lq = t.load_block_at(
14392        "load_q",
14393        qt,
14394        0,
14395        row_steps(per_row * head_t, vec![(LV_ITEM_X, head_t)]),
14396        1,
14397        head_t,
14398        head_t,
14399        qin,
14400    );
14401    let position = match spec.from {
14402        Some((RowIndex::Named(_), cache)) => {
14403            let (pn, pf) = names.operand(3);
14404            let ix = index_type(target, index_space(cache as u64, 1, pf))?;
14405            Some((
14406                t.tensor(
14407                    pn,
14408                    TensorRole::Input,
14409                    1,
14410                    1,
14411                    ix,
14412                    super::from_tensor_op::init_of(pf, 1),
14413                ),
14414                ix,
14415                cache,
14416            ))
14417        }
14418        _ => None,
14419    };
14420    let base = match spec.from {
14421        Some((RowIndex::Literal(p), cache)) => {
14422            if p < 1 || p as usize + rows - 1 > cache {
14423                return Err(SpatialReject::NotLowered("a literal rope position whose block leaves the cache — the program fails at run time on every tier, so there is no section to lower"));
14424            }
14425            (p as usize - 1) * half_t
14426        }
14427        _ => 0,
14428    };
14429    let table_load = |t: &mut Template, name: &str, tensor: TensorId, out: BufferId| -> StageId {
14430        let Some((pt, ix, cache)) = position else {
14431            return t.load_block_at(
14432                name,
14433                tensor,
14434                base,
14435                row_steps(half_t, vec![]),
14436                1,
14437                half_t,
14438                half_t,
14439                out,
14440            );
14441        };
14442        let w1 = wide(1, target);
14443        let pin = t.buffer(&format!("{name}_position"), 1, ix);
14444        let lp = t.load(&format!("{name}_position"), pt, vec![], 1, pin);
14445        let biased = t.scratch_deep(&format!("{name}_biased"), 1, 1, ix);
14446        let sb = t.stage(
14447            &format!("{name}_biased"),
14448            StageKind::Compute(ComputeOp::Generic {
14449                iters: flat_space(1),
14450                ins: vec![fixed(0)],
14451                out: ident_map(w1),
14452                body: *sub(op(0), konst(1.0)),
14453            }),
14454            vec![pin],
14455            vec![biased],
14456        );
14457        t.edge(lp, pin, sb);
14458        let row0 = t.buffer(&format!("{name}_row"), 1, ix);
14459        let guard = ElemAccess {
14460            base: 0,
14461            terms: Vec::new(),
14462            indirect: Some(IndirectTerm {
14463                operand: 0,
14464                stride: 0,
14465                modulus: cache - rows + 1,
14466            }),
14467        };
14468        let s = t.stage(
14469            &format!("{name}_row"),
14470            StageKind::Compute(ComputeOp::Generic {
14471                iters: flat_space(1),
14472                ins: vec![guard],
14473                out: ident_map(w1),
14474                body: TensorBody::Operand(0),
14475            }),
14476            vec![biased],
14477            vec![row0],
14478        );
14479        t.edge(sb, biased, s);
14480        let origin = TileAddr {
14481            base: 0,
14482            steps: row_steps(half_t, vec![]),
14483            indirect: Some(IndirectTile {
14484                index: row0,
14485                tile_stride: half_t,
14486                elems_per_tile: 1,
14487                bound: (0, cache - rows),
14488            }),
14489        };
14490        let l = t.stage(
14491            name,
14492            StageKind::Load {
14493                tensor,
14494                pattern: AccessPattern {
14495                    origin,
14496                    block_h: 1,
14497                    block_w: half_t,
14498                    stride_w: 1,
14499                    stride_h: half_t,
14500                },
14501                num_blocks: 1,
14502                multicast: None,
14503            },
14504            vec![row0],
14505            vec![out],
14506        );
14507        t.edge(s, row0, l);
14508        l
14509    };
14510    // The windowed table read: item `(y, ·)`'s row starts at element `e = (P₀ + y)·half` — `P₀` the position less one
14511    // (checked, as a row index is) or the literal's, `y·half` from a compile-time table one tile per query row — its
14512    // window at tile `min(⌊e/32⌋, T − W)`, so no window reads past the table, and the row at `e − 32·start` in the page.
14513    let ix_rows = index_type(target, (table_rows * half + tw) as u64)?;
14514    let rows_table = windowed.then(|| {
14515        let literal = match spec.from {
14516            Some((RowIndex::Literal(p), _)) => p as usize - 1,
14517            _ => 0,
14518        };
14519        let tiles: Vec<i64> = (0..rows).flat_map(|y| std::iter::once(((literal + y) * half) as i64).chain(std::iter::repeat_n(0, tw - 1))).collect();
14520        t.tensor(&format!("{}__rope_rows", names.out), TensorRole::Param, 1, rows * tw, ix_rows, TensorInit::Values(tiles))
14521    });
14522    let first_element = |t: &mut Template, name: &str| -> Result<(Vec<StageId>, BufferId, Precision), SpatialReject> {
14523        let ix = ix_rows;
14524        let rows_t = rows_table.expect("a windowed rope declares its rows");
14525        let at_row = t.buffer_deep(&format!("{name}_row_start"), 1, 1, ix);
14526        let steps = if rows > 1 { vec![(LV_ITEM_Y, 1)] } else { Vec::new() };
14527        let l_row = t.load_block_at(&format!("{name}_row_start"), rows_t, 0, steps, 1, 1, 1, at_row);
14528        let w1 = wide(1, target);
14529        let e = t.scratch_deep(&format!("{name}_first"), 1, 1, ix);
14530        let mut stages = vec![l_row];
14531        match position {
14532            None => {
14533                let s = t.stage(&format!("{name}_first"), generic_op(flat_space(1), vec![fixed(0)], ident_map(w1), TensorBody::Operand(0)), vec![at_row], vec![e]);
14534                t.edge(l_row, at_row, s);
14535                stages.push(s);
14536            }
14537            Some((pt, pix, cache)) => {
14538                let pin = t.buffer(&format!("{name}_position"), 1, pix);
14539                let lp = t.load(&format!("{name}_position"), pt, vec![], 1, pin);
14540                let biased = t.scratch_deep(&format!("{name}_biased"), 1, 1, pix);
14541                let s_b = t.stage(&format!("{name}_biased"), generic_op(flat_space(1), vec![fixed(0)], ident_map(w1), *sub(op(0), konst(1.0))), vec![pin], vec![biased]);
14542                t.edge(lp, pin, s_b);
14543                // `P − 1`, inside the rows a block of `rows` may start at — the stride-zero guard a row index gets.
14544                let guard = ElemAccess { base: 0, terms: Vec::new(), indirect: Some(IndirectTerm { operand: 0, stride: 0, modulus: cache - rows + 1 }) };
14545                let s = t.stage(
14546                    &format!("{name}_first"),
14547                    generic_op(flat_space(1), vec![guard, fixed(0)], ident_map(w1), *add(mul(op(0), konst(half as f64)), op(1))),
14548                    vec![biased, at_row],
14549                    vec![e],
14550                );
14551                t.edge(s_b, biased, s);
14552                t.edge(l_row, at_row, s);
14553                stages.extend([lp, s_b, s]);
14554            }
14555        }
14556        Ok((stages, e, ix))
14557    };
14558    let window_start = |t: &mut Template, name: &str, e: BufferId, ix: Precision| -> (StageId, BufferId) {
14559        let start = t.scratch_deep(&format!("{name}_window"), 1, 1, ix);
14560        let last = (table_t - window_t) as f64;
14561        let body = TensorBody::Min(Box::new(TensorBody::FloorDiv(op(0), konst(tw as f64))), konst(last));
14562        let s = t.stage(&format!("{name}_window"), generic_op(flat_space(1), vec![fixed(0)], ident_map(wide(1, target)), body), vec![e], vec![start]);
14563        (s, start)
14564    };
14565    let windowed_load = |t: &mut Template, name: &str, tensor: TensorId, out: BufferId| -> Result<(StageId, Vec<StageId>), SpatialReject> {
14566        let (mut stages, e, ix) = first_element(t, name)?;
14567        let (s_start, start) = window_start(t, name, e, ix);
14568        t.edge(*stages.last().expect("a first element"), e, s_start);
14569        stages.push(s_start);
14570        let l = t.stage(
14571            name,
14572            StageKind::Load {
14573                tensor,
14574                pattern: AccessPattern {
14575                    origin: TileAddr {
14576                        base: 0,
14577                        steps: Vec::new(),
14578                        indirect: Some(IndirectTile { index: start, tile_stride: 1, elems_per_tile: 1, bound: (0, table_t - window_t) }),
14579                    },
14580                    block_h: 1,
14581                    block_w: window_t,
14582                    stride_w: 1,
14583                    stride_h: window_t,
14584                },
14585                num_blocks: 1,
14586                multicast: None,
14587            },
14588            vec![start],
14589            vec![out],
14590        );
14591        t.edge(s_start, start, l);
14592        Ok((l, stages))
14593    };
14594    let (lc, ls, offset) = if windowed {
14595        let (lc, _) = windowed_load(&mut t, "load_rc", ct, cin)?;
14596        let (ls, _) = windowed_load(&mut t, "load_rs", st, sin)?;
14597        // The row's offset in its window: `e − 32·start`, from a first element of its own.
14598        let (stages, e, ix) = first_element(&mut t, "rope_offset")?;
14599        let (s_start, start) = window_start(&mut t, "rope_offset", e, ix);
14600        t.edge(*stages.last().expect("a first element"), e, s_start);
14601        let off = t.scratch_deep("rope_offset", 1, 1, ix);
14602        let s_off = t.stage(
14603            "rope_offset",
14604            generic_op(flat_space(1), vec![fixed(0), fixed(0)], ident_map(wide(1, target)), *sub(op(0), mul(op(1), konst(tw as f64)))),
14605            vec![e, start],
14606            vec![off],
14607        );
14608        t.edge(s_start, start, s_off);
14609        (lc, ls, Some((s_off, off)))
14610    } else {
14611        (table_load(&mut t, "load_rc", ct, cin), table_load(&mut t, "load_rs", st, sin), None)
14612    };
14613    let w = wide(head_t, target);
14614    // Inside one head the three regions are the rotation's own: rotate the low half against the high,
14615    // the high half against the low, and pass through whatever the rotary does not cover.
14616    let mut slots: Vec<ElemAccess> = vec![
14617        at(0),              // 0: q[c]
14618        at(half as i64),    // 1: q[c + half]
14619        at(-(half as i64)), // 2: q[c - half]
14620        at(0),              // 3: rc[c]
14621        at(0),              // 4: rs[c]
14622        at(-(half as i64)), // 5: rs[c - half]
14623        at(-(half as i64)), // 6: rc[c - half]
14624    ];
14625    let mut ins = vec![qin, qin, qin, cin, sin, sin, cin];
14626    // A windowed table is read at its row's offset in the page: operand 7.
14627    if let Some((_, off)) = offset {
14628        for slot in &mut slots[3..7] {
14629            slot.indirect = Some(IndirectTerm { operand: 7, stride: 1, modulus: window_t * tw });
14630        }
14631        slots.push(fixed(0));
14632        ins.push(off);
14633    }
14634    let mut bands = vec![
14635        (half, *sub(mul(op(0), op(3)), mul(op(1), op(4)))),
14636        (2 * half, *add(mul(op(2), op(5)), mul(op(0), op(6)))),
14637    ];
14638    if 2 * half < hd {
14639        bands.push((hd, TensorBody::Operand(0)));
14640    }
14641    let s = t.stage(
14642        "rope",
14643        StageKind::Compute(ComputeOp::Generic {
14644            iters: flat_space(w),
14645            ins: slots,
14646            out: ident_map(w),
14647            body: banded(bands),
14648        }),
14649        ins,
14650        vec![ob],
14651    );
14652    let store = t.store_block(
14653        "store_y",
14654        yt,
14655        0,
14656        row_steps(per_row * head_t, vec![(LV_ITEM_X, head_t)]),
14657        1,
14658        head_t,
14659        head_t,
14660        ob,
14661    );
14662    for l in [lq, lc, ls] {
14663        let b = match l {
14664            x if x == lq => qin,
14665            x if x == lc => cin,
14666            _ => sin,
14667        };
14668        t.edge(l, b, s);
14669    }
14670    if let Some((s_off, off)) = offset {
14671        t.edge(s_off, off, s);
14672    }
14673    t.edge(s, ob, store);
14674    let items: Vec<WorkItem> = (0..heads)
14675        .map(|h| {
14676            let coord = if rows > 1 {
14677                (h / per_row, h % per_row)
14678            } else {
14679                (0, h)
14680            };
14681            WorkItem {
14682                id: WorkItemId::new(h),
14683                coord,
14684                out: TileRect {
14685                    row0_t: 0,
14686                    col0_t: h * head_t,
14687                    h_t: 1,
14688                    w_t: head_t,
14689                },
14690            }
14691        })
14692        .collect();
14693    t.finish(
14694        rope_name(spec),
14695        items,
14696        AccPolicy::F32,
14697        approx("rotary embedding", precision),
14698        target,
14699    )
14700}
14701
14702/// [`from_row_norm_flat`] with the HEAD as the work item: `y[h, i] = x[h, i]·w[i]·rsqrt(mean_i x[h, i]²)`.
14703/// The weight is a per-COLUMN table shared by every head, so its tile address does not move with the item.
14704fn blocked_head_rms(
14705    spec: &FusedRowSpec,
14706    names: &StageNames,
14707    target: &SpatialTarget,
14708) -> Result<SpatialProgram, SpatialReject> {
14709    let (rows, cols) = (spec.rows, spec.cols);
14710    let n = rows * cols;
14711    let head_t = tiles_of(cols, target);
14712    let precision = elem_precision(target)?;
14713    let acc = AccPolicy::F32.interm_precision();
14714    let mut t = Template::new(target)?;
14715    let (xn, xf) = names.operand(0);
14716    let xt = t.tensor(
14717        xn,
14718        TensorRole::Input,
14719        1,
14720        n,
14721        precision,
14722        super::from_tensor_op::init_of(xf, n),
14723    );
14724    let yt = t.tensor(
14725        &names.out,
14726        TensorRole::Output,
14727        1,
14728        n,
14729        precision,
14730        TensorInit::Zero,
14731    );
14732    let weight = spec.weight.map(|_| names.operand(1));
14733    let ones = t.tensor(
14734        &format!("{}__rms_scaler", names.out),
14735        TensorRole::Param,
14736        target.tile.rows,
14737        target.tile.cols,
14738        precision,
14739        TensorInit::ConstF(1.0f64.to_bits()),
14740    );
14741    let xin = t.buffer_deep("in_x", 1, head_t, precision);
14742    let oin = t.buffer_deep("in_scaler", 1, 1, precision);
14743    let sq = t.scratch_deep("head_sq", 1, head_t, acc);
14744    let ssq = t.scratch_deep("head_ssq", 1, 1, acc);
14745    let inv = t.scratch_deep("head_inv", 1, 1, acc);
14746    let ob = t.buffer_deep("out", 1, head_t, precision);
14747    let step = vec![(LV_ITEM_X, head_t)];
14748    let lx = t.load_block_at("load_head", xt, 0, step.clone(), 1, head_t, head_t, xin);
14749    let lo = t.load_block_at("load_scaler", ones, 0, vec![], 1, 1, 1, oin);
14750    let w = wide(head_t, target);
14751    let one = wide(1, target);
14752    let s_sq = t.stage(
14753        "square",
14754        StageKind::Compute(ComputeOp::Generic {
14755            iters: head_space(1, cols, None),
14756            ins: vec![map2(0, cols as i64, 1), map2(0, cols as i64, 1)],
14757            out: ident_map(w),
14758            body: *mul(op(0), op(1)),
14759        }),
14760        vec![xin, xin],
14761        vec![sq],
14762    );
14763    let s_sum = t.stage(
14764        "head_sum",
14765        StageKind::Compute(ComputeOp::Generic {
14766            iters: head_space(1, cols, Some(ReducePool::Sum)),
14767            ins: vec![map2(0, w as i64, 1), const_map()],
14768            out: per_row_map(one),
14769            body: TensorBody::Operand(0),
14770        }),
14771        vec![sq, oin],
14772        vec![ssq],
14773    );
14774    let s_inv = t.stage(
14775        "head_rsqrt",
14776        StageKind::Compute(ComputeOp::Generic {
14777            iters: head_space(1, target.tile.cols, None),
14778            ins: vec![per_row_map(one)],
14779            out: per_row_map(one),
14780            // `over: cols` — the divide by the head width happens ONCE, inside the transform, which is the
14781            // arrangement `from_fused_row` documents and measured a bf16 ulp on when it did not.
14782            body: TensorBody::Unary {
14783                kind: SfpuKind::Rsqrt {
14784                    eps_nano: eps_nano(spec.eps)?,
14785                    over: cols,
14786                },
14787                arg: op(0),
14788            },
14789        }),
14790        vec![ssq],
14791        vec![inv],
14792    );
14793    let (wbuf, lw) = match weight {
14794        Some((wn, wf)) => {
14795            let wt = t.tensor(
14796                wn,
14797                TensorRole::Param,
14798                1,
14799                cols,
14800                precision,
14801                super::from_tensor_op::init_of(wf, cols),
14802            );
14803            let b = t.buffer_deep("in_w", 1, head_t, precision);
14804            let l = t.load_block_at("load_w", wt, 0, vec![], 1, head_t, head_t, b);
14805            (Some(b), Some(l))
14806        }
14807        None => (None, None),
14808    };
14809    let mut ins = vec![xin, inv];
14810    let mut maps = vec![map2(0, cols as i64, 1), fixed(0)];
14811    let mut body = mul(op(0), op(1));
14812    if let Some(b) = wbuf {
14813        ins.push(b);
14814        maps.push(map2(0, cols as i64, 1));
14815        body = mul(body, op(ins.len() - 1));
14816    }
14817    let s_out = t.stage(
14818        "rescale",
14819        StageKind::Compute(ComputeOp::Generic {
14820            iters: head_space(1, cols, None),
14821            ins: maps,
14822            out: ident_map(w),
14823            body: *body,
14824        }),
14825        ins,
14826        vec![ob],
14827    );
14828    let store = t.store_block("store_y", yt, 0, step, 1, head_t, head_t, ob);
14829    t.edge(lx, xin, s_sq);
14830    t.edge(lo, oin, s_sum);
14831    t.edge(s_sq, sq, s_sum);
14832    t.edge(s_sum, ssq, s_inv);
14833    t.edge(s_inv, inv, s_out);
14834    if let (Some(l), Some(b)) = (lw, wbuf) {
14835        t.edge(l, b, s_out);
14836    }
14837    t.edge(s_out, ob, store);
14838    let items: Vec<WorkItem> = (0..rows)
14839        .map(|h| WorkItem {
14840            id: WorkItemId::new(h),
14841            coord: (0, h),
14842            out: TileRect {
14843                row0_t: 0,
14844                col0_t: h * head_t,
14845                h_t: 1,
14846                w_t: head_t,
14847            },
14848        })
14849        .collect();
14850    t.finish(
14851        format!("rmsnorm_heads_{rows}x{cols}"),
14852        items,
14853        AccPolicy::F32,
14854        approx("per-head rmsnorm", precision),
14855        target,
14856    )
14857}