Skip to main content

logicaffeine_compile/spatial/lower/
spatial_lang.rs

1//! **Spatial source emission** for a VERIFIED `SpatialProgram` — the RDU's own front-door language (Spatial,
2//! Koeplinger et al. PLDI 2018; the Plasticine/RDU lineage; the public Spatial → PIR → plastiroute → Tungsten
3//! chain built by `scripts/plasticine/`). The `SpatialProgram` twin of `codegen_sva::spatial_emit`,
4//! which prints a `LoopFabric`.
5//!
6//! Every construct printed is the image of one IR object:
7//!   * a tensor is a `DRAM[T]`, filled on the host from its recorded fill;
8//!   * a buffer is an `SRAM[T]` sized by the page its producer fills (Spatial N-buffers a memory written in
9//!     one controller and read in the next, so the IR's depth is recorded beside the declaration rather
10//!     than spelled), or a `FIFO[T]` when its edge is a `Fifo` handshake;
11//!   * a stage is one controller: a Load is a `load` (DRAM → SRAM), a Store a `store`, a `MatmulTiles`
12//!     the row/column `Foreach` nest around a lane-parallel `Reduce` dot product — inside a `MemReduce`
13//!     over the K blocks when the IR spills partials through an `Interm` buffer — and a `Generic` an
14//!     element-wise `Foreach` whose reduce axis is a `Reduce`, its operands read through their access maps
15//!     with an absent axis reading coordinate 0 (the broadcast), exactly as `spatial::eval` reads them;
16//!   * the host checks the output's CHECKSUM against the gold: the spatial oracle's answer
17//!     (`eval_choreography` over the recorded fills) unless the caller hands over the VM's.
18//! The handshake ops (`Reserve`/`Push`/`Wait`/`Pop`, the accumulator acquire/release, the barriers) have no
19//! spelling of their own: Spatial's compiler derives the buffering and the pipeline from the controller
20//! nest, which is why a depth is a comment and not a construct.
21//!
22//! What is not spelled is refused BY NAME (`SpatialLowerError::Unsupported`): multicast groups and `Recv`
23//! stages, `Scalar` bodies, the fabric's fused epilogues, external or vendor-random fills, weight-stationary
24//! dataflow, cross-site handshakes, non-rectangular tile blocks, aliased buffers, sites carrying several
25//! work items. A multi-site program prints as a `Parallel` of per-site `Pipe`s — spelled, not yet gated
26//! through the chain (the gate runs the single-site RDU GEMM).
27
28use std::collections::{BTreeMap, BTreeSet};
29
30use super::{require_verified, SpatialLowerError};
31use crate::codegen_sva::fabric_target::Precision;
32use crate::spatial::eval::{eval_choreography, Inputs, Value};
33use crate::spatial::ir::*;
34use crate::spatial::passes::{LV_ITEM, LV_ITEM_X, LV_ITEM_Y, LV_KBLOCK, LV_SUB_H, LV_SUB_W};
35use crate::spatial::target::{DataflowPolicy, SpatialTarget, TileShape};
36
37/// The Spatial scalar the whole app computes in. `Int` when every value the program touches provably fits
38/// 32 bits (what the chain's GEMM regression apps use), `Long` otherwise, `Float` for any float precision.
39#[derive(Clone, Copy, Debug, PartialEq, Eq)]
40enum Scalar {
41    Int,
42    Long,
43    Float,
44}
45
46impl Scalar {
47    fn spatial(self) -> &'static str {
48        match self {
49            Scalar::Int => "Int",
50            Scalar::Long => "Long",
51            Scalar::Float => "Float",
52        }
53    }
54    fn is_float(self) -> bool {
55        self == Scalar::Float
56    }
57}
58
59/// Which half of the public chain an app is printed for — the two differ only in how a fold is spelled.
60#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub enum SpatialFlavor {
62    /// PIR and the fabric: every fold a `Reduce`, which PIR maps onto a PCU's reduction tree.
63    Fabric,
64    /// Spatial's Scala simulator: every fold an explicit register scan — seeded by the first point, combined in
65    /// order, as the oracle folds. This Spatial's Scala backend drops the first-iteration flags of an unrolled
66    /// `Reduce` when it splits the node's kernel out (the kernel reads flags it is never passed).
67    Simulator,
68}
69
70/// The Spatial type a memory is declared with: `T`, the app's scalar, or `I` — the integer scalar of a program that
71/// holds floats and integers both (its indices, positions and tokens). A program of one class declares everything `T`.
72#[derive(Clone, Copy, Debug, PartialEq, Eq)]
73enum Ty {
74    T,
75    I,
76}
77
78impl Ty {
79    fn name(self) -> &'static str {
80        match self {
81            Ty::T => "T",
82            Ty::I => "I",
83        }
84    }
85}
86
87#[derive(Clone, Copy, Debug, PartialEq)]
88enum Gold {
89    Int(i64),
90    Float(f64),
91}
92
93/// A stage's shape, read ONCE from the IR (this is the one place the compute enum is spelled): every
94/// printer below matches on this view, so the case analysis over compute shapes has exactly one home.
95enum Shape<'a> {
96    Load,
97    Store,
98    Recv,
99    Matmul {
100        block: &'a BlockShape,
101        sub_blocks: (usize, usize),
102        num_blocks: usize,
103        operand: Precision,
104        acc: AccPolicy,
105        dataflow: DataflowPolicy,
106        interm: Option<BufferId>,
107    },
108    Generic {
109        iters: &'a [IterSpec],
110        ins: &'a [ElemAccess],
111        out: &'a ElemAccess,
112        body: &'a TensorBody,
113    },
114    /// The fabric's fused epilogues (requant, bias, relu).
115    Epilogue,
116    Scalar,
117    /// A cross-chip reduction of the partials of a tensor-parallel shard.
118    AllReduce,
119}
120
121fn shape(s: &Stage) -> Shape<'_> {
122    match &s.kind {
123        StageKind::Load { .. } => Shape::Load,
124        StageKind::Store { .. } => Shape::Store,
125        StageKind::Recv { .. } => Shape::Recv,
126        StageKind::AllReduce { .. } => Shape::AllReduce,
127        StageKind::Scalar(_) => Shape::Scalar,
128        StageKind::Compute(ComputeOp::MatmulTiles {
129            block,
130            sub_blocks,
131            num_blocks,
132            operand,
133            acc,
134            dataflow,
135            interm,
136            ..
137        }) => Shape::Matmul {
138            block,
139            sub_blocks: *sub_blocks,
140            num_blocks: *num_blocks,
141            operand: *operand,
142            acc: *acc,
143            dataflow: *dataflow,
144            interm: *interm,
145        },
146        StageKind::Compute(ComputeOp::Generic {
147            iters,
148            ins,
149            out,
150            body,
151        }) => Shape::Generic {
152            iters,
153            ins,
154            out,
155            body,
156        },
157        StageKind::Compute(_) => Shape::Epilogue,
158    }
159}
160
161fn ident(s: &str) -> String {
162    s.chars()
163        .map(|c| if c.is_ascii_alphanumeric() { c } else { '_' })
164        .collect()
165}
166
167fn dram_name(t: &TensorDecl) -> String {
168    format!("d_{}", ident(&t.name))
169}
170
171/// Every buffer's Scala identifier: its name, and its id too where two buffers share a name — a section whose
172/// work items share a site clones its buffers under one name, and one Scala scope holds them all.
173fn buffer_names(p: &SpatialProgram) -> BTreeMap<BufferId, String> {
174    let mut uses: BTreeMap<String, usize> = BTreeMap::new();
175    for b in &p.buffers {
176        *uses.entry(ident(&b.name)).or_default() += 1;
177    }
178    p.buffers
179        .iter()
180        .map(|b| {
181            (
182                b.id,
183                if uses[&ident(&b.name)] > 1 {
184                    format!("{}_b{}", ident(&b.name), b.id.idx())
185                } else {
186                    ident(&b.name)
187                },
188            )
189        })
190        .collect()
191}
192
193fn is_float_precision(p: Precision) -> bool {
194    matches!(
195        p,
196        Precision::Fp8 | Precision::Bf16 | Precision::Fp16 | Precision::Fp32
197    )
198}
199
200/// Whether any constant of `body` satisfies `pred`.
201fn body_has_const(body: &TensorBody, pred: &dyn Fn(u64) -> bool) -> bool {
202    match body {
203        TensorBody::Const(bits) => pred(*bits),
204        TensorBody::Operand(_) => false,
205        TensorBody::Add(x, y)
206        | TensorBody::Sub(x, y)
207        | TensorBody::Mul(x, y)
208        | TensorBody::Min(x, y)
209        | TensorBody::Max(x, y)
210        | TensorBody::FloorDiv(x, y) => body_has_const(x, pred) || body_has_const(y, pred),
211        TensorBody::Unary { arg, .. } | TensorBody::Shr { arg, .. } => body_has_const(arg, pred),
212        TensorBody::Select {
213            then, otherwise, ..
214        } => body_has_const(then, pred) || body_has_const(otherwise, pred),
215    }
216}
217
218/// A literal of type `ty` — a float Scala parses as a `Double` when `float`, else an integer — lifted to it.
219fn lit_of(v: f64, ty: Ty, float: bool) -> String {
220    let s = if float {
221        format!("{v:?}")
222    } else {
223        (v as i64).to_string()
224    };
225    if v < 0.0 {
226        format!("({s}).to[{}]", ty.name())
227    } else {
228        format!("{s}.to[{}]", ty.name())
229    }
230}
231
232/// `expr`, of type `from`, as a `to`.
233fn cast(expr: String, from: Ty, to: Ty) -> String {
234    if from == to {
235        expr
236    } else {
237        format!("{expr}.to[{}]", to.name())
238    }
239}
240
241/// A float literal Scala parses as a `Double`, lifted to `T`.
242fn float_lit(v: f64) -> String {
243    let s = format!("{v:?}");
244    if v < 0.0 {
245        format!("({s}).to[T]")
246    } else {
247        format!("{s}.to[T]")
248    }
249}
250
251/// A 64-bit host literal for the checksum comparison: `(vL: Long)` is legal for every `i64` and lifted by
252/// Spatial's `FixFromLong` (a bare `.to[Long]` on a wide literal is "integer number too large").
253fn long_lit(v: i64) -> String {
254    format!("({v}L: Long)")
255}
256
257/// The largest lane count that divides `n`, at most the PCU's 16 SIMD lanes — the `par` of the dot product.
258fn lanes(n: usize) -> usize {
259    (1..=16.min(n.max(1)))
260        .rev()
261        .find(|l| n % l == 0)
262        .unwrap_or(1)
263}
264
265// ── Refusals ──────────────────────────────────────────────────────────────────────────────────────
266
267/// Every construct the emitter does not spell, refused by name before anything is printed.
268fn refuse_unspelled(p: &SpatialProgram) -> Result<(), SpatialLowerError> {
269    let pl = p.placement.as_ref().expect("a verified program is placed");
270    let non_finite = |bits: u64| !f64::from_bits(bits).is_finite();
271    let fills_non_finite = |t: &TensorDecl| match &t.init {
272        TensorInit::ConstF(b) => non_finite(*b),
273        TensorInit::FloatValues(v) => v
274            .iter()
275            .take(t.elems.rows * t.elems.cols)
276            .any(|b| non_finite(*b)),
277        _ => false,
278    };
279    let body_non_finite = |s: &Stage| match shape(s) {
280        Shape::Generic { body, .. } => body_has_const(body, &non_finite),
281        _ => false,
282    };
283    if p.tensors.iter().any(fills_non_finite) || p.stages.iter().any(body_non_finite) {
284        return Err(SpatialLowerError::Unsupported("a non-finite float constant (±∞ or NaN, in a body or a fill) is not spelled — Spatial has no verified FltPt literal for one"));
285    }
286    let multicast = !pl.groups.is_empty()
287        || p.stages.iter().any(|s| {
288            matches!(s.kind, StageKind::Recv { .. })
289                || matches!(
290                    s.kind,
291                    StageKind::Load {
292                        multicast: Some(_),
293                        ..
294                    }
295                )
296        });
297    if multicast {
298        return Err(SpatialLowerError::Unsupported("multicast groups (a Load that broadcasts a block to a rectangle of sites, and the Recv stages that receive it) are not spelled — Spatial has no group broadcast; place the program without reader fusion"));
299    }
300    for s in &p.stages {
301        match shape(s) {
302            Shape::Scalar => {
303                return Err(SpatialLowerError::Unsupported("a Scalar stage (the fabric's LoopCell datapath) is not spelled — print it through codegen_sva::spatial_emit's counted-loop form"));
304            }
305            Shape::Matmul {
306                dataflow,
307                num_blocks,
308                interm,
309                ..
310            } => {
311                if dataflow == DataflowPolicy::WeightStationary {
312                    return Err(SpatialLowerError::Unsupported("a weight-stationary MatmulTiles dataflow (the weight block resident across output blocks) is not spelled — only the output-stationary nest is"));
313                }
314                if num_blocks > 1 && interm.is_none() {
315                    return Err(SpatialLowerError::Unsupported("a MatmulTiles reduction over several K blocks with no Interm buffer (the accumulator held in registers across blocks) is not spelled — the MemReduce needs the spill buffer"));
316                }
317            }
318            Shape::Generic {
319                iters, ins, out, ..
320            } => {
321                if iters.len() != 2 {
322                    return Err(SpatialLowerError::Unsupported("a Generic compute stage with a contraction axis (LV_GEN_K, a third folded iteration variable) — no Spatial controller spelled for it"));
323                }
324                if ins.len() != s.ins.len() {
325                    return Err(SpatialLowerError::Unsupported(
326                        "a Generic compute whose access maps do not match its operand buffers",
327                    ));
328                }
329                if out.indirect.is_some() {
330                    return Err(SpatialLowerError::Unsupported("an indirect (value-dependent) OUTPUT map — an element scatter — is not spelled; the oracle gives it no semantics either"));
331                }
332            }
333            Shape::Epilogue => {
334                return Err(SpatialLowerError::Unsupported("the fabric's fused epilogues (requant, bias, relu) are not spelled — Spatial's compute unit has no packer epilogue slot; legalize them into a Generic"));
335            }
336            Shape::AllReduce => {
337                return Err(SpatialLowerError::Unsupported("a cross-chip AllReduce (the sum of a tensor-parallel shard's partials) is not spelled — one Spatial app describes ONE RDU; emit the per-chip program and let the host collective combine them"));
338            }
339            Shape::Load | Shape::Store => {
340                let (tensor, pattern) = match &s.kind {
341                    StageKind::Load {
342                        tensor, pattern, ..
343                    }
344                    | StageKind::Store {
345                        tensor, pattern, ..
346                    } => (p.tensor(*tensor), pattern),
347                    _ => unreachable!(),
348                };
349                if pattern.stride_w != 1 || pattern.stride_h != tensor.tiles.cols_t {
350                    return Err(SpatialLowerError::Unsupported("a Load or Store whose tile block is not a contiguous rectangle (stride_w other than 1, stride_h other than the tensor's tile row) is not spelled as one DRAM range"));
351                }
352                if pattern.origin.indirect.is_some() {
353                    return Err(SpatialLowerError::Unsupported("an indirect (value-dependent) tile address — a gathered Load or scattered Store — is not spelled; a DRAM range needs an affine origin"));
354                }
355            }
356            Shape::Recv => unreachable!("refused above"),
357        }
358    }
359    if !p.trips.is_empty() {
360        return Err(SpatialLowerError::Unsupported("a loop whose count the program reads from its data (a trip) is not spelled — the emitter prints every K-block loop with a fixed count"));
361    }
362    for e in &p.edges {
363        if !matches!(e.handshake, Handshake::Fifo | Handshake::ValidReady) {
364            return Err(SpatialLowerError::Unsupported("a cross-site handshake edge (Token, Semaphore or Barrier) is not spelled — Spatial has no inter-controller semaphore; every edge must be a same-site Fifo or ValidReady hand-off"));
365        }
366    }
367    for b in &p.buffers {
368        if b.alias_of.is_some() {
369            return Err(SpatialLowerError::Unsupported("an aliased buffer (a spill sharing its memory region with the output) is not spelled — Spatial memories do not alias"));
370        }
371    }
372    for t in &p.tensors {
373        match t.init {
374            TensorInit::External => {
375                return Err(SpatialLowerError::Unsupported("a tensor whose fill is External has no host data to fill its DRAM with — the regression harness needs a deterministic fill"));
376            }
377            TensorInit::RandomBf16 { .. } => {
378                return Err(SpatialLowerError::Unsupported("a tensor with a RandomBf16 fill — the vendor's generator is not reproduced on the host"));
379            }
380            _ => {}
381        }
382    }
383    // ★ SEVERAL ITEMS ON ONE SITE ONLY MATTERS IF SOMETHING INDEXES BY THE ITEM.
384    //
385    // Every stage carries its OWN work item (`Stage::item`), and `resolve_origin` already reads that
386    // item's coordinate for `LV_ITEM_Y` / `LV_ITEM_X`. So a site holding one stage per section resolves
387    // each of them against the right item and needs no loop printed at all. The single thing this
388    // emitter cannot spell is `LV_ITEM` — the item's LINEAR index — which `resolve_origin` pins to zero
389    // because it had been guaranteed one item per site.
390    //
391    // Refusing every multi-item site to protect that one case refused the whole composed Qwen3.8
392    // forward, whose 117 sections stack onto shared sites. Measured on that program: the schedule
393    // contains 39 counted loops and every one of them is `LV_KBLOCK`; `LV_ITEM` does not appear. So the
394    // guard was turning away a program it could spell correctly. It now refuses exactly what it cannot
395    // resolve — an address that reads the linear item index on a site carrying more than one, or an
396    // `LV_ITEM` loop anywhere in the schedule, whose body it does not print.
397    let crowded: std::collections::BTreeSet<Site> = pl
398        .items_of_site
399        .iter()
400        .filter(|(_, v)| v.len() != 1)
401        .map(|(s, _)| *s)
402        .collect();
403    if !crowded.is_empty() {
404        let reads_item = |at: &TileAddr| at.steps.iter().any(|(v, s)| *v == LV_ITEM && *s != 0);
405        for st in &p.stages {
406            let on_crowded = st.site.map(|s| crowded.contains(&s)).unwrap_or(false);
407            if !on_crowded {
408                continue;
409            }
410            let at = match &st.kind {
411                StageKind::Load { pattern, .. } | StageKind::Store { pattern, .. } => {
412                    Some(&pattern.origin)
413                }
414                _ => None,
415            };
416            if at.map(reads_item).unwrap_or(false) {
417                return Err(SpatialLowerError::Unsupported("a tile address that reads the linear work-item index on a site carrying several items — the emitter resolves an item's coordinate but has no counted loop to give that index a value"));
418            }
419        }
420        if let Some(sched) = p.schedule.as_ref() {
421            let mut loops_item = false;
422            for prog in &sched.programs {
423                walk_tagged_ops(&prog.ops, 1, &mut |op, _| {
424                    if let ChoreoOp::Loop { var, .. } = op {
425                        loops_item |= *var == LV_ITEM;
426                    }
427                });
428            }
429            if loops_item {
430                return Err(SpatialLowerError::Unsupported("a counted loop over work items — the emitter prints the K-block loop only, so an item loop's body would be spelled once instead of once per item"));
431            }
432        }
433    }
434    Ok(())
435}
436
437// ── The oracle: fills in, checksum out ────────────────────────────────────────────────────────────
438
439/// Every Input/Param tensor materialized from its recorded fill.
440fn fill_inputs(p: &SpatialProgram) -> Result<Inputs, SpatialLowerError> {
441    let mut inputs = Inputs::default();
442    for t in &p.tensors {
443        if t.role == TensorRole::Output {
444            continue;
445        }
446        let n = t.elems.rows * t.elems.cols;
447        if t.init.is_float() {
448            let vals = (0..n)
449                .map(|i| t.init.element_f64(i))
450                .collect::<Option<Vec<f64>>>();
451            inputs.tensors.insert(
452                t.id,
453                Value::Float(vals.ok_or(SpatialLowerError::Unsupported(
454                    "a tensor whose value fill is shorter than the tensor",
455                ))?),
456            );
457        } else {
458            let vals = (0..n)
459                .map(|i| t.init.element(i))
460                .collect::<Option<Vec<i64>>>();
461            inputs.tensors.insert(
462                t.id,
463                Value::Int(vals.ok_or(SpatialLowerError::Unsupported(
464                    "a tensor whose value fill is shorter than the tensor",
465                ))?),
466            );
467        }
468    }
469    Ok(inputs)
470}
471
472/// The checksum the host compares: the wrapping sum of an integer output, the plain sum of a float one.
473fn checksum_gold(out: &Value) -> Gold {
474    match out {
475        Value::Int(v) => Gold::Int(v.iter().fold(0i64, |s, x| s.wrapping_add(*x))),
476        Value::Float(v) => Gold::Float(v.iter().sum()),
477    }
478}
479
480fn parse_gold(s: &str, scalar: Scalar) -> Result<Gold, SpatialLowerError> {
481    let s = s.trim();
482    if scalar.is_float() {
483        s.parse::<f64>().map(Gold::Float).map_err(|_| {
484            SpatialLowerError::Unsupported(
485                "a gold checksum that is not a number of the program's float output precision",
486            )
487        })
488    } else {
489        s.parse::<i64>().map(Gold::Int).map_err(|_| {
490            SpatialLowerError::Unsupported(
491                "a gold checksum that is not an integer of the program's output precision",
492            )
493        })
494    }
495}
496
497/// `Int` only when every value the app will hold — the fills, the output, and every partial sum of every
498/// matmul — provably fits 32 bits; `Long` otherwise.
499fn scalar_type(p: &SpatialProgram, ix: &Index, inputs: &Inputs, out: &Value) -> Scalar {
500    let any_float = p.tensors.iter().any(|t| is_float_precision(t.dtype))
501        || p.buffers.iter().any(|b| is_float_precision(b.dtype))
502        || matches!(out, Value::Float(_));
503    if any_float {
504        return Scalar::Float;
505    }
506    let limit = i32::MAX as i64;
507    let mut bound: i64 = 0;
508    for v in inputs.tensors.values() {
509        if let Value::Int(xs) = v {
510            bound = bound.max(xs.iter().map(|x| x.abs()).max().unwrap_or(0));
511        }
512    }
513    if let Value::Int(xs) = out {
514        bound = bound.max(xs.iter().map(|x| x.abs()).max().unwrap_or(0));
515    }
516    for s in &p.stages {
517        if let Shape::Matmul { .. } = shape(s) {
518            let operand_bound = |buf: BufferId| -> Option<i64> {
519                ix.producer
520                    .get(&buf)
521                    .map(|id| p.stage(*id))
522                    .and_then(|st| match &st.kind {
523                        StageKind::Load { tensor, .. } => p.tensor(*tensor).abs_bound(),
524                        _ => None,
525                    })
526            };
527            let (a, b) = (
528                s.ins.first().and_then(|b| operand_bound(*b)),
529                s.ins.get(1).and_then(|b| operand_bound(*b)),
530            );
531            match (a, b) {
532                (Some(a), Some(b)) => {
533                    let k = (p.work.k_tiles * p.tile.cols) as i64;
534                    bound = bound.max(k.saturating_mul(a).saturating_mul(b));
535                }
536                _ => return Scalar::Long,
537            }
538        }
539    }
540    if bound <= limit {
541        Scalar::Int
542    } else {
543        Scalar::Long
544    }
545}
546
547// ── Geometry ──────────────────────────────────────────────────────────────────────────────────────
548
549/// A buffer's memory as (rows, cols) of elements: the page its producer fills.
550fn buffer_geom(ctx: &Ctx, b: &Buffer) -> (usize, usize) {
551    let p = ctx.p;
552    let (tr, tc) = (p.tile.rows, p.tile.cols);
553    let by_matmul = |s: &Stage| match shape(s) {
554        Shape::Matmul {
555            block, sub_blocks, ..
556        } => Some((sub_blocks.0 * block.osh * tr, sub_blocks.1 * block.osw * tc)),
557        _ => None,
558    };
559    if let Some(s) = ctx.ix.producer.get(&b.id).map(|id| p.stage(*id)) {
560        if let StageKind::Load { pattern, .. } = &s.kind {
561            return (pattern.block_h * tr, pattern.block_w * tc);
562        }
563        if let Some(g) = by_matmul(s) {
564            return g;
565        }
566    }
567    // A matmul's spill buffer is read and written by the same stage, so it is not in `outs`.
568    if let Some(s) = ctx.ix.spill.get(&b.id).map(|id| p.stage(*id)) {
569        return by_matmul(s).unwrap();
570    }
571    (tr, b.kind.page_tiles() * tc)
572}
573
574/// A tile address with the work item's coordinates folded in: the remaining runtime step is the K block's.
575struct Origin {
576    row_t: usize,
577    col_t: usize,
578    /// Tile rows / tile columns the block advances per K block.
579    k_row_step: usize,
580    k_col_step: usize,
581}
582
583fn resolve_origin(
584    p: &SpatialProgram,
585    at: &TileAddr,
586    tensor: &TensorDecl,
587    item: &WorkItem,
588    k_blocks: usize,
589    sub: Option<(usize, usize)>,
590) -> Result<Origin, SpatialLowerError> {
591    let cols_t = tensor.tiles.cols_t.max(1);
592    let mut base = at.base;
593    let (mut k_row_step, mut k_col_step) = (0, 0);
594    for (var, stride) in &at.steps {
595        if *var == LV_ITEM_Y {
596            base += item.coord.0 * stride;
597        } else if *var == LV_ITEM_X {
598            base += item.coord.1 * stride;
599        } else if *var == LV_ITEM {
600            // One item per site (refused otherwise): the item index is 0.
601        } else if *var == LV_KBLOCK {
602            if k_blocks > 1 {
603                if *stride < cols_t {
604                    k_col_step = *stride;
605                } else if stride % cols_t == 0 {
606                    k_row_step = stride / cols_t;
607                } else {
608                    return Err(SpatialLowerError::Unsupported("a K-block tile step that is neither a whole tile row nor within one (the block would wrap the tensor's tile grid)"));
609                }
610            }
611        } else if *var == LV_SUB_H || *var == LV_SUB_W {
612            // A store's sub-block steps are folded into one whole-block range by the caller.
613            let (block_h, block_w) = sub.ok_or(SpatialLowerError::Unsupported(
614                "a Load tile address stepping on a sub-block loop variable",
615            ))?;
616            let expected = if *var == LV_SUB_H {
617                block_h * cols_t
618            } else {
619                block_w
620            };
621            if *stride != expected {
622                return Err(SpatialLowerError::Unsupported("a Store whose sub-blocks do not tile a contiguous rectangle of the output block"));
623            }
624        } else {
625            let _ = p;
626            return Err(SpatialLowerError::Unsupported("a tile address stepping on a loop variable the emitter does not spell (only the work item, the sub-block and the K block)"));
627        }
628    }
629    Ok(Origin {
630        row_t: base / cols_t,
631        col_t: base % cols_t,
632        k_row_step,
633        k_col_step,
634    })
635}
636
637/// A DRAM range `r0::r1, c0::c1` in elements — with the K-block iterator `kb` folded in when the block
638/// advances per K block.
639fn dram_range(
640    o: &Origin,
641    block_h_t: usize,
642    block_w_t: usize,
643    tile: TileShape,
644    kb: Option<&str>,
645) -> String {
646    let (r0, c0) = (o.row_t * tile.rows, o.col_t * tile.cols);
647    let (h, w) = (block_h_t * tile.rows, block_w_t * tile.cols);
648    let axis = |base: usize, len: usize, step_t: usize, elems: usize| -> String {
649        match kb {
650            Some(kb) if step_t > 0 => {
651                let step = step_t * elems;
652                format!(
653                    "(({kb} * {step}) + {base})::(({kb} * {step}) + {})",
654                    base + len
655                )
656            }
657            _ => format!("{base}::{}", base + len),
658        }
659    };
660    format!(
661        "{}, {}",
662        axis(r0, h, o.k_row_step, tile.rows),
663        axis(c0, w, o.k_col_step, tile.cols)
664    )
665}
666
667// ── The program, read for structure ───────────────────────────────────────────────────────────────
668
669/// What the printer asks of the program per section, site, stage and buffer, each derived in one pass. A composed
670/// program is printed section by section and site by site, so every question is answered here by a lookup, and the
671/// print stays linear in the program.
672struct Index {
673    /// Each site's stages of each section, in dataflow order — a producer before every consumer on the site,
674    /// declaration order breaking ties — keyed by `(site, section)`; `None` is the whole site.
675    order: BTreeMap<(Site, Option<usize>), Vec<StageId>>,
676    /// Each site's buffers, in id order.
677    site_buffers: BTreeMap<Site, Vec<BufferId>>,
678    /// The first stage whose outputs hold each buffer: its producer.
679    producer: BTreeMap<BufferId, StageId>,
680    /// The first matmul spilling into each buffer.
681    spill: BTreeMap<BufferId, StageId>,
682    /// How many K blocks the schedule loops each stage over, for a stage some op of which sits inside a K-block loop
683    /// that runs more than once (the last such op read, in program then op order).
684    k_blocks: BTreeMap<StageId, usize>,
685    /// The section of the first Store of each tensor.
686    stored_in: BTreeMap<TensorId, Option<usize>>,
687    /// The sites each section runs on.
688    section_sites: Vec<BTreeSet<Site>>,
689}
690
691impl Index {
692    fn of(p: &SpatialProgram) -> Index {
693        let (mut producer, mut spill, mut stored_in) =
694            (BTreeMap::new(), BTreeMap::new(), BTreeMap::new());
695        let mut section_sites = vec![BTreeSet::new(); p.sections.len()];
696        let mut on_site: BTreeMap<Site, Vec<StageId>> = BTreeMap::new();
697        for st in &p.stages {
698            for b in &st.outs {
699                producer.entry(*b).or_insert(st.id);
700            }
701            if let Shape::Matmul {
702                interm: Some(i), ..
703            } = shape(st)
704            {
705                spill.entry(i).or_insert(st.id);
706            }
707            if let StageKind::Store { tensor, .. } = &st.kind {
708                stored_in
709                    .entry(*tensor)
710                    .or_insert_with(|| p.section_of_stage(st.id));
711            }
712            if let Some(site) = st.site {
713                on_site.entry(site).or_default().push(st.id);
714                if let Some(set) = p
715                    .section_of_stage(st.id)
716                    .and_then(|k| section_sites.get_mut(k))
717                {
718                    set.insert(site);
719                }
720            }
721        }
722        // The dataflow among each site's own stages.
723        let mut succ: BTreeMap<StageId, Vec<StageId>> = BTreeMap::new();
724        let mut indeg: BTreeMap<StageId, usize> = BTreeMap::new();
725        for e in &p.edges {
726            let from_site = p.stage(e.from).site;
727            if from_site.is_none() {
728                continue;
729            }
730            for to in &e.to {
731                if p.stage(*to).site == from_site && *to != e.from {
732                    succ.entry(e.from).or_default().push(*to);
733                    *indeg.entry(*to).or_default() += 1;
734                }
735            }
736        }
737        let mut order: BTreeMap<(Site, Option<usize>), Vec<StageId>> = BTreeMap::new();
738        for (site, stages) in &on_site {
739            let mut ready: BTreeSet<StageId> = stages
740                .iter()
741                .copied()
742                .filter(|s| indeg.get(s).copied().unwrap_or(0) == 0)
743                .collect();
744            let mut whole = Vec::with_capacity(stages.len());
745            while let Some(next) = ready.pop_first() {
746                whole.push(next);
747                for s in succ.get(&next).into_iter().flatten() {
748                    let d = indeg
749                        .get_mut(s)
750                        .expect("a successor on the site has an in-degree");
751                    *d -= 1;
752                    if *d == 0 {
753                        ready.insert(*s);
754                    }
755                }
756            }
757            for id in &whole {
758                if let Some(k) = p.section_of_stage(*id) {
759                    order.entry((*site, Some(k))).or_default().push(*id);
760                }
761            }
762            order.insert((*site, None), whole);
763        }
764        let mut site_buffers: BTreeMap<Site, Vec<BufferId>> = BTreeMap::new();
765        for b in &p.buffers {
766            if let Some(site) = b.site {
767                site_buffers.entry(site).or_default().push(b.id);
768            }
769        }
770        let mut k_blocks = BTreeMap::new();
771        fn walk(ops: &[TaggedOp], under: Option<usize>, found: &mut BTreeMap<StageId, usize>) {
772            for t in ops {
773                if let ChoreoOp::Loop { var, count, body } = &t.op {
774                    let inner = if *var == LV_KBLOCK {
775                        Some(count.bound())
776                    } else {
777                        under
778                    };
779                    walk(body, inner, found);
780                } else if let Some(n) = under.filter(|n| *n > 1) {
781                    found.insert(t.stage, n);
782                }
783            }
784        }
785        for prog in p.schedule.iter().flat_map(|sch| sch.programs.iter()) {
786            walk(&prog.ops, None, &mut k_blocks);
787        }
788        Index {
789            order,
790            site_buffers,
791            producer,
792            spill,
793            k_blocks,
794            stored_in,
795            section_sites,
796        }
797    }
798
799    fn k_blocks_of(&self, stage: StageId) -> Option<usize> {
800        self.k_blocks.get(&stage).copied()
801    }
802
803    /// The section whose Store writes `t`, if any.
804    fn storing_section(&self, t: TensorId) -> Option<usize> {
805        self.stored_in.get(&t).copied().flatten()
806    }
807}
808
809// ── Printing ──────────────────────────────────────────────────────────────────────────────────────
810
811struct Emit {
812    s: String,
813    ind: usize,
814}
815
816impl Emit {
817    fn l(&mut self, line: impl AsRef<str>) {
818        for _ in 0..self.ind {
819            self.s.push_str("  ");
820        }
821        self.s.push_str(line.as_ref());
822        self.s.push('\n');
823    }
824    fn open(&mut self, line: impl AsRef<str>) {
825        self.l(line);
826        self.ind += 1;
827    }
828    fn close(&mut self, line: &str) {
829        self.ind -= 1;
830        self.l(line);
831    }
832}
833
834struct Ctx<'a> {
835    p: &'a SpatialProgram,
836    scalar: Scalar,
837    /// The integer scalar `I` of a program holding floats and integers both; `None` when every memory is `T`.
838    int_scalar: Option<Scalar>,
839    /// Buffers whose edge is a `Fifo` handshake: spelled `FIFO[T]`, read and written in stream order.
840    fifo: BTreeSet<BufferId>,
841    names: BTreeMap<BufferId, String>,
842    flavor: SpatialFlavor,
843    ix: Index,
844}
845
846impl<'a> Ctx<'a> {
847    fn ty(&self, dtype: Precision) -> Ty {
848        if self.int_scalar.is_some() && !is_float_precision(dtype) {
849            Ty::I
850        } else {
851            Ty::T
852        }
853    }
854    fn name(&self, b: &Buffer) -> String {
855        self.names[&b.id].clone()
856    }
857    fn buf_ty(&self, b: &Buffer) -> Ty {
858        self.ty(b.dtype)
859    }
860    fn is_float(&self, ty: Ty) -> bool {
861        ty == Ty::T && self.scalar.is_float()
862    }
863    fn buf(&self, b: BufferId) -> &'a Buffer {
864        self.p.buffer(b)
865    }
866    fn tensor_of_load(&self, s: &Stage) -> &'a TensorDecl {
867        match &s.kind {
868            StageKind::Load { tensor, .. } | StageKind::Store { tensor, .. } => {
869                self.p.tensor(*tensor)
870            }
871            _ => unreachable!(),
872        }
873    }
874    fn item_of(&self, s: &Stage) -> &'a WorkItem {
875        let id = s.item.expect("a placed stage carries its work item");
876        &self.p.work.items[id.idx()]
877    }
878}
879
880/// The iteration variables a body or a predicate may name: the Generic's (row, col) axes.
881struct Axes {
882    row: LoopVar,
883    col: LoopVar,
884}
885
886impl Axes {
887    fn name(&self, v: LoopVar) -> Result<&'static str, SpatialLowerError> {
888        if v == self.row {
889            Ok("r")
890        } else if v == self.col {
891            Ok("c")
892        } else {
893            Err(SpatialLowerError::Unsupported(
894                "a body predicate on a loop variable outside the (rows, cols) iteration space",
895            ))
896        }
897    }
898}
899
900/// Whether `b` is positive at every point by its structure alone — a positive constant, an exponential, and their
901/// sums and products.
902fn positive(b: &TensorBody) -> bool {
903    match b {
904        TensorBody::Const(bits) => f64::from_bits(*bits) > 0.0,
905        TensorBody::Unary {
906            kind: SfpuKind::Exp,
907            ..
908        } => true,
909        TensorBody::Add(x, y) | TensorBody::Mul(x, y) => positive(x) && positive(y),
910        _ => false,
911    }
912}
913
914/// `(a, b)` when `body` is `e^(a·ln b)` with `b` positive by structure.
915fn as_power(body: &TensorBody) -> Option<(&TensorBody, &TensorBody)> {
916    let TensorBody::Unary {
917        kind: SfpuKind::Exp,
918        arg,
919    } = body
920    else {
921        return None;
922    };
923    let TensorBody::Mul(x, y) = &**arg else {
924        return None;
925    };
926    let (a, b) = match (&**x, &**y) {
927        (
928            a,
929            TensorBody::Unary {
930                kind: SfpuKind::Ln,
931                arg: b,
932            },
933        )
934        | (
935            TensorBody::Unary {
936                kind: SfpuKind::Ln,
937                arg: b,
938            },
939            a,
940        ) => (a, &**b),
941        _ => return None,
942    };
943    positive(b).then_some((a, b))
944}
945
946/// A body, computed in `ty` (every operand already read as a `ty`).
947fn body_expr(
948    ctx: &Ctx,
949    ty: Ty,
950    body: &TensorBody,
951    axes: &Axes,
952    operand: &dyn Fn(usize) -> String,
953) -> Result<String, SpatialLowerError> {
954    // ★ `e^(a·ln b)` IS `b^a` FOR b > 0, AND IS SPELLED SO. The fabric chain's Tungsten backend prints `ln` as a bare
955    // `log(…)` that its contexts shadow — the generated C++ does not compile — while `pow` reaches the library.
956    // Gated DeltaNet's decay `e^(A·ln(1 + e^u))` is the case: `1 + e^u` is positive by structure.
957    if let Some((a, b)) = as_power(body) {
958        return Ok(format!(
959            "pow({}, {})",
960            body_expr(ctx, ty, b, axes, operand)?,
961            body_expr(ctx, ty, a, axes, operand)?
962        ));
963    }
964    Ok(match body {
965        TensorBody::Operand(i) => operand(*i),
966        TensorBody::Const(bits) => lit_of(f64::from_bits(*bits), ty, ctx.is_float(ty)),
967        TensorBody::Add(x, y) => format!(
968            "({} + {})",
969            body_expr(ctx, ty, x, axes, operand)?,
970            body_expr(ctx, ty, y, axes, operand)?
971        ),
972        TensorBody::Sub(x, y) => format!(
973            "({} - {})",
974            body_expr(ctx, ty, x, axes, operand)?,
975            body_expr(ctx, ty, y, axes, operand)?
976        ),
977        TensorBody::Mul(x, y) => format!(
978            "({} * {})",
979            body_expr(ctx, ty, x, axes, operand)?,
980            body_expr(ctx, ty, y, axes, operand)?
981        ),
982        TensorBody::Min(..) | TensorBody::Max(..) | TensorBody::FloorDiv(..) | TensorBody::Shr { .. } => {
983            return Err(SpatialLowerError::Unsupported(
984                "an integer model's min, max, floor division or arithmetic shift — the Spatial-language printer has no lock holding them to the VM yet",
985            ))
986        }
987        TensorBody::Unary { kind, arg } => {
988            let x = body_expr(ctx, ty, arg, axes, operand)?;
989            let t = ty.name();
990            match kind {
991                SfpuKind::Exp => format!("exp({x})"),
992                SfpuKind::Recip => format!("(1.to[{t}] / {x})"),
993                SfpuKind::Rsqrt { eps_nano, over } => {
994                    let eps = *eps_nano as f64 / 1e9;
995                    format!(
996                        "(1.to[{t}] / sqrt(({x} / {}) + {}))",
997                        lit_of((*over).max(1) as f64, ty, true),
998                        lit_of(eps, ty, true)
999                    )
1000                }
1001                SfpuKind::Silu => format!("({x} / (1.to[{t}] + exp(0.to[{t}] - {x})))"),
1002                SfpuKind::Relu => format!("max({x}, 0.to[{t}])"),
1003                SfpuKind::Ln if ctx.flavor == SpatialFlavor::Fabric => {
1004                    return Err(SpatialLowerError::Unsupported("a natural log outside `e^(a·ln b)` with b positive, on the fabric flavor — Tungsten prints `ln` as a bare `log(…)` its contexts shadow"));
1005                }
1006                SfpuKind::Ln => format!("ln({x})"),
1007            }
1008        }
1009        // Decided by POSITION: a `mux` on the axis, both branches computed (branch-free), as the oracle does.
1010        TensorBody::Select {
1011            cond,
1012            then,
1013            otherwise,
1014        } => {
1015            let axis = axes.name(cond.var())?;
1016            let op = match cond {
1017                Pred::AxisLt(..) => "<",
1018                Pred::AxisGe(..) => ">=",
1019            };
1020            let test = match cond.rhs() {
1021                PredRhs::Const(k) => format!("{axis} {op} {k}"),
1022                PredRhs::Operand(i) => format!("{axis}.to[{}] {op} {}", ty.name(), operand(i)),
1023            };
1024            format!(
1025                "mux({test}, {}, {})",
1026                body_expr(ctx, ty, then, axes, operand)?,
1027                body_expr(ctx, ty, otherwise, axes, operand)?
1028            )
1029        }
1030    })
1031}
1032
1033/// A map's strides along the Generic's `(row, col)` axes; a step on any other loop variable is refused.
1034fn map_strides(m: &ElemAccess, axes: &Axes) -> Result<(i64, i64), SpatialLowerError> {
1035    let (mut sr, mut sc) = (0i64, 0i64);
1036    for (v, st) in &m.terms {
1037        if *v == axes.row {
1038            sr += st;
1039        } else if *v == axes.col {
1040            sc += st;
1041        } else if *st != 0 {
1042            return Err(SpatialLowerError::Unsupported("an access map stepping on a loop variable outside the (rows, cols) iteration space"));
1043        }
1044    }
1045    Ok((sr, sc))
1046}
1047
1048/// The affine `(row, col)` of a map in a page `width` elements wide and `plane_rows` high — the plane the oracle
1049/// reads: `off = base + r·sr + c·sc`, `(off / width, off % width)`, with the division folded away. `Some` exactly
1050/// when every offset of the space is on the page and its column part never carries into the next page row;
1051/// [`flat_read`] reads every other map through its offset.
1052fn affine_rc(
1053    base: i64,
1054    sr: i64,
1055    sc: i64,
1056    width: usize,
1057    plane_rows: usize,
1058    r_ext: usize,
1059    c_ext: usize,
1060) -> Option<(String, String)> {
1061    if base < 0 || sr < 0 || sc < 0 {
1062        return None;
1063    }
1064    let w = width.max(1) as i64;
1065    let (rm, cm) = ((r_ext as i64 - 1).max(0), (c_ext as i64 - 1).max(0));
1066    if base % w + rm * (sr % w) + cm * (sc % w) >= w
1067        || base / w + rm * (sr / w) + cm * (sc / w) >= plane_rows as i64
1068    {
1069        return None;
1070    }
1071    let term = |coef: i64, var: &str| -> Option<String> {
1072        match coef {
1073            0 => None,
1074            1 => Some(var.to_string()),
1075            k => Some(format!("{var} * {k}")),
1076        }
1077    };
1078    let join = |konst: i64, a: Option<String>, b: Option<String>| -> String {
1079        let mut parts: Vec<String> = Vec::new();
1080        if konst != 0 {
1081            parts.push(konst.to_string());
1082        }
1083        parts.extend(a);
1084        parts.extend(b);
1085        match parts.len() {
1086            0 => "0".to_string(),
1087            1 => parts.remove(0),
1088            _ => format!("({})", parts.join(" + ")),
1089        }
1090    };
1091    Some((
1092        join(base / w, term(sr / w, "r"), term(sc / w, "c")),
1093        join(base % w, term(sr % w, "r"), term(sc % w, "c")),
1094    ))
1095}
1096
1097/// Operand `i`'s read at the point `(r, c)`, in its own memory's type, and the `val`s it needs first.
1098///
1099/// A map the affine split covers reads `buf(row, col)`. Any other — a signed window that starts before the page
1100/// (rope's partner, the conv ring's shift), a flat offset that carries into the next page row, a GATHER whose
1101/// `IndirectTerm` adds `index · stride` — reads through its flat offset `o`, split at runtime. An offset the space
1102/// cannot drive off the page reads directly; one it can (or one that carries data) is clamped onto the page for the
1103/// address and the read muxed to 0 off it — the oracle's reading of an offset outside the plane, exactly.
1104fn flat_read(
1105    ctx: &Ctx,
1106    b: &Buffer,
1107    m: &ElemAccess,
1108    axes: &Axes,
1109    r_ext: usize,
1110    c_ext: usize,
1111    gather: Option<(&str, i64)>,
1112    tag: usize,
1113) -> Result<(Vec<String>, String), SpatialLowerError> {
1114    let ty = ctx.buf_ty(b);
1115    let (rows, width) = plane_of(ctx, b)?;
1116    let (sr, sc) = map_strides(m, axes)?;
1117    let name = ctx.name(b);
1118    if gather.is_none() {
1119        if let Some((r, c)) = affine_rc(m.base, sr, sc, width, rows, r_ext, c_ext) {
1120            return Ok((Vec::new(), format!("{name}({r}, {c})")));
1121        }
1122    }
1123    let n = (rows * width) as i64;
1124    let (rm, cm) = ((r_ext as i64 - 1).max(0), (c_ext as i64 - 1).max(0));
1125    let lo = m.base + (rm * sr).min(0) + (cm * sc).min(0);
1126    let hi = m.base + (rm * sr).max(0) + (cm * sc).max(0);
1127    let mut parts: Vec<String> = Vec::new();
1128    if m.base != 0 {
1129        parts.push(m.base.to_string());
1130    }
1131    for (coef, var) in [(sr, "r"), (sc, "c")] {
1132        match coef {
1133            0 => {}
1134            1 => parts.push(var.to_string()),
1135            k => parts.push(format!("{var} * {k}")),
1136        }
1137    }
1138    if let Some((index, stride)) = gather {
1139        parts.push(if stride == 1 {
1140            index.to_string()
1141        } else {
1142            format!("{index} * {stride}")
1143        });
1144    }
1145    let off = if parts.is_empty() {
1146        "0.to[Int]".to_string()
1147    } else {
1148        parts.join(" + ")
1149    };
1150    let o = format!("o{tag}");
1151    let mut pre = vec![format!("val {o}: Int = {off}")];
1152    if gather.is_none() && lo >= 0 && hi < n {
1153        return Ok((pre, format!("{name}({o} / {width}, {o} % {width})")));
1154    }
1155    let a = format!("a{tag}");
1156    pre.push(format!(
1157        "val {a} = mux({o} >= 0, mux({o} < {n}, {o}, 0.to[Int]), 0.to[Int])"
1158    ));
1159    let zero = lit_of(0.0, ty, ctx.is_float(ty));
1160    Ok((
1161        pre,
1162        format!(
1163            "mux({o} >= 0, mux({o} < {n}, {name}({a} / {width}, {a} % {width}), {zero}), {zero})"
1164        ),
1165    ))
1166}
1167
1168/// Every operand's read at the point `(r, c)` — each in its own memory's type — and the `val`s they need first.
1169/// A gather's index is its index operand read through that operand's AFFINE map (the oracle's `offset_at`), as
1170/// an `Int`; a stride-0 gather is a bounds check that moves nothing, and reads its affine offset.
1171fn operand_reads(
1172    ctx: &Ctx,
1173    s: &Stage,
1174    maps: &[ElemAccess],
1175    axes: &Axes,
1176    r_ext: usize,
1177    c_ext: usize,
1178    folds: bool,
1179) -> Result<(Vec<String>, Vec<String>), SpatialLowerError> {
1180    let mut prelude: Vec<String> = Vec::new();
1181    let mut indices: BTreeMap<usize, String> = BTreeMap::new();
1182    for m in maps {
1183        if let Some(ind) = m.indirect.filter(|ind| ind.stride != 0) {
1184            if ind.operand >= maps.len() {
1185                return Err(SpatialLowerError::Unsupported(
1186                    "a gather whose index names an absent operand",
1187                ));
1188            }
1189            indices
1190                .entry(ind.operand)
1191                .or_insert_with(|| format!("ix{}", ind.operand));
1192        }
1193    }
1194    for (&k, index) in &indices {
1195        let b = ctx.buf(s.ins[k]);
1196        if ctx.fifo.contains(&b.id) {
1197            return Err(SpatialLowerError::Unsupported("a gather whose index is behind a Fifo handshake — the index would be dequeued once for the address and again as an operand"));
1198        }
1199        let affine = ElemAccess {
1200            indirect: None,
1201            ..maps[k].clone()
1202        };
1203        let (pre, read) = flat_read(ctx, b, &affine, axes, r_ext, c_ext, None, maps.len() + k)?;
1204        prelude.extend(pre);
1205        prelude.push(format!("val {index} = {read}.to[Int]"));
1206    }
1207    let mut reads: Vec<String> = Vec::new();
1208    for (i, m) in maps.iter().enumerate() {
1209        let b = ctx.buf(s.ins[i]);
1210        if ctx.fifo.contains(&b.id) {
1211            let (rows, width) = plane_of(ctx, b)?;
1212            let (sr, sc) = map_strides(m, axes)?;
1213            let identity = !folds
1214                && m.indirect.is_none()
1215                && affine_rc(m.base, sr, sc, width, rows, r_ext, c_ext)
1216                    == Some(("r".to_string(), "c".to_string()));
1217            if !identity {
1218                return Err(SpatialLowerError::Unsupported("a Fifo-handshake buffer read out of stream order (a broadcast, window, gather or reduce map) — Spatial's FIFO dequeues in order; place the program on a target with ValidReady handshakes"));
1219            }
1220            reads.push(format!("{}.deq()", ctx.name(b)));
1221            continue;
1222        }
1223        let gather = m
1224            .indirect
1225            .filter(|ind| ind.stride != 0)
1226            .map(|ind| (indices[&ind.operand].as_str(), ind.stride));
1227        let (pre, read) = flat_read(ctx, b, m, axes, r_ext, c_ext, gather, i)?;
1228        prelude.extend(pre);
1229        reads.push(read);
1230    }
1231    Ok((prelude, reads))
1232}
1233
1234/// The plane a Generic operand's page is read as: `(rows, width)` in elements — one tile high, `page_tiles`
1235/// tiles wide, which is the SRAM a one-row Load or a Generic fills.
1236fn plane_of(ctx: &Ctx, b: &Buffer) -> Result<(usize, usize), SpatialLowerError> {
1237    let (rows, cols) = buffer_geom(ctx, b);
1238    if rows != ctx.p.tile.rows {
1239        return Err(SpatialLowerError::Unsupported("a Generic operand or output whose page is a multi-row tile block — the plane a flat map addresses is one tile row of pages"));
1240    }
1241    Ok((rows, cols))
1242}
1243
1244fn emit_generic(
1245    ctx: &Ctx,
1246    e: &mut Emit,
1247    s: &Stage,
1248    iters: &[IterSpec],
1249    maps: &[ElemAccess],
1250    out: &ElemAccess,
1251    body: &TensorBody,
1252) -> Result<(), SpatialLowerError> {
1253    let (row_ax, col_ax) = (iters[0], iters[1]);
1254    let axes = Axes {
1255        row: row_ax.var,
1256        col: col_ax.var,
1257    };
1258    let out_buf = ctx.buf(*s.outs.first().ok_or(SpatialLowerError::Unsupported(
1259        "a Generic compute stage without an output buffer",
1260    ))?);
1261    let pool_of = |k: IterKind| match k {
1262        IterKind::Reduce { pool } => Some(pool),
1263        IterKind::Parallel => None,
1264    };
1265    let (row_pool, col_pool) = (pool_of(row_ax.kind), pool_of(col_ax.kind));
1266    let fold_of = |pool: ReducePool| -> Result<&'static str, SpatialLowerError> {
1267        match pool {
1268            ReducePool::Sum => Ok("{ _ + _ }"),
1269            ReducePool::Max => Ok("{ (a, b) => max(a, b) }"),
1270            ReducePool::ArgMax => Err(SpatialLowerError::Unsupported("an ArgMax reduce pool over the row axis, or over both — the index scan is spelled along the columns")),
1271        }
1272    };
1273    let (r_ext, c_ext) = (row_ax.extent, col_ax.extent);
1274    let out_ty = ctx.buf_ty(out_buf);
1275    let in_ty: Vec<Ty> = s.ins.iter().map(|b| ctx.buf_ty(ctx.buf(*b))).collect();
1276    // The body computes in its output's type — except an index scan's, which compares the VALUES it reads.
1277    let index_scan = col_pool == Some(ReducePool::ArgMax) && row_pool.is_none();
1278    let body_ty = if index_scan {
1279        in_ty
1280            .iter()
1281            .copied()
1282            .find(|t| ctx.is_float(*t))
1283            .or(in_ty.first().copied())
1284            .unwrap_or(out_ty)
1285    } else {
1286        out_ty
1287    };
1288    let (prelude, reads) = operand_reads(
1289        ctx,
1290        s,
1291        maps,
1292        &axes,
1293        r_ext,
1294        c_ext,
1295        row_pool.is_some() || col_pool.is_some(),
1296    )?;
1297    let operand = |i: usize| -> String { cast(reads[i].clone(), in_ty[i], body_ty) };
1298    let expr = body_expr(ctx, body_ty, body, &axes, &operand)?;
1299    let (out_rows, out_width) = plane_of(ctx, out_buf)?;
1300    let out_name = ctx.name(out_buf);
1301    let out_fifo = ctx.fifo.contains(&out_buf.id);
1302    if out_fifo && (row_pool.is_some() || col_pool.is_some()) {
1303        return Err(SpatialLowerError::Unsupported("a Fifo-handshake buffer written by a reduced Generic (one scalar broadcast over a tile) — Spatial's FIFO enqueues one value per point"));
1304    }
1305    for it in [row_ax, col_ax] {
1306        if matches!(it.kind, IterKind::Reduce { .. }) && out.varies_along(it.var) {
1307            return Err(SpatialLowerError::Unsupported(
1308                "a Generic compute whose output map varies along a reduced axis",
1309            ));
1310        }
1311    }
1312    let (bt, ot) = (body_ty.name(), out_ty.name());
1313    // `val fold` over `space` (bound to `binders`, whose axes `axes_folded` are the ones folded): a `Reduce` for the
1314    // fabric, or a register scan seeded at the first point for the simulator.
1315    let fold_line = |e: &mut Emit,
1316                     space: String,
1317                     binders: &str,
1318                     axes_folded: &[&str],
1319                     pool: ReducePool|
1320     -> Result<(), SpatialLowerError> {
1321        match ctx.flavor {
1322            SpatialFlavor::Fabric => {
1323                let head = format!("val fold = Reduce(Reg[{bt}])({space}) {{ {binders} =>");
1324                let tail = fold_of(pool)?;
1325                if prelude.is_empty() {
1326                    e.l(format!("{head} {expr} }} {tail}"));
1327                } else {
1328                    e.open(head);
1329                    for l in &prelude {
1330                        e.l(l);
1331                    }
1332                    e.l(&expr);
1333                    e.close(&format!("}} {tail}"));
1334                }
1335            }
1336            SpatialFlavor::Simulator => {
1337                let combined = match pool {
1338                    ReducePool::Sum => "fold.value + v",
1339                    ReducePool::Max => "max(fold.value, v)",
1340                    ReducePool::ArgMax => return Err(SpatialLowerError::Unsupported("an ArgMax reduce on a Generic's row axis in the simulator flavor — the index scan is spelled along the columns")),
1341                };
1342                let seeded = axes_folded.iter().rev().fold("v".to_string(), |inner, ax| {
1343                    format!("mux({ax} < 1, {inner}, {combined})")
1344                });
1345                e.l(format!("val fold = Reg[{bt}]"));
1346                e.open(format!("Sequential.Foreach({space}) {{ {binders} =>"));
1347                for l in &prelude {
1348                    e.l(l);
1349                }
1350                e.l(format!("val v = {expr}"));
1351                e.l(format!("fold := {seeded}"));
1352                e.close("}");
1353            }
1354        }
1355        Ok(())
1356    };
1357    // ★ THE ORACLE'S PAGE STARTS AT ZERO. A stage writes the points of its space, and a space narrower than the page
1358    // (a 1-row space over a 32-row tile, a row of 8 over 32 columns) leaves the rest as the oracle leaves it: 0. An
1359    // SRAM starts undefined, and a later section that folds over a whole tile — a matmul's K — reads that rest.
1360    let covered = match (row_pool, col_pool) {
1361        (None, None) => r_ext >= out_rows && c_ext >= out_width,
1362        (None, Some(_)) => r_ext >= out_rows,
1363        (Some(_), None) => c_ext >= out_width,
1364        (Some(_), Some(_)) => true,
1365    };
1366    if !covered && !out_fifo {
1367        e.l(format!("// zero page: {} writes {r_ext}x{c_ext} of its {out_rows}x{out_width} page, and the rest is 0 as the oracle's page starts", s.name));
1368        e.l(format!("Foreach(0 until {out_rows}, 0 until {out_width}) {{ (r, c) => {out_name}(r, c) = 0.to[{ot}] }}"));
1369    }
1370    match (row_pool, col_pool) {
1371        // Every point of the space is written at (r, c) — the oracle writes a parallel space that way whatever
1372        // the output map says, so a per-row value lands across its whole row and a per-row reader finds it at
1373        // column 0.
1374        (None, None) => {
1375            if c_ext > out_width || r_ext > out_rows {
1376                return Err(SpatialLowerError::Unsupported("a Generic compute whose (rows, cols) space is larger than its output page plane"));
1377            }
1378            e.l(format!(
1379                "// stage {} (Generic: {r_ext} rows x {c_ext} cols, body {body:?})",
1380                s.name
1381            ));
1382            e.open(format!(
1383                "Foreach(0 until {r_ext}, 0 until {c_ext}) {{ (r, c) =>"
1384            ));
1385            for l in &prelude {
1386                e.l(l);
1387            }
1388            if out_fifo {
1389                e.l(format!(
1390                    "{out_name}.enq({})",
1391                    cast(expr.clone(), body_ty, out_ty)
1392                ));
1393            } else {
1394                e.l(format!(
1395                    "{out_name}(r, c) = {}",
1396                    cast(expr.clone(), body_ty, out_ty)
1397                ));
1398            }
1399            e.close("}");
1400        }
1401        // The first maximum of each row, written ACROSS the output tile row: the VM's `>` scan, exactly as the
1402        // `ArgMax` fold is — column 0 seeds the running best, and a later column replaces it only when strictly
1403        // greater. A scan carries its best from column to column, so it runs `Sequential`.
1404        // ★ ON THE FABRIC THE SAME SCAN IS TWO PLAIN FOLDS. PIR maps a register only as an accumulation (`acc = f(acc,
1405        // x)`), and refuses the scan's pair of mux-updated registers. With `v0` the seed and `w_c = v_c`, or `v0` where
1406        // `v_c` is NaN: `m = max w`, `first = min { c : w_c = m }`, and the answer is `first` — or 0 when `v0` is NaN.
1407        // A NaN replaced by `v0` can create no new maximum and ties go to the lowest index, so this IS the scan.
1408        (None, Some(ReducePool::ArgMax)) if ctx.flavor == SpatialFlavor::Fabric => {
1409            e.l(format!("// stage {} (Generic: {r_ext} rows, the first maximum of {c_ext} cols — the scan as a max fold and a min-index fold; body {body:?})", s.name));
1410            e.open(format!("Foreach(0 until {r_ext}) {{ r =>"));
1411            e.open("val v0 = {");
1412            e.l("val c = 0.to[I32]");
1413            for l in &prelude {
1414                e.l(l);
1415            }
1416            e.l(&expr);
1417            e.close("}");
1418            e.open(format!(
1419                "val m = Reduce(Reg[{bt}])(0 until {c_ext}) {{ c =>"
1420            ));
1421            for l in &prelude {
1422                e.l(l);
1423            }
1424            e.l(format!("val v = {expr}"));
1425            e.l("mux(v === v, v, v0)");
1426            e.close("} { (a, b) => max(a, b) }");
1427            e.open(format!(
1428                "val first = Reduce(Reg[I32])(0 until {c_ext}) {{ c =>"
1429            ));
1430            for l in &prelude {
1431                e.l(l);
1432            }
1433            e.l(format!("val v = {expr}"));
1434            e.l(format!(
1435                "mux(mux(v === v, v, v0) === m.value, c, {c_ext}.to[I32])"
1436            ));
1437            e.close("} { (a, b) => min(a, b) }");
1438            e.open(format!("Foreach(0 until {out_width}) {{ cc =>"));
1439            e.l(format!(
1440                "{out_name}(r, cc) = mux(v0 === v0, first.value, 0.to[I32]).to[{ot}]"
1441            ));
1442            e.close("}");
1443            e.close("}");
1444        }
1445        (None, Some(ReducePool::ArgMax)) => {
1446            e.l(format!("// stage {} (Generic: {r_ext} rows, the first maximum of {c_ext} cols — seeded by col 0, replaced only on `>`; body {body:?})", s.name));
1447            e.open(format!("Foreach(0 until {r_ext}) {{ r =>"));
1448            e.l(format!("val best = Reg[{bt}]"));
1449            e.l(format!("val at = Reg[{ot}]"));
1450            e.open(format!("Sequential.Foreach(0 until {c_ext}) {{ c =>"));
1451            for l in &prelude {
1452                e.l(l);
1453            }
1454            e.l(format!("val v = {expr}"));
1455            e.l("val ahead = v > best.value");
1456            e.l("best := mux(c < 1, v, mux(ahead, v, best.value))");
1457            e.l(format!(
1458                "at := mux(c < 1, c.to[{ot}], mux(ahead, c.to[{ot}], at.value))"
1459            ));
1460            e.close("}");
1461            e.open(format!("Foreach(0 until {out_width}) {{ cc =>"));
1462            e.l(format!("{out_name}(r, cc) = at.value"));
1463            e.close("}");
1464            e.close("}");
1465        }
1466        // One scalar per row, written ACROSS the output tile row (a per-row broadcast reads column 0).
1467        (None, Some(pool)) => {
1468            e.l(format!(
1469                "// stage {} (Generic: {r_ext} rows, {c_ext} cols reduced {pool:?}, body {body:?})",
1470                s.name
1471            ));
1472            e.open(format!("Foreach(0 until {r_ext}) {{ r =>"));
1473            fold_line(e, format!("0 until {c_ext}"), "c", &["c"], pool)?;
1474            e.open(format!("Foreach(0 until {out_width}) {{ cc =>"));
1475            e.l(format!(
1476                "{out_name}(r, cc) = {}",
1477                cast("fold".into(), body_ty, out_ty)
1478            ));
1479            e.close("}");
1480            e.close("}");
1481        }
1482        // One scalar per column, written DOWN the output tile (a per-column broadcast reads row 0).
1483        (Some(pool), None) => {
1484            e.l(format!(
1485                "// stage {} (Generic: {c_ext} cols, {r_ext} rows reduced {pool:?}, body {body:?})",
1486                s.name
1487            ));
1488            e.open(format!("Foreach(0 until {c_ext}) {{ c =>"));
1489            fold_line(e, format!("0 until {r_ext}"), "r", &["r"], pool)?;
1490            e.open(format!("Foreach(0 until {out_rows}) {{ rr =>"));
1491            e.l(format!(
1492                "{out_name}(rr, c) = {}",
1493                cast("fold".into(), body_ty, out_ty)
1494            ));
1495            e.close("}");
1496            e.close("}");
1497        }
1498        // One scalar, on the whole output tile.
1499        (Some(rp), Some(cp)) => {
1500            if rp != cp {
1501                return Err(SpatialLowerError::Unsupported(
1502                    "a Generic compute folding its two axes with two different reduce pools",
1503                ));
1504            }
1505            fold_of(rp)?;
1506            e.l(format!("// stage {} (Generic: {r_ext} rows x {c_ext} cols both reduced {rp:?}, body {body:?})", s.name));
1507            fold_line(
1508                e,
1509                format!("0 until {r_ext}, 0 until {c_ext}"),
1510                "(r, c)",
1511                &["r", "c"],
1512                rp,
1513            )?;
1514            e.open(format!(
1515                "Foreach(0 until {out_rows}, 0 until {out_width}) {{ (rr, cc) =>"
1516            ));
1517            e.l(format!(
1518                "{out_name}(rr, cc) = {}",
1519                cast("fold".into(), body_ty, out_ty)
1520            ));
1521            e.close("}");
1522        }
1523    }
1524    Ok(())
1525}
1526
1527fn emit_load(
1528    ctx: &Ctx,
1529    e: &mut Emit,
1530    s: &Stage,
1531    kb: Option<&str>,
1532) -> Result<(), SpatialLowerError> {
1533    let StageKind::Load {
1534        pattern,
1535        num_blocks,
1536        ..
1537    } = &s.kind
1538    else {
1539        unreachable!()
1540    };
1541    let t = ctx.tensor_of_load(s);
1542    let item = ctx.item_of(s);
1543    let k_blocks = if kb.is_some() { *num_blocks } else { 1 };
1544    let o = resolve_origin(ctx.p, &pattern.origin, t, item, k_blocks, None)?;
1545    if o.row_t + (k_blocks - 1) * o.k_row_step + pattern.block_h > t.tiles.rows_t
1546        || o.col_t + (k_blocks - 1) * o.k_col_step + pattern.block_w > t.tiles.cols_t
1547    {
1548        return Err(SpatialLowerError::Unsupported(
1549            "a Load whose K blocks run past the tensor's tile grid",
1550        ));
1551    }
1552    let buf = ctx.buf(*s.outs.first().ok_or(SpatialLowerError::Unsupported(
1553        "a Load stage without an output buffer",
1554    ))?);
1555    let step = if k_blocks > 1 {
1556        format!(
1557            ", stepping {} tile rows / {} tile cols per K block over {} blocks",
1558            o.k_row_step, o.k_col_step, k_blocks
1559        )
1560    } else {
1561        String::new()
1562    };
1563    e.l(format!(
1564        "// stage {} (Load {}: {}x{} tiles at tile ({},{}){step})",
1565        s.name, t.name, pattern.block_h, pattern.block_w, o.row_t, o.col_t
1566    ));
1567    if ctx.fifo.contains(&buf.id) {
1568        // A FIFO fills in stream order: one DRAM row of the block per enqueue burst.
1569        let (r0, c0) = (o.row_t * ctx.p.tile.rows, o.col_t * ctx.p.tile.cols);
1570        let (h, w) = (
1571            pattern.block_h * ctx.p.tile.rows,
1572            pattern.block_w * ctx.p.tile.cols,
1573        );
1574        e.open(format!("Foreach(0 until {h}) {{ row =>"));
1575        e.l(format!(
1576            "{} load {}({r0} + row, {c0}::{})",
1577            ctx.name(buf),
1578            dram_name(t),
1579            c0 + w
1580        ));
1581        e.close("}");
1582    } else {
1583        e.l(format!(
1584            "{} load {}({})",
1585            ctx.name(buf),
1586            dram_name(t),
1587            dram_range(&o, pattern.block_h, pattern.block_w, ctx.p.tile, kb)
1588        ));
1589    }
1590    Ok(())
1591}
1592
1593fn emit_store(ctx: &Ctx, e: &mut Emit, s: &Stage) -> Result<(), SpatialLowerError> {
1594    let StageKind::Store {
1595        pattern,
1596        sub_blocks,
1597        ..
1598    } = &s.kind
1599    else {
1600        unreachable!()
1601    };
1602    let t = ctx.tensor_of_load(s);
1603    let item = ctx.item_of(s);
1604    let o = resolve_origin(
1605        ctx.p,
1606        &pattern.origin,
1607        t,
1608        item,
1609        1,
1610        Some((pattern.block_h, pattern.block_w)),
1611    )?;
1612    let (h_t, w_t) = (
1613        sub_blocks.0 * pattern.block_h,
1614        sub_blocks.1 * pattern.block_w,
1615    );
1616    if o.row_t + h_t > t.tiles.rows_t || o.col_t + w_t > t.tiles.cols_t {
1617        return Err(SpatialLowerError::Unsupported(
1618            "a Store whose block runs past the tensor's tile grid",
1619        ));
1620    }
1621    let buf = ctx.buf(*s.ins.first().ok_or(SpatialLowerError::Unsupported(
1622        "a Store stage without an input buffer",
1623    ))?);
1624    let (rows, cols) = buffer_geom(ctx, buf);
1625    if (rows, cols) != (h_t * ctx.p.tile.rows, w_t * ctx.p.tile.cols) {
1626        return Err(SpatialLowerError::Unsupported(
1627            "a Store whose block does not match the buffer it drains (the compute's output block)",
1628        ));
1629    }
1630    e.l(format!(
1631        "// stage {} (Store {}: {}x{} sub-blocks of {}x{} tiles at tile ({},{}))",
1632        s.name,
1633        t.name,
1634        sub_blocks.0,
1635        sub_blocks.1,
1636        pattern.block_h,
1637        pattern.block_w,
1638        o.row_t,
1639        o.col_t
1640    ));
1641    if ctx.fifo.contains(&buf.id) {
1642        let (r0, c0) = (o.row_t * ctx.p.tile.rows, o.col_t * ctx.p.tile.cols);
1643        e.open(format!("Foreach(0 until {rows}) {{ row =>"));
1644        e.l(format!(
1645            "{}({r0} + row, {c0}::{}) store {}",
1646            dram_name(t),
1647            c0 + cols,
1648            ctx.name(buf)
1649        ));
1650        e.close("}");
1651    } else {
1652        e.l(format!(
1653            "{}({}) store {}",
1654            dram_name(t),
1655            dram_range(&o, h_t, w_t, ctx.p.tile, None),
1656            ctx.name(buf)
1657        ));
1658    }
1659    Ok(())
1660}
1661
1662/// The dot-product nest of one `MatmulTiles` stage, writing into `dst` (its output buffer, or its spill
1663/// buffer inside a `MemReduce`).
1664fn emit_matmul_nest(
1665    ctx: &Ctx,
1666    e: &mut Emit,
1667    s: &Stage,
1668    dst: &Buffer,
1669) -> Result<(), SpatialLowerError> {
1670    let Shape::Matmul { block, .. } = shape(s) else {
1671        unreachable!()
1672    };
1673    if s.ins.len() != 2 {
1674        return Err(SpatialLowerError::Unsupported(
1675            "a MatmulTiles stage without exactly two operand buffers (the A block and the B block)",
1676        ));
1677    }
1678    let (a, b) = (ctx.buf(s.ins[0]), ctx.buf(s.ins[1]));
1679    if ctx.fifo.contains(&a.id) || ctx.fifo.contains(&b.id) {
1680        return Err(SpatialLowerError::Unsupported("a MatmulTiles operand behind a Fifo handshake — the dot product needs random access to the block"));
1681    }
1682    let (rows, cols) = buffer_geom(ctx, dst);
1683    let kb = block.ibw * ctx.p.tile.cols;
1684    if buffer_geom(ctx, a) != (rows, kb) || buffer_geom(ctx, b) != (kb, cols) {
1685        return Err(SpatialLowerError::Unsupported("MatmulTiles operand blocks that do not match the output block (A must be rows x K block, B must be K block x cols)"));
1686    }
1687    let d = ctx.buf_ty(dst);
1688    let (ra, rb) = (
1689        cast(format!("{}(i, k)", ctx.name(a)), ctx.buf_ty(a), d),
1690        cast(format!("{}(k, j)", ctx.name(b)), ctx.buf_ty(b), d),
1691    );
1692    e.open(format!("Foreach(0 until {rows}) {{ i =>"));
1693    e.open(format!("Foreach(0 until {cols}) {{ j =>"));
1694    match ctx.flavor {
1695        SpatialFlavor::Fabric => e.l(format!(
1696            "val dot = Reduce(Reg[{}])(0 until {kb} par {}) {{ k => {ra} * {rb} }} {{ _ + _ }}",
1697            d.name(),
1698            lanes(kb)
1699        )),
1700        SpatialFlavor::Simulator => {
1701            e.l(format!("val dot = Reg[{}]", d.name()));
1702            e.l(format!("Sequential.Foreach(0 until {kb}) {{ k => dot := mux(k < 1, {ra} * {rb}, dot.value + {ra} * {rb}) }}"));
1703        }
1704    }
1705    e.l(format!("Pipe {{ {}(i, j) = dot }}", ctx.name(dst)));
1706    e.close("}");
1707    e.close("}");
1708    Ok(())
1709}
1710
1711fn emit_matmul(
1712    ctx: &Ctx,
1713    e: &mut Emit,
1714    s: &Stage,
1715    k_loads: &[&Stage],
1716) -> Result<(), SpatialLowerError> {
1717    let Shape::Matmul {
1718        block,
1719        sub_blocks,
1720        num_blocks,
1721        operand,
1722        acc,
1723        interm,
1724        ..
1725    } = shape(s)
1726    else {
1727        unreachable!()
1728    };
1729    let out = ctx.buf(*s.outs.first().ok_or(SpatialLowerError::Unsupported(
1730        "a MatmulTiles stage without an output buffer",
1731    ))?);
1732    let k_blocks = ctx.ix.k_blocks_of(s.id).unwrap_or(1);
1733    if k_blocks != num_blocks && k_blocks > 1 {
1734        return Err(SpatialLowerError::Unsupported("a MatmulTiles stage whose schedule loops over a different number of K blocks than its plan"));
1735    }
1736    e.l(format!(
1737        "// stage {} (MatmulTiles: {}x{} sub-blocks of {}x{} tiles, K block {} tiles, {:?} operands accumulated {:?}, output-stationary, {} K block(s))",
1738        s.name, sub_blocks.0, sub_blocks.1, block.osh, block.osw, block.ibw, operand, acc, num_blocks
1739    ));
1740    if k_blocks > 1 {
1741        if ctx.flavor == SpatialFlavor::Simulator {
1742            return Err(SpatialLowerError::Unsupported("a MatmulTiles reduction over several K blocks (a MemReduce) in the simulator flavor — its K-block fold is a Reduce the simulator cannot split"));
1743        }
1744        let spill = ctx.buf(interm.expect("refused above when absent"));
1745        e.open(format!(
1746            "MemReduce({})(0 until {k_blocks}) {{ kb =>",
1747            ctx.name(out)
1748        ));
1749        for l in k_loads {
1750            emit_load(ctx, e, l, Some("kb"))?;
1751        }
1752        emit_matmul_nest(ctx, e, s, spill)?;
1753        e.l(ctx.name(spill));
1754        e.close("} { _ + _ }");
1755    } else {
1756        emit_matmul_nest(ctx, e, s, out)?;
1757    }
1758    Ok(())
1759}
1760
1761/// One site's memories and controllers — of section `section` alone in a composed program, whose sections run one
1762/// after another.
1763fn emit_site(
1764    ctx: &Ctx,
1765    e: &mut Emit,
1766    site: Site,
1767    section: Option<usize>,
1768) -> Result<(), SpatialLowerError> {
1769    let p = ctx.p;
1770    let pl = p.placement.as_ref().unwrap();
1771    let order: &[StageId] = ctx
1772        .ix
1773        .order
1774        .get(&(site, section))
1775        .map(Vec::as_slice)
1776        .unwrap_or(&[]);
1777    let item = match section {
1778        None => &p.work.items[pl.items_at(site)[0].idx()],
1779        Some(_) => ctx.item_of(p.stage(*order.first().ok_or(SpatialLowerError::Unsupported(
1780            "a section with no stage on a site it is placed on",
1781        ))?)),
1782    };
1783    e.l(format!(
1784        "// site ({},{}): work item {} -> output block tiles ({},{}) {}x{}",
1785        site.x,
1786        site.y,
1787        item.id.idx(),
1788        item.out.row0_t,
1789        item.out.col0_t,
1790        item.out.h_t,
1791        item.out.w_t
1792    ));
1793    let used: BTreeSet<BufferId> = order
1794        .iter()
1795        .flat_map(|id| {
1796            let st = p.stage(*id);
1797            let interm = match shape(st) {
1798                Shape::Matmul { interm, .. } => interm,
1799                _ => None,
1800            };
1801            st.ins
1802                .iter()
1803                .chain(st.outs.iter())
1804                .copied()
1805                .chain(interm)
1806                .collect::<Vec<_>>()
1807        })
1808        .collect();
1809    let buffers: Vec<&Buffer> = match section {
1810        None => ctx
1811            .ix
1812            .site_buffers
1813            .get(&site)
1814            .into_iter()
1815            .flatten()
1816            .map(|b| p.buffer(*b))
1817            .collect(),
1818        Some(_) => used
1819            .iter()
1820            .map(|b| p.buffer(*b))
1821            .filter(|b| b.site == Some(site))
1822            .collect(),
1823    };
1824    for b in buffers {
1825        let (rows, cols) = buffer_geom(ctx, b);
1826        let tier = format!("{:?}", b.space);
1827        let depth = match b.kind {
1828            BufferKind::Fifo {
1829                depth_pages,
1830                page_tiles,
1831            } => format!("Fifo depth {depth_pages} pages x {page_tiles} tiles"),
1832            BufferKind::Interm {
1833                depth_pages,
1834                page_tiles,
1835            } => format!("Interm depth {depth_pages} pages x {page_tiles} tiles"),
1836            BufferKind::Accumulator { tiles } => format!("Accumulator {tiles} tiles"),
1837        };
1838        if ctx.fifo.contains(&b.id) {
1839            let cap = b.kind.tiles() * p.tile.rows * p.tile.cols;
1840            e.l(format!(
1841                "// buffer {} {tier} {depth} ({:?}) — Fifo handshake: a {cap}-element FIFO",
1842                b.name, b.dtype
1843            ));
1844            e.l(format!(
1845                "val {} = FIFO[{}]({cap})",
1846                ctx.name(b),
1847                ctx.buf_ty(b).name()
1848            ));
1849        } else {
1850            e.l(format!("// buffer {} {tier} {depth} ({:?}) — one page is {rows}x{cols}; Spatial N-buffers it across the controllers", b.name, b.dtype));
1851            e.l(format!(
1852                "val {} = SRAM[{}]({rows}, {cols})",
1853                ctx.name(b),
1854                ctx.buf_ty(b).name()
1855            ));
1856        }
1857    }
1858    // Loads the schedule runs once per K block belong inside the matmul's MemReduce.
1859    let k_loads: Vec<&Stage> = order
1860        .iter()
1861        .map(|id| p.stage(*id))
1862        .filter(|s| matches!(s.kind, StageKind::Load { .. }) && ctx.ix.k_blocks_of(s.id).is_some())
1863        .collect();
1864    for id in order {
1865        let s = p.stage(*id);
1866        match shape(s) {
1867            Shape::Load => {
1868                if ctx.ix.k_blocks_of(s.id).is_none() {
1869                    emit_load(ctx, e, s, None)?;
1870                }
1871            }
1872            Shape::Store => emit_store(ctx, e, s)?,
1873            Shape::Matmul { .. } => emit_matmul(ctx, e, s, &k_loads)?,
1874            Shape::Generic {
1875                iters,
1876                ins,
1877                out,
1878                body,
1879            } => {
1880                if ctx.ix.k_blocks_of(s.id).is_some() {
1881                    return Err(SpatialLowerError::Unsupported("a Generic compute stage inside the K-block loop — only the matmul's operand loads run per K block"));
1882                }
1883                emit_generic(ctx, e, s, iters, ins, out, body)?;
1884            }
1885            Shape::Recv | Shape::Epilogue | Shape::Scalar | Shape::AllReduce => {
1886                unreachable!("refused by refuse_unspelled")
1887            }
1888        }
1889    }
1890    Ok(())
1891}
1892
1893/// The host's element checks: each slice read back from its DRAM and compared with its gold — exactly for an integer
1894/// tensor, within its tolerance for a float one — every result printed, and `cksum` their conjunction.
1895fn emit_element_checks(
1896    e: &mut Emit,
1897    ctx: &Ctx,
1898    checks: &[SpatialCheck],
1899) -> Result<(), SpatialLowerError> {
1900    let mut oks: Vec<String> = Vec::new();
1901    for (k, c) in checks.iter().enumerate() {
1902        let t = ctx.p.tensors.iter().find(|t| t.name == c.tensor).ok_or(
1903            SpatialLowerError::Unsupported("a check naming a tensor the program does not have"),
1904        )?;
1905        let n = c.gold.len();
1906        if n == 0 || c.offset + n > t.elems.rows * t.elems.cols {
1907            return Err(SpatialLowerError::Unsupported(
1908                "a check whose slice is empty or runs past its tensor",
1909            ));
1910        }
1911        let (ty, float) = (ctx.ty(t.dtype), ctx.is_float(ctx.ty(t.dtype)));
1912        let (name, cols, off) = (dram_name(t), t.elems.cols, c.offset);
1913        e.l(format!(
1914            "// check {k}: {} elements {off}..{} of {} against the caller's gold",
1915            n,
1916            off + n,
1917            t.name
1918        ));
1919        // Each check is its own method: a gold table costs JVM code per literal, and `main` holds one method's worth.
1920        e.open(format!("def check{k}(): Bit = {{"));
1921        e.l(format!("val {name}_check{k} = getMatrix({name})"));
1922        let lits: Vec<String> = c.gold.iter().map(|v| lit_of(*v, ty, float)).collect();
1923        let table = format!("gold{k}");
1924        emit_table(e, &table, ty.name(), &lits);
1925        let got = format!("{name}_check{k}(({off} + q) / {cols}, ({off} + q) % {cols})");
1926        if float {
1927            let tol = c.rel_tol * c.gold.iter().fold(0.0f64, |m, v| m.max(v.abs())) + 1e-6;
1928            e.l(format!(
1929                "val err{k} = Array.tabulate({n}) {{ q => abs({got} - {table}_vals(q)) }}.max"
1930            ));
1931            e.l(format!("val ok{k} = err{k} <= {}", lit_of(tol, ty, true)));
1932            // The Scala backend's harness reads any output line holding "error" as a failure, so the deviation is
1933            // printed as a deviation.
1934            e.l(format!(
1935                "println(r\"check {k} {}: max deviation $err{k} tolerance {tol:e} ok $ok{k}\")",
1936                t.name
1937            ));
1938        } else {
1939            e.l(format!("val ok{k} = Array.tabulate({n}) {{ q => {got} == {table}_vals(q) }}.reduce {{ _ && _ }}"));
1940            e.l(format!(
1941                "println(r\"check {k} {}: exact ok $ok{k}\")",
1942                t.name
1943            ));
1944        }
1945        e.l(format!("ok{k}"));
1946        e.close("}");
1947        e.l(format!("val ok{k} = check{k}()"));
1948        oks.push(format!("ok{k}"));
1949    }
1950    e.l(format!("val cksum = {}", oks.join(" && ")));
1951    Ok(())
1952}
1953
1954/// The host-side fill of one tensor from its recorded init.
1955/// The literals of a table fill, as `val <name>_vals: Array[ty]`. A JVM method holds at most 64 KB of code and every
1956/// staged literal costs dozens of bytes of it, so a table past [`TABLE_CHUNK`] literals is built from chunks, each
1957/// its own `def`, joined with `++`.
1958fn emit_table(e: &mut Emit, name: &str, tn: &str, lits: &[String]) {
1959    if lits.len() <= TABLE_CHUNK {
1960        e.l(format!(
1961            "val {name}_vals = Array[{tn}]({})",
1962            lits.join(", ")
1963        ));
1964        return;
1965    }
1966    let chunks: Vec<&[String]> = lits.chunks(TABLE_CHUNK).collect();
1967    for (k, chunk) in chunks.iter().enumerate() {
1968        e.l(format!(
1969            "def {name}_part{k}(): Array[{tn}] = Array[{tn}]({})",
1970            chunk.join(", ")
1971        ));
1972    }
1973    let parts: Vec<String> = (0..chunks.len())
1974        .map(|k| format!("{name}_part{k}()"))
1975        .collect();
1976    e.l(format!("val {name}_vals = {}", parts.join(" ++ ")));
1977}
1978
1979/// Literals per table chunk: well inside a JVM method's code limit at every literal width.
1980const TABLE_CHUNK: usize = 256;
1981
1982/// With [`HostFill::OwnMethod`], the data is built inside its own `def` — a composed program's fills are many, and one
1983/// JVM method holding them all is past what the Scala backend can analyze; with [`HostFill::MethodOnly`] the DRAM is
1984/// a registry's ([`emit_dram_registry`]) and the caller calls the `def`.
1985fn emit_host_fill(e: &mut Emit, ctx: &Ctx, t: &TensorDecl, mode: HostFill) {
1986    let (rows, cols) = (t.elems.rows, t.elems.cols);
1987    let (pr, pc) = padded_dims(ctx.p, t);
1988    let name = dram_name(t);
1989    let (ty, float) = (ctx.ty(t.dtype), ctx.is_float(ctx.ty(t.dtype)));
1990    let tn = ty.name();
1991    e.l(format!(
1992        "// tensor {} {:?} {}x{} {:?} {:?}",
1993        t.name, t.role, rows, cols, t.dtype, t.space
1994    ));
1995    if mode == HostFill::OwnMethod {
1996        e.l(format!("val {name} = DRAM[{tn}]({pr}, {pc})"));
1997    }
1998    if mode != HostFill::Inline {
1999        e.open(format!("def fill_{name}(): Unit = {{"));
2000    }
2001    let lit = |v: f64| lit_of(v, ty, float);
2002    // A tensor is tiled, and a tile past its elements holds zeros — the oracle's padding, which a tile-block load
2003    // reads. `inside` keeps a value on the elements and zeroes the padding; a tile-aligned tensor has none.
2004    let padded = (pr, pc) != (rows, cols);
2005    let zero = format!("0.to[{tn}]");
2006    let inside = |v: String| {
2007        if padded {
2008            format!("mux(i < {rows}, mux(j < {cols}, {v}, {zero}), {zero})")
2009        } else {
2010            v
2011        }
2012    };
2013    let index = if padded {
2014        format!("mux(i < {rows}, mux(j < {cols}, i * {cols} + j, 0), 0)")
2015    } else {
2016        format!("i * {cols} + j")
2017    };
2018    match &t.init {
2019        TensorInit::Zero => e.l(format!(
2020            "val {name}_data = (0::{pr}, 0::{pc}) {{ (i, j) => 0.to[{tn}] }}"
2021        )),
2022        TensorInit::ConstF(b) => e.l(format!(
2023            "val {name}_data = (0::{pr}, 0::{pc}) {{ (i, j) => {} }}",
2024            inside(format!("{}.to[{tn}]", lit(f64::from_bits(*b))))
2025        )),
2026        TensorInit::Pattern { modulus, shift } => {
2027            let shift = if *shift < 0 {
2028                format!("- {}", -shift)
2029            } else {
2030                format!("+ {shift}")
2031            };
2032            e.l(format!(
2033                "val {name}_data = (0::{pr}, 0::{pc}) {{ (i, j) => {} }}",
2034                inside(format!("(((i * {cols} + j) % {modulus}) {shift}).to[{tn}]"))
2035            ));
2036        }
2037        TensorInit::Values(_) | TensorInit::FloatValues(_) => {
2038            let vals: Vec<f64> = match &t.init {
2039                TensorInit::Values(v) => v.iter().take(rows * cols).map(|x| *x as f64).collect(),
2040                TensorInit::FloatValues(v) => v
2041                    .iter()
2042                    .take(rows * cols)
2043                    .map(|b| f64::from_bits(*b))
2044                    .collect(),
2045                _ => unreachable!(),
2046            };
2047            if vals.len() == rows * cols
2048                && vals.windows(2).all(|w| w[0].to_bits() == w[1].to_bits())
2049                && !vals.is_empty()
2050            {
2051                // One value everywhere (a scaler table): a formula, not a table.
2052                e.l(format!(
2053                    "val {name}_data = (0::{pr}, 0::{pc}) {{ (i, j) => {} }}",
2054                    inside(lit(vals[0]))
2055                ));
2056            } else {
2057                let lits: Vec<String> = vals.iter().map(|v| lit(*v)).collect();
2058                emit_table(e, &name, tn, &lits);
2059                e.l(format!(
2060                    "val {name}_data = (0::{pr}, 0::{pc}) {{ (i, j) => {} }}",
2061                    inside(format!("{name}_vals({index})"))
2062                ));
2063            }
2064        }
2065        TensorInit::External | TensorInit::RandomBf16 { .. } => {
2066            unreachable!("refused by refuse_unspelled")
2067        }
2068    }
2069    match mode {
2070        HostFill::Inline => {
2071            e.l(format!("val {name} = DRAM[{tn}]({pr}, {pc})"));
2072            e.l(format!("setMem({name}, {name}_data)"));
2073        }
2074        HostFill::OwnMethod => {
2075            e.l(format!("setMem({name}, {name}_data)"));
2076            e.close("}");
2077            e.l(format!("fill_{name}()"));
2078        }
2079        HostFill::MethodOnly => {
2080            e.l(format!("setMem({name}, {name}_data)"));
2081            e.close("}");
2082        }
2083    }
2084}
2085
2086/// How the host declares and fills one tensor.
2087#[derive(Clone, Copy, PartialEq, Eq)]
2088enum HostFill {
2089    /// `val d = DRAM…`, its data, `setMem` — all in `main` (a one-section program).
2090    Inline,
2091    /// `val d = DRAM…` in `main`, the data and `setMem` in `def fill_d()`, called at once.
2092    OwnMethod,
2093    /// Only `def fill_d()`: the DRAM is a registry's, and the fills are called in chunks.
2094    MethodOnly,
2095}
2096
2097/// ★ THE JVM BOUNDS ONE METHOD'S CODE AT 64 KB, AND THE HOST IS ONE METHOD. Every staged `DRAM[T](r, c)` carries
2098/// the macro-generated source context Spatial threads through it — a few hundred bytes of `main` each — so a
2099/// composed program's host holds [`HOST_INLINE`] tensors at most (measured: 201 compile, 398 do not). Past that, the
2100/// DRAMs are made in [`DRAM_CHUNK`]s by methods of their own into registries (a Scala `Vector` — `Array` is Spatial's
2101/// staged host array), each tensor's name an accessor into one (`def d_x: DRAM2[T] = dram_T_0_all(5)` — a method, so
2102/// `main` holds none of it), and the fills are called in chunks the same way.
2103const HOST_INLINE: usize = 128;
2104const DRAM_CHUNK: usize = 64;
2105
2106/// The DRAM registries of a host past [`HOST_INLINE`] tensors: per element type, [`DRAM_CHUNK`] allocations to a
2107/// method, the method's `Seq` held once in `main`, and one accessor per tensor.
2108fn emit_dram_registry(e: &mut Emit, ctx: &Ctx, host: &[TensorDecl]) {
2109    let mut by_type: BTreeMap<&'static str, Vec<&TensorDecl>> = BTreeMap::new();
2110    for t in host {
2111        by_type.entry(ctx.ty(t.dtype).name()).or_default().push(t);
2112    }
2113    for (tn, tensors) in &by_type {
2114        for (c, chunk) in tensors.chunks(DRAM_CHUNK).enumerate() {
2115            e.open(format!("def dram_{tn}_{c}(): scala.collection.immutable.Vector[DRAM2[{tn}]] = scala.collection.immutable.Vector("));
2116            for (i, t) in chunk.iter().enumerate() {
2117                let (pr, pc) = padded_dims(ctx.p, t);
2118                let sep = if i + 1 < chunk.len() { "," } else { "" };
2119                e.l(format!("DRAM[{tn}]({pr}, {pc}){sep} // {}", t.name));
2120            }
2121            e.close(")");
2122            e.l(format!("val dram_{tn}_{c}_all = dram_{tn}_{c}()"));
2123        }
2124    }
2125    for (tn, tensors) in &by_type {
2126        for (c, chunk) in tensors.chunks(DRAM_CHUNK).enumerate() {
2127            for (i, t) in chunk.iter().enumerate() {
2128                e.l(format!(
2129                    "def {}: DRAM2[{tn}] = dram_{tn}_{c}_all({i})",
2130                    dram_name(t)
2131                ));
2132            }
2133        }
2134    }
2135}
2136
2137/// ★ SPATIAL WRITES A CONTROLLER'S BLOCK AS ONE PARAMETER LIST, AND THE JVM BOUNDS ONE AT 254. A sequence past
2138/// [`SEQUENCE_CHUNK`] calls is nested — each chunk its own `def` running its calls in a `Sequential` — until the
2139/// outermost holds at most that many; the order is the calls' order at every level. Returns the outermost calls.
2140const SEQUENCE_CHUNK: usize = 128;
2141
2142fn chunk_sequence(
2143    e: &mut Emit,
2144    what: &str,
2145    mut calls: Vec<String>,
2146    controller: Option<&str>,
2147) -> Vec<String> {
2148    let mut level = 0;
2149    while calls.len() > SEQUENCE_CHUNK {
2150        let mut next = Vec::new();
2151        for (g, chunk) in calls.chunks(SEQUENCE_CHUNK).enumerate() {
2152            let name = format!("{what}_{level}_{g}");
2153            e.open(format!("def {name}(): Unit = {{"));
2154            if let Some(c) = controller {
2155                e.open(format!("{c} {{"));
2156            }
2157            for call in chunk {
2158                e.l(call);
2159            }
2160            if controller.is_some() {
2161                e.close("}");
2162            }
2163            e.close("}");
2164            next.push(format!("{name}()"));
2165        }
2166        calls = next;
2167        level += 1;
2168    }
2169    calls
2170}
2171
2172/// A tensor's DRAM extent: its tile grid, in elements — every tile a Load or Store moves lies inside it.
2173fn padded_dims(p: &SpatialProgram, t: &TensorDecl) -> (usize, usize) {
2174    (t.tiles.rows_t * p.tile.rows, t.tiles.cols_t * p.tile.cols)
2175}
2176
2177/// Emit the Spatial regression app (`@spatial class Logos_<name> extends SpatialTest`) for a verified spatial
2178/// program: every tensor a `DRAM[T]` filled from its recorded fill, every buffer an `SRAM[T]`/`FIFO[T]`, every
2179/// stage one controller, and the host checking the output tensor's checksum against `gold` — the VM's answer
2180/// when the caller has it (`Show checksum` of the output), the spatial oracle's (`eval_choreography`)
2181/// otherwise. `target` is the descriptor the program was compiled for (its name is recorded in the header).
2182/// The harness's `PASS: true (Logos_<name>)` line is what the toolchain ledger parses.
2183pub fn emit_spatial_program(
2184    p: &SpatialProgram,
2185    target: &SpatialTarget,
2186    gold: Option<&str>,
2187) -> Result<String, SpatialLowerError> {
2188    emit_app(p, target, AppCheck::Checksum(gold), SpatialFlavor::Fabric)
2189}
2190
2191/// A slice of a tensor the app checks, element by element, after the accelerator has run: `gold` holds the values
2192/// the caller vouches for (the VM's) for the row-major elements `offset..offset + gold.len()`. An integer tensor
2193/// must match exactly; a float one within `rel_tol · max|gold|`.
2194#[derive(Clone, Debug, PartialEq)]
2195pub struct SpatialCheck {
2196    pub tensor: String,
2197    pub offset: usize,
2198    pub gold: Vec<f64>,
2199    pub rel_tol: f64,
2200}
2201
2202/// [`emit_spatial_program`] with the host checking every slice of `checks` — the output, and any activation or
2203/// carried state the caller holds a gold for — rather than one checksum of the output.
2204pub fn emit_spatial_program_checked(
2205    p: &SpatialProgram,
2206    target: &SpatialTarget,
2207    checks: &[SpatialCheck],
2208    flavor: SpatialFlavor,
2209) -> Result<String, SpatialLowerError> {
2210    if checks.is_empty() {
2211        return Err(SpatialLowerError::Unsupported(
2212            "a checked program with no tensor to check — its host would pass against nothing",
2213        ));
2214    }
2215    emit_app(p, target, AppCheck::Elements(checks), flavor)
2216}
2217
2218enum AppCheck<'a> {
2219    Checksum(Option<&'a str>),
2220    Elements(&'a [SpatialCheck]),
2221    /// A run of sections, entered with the oracle's values and checked against them.
2222    Segment {
2223        sections: std::ops::Range<usize>,
2224        values: &'a BTreeMap<TensorId, Value>,
2225        rel_tol: f64,
2226    },
2227}
2228
2229/// A run of consecutive sections of a composed program, printed as ITS OWN app — the unit PIR can map, since PIR
2230/// writes an accelerator as one JVM class and a whole composed program is past what the JVM allows one. The app
2231/// holds the tensors the segment's stages touch, each filled with its value at the segment's ENTRY; runs the
2232/// sections in order; and checks every tensor they store against its value at the segment's EXIT — exactly for an
2233/// integer tensor, within `rel_tol · max|value|` for a float one.
2234///
2235/// `values` is every tensor at the program's end ([`crate::spatial::eval::eval_choreography_every_tensor`]). A
2236/// composed program stores each tensor in one section at most, so a tensor stored before the segment enters with
2237/// that value, and any other enters with its fill.
2238pub fn emit_spatial_segment(
2239    p: &SpatialProgram,
2240    target: &SpatialTarget,
2241    sections: std::ops::Range<usize>,
2242    values: &BTreeMap<TensorId, Value>,
2243    rel_tol: f64,
2244    flavor: SpatialFlavor,
2245) -> Result<String, SpatialLowerError> {
2246    if sections.is_empty() || sections.end > p.sections.len() {
2247        return Err(SpatialLowerError::Unsupported(
2248            "a segment that is empty or runs past the program's sections",
2249        ));
2250    }
2251    emit_app(
2252        p,
2253        target,
2254        AppCheck::Segment {
2255            sections,
2256            values,
2257            rel_tol,
2258        },
2259        flavor,
2260    )
2261}
2262
2263/// A tensor's elements as a fill: the oracle's value, carried into a segment's app.
2264fn fill_of_value(v: &Value) -> TensorInit {
2265    match v {
2266        Value::Float(f) => TensorInit::FloatValues(f.iter().map(|x| x.to_bits()).collect()),
2267        Value::Int(i) => TensorInit::Values(i.clone()),
2268    }
2269}
2270
2271fn emit_app(
2272    p: &SpatialProgram,
2273    target: &SpatialTarget,
2274    check: AppCheck,
2275    flavor: SpatialFlavor,
2276) -> Result<String, SpatialLowerError> {
2277    let gold = match check {
2278        AppCheck::Checksum(g) => g,
2279        AppCheck::Elements(_) | AppCheck::Segment { .. } => None,
2280    };
2281    // The sections this app runs, and the tensors it holds — every one, or a segment's.
2282    let run: std::ops::Range<usize> = match &check {
2283        AppCheck::Segment { sections, .. } => sections.clone(),
2284        _ => 0..p.sections.len(),
2285    };
2286    let touched: BTreeSet<TensorId> = p
2287        .stages
2288        .iter()
2289        .filter(|st| {
2290            p.sections.is_empty() || p.section_of_stage(st.id).is_some_and(|k| run.contains(&k))
2291        })
2292        .filter_map(|st| match &st.kind {
2293            StageKind::Load { tensor, .. } | StageKind::Store { tensor, .. } => Some(*tensor),
2294            _ => None,
2295        })
2296        .collect();
2297    require_verified(p)?;
2298    refuse_unspelled(p)?;
2299    let inputs = fill_inputs(p)?;
2300    let outputs = eval_choreography(p, &inputs).map_err(SpatialLowerError::Uncertified)?;
2301    let out_t = p.output_tensor().ok_or(SpatialLowerError::Unsupported(
2302        "a program without an Output tensor has nothing for the host to check",
2303    ))?;
2304    let out_v = outputs
2305        .tensors
2306        .get(&out_t.id)
2307        .ok_or(SpatialLowerError::Unsupported(
2308            "the oracle produced no value for the Output tensor",
2309        ))?;
2310    let ix = Index::of(p);
2311    let scalar = scalar_type(p, &ix, &inputs, out_v);
2312    // A program holding floats and integers both computes its floats in `T` and declares its integer memories —
2313    // indices, positions, tokens — `I`.
2314    let dtypes: Vec<Precision> = p
2315        .tensors
2316        .iter()
2317        .map(|t| t.dtype)
2318        .chain(p.buffers.iter().map(|b| b.dtype))
2319        .collect();
2320    let int_scalar = (dtypes.iter().any(|d| is_float_precision(*d))
2321        && dtypes.iter().any(|d| !is_float_precision(*d)))
2322    .then(|| {
2323        if dtypes.contains(&Precision::Int64) {
2324            Scalar::Long
2325        } else {
2326            Scalar::Int
2327        }
2328    });
2329    let gold_source = if gold.is_some() {
2330        "the VM's checksum, handed over by the caller"
2331    } else {
2332        "the spatial oracle's checksum (eval_choreography over the recorded fills)"
2333    };
2334    let gold_note = format!(" // {gold_source}");
2335    // The gold is of the output's class: a program of floats and integers may emit an integer (a token).
2336    let out_scalar = if int_scalar.is_some() && !is_float_precision(out_t.dtype) {
2337        int_scalar.unwrap()
2338    } else {
2339        scalar
2340    };
2341    let gold = match gold {
2342        Some(s) => parse_gold(s, out_scalar)?,
2343        None => checksum_gold(out_v),
2344    };
2345    if let Gold::Float(v) = gold {
2346        if !v.is_finite() {
2347            return Err(SpatialLowerError::Unsupported("a float gold checksum that is not finite — the output tensor overflowed its precision"));
2348        }
2349    }
2350    let fifo: BTreeSet<BufferId> = p
2351        .edges
2352        .iter()
2353        .filter(|e| matches!(e.handshake, Handshake::Fifo))
2354        .map(|e| e.buffer)
2355        .collect();
2356    let ctx = Ctx {
2357        p,
2358        scalar,
2359        int_scalar,
2360        fifo,
2361        names: buffer_names(p),
2362        flavor,
2363        ix,
2364    };
2365    let class = match &check {
2366        AppCheck::Segment { sections, .. } => format!(
2367            "Logos_{}_s{}_{}",
2368            ident(&p.name),
2369            sections.start,
2370            sections.end
2371        ),
2372        _ => format!("Logos_{}", ident(&p.name)),
2373    };
2374    let sites: Vec<Site> = p.sites().into_iter().collect();
2375
2376    let mut e = Emit {
2377        s: String::new(),
2378        ind: 0,
2379    };
2380    e.l(format!(
2381        "// Generated by LOGOS (logicaffeine_compile::spatial::lower::spatial_lang) — Spatial regression app for the verified\n// spatial program `{}` on {} ({} site(s), tile {}x{}, numerics {:?}). Every controller and memory below is the\n// image of one IR object (stage / buffer / tensor); the host fills every tensor from its recorded fill and checks the\n// output checksum against the gold (the VM's checksum when the caller hands it over, else the spatial oracle's).\n// Values are Spatial `{}`.",
2382        p.name, target.name, sites.len(), p.tile.rows, p.tile.cols, p.claim, scalar.spatial()
2383    ));
2384    e.l("import spatial.dsl._");
2385    e.l("");
2386    e.open(format!("@spatial class {class} extends SpatialTest {{"));
2387    e.l(format!("type T = {}", scalar.spatial()));
2388    if let Some(i) = int_scalar {
2389        e.l(format!("type I = {}", i.spatial()));
2390    }
2391    e.open("def main(args: Array[String]): Unit = {");
2392    // A segment enters with the oracle's value of every tensor an earlier section stored.
2393    let mut host: Vec<TensorDecl> = Vec::new();
2394    for t in p
2395        .tensors
2396        .iter()
2397        .filter(|t| touched.contains(&t.id) || !matches!(check, AppCheck::Segment { .. }))
2398    {
2399        host.push(match &check {
2400            AppCheck::Segment {
2401                sections, values, ..
2402            } if ctx
2403                .ix
2404                .storing_section(t.id)
2405                .is_some_and(|k| k < sections.start) =>
2406            {
2407                let v = values.get(&t.id).ok_or(SpatialLowerError::Unsupported(
2408                    "a segment entering with a tensor the oracle's values do not hold",
2409                ))?;
2410                TensorDecl {
2411                    init: fill_of_value(v),
2412                    role: TensorRole::Activation,
2413                    ..t.clone()
2414                }
2415            }
2416            _ => t.clone(),
2417        });
2418    }
2419    if !p.sections.is_empty() && host.len() > HOST_INLINE {
2420        emit_dram_registry(&mut e, &ctx, &host);
2421        for t in host.iter().filter(|t| t.role != TensorRole::Output) {
2422            emit_host_fill(&mut e, &ctx, t, HostFill::MethodOnly);
2423        }
2424        let fills: Vec<String> = host
2425            .iter()
2426            .filter(|t| t.role != TensorRole::Output)
2427            .map(|t| format!("fill_{}()", dram_name(t)))
2428            .collect();
2429        let mut chunked = Vec::new();
2430        for (g, chunk) in fills.chunks(DRAM_CHUNK).enumerate() {
2431            e.open(format!("def fill_all_{g}(): Unit = {{"));
2432            for call in chunk {
2433                e.l(call);
2434            }
2435            e.close("}");
2436            chunked.push(format!("fill_all_{g}()"));
2437        }
2438        for call in chunked {
2439            e.l(call);
2440        }
2441    } else {
2442        for t in &host {
2443            if t.role == TensorRole::Output {
2444                e.l(format!(
2445                    "// tensor {} {:?} {}x{} {:?} {:?}",
2446                    t.name, t.role, t.elems.rows, t.elems.cols, t.dtype, t.space
2447                ));
2448                let (pr, pc) = padded_dims(p, t);
2449                e.l(format!(
2450                    "val {} = DRAM[{}]({pr}, {pc})",
2451                    dram_name(t),
2452                    ctx.ty(t.dtype).name()
2453                ));
2454            } else {
2455                emit_host_fill(
2456                    &mut e,
2457                    &ctx,
2458                    t,
2459                    if p.sections.is_empty() {
2460                        HostFill::Inline
2461                    } else {
2462                        HostFill::OwnMethod
2463                    },
2464                );
2465            }
2466        }
2467    }
2468    // A composed program's sections are each their own method, for the same reason as its fills.
2469    for (k, section) in p
2470        .sections
2471        .iter()
2472        .enumerate()
2473        .filter(|(k, _)| run.contains(k))
2474    {
2475        let on = &ctx.ix.section_sites[k];
2476        e.l(format!(
2477            "// section {k}: {} ({} site(s))",
2478            section.name,
2479            on.len()
2480        ));
2481        e.open(format!("def section_{k}(): Unit = {{"));
2482        e.open(if on.len() == 1 {
2483            "Sequential {"
2484        } else {
2485            "Parallel {"
2486        });
2487        for site in on {
2488            if on.len() > 1 {
2489                e.open("Pipe {");
2490            }
2491            emit_site(&ctx, &mut e, *site, Some(k))?;
2492            if on.len() > 1 {
2493                e.close("}");
2494            }
2495        }
2496        e.close("}");
2497        e.close("}");
2498    }
2499    let sequence = chunk_sequence(
2500        &mut e,
2501        "run",
2502        run.clone().map(|k| format!("section_{k}()")).collect(),
2503        Some("Sequential"),
2504    );
2505    e.open("Accel {");
2506    if p.sections.is_empty() {
2507        if sites.len() == 1 {
2508            emit_site(&ctx, &mut e, sites[0], None)?;
2509        } else {
2510            e.open("Parallel {");
2511            for site in &sites {
2512                e.open("Pipe {");
2513                emit_site(&ctx, &mut e, *site, None)?;
2514                e.close("}");
2515            }
2516            e.close("}");
2517        }
2518    } else {
2519        // ★ THE RENDEZVOUS IS THE SEQUENCE. A composed program's sections hand over through DRAM, and its gather /
2520        // release barrier holds every site of section k+1 until every site of section k is done. Spatial orders no
2521        // DRAM store before a DRAM load in a parallel controller, so the barrier is spelled as what it means: a
2522        // `Sequential` over the sections, each a `Parallel` of its sites.
2523        e.open("Sequential {");
2524        for call in &sequence {
2525            e.l(call);
2526        }
2527        e.close("}");
2528    }
2529    e.close("}");
2530    let segment_checks: Vec<SpatialCheck>;
2531    let element_checks = match &check {
2532        AppCheck::Elements(checks) => Some(*checks),
2533        AppCheck::Segment {
2534            sections,
2535            values,
2536            rel_tol,
2537        } => {
2538            segment_checks = p
2539                .tensors
2540                .iter()
2541                .filter(|t| {
2542                    ctx.ix
2543                        .storing_section(t.id)
2544                        .is_some_and(|k| sections.contains(&k))
2545                })
2546                .map(|t| {
2547                    let v = values.get(&t.id).ok_or(SpatialLowerError::Unsupported(
2548                        "a segment storing a tensor the oracle's values do not hold",
2549                    ))?;
2550                    let gold: Vec<f64> = match v {
2551                        Value::Float(f) => f.clone(),
2552                        Value::Int(i) => i.iter().map(|x| *x as f64).collect(),
2553                    };
2554                    Ok(SpatialCheck {
2555                        tensor: t.name.clone(),
2556                        offset: 0,
2557                        gold,
2558                        rel_tol: *rel_tol,
2559                    })
2560                })
2561                .collect::<Result<_, SpatialLowerError>>()?;
2562            if segment_checks.is_empty() {
2563                return Err(SpatialLowerError::Unsupported(
2564                    "a segment that stores nothing — it would pass against nothing",
2565                ));
2566            }
2567            Some(&segment_checks[..])
2568        }
2569        AppCheck::Checksum(_) => None,
2570    };
2571    if let Some(checks) = element_checks {
2572        emit_element_checks(&mut e, &ctx, checks)?;
2573        e.l(format!("println(\"PASS: \" + cksum + \" ({class})\")"));
2574        e.l("assert(cksum)");
2575        e.close("}");
2576        e.close("}");
2577        return Ok(e.s);
2578    }
2579    let out_name = dram_name(out_t);
2580    let (rows, cols) = (out_t.elems.rows, out_t.elems.cols);
2581    e.l(format!("val {out_name}_result = getMatrix({out_name})"));
2582    match gold {
2583        Gold::Int(v) => {
2584            e.l(format!("val checksum = Array.tabulate({}) {{ q => {out_name}_result(q / {cols}, q % {cols}).to[Long] }}.reduce {{ _ + _ }}", rows * cols));
2585            e.l(format!("val gold = {}{gold_note}", long_lit(v)));
2586            e.l("val cksum = checksum == gold");
2587        }
2588        Gold::Float(v) => {
2589            let tol = v.abs() * 1e-2 + 1e-3;
2590            e.l(format!("val checksum = Array.tabulate({}) {{ q => {out_name}_result(q / {cols}, q % {cols}) }}.reduce {{ _ + _ }}", rows * cols));
2591            e.l(format!("val gold = {}{gold_note}", float_lit(v)));
2592            e.l(format!("val tol = {}", float_lit(tol)));
2593            e.l("val cksum = abs(checksum - gold) <= tol");
2594        }
2595    }
2596    e.l("println(r\"checksum $checksum gold $gold\")");
2597    e.l(format!("println(\"PASS: \" + cksum + \" ({class})\")"));
2598    e.l("assert(cksum)");
2599    e.close("}");
2600    e.close("}");
2601    Ok(e.s)
2602}