Skip to main content

jay/
fuse.rs

1//! Fusing chains of elementwise verbs into one blockwise pass.
2//!
3//! A chain like `+/ w * x` runs one pass per verb: the product is written to
4//! memory in full and read back to be reduced. This pass finds maximal
5//! subtrees of elementwise primitives at compile time and replaces them with
6//! [`Expr::Fused`], which evaluates the whole chain a block at a time — the
7//! block stays in cache, so the arrays at the leaves are read once and the
8//! result is written once.
9//!
10//! The kernel is a postfix program over a small stack of block buffers. It
11//! covers only what it can compute exactly as the unfused pipeline would;
12//! everything else — a shape that needs broadcasting, a dtype the chain
13//! would narrow, an integer overflow — declines at run time and the original
14//! subtree, kept inside the node, evaluates instead. Fusion therefore cannot
15//! change a result or an error message.
16//!
17//! A chain does not have to be written as one sentence. `d =. {x} - m`
18//! followed by `+/ d * d` names a value that nothing needs as an array, and
19//! the pass moves such a value into the sentences that read it — see
20//! `inline_once` for the rules that keep that sound.
21
22use std::sync::atomic::{AtomicU64, Ordering};
23
24use crate::array::{Array, Data, Layout};
25use crate::dtype::DType;
26use crate::error::Span;
27use crate::ir::{Expr, Program, Scope};
28use crate::par;
29use crate::simd::multiversioned;
30use crate::verb::{
31    tol_cmp, windows_into, DyadOp, MonadOp, ScalarDyad, ScalarMonad, Tol, Verb, WindowKind,
32    RANK_INF,
33};
34
35/// Elements a block buffer holds.
36///
37/// The working set is `slots` buffers of this size — two or three for the
38/// benchmark kernels — so 8,192 f64 is 128 to 192 KB and stays inside a
39/// 256 KB L2. The value is not delicate: measured at 2,048 / 4,096 / 8,192 /
40/// 16,384 / 32,768 on `+/ w * x` and `+/ ^ x` over 20M rows, the whole range
41/// lands within a few per cent of the best, because what the kernel is
42/// really bounded by is streaming the leaves in from memory once.
43pub const BLOCK: usize = 8_192;
44
45/// The largest window a kernel absorbs.
46///
47/// A block computes the wide axis its own windows need, which is the block
48/// plus a halo of about three window lengths, and holds it in the same
49/// buffers the arithmetic uses. Past this size the halo is most of the work
50/// and the buffers are past any cache worth staying in, so a longer window
51/// stays outside the kernel and takes the pass it has always taken.
52pub const MAX_WINDOW: usize = 1_024;
53
54/// One step of a kernel: postfix, so operands are already on the stack.
55#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub enum Instr {
57    /// Push input `k`.
58    Load(usize),
59    /// Replace the top of the stack.
60    Monad(ScalarMonad),
61    /// Replace the top two, left below right.
62    Dyad(ScalarDyad),
63    /// Keep the top of the stack as let `k`, a value the rest of the
64    /// program reads more than once. It holds its block buffer until the
65    /// block is finished; nothing pops it.
66    Store(usize),
67    /// Push let `k` again.
68    Let(usize),
69    /// Fold every window of `k` consecutive items of the top of the stack
70    /// into one item. The operand stands on the wide axis and the result on
71    /// the kernel's own, which is `k - 1` items shorter.
72    Window(ScalarDyad, usize),
73    /// Replace the top of the stack with its running fold: item `i` becomes
74    /// the fold of items `0 .. i`. Both stand on the same axis.
75    Scan(ScalarDyad),
76}
77
78/// Which axis a value inside a kernel stands on.
79///
80/// A kernel that folds windows reads two: the one its result stands on, and
81/// the wider one every window step reads, which is `k - 1` items longer.
82/// Where a value stands is decided by the chain — everything under a window
83/// step is wide — so the two never have to be told apart at run time.
84#[derive(Clone, Copy, Debug, PartialEq, Eq)]
85enum Dom {
86    Result,
87    Wide,
88}
89
90/// The stages a chain absorbs.
91///
92/// A chain takes moving windows or running folds, not both, and every
93/// window step in one kernel folds windows of the same length: that is what
94/// leaves exactly two axes to align, which shapes alone can then decide.
95/// Anything else — a second window length, a window inside a window, a
96/// running fold beside a window — is read as a leaf and runs as the pass it
97/// was.
98#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
99struct Plan {
100    window: Option<usize>,
101    scan: bool,
102}
103
104/// What one evaluation of a kernel produces.
105#[derive(Clone, Copy, Debug, PartialEq, Eq)]
106pub enum Yield {
107    /// The mapped values, as an array of the chain's own shape.
108    Values,
109    /// The mapped values folded into one by an absorbed reduction.
110    Reduce(ScalarDyad),
111    /// How many items the mapped values would have — `#` over a chain. The
112    /// shapes answer that before any arithmetic runs, so none runs.
113    Tally,
114}
115
116/// A fused elementwise chain and what is made of its values.
117#[derive(Clone, Debug)]
118pub struct FusedKernel {
119    code: Vec<Instr>,
120    /// Block buffers one evaluation needs at once.
121    slots: usize,
122    yields: Yield,
123    /// The input each leaf of the chain reads, in the order the chain
124    /// reaches them. Two leaves that are the same subtree share one input,
125    /// so this is not the identity, and the fallback needs it to give every
126    /// leaf back the value it was given.
127    leaves: Vec<usize>,
128    /// The axis each input is read on. None is an input the chain reads on
129    /// both, which only a scalar can satisfy.
130    doms: Vec<Option<Dom>>,
131    /// The axis each let stands on, which is the axis every repeat it was
132    /// taken for was written on.
133    let_doms: Vec<Dom>,
134    /// The stages the chain was built with, so that the fallback walks it
135    /// exactly as the pass walked it.
136    plan: Plan,
137    /// The window every window step folds, when the code holds one.
138    window: Option<usize>,
139    /// Running folds in the code. Each carries an accumulator from block to
140    /// block, so a kernel that has any runs its blocks in order.
141    scans: usize,
142    /// The dialect's comparison tolerance, so that a comparison inside the
143    /// kernel answers exactly as the same comparison outside it does.
144    tol: Tol,
145}
146
147impl FusedKernel {
148    pub fn code(&self) -> &[Instr] {
149        &self.code
150    }
151
152    pub fn yields(&self) -> Yield {
153        self.yields
154    }
155
156    pub fn reduce(&self) -> Option<ScalarDyad> {
157        match self.yields {
158            Yield::Reduce(op) => Some(op),
159            _ => None,
160        }
161    }
162
163    /// The comparison tolerance the program was compiled with. A backend
164    /// that generates its own code for this kernel needs it, so that a
165    /// comparison answers there as it answers everywhere else.
166    pub fn tol(&self) -> Tol {
167        self.tol
168    }
169}
170
171/// How often a fused node has handed its work back to the original subtree.
172/// A counter rather than a log: the fallback is correct, only slower, and
173/// what a caller wants to know is whether it is happening at all.
174static FALLBACKS: AtomicU64 = AtomicU64::new(0);
175
176/// Number of fallbacks since the process started.
177pub fn fallback_count() -> u64 {
178    FALLBACKS.load(Ordering::Relaxed)
179}
180
181fn note_fallback() {
182    FALLBACKS.fetch_add(1, Ordering::Relaxed);
183}
184
185// ------------------------------------------------------------- the op set
186//
187// A verb may join a kernel only if it cannot fail on numeric data: the
188// kernel reports no errors of its own, so anything that could raise one
189// (APL's `÷` by zero, `%:` and `^.` of a negative, `^`'s zero to a negative
190// power, APL's `~` off 0/1) stays outside and breaks the chain there.
191
192/// The elementwise monad this verb performs, if the kernel covers it.
193fn fusable_monad(v: &Verb) -> Option<ScalarMonad> {
194    use ScalarMonad::*;
195    let Verb::Prim(p) = v else { return None };
196    let MonadOp::Scalar(op) = p.monad else { return None };
197    matches!(
198        op,
199        Conj | Neg | Abs | Signum | Recip | Floor | Ceil | Inc | Dec | Double | Halve | Square
200            | OneMinus | Exp
201    )
202    .then_some(op)
203}
204
205/// The elementwise dyad this verb performs, if the kernel covers it.
206fn fusable_dyad(v: &Verb) -> Option<ScalarDyad> {
207    use ScalarDyad::*;
208    let Verb::Prim(p) = v else { return None };
209    let DyadOp::Scalar(op) = p.dyad else { return None };
210    matches!(op, Add | Sub | Mul | DivJ | Min | Max | Residue | Eq | Ne | Lt | Le | Gt | Ge)
211        .then_some(op)
212}
213
214/// The reduction this verb performs over the leading axis, if the kernel can
215/// absorb it: an associative arithmetic primitive, applied at full rank.
216/// APL's `+/` is the same thing under a rank wrapper.
217fn absorbable_reduce(v: &Verb) -> Option<ScalarDyad> {
218    use ScalarDyad::*;
219    let inner = match v {
220        // APL's `f/` is the same fold monadically; only its dyad differs.
221        Verb::Reduce(u) | Verb::NWise(u) => u,
222        // The wrapper applies the reduction to cells of rank >= 1; over the
223        // rank-1 argument this kernel insists on, that is the whole array.
224        Verb::Rank(u, r) if r[0] >= 1 => match &**u {
225            Verb::Reduce(inner) | Verb::NWise(inner) => inner,
226            _ => return None,
227        },
228        _ => return None,
229    };
230    let Verb::Prim(p) = &**inner else { return None };
231    let DyadOp::Scalar(op) = p.dyad else { return None };
232    matches!(op, Add | Mul | Min | Max).then_some(op)
233}
234
235/// The moving fold this dyad performs, if the kernel can absorb it: `k u/\ y`
236/// over an associative arithmetic `u` and a window the compiler knows the
237/// length of. APL's n-wise reduction `n f/ y` is the same fold over the same
238/// windows, so it takes the same path. A left argument of more than one
239/// number is a frame in J — several window lengths, several results — and
240/// stays outside; in APL it is not a window length at all.
241fn absorbable_window(e: &Expr) -> Option<(ScalarDyad, usize)> {
242    let (op, x) = match e {
243        Expr::Dyad { verb: Verb::Windowed(u, WindowKind::Prefix), x, .. } => {
244            (absorbable_reduce(u)?, x)
245        }
246        // `absorbable_reduce` also unwraps the rank wrapper APL's `/` wears;
247        // over the rank-1 argument this kernel insists on it changes nothing.
248        Expr::Dyad { verb: v @ (Verb::NWise(_) | Verb::Rank(..)), x, .. }
249            if is_nwise(v) =>
250        {
251            (absorbable_reduce(v)?, x)
252        }
253        _ => return None,
254    };
255    let Expr::Const(a, _) = &**x else { return None };
256    if a.rank() != 0 {
257        return None;
258    }
259    let k = *a.to_i64_vec()?.first()?;
260    // A negative left argument reverses each window in APL and cuts the
261    // argument into chunks in J, and a zero takes the empty runs between the
262    // items: none of the three is a plain moving window.
263    (1..=MAX_WINDOW as i64).contains(&k).then_some((op, k as usize))
264}
265
266/// Is this verb APL's `f/` or `f⌿` — the one whose dyad is the n-wise
267/// reduction — rather than J's `u/`, whose dyad is the table?
268fn is_nwise(v: &Verb) -> bool {
269    match v {
270        Verb::NWise(_) => true,
271        Verb::Rank(u, r) => r[0] >= 1 && matches!(&**u, Verb::NWise(_)),
272        _ => false,
273    }
274}
275
276/// The running fold this monad performs, if the kernel can absorb it. J's
277/// `u\` and APL's `f\` are the same scan; `u\.` folds from the far end,
278/// where an accumulator cannot be handed from one block to the next.
279fn absorbable_scan(e: &Expr) -> Option<ScalarDyad> {
280    let Expr::Monad { verb, .. } = e else { return None };
281    // APL's `f\` scans the last axis, which over the vector this stage
282    // insists on is the whole argument — the same wrapper `+/` wears.
283    let inner = match verb {
284        Verb::Rank(u, r) if r[0] >= 1 => &**u,
285        v => v,
286    };
287    let Verb::Windowed(u, kind) = inner else { return None };
288    if *kind == WindowKind::Suffix {
289        return None;
290    }
291    absorbable_reduce(u)
292}
293
294/// Is this the tally, applied to the array as a whole? `#"1` and its like
295/// count the items of cells instead, which is not what the shape says.
296fn is_tally(v: &Verb) -> bool {
297    matches!(v, Verb::Prim(p) if p.monad == MonadOp::Tally && p.ranks[0] == RANK_INF)
298}
299
300// ------------------------------------------------------------- the pass
301
302/// The chain as a tree, before it becomes postfix code.
303#[derive(Clone, PartialEq)]
304enum Node {
305    /// A subtree the kernel does not cover: an input, with its index.
306    Leaf(usize),
307    Monad(ScalarMonad, Box<Node>),
308    Dyad(ScalarDyad, Box<Node>, Box<Node>),
309    /// A moving fold: its operand stands on the wide axis, it on the
310    /// kernel's own.
311    Window(ScalarDyad, usize, Box<Node>),
312    Scan(ScalarDyad, Box<Node>),
313}
314
315/// The subtrees a chain reads.
316///
317/// Inputs are numbered in the order the evaluator would reach them — a
318/// dyad's right argument first — so that a fused node evaluates its leaves
319/// exactly when and where the unfused tree does. Two leaves that are the
320/// same subtree take the same input: nothing inside a chain can assign, so
321/// the second writing of `+/ {x}` reads what the first one read, and
322/// evaluating it once is what the sentence means either way.
323#[derive(Default)]
324struct Leaves<'a> {
325    inputs: Vec<&'a Expr>,
326    /// The input each leaf position reads, in chain order.
327    order: Vec<usize>,
328    /// The axis each input is read on, None where the chain reads it on
329    /// both — which only a scalar can be.
330    doms: Vec<Option<Dom>>,
331}
332
333impl<'a> Leaves<'a> {
334    fn push(&mut self, e: &'a Expr, dom: Dom) -> usize {
335        let i = match self.inputs.iter().position(|&p| same(p, e)) {
336            Some(i) => i,
337            None => {
338                self.inputs.push(e);
339                self.doms.push(Some(dom));
340                self.inputs.len() - 1
341            }
342        };
343        if self.doms[i] != Some(dom) {
344            self.doms[i] = None;
345        }
346        self.order.push(i);
347        i
348    }
349}
350
351/// A name the chain reads through to the value assigned to it, as inlining
352/// that assignment would; `hits` counts the uses it absorbed.
353struct Inline<'a> {
354    name: &'a str,
355    def: &'a Expr,
356    hits: usize,
357}
358
359/// The name a chain reads through, where the pass is moving one.
360fn read_through<'a>(e: &Expr, sub: Option<&Inline<'a>>) -> Option<&'a Expr> {
361    match (e, sub) {
362        (Expr::Name(n, _), Some(s)) if n == s.name => Some(s.def),
363        _ => None,
364    }
365}
366
367/// The stages the chain rooted at `e` may absorb.
368///
369/// Decided before the chain is built and then consulted by everything that
370/// walks it, so the pass, the fallback and the inliner all read the same
371/// tree. Window lengths are collected from the positions a window could be
372/// absorbed at; where they do not all agree there is more than one wide
373/// axis, and none is taken.
374fn plan_of(e: &Expr, sub: Option<&Inline<'_>>) -> Plan {
375    fn walk(e: &Expr, sub: Option<&Inline<'_>>, inside: bool, ks: &mut Vec<usize>, s: &mut bool) {
376        if let Some(def) = read_through(e, sub) {
377            return walk(def, sub, inside, ks, s);
378        }
379        match e {
380            Expr::Monad { verb, y, .. } if fusable_monad(verb).is_some() => {
381                walk(y, sub, inside, ks, s)
382            }
383            Expr::Dyad { verb, x, y, .. } if fusable_dyad(verb).is_some() => {
384                walk(y, sub, inside, ks, s);
385                walk(x, sub, inside, ks, s);
386            }
387            Expr::Dyad { y, .. } if !inside && absorbable_window(e).is_some() => {
388                ks.push(absorbable_window(e).expect("just matched").1);
389                walk(y, sub, true, ks, s);
390            }
391            Expr::Monad { y, .. } if absorbable_scan(e).is_some() => {
392                *s = true;
393                walk(y, sub, inside, ks, s);
394            }
395            _ => {}
396        }
397    }
398    let (mut ks, mut scan) = (Vec::new(), false);
399    walk(e, sub, false, &mut ks, &mut scan);
400    let window = match ks.split_first() {
401        Some((k, rest)) if rest.iter().all(|r| r == k) => Some(*k),
402        _ => None,
403    };
404    Plan { window, scan: window.is_none() && scan }
405}
406
407/// Build the chain rooted at `e`, collecting the subtrees that feed it.
408fn chain<'a>(
409    e: &'a Expr,
410    lv: &mut Leaves<'a>,
411    sub: &mut Option<Inline<'a>>,
412    plan: Plan,
413    dom: Dom,
414) -> Node {
415    if read_through(e, sub.as_ref()).is_some() {
416        let def = read_through(e, sub.as_ref()).expect("just matched");
417        if let Some(s) = sub.as_mut() {
418            s.hits += 1;
419        }
420        return chain(def, lv, sub, plan, dom);
421    }
422    match e {
423        Expr::Monad { verb, y, .. } => {
424            if let Some(op) = fusable_monad(verb) {
425                return Node::Monad(op, Box::new(chain(y, lv, sub, plan, dom)));
426            }
427            if plan.scan && let Some(op) = absorbable_scan(e) {
428                return Node::Scan(op, Box::new(chain(y, lv, sub, plan, dom)));
429            }
430            Node::Leaf(lv.push(e, dom))
431        }
432        Expr::Dyad { verb, x, y, .. } => {
433            if let Some(op) = fusable_dyad(verb) {
434                let ry = chain(y, lv, sub, plan, dom);
435                let rx = chain(x, lv, sub, plan, dom);
436                return Node::Dyad(op, Box::new(rx), Box::new(ry));
437            }
438            // A window inside a window would want a third axis; only the
439            // outer one is taken, and the inner reads as the leaf it is.
440            if dom == Dom::Result
441                && let Some((op, k)) = absorbable_window(e)
442                && plan.window == Some(k)
443            {
444                return Node::Window(op, k, Box::new(chain(y, lv, sub, plan, Dom::Wide)));
445            }
446            Node::Leaf(lv.push(e, dom))
447        }
448        _ => Node::Leaf(lv.push(e, dom)),
449    }
450}
451
452fn ops(n: &Node) -> usize {
453    match n {
454        Node::Leaf(_) => 0,
455        Node::Monad(_, y) | Node::Window(_, _, y) | Node::Scan(_, y) => 1 + ops(y),
456        Node::Dyad(_, x, y) => 1 + ops(x) + ops(y),
457    }
458}
459
460/// Every subtree of the chain that computes something, with the axis it
461/// stands on.
462fn subtrees<'a>(n: &'a Node, dom: Dom, out: &mut Vec<(&'a Node, Dom)>) {
463    if ops(n) == 0 {
464        return;
465    }
466    out.push((n, dom));
467    match n {
468        Node::Leaf(_) => {}
469        Node::Monad(_, y) | Node::Scan(_, y) => subtrees(y, dom, out),
470        Node::Window(_, _, y) => subtrees(y, Dom::Wide, out),
471        Node::Dyad(_, x, y) => {
472            subtrees(x, dom, out);
473            subtrees(y, dom, out);
474        }
475    }
476}
477
478/// The values the chain computes more than once, largest first.
479///
480/// `+/ d * d` over an inlined `d` writes the same arithmetic twice, and a
481/// block-at-a-time kernel can do what the assignment did: compute it once
482/// and read it twice. Each of these becomes a let — a block buffer of its
483/// own, held for the length of the block. Only maximal repeats are taken,
484/// so a repeat inside a let is part of that let rather than one more.
485///
486/// A value written on both axes of a windowed chain is not one value: the
487/// two are different lengths and read different items, so a repeat counts
488/// only against the repeats on its own axis, and only where the other axis
489/// holds none.
490fn lets_of(n: &Node) -> Vec<(Node, Dom)> {
491    let mut all = Vec::new();
492    subtrees(n, Dom::Result, &mut all);
493    let mut out = Vec::new();
494    fn walk(n: &Node, dom: Dom, all: &[(&Node, Dom)], out: &mut Vec<(Node, Dom)>) {
495        let count = |d: Dom| all.iter().filter(|(m, md)| *m == n && *md == d).count();
496        if ops(n) >= 1 && count(dom) >= 2 && count(other(dom)) == 0 {
497            if !out.iter().any(|(m, _)| m == n) {
498                out.push((n.clone(), dom));
499            }
500            return;
501        }
502        match n {
503            Node::Leaf(_) => {}
504            Node::Monad(_, y) | Node::Scan(_, y) => walk(y, dom, all, out),
505            Node::Window(_, _, y) => walk(y, Dom::Wide, all, out),
506            Node::Dyad(_, x, y) => {
507                walk(x, dom, all, out);
508                walk(y, dom, all, out);
509            }
510        }
511    }
512    walk(n, Dom::Result, &all, &mut out);
513    out
514}
515
516fn other(d: Dom) -> Dom {
517    match d {
518        Dom::Result => Dom::Wide,
519        Dom::Wide => Dom::Result,
520    }
521}
522
523/// Postfix code for the chain: the lets first, each into a slot of its own,
524/// then the chain that reads them.
525fn emit_all(n: &Node, lets: &[(Node, Dom)], code: &mut Vec<Instr>) {
526    for (k, (l, _)) in lets.iter().enumerate() {
527        // A let is emitted from the lets before it, so it cannot read
528        // itself; maximal repeats never nest, so there is nothing else.
529        emit(l, &lets[..k], code);
530        code.push(Instr::Store(k));
531    }
532    emit(n, lets, code);
533}
534
535/// Postfix code for the chain: a dyad's left operand is pushed first.
536fn emit(n: &Node, lets: &[(Node, Dom)], code: &mut Vec<Instr>) {
537    if let Some(k) = lets.iter().position(|(l, _)| l == n) {
538        code.push(Instr::Let(k));
539        return;
540    }
541    match n {
542        Node::Leaf(i) => code.push(Instr::Load(*i)),
543        Node::Monad(op, y) => {
544            emit(y, lets, code);
545            code.push(Instr::Monad(*op));
546        }
547        Node::Window(op, k, y) => {
548            emit(y, lets, code);
549            code.push(Instr::Window(*op, *k));
550        }
551        Node::Scan(op, y) => {
552            emit(y, lets, code);
553            code.push(Instr::Scan(*op));
554        }
555        Node::Dyad(op, x, y) => {
556            emit(x, lets, code);
557            emit(y, lets, code);
558            code.push(Instr::Dyad(*op));
559        }
560    }
561}
562
563/// Block buffers the postfix program needs at once.
564///
565/// Only a computed value holds one — an input is read where it lies — and
566/// the buffer being written is allocated before the operands are released,
567/// so the peak is the live count at some operation plus one.
568fn slots(code: &[Instr]) -> usize {
569    let mut stack: Vec<bool> = Vec::new();
570    let mut live = 0usize;
571    let mut max = 1usize;
572    for ins in code {
573        let operands = match ins {
574            Instr::Load(_) => {
575                stack.push(false);
576                continue;
577            }
578            // A let holds its buffer for the whole block: it is never
579            // released, so the count it added when it was computed stands
580            // and reading it takes nothing.
581            Instr::Let(_) => {
582                stack.push(false);
583                continue;
584            }
585            Instr::Store(_) => {
586                stack.pop();
587                continue;
588            }
589            Instr::Monad(_) | Instr::Window(..) | Instr::Scan(_) => 1,
590            Instr::Dyad(_) => 2,
591        };
592        max = max.max(live + 1);
593        for _ in 0..operands {
594            if stack.pop().unwrap_or(false) {
595                live -= 1;
596            }
597        }
598        live += 1;
599        stack.push(true);
600    }
601    max
602}
603
604/// Is this subtree free of effects?
605///
606/// A fused node evaluates its leaves in the order the unfused tree would,
607/// so an effect in one would still happen exactly once — but a node that
608/// can fall back is easier to be sure of when nothing inside it can act on
609/// the world, and a chain with `echo` in it is not the kind worth fusing.
610fn replayable(e: &Expr) -> bool {
611    match e {
612        Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => true,
613        Expr::Assign { .. }
614        | Expr::PrintPass { .. }
615        | Expr::Input { .. }
616        | Expr::Elided { .. }
617        | Expr::Control(..)
618        | Expr::AmendIndex { .. }
619        | Expr::VerbDef { .. }
620        | Expr::ModDef { .. } => false,
621        Expr::Monad { verb, y, .. } => verb.is_pure() && replayable(y),
622        Expr::Dyad { verb, x, y, .. } => verb.is_pure() && replayable(x) && replayable(y),
623        Expr::Fused { inputs, .. } => inputs.iter().all(replayable),
624    }
625}
626
627/// Are these the same computation? Two writings of one subexpression differ
628/// in their spans, which are positions in the source and mean nothing to
629/// the value, so spans are not compared. Assignments, output and fused
630/// nodes are never the same as anything: only leaves of a chain reach here,
631/// and a chain holds none of those.
632fn same(a: &Expr, b: &Expr) -> bool {
633    match (a, b) {
634        (Expr::Const(p, _), Expr::Const(q, _)) => p == q,
635        (Expr::Param(p, _), Expr::Param(q, _)) => p == q,
636        (Expr::Name(p, _), Expr::Name(q, _)) => p == q,
637        (Expr::Monad { verb: u, y: p, .. }, Expr::Monad { verb: v, y: q, .. }) => {
638            same_verb(u, v) && same(p, q)
639        }
640        (
641            Expr::Dyad { verb: u, x: px, y: py, .. },
642            Expr::Dyad { verb: v, x: qx, y: qy, .. },
643        ) => same_verb(u, v) && same(px, qx) && same(py, qy),
644        _ => false,
645    }
646}
647
648fn same_verb(a: &Verb, b: &Verb) -> bool {
649    match (a, b) {
650        (Verb::Prim(p), Verb::Prim(q)) => p == q,
651        (Verb::Rank(u, r), Verb::Rank(v, s)) => r == s && same_verb(u, v),
652        (Verb::Reduce(u), Verb::Reduce(v))
653        | (Verb::NWise(u), Verb::NWise(v))
654        | (Verb::Commute(u), Verb::Commute(v)) => same_verb(u, v),
655        (Verb::Windowed(u, j), Verb::Windowed(v, k)) => j == k && same_verb(u, v),
656        (Verb::PowerN(u, m), Verb::PowerN(v, n)) => m == n && same_verb(u, v),
657        (Verb::Fork(f, g, h), Verb::Fork(f2, g2, h2)) => {
658            same_verb(f, f2) && same_verb(g, g2) && same_verb(h, h2)
659        }
660        (Verb::NounFork(m, g, h), Verb::NounFork(n, g2, h2)) => {
661            m == n && same_verb(g, g2) && same_verb(h, h2)
662        }
663        (Verb::Hook(g, h), Verb::Hook(g2, h2))
664        | (Verb::Atop(g, h), Verb::Atop(g2, h2))
665        | (Verb::Compose(g, h), Verb::Compose(g2, h2)) => same_verb(g, g2) && same_verb(h, h2),
666        (Verb::BondLeft(m, u), Verb::BondLeft(n, v)) => m == n && same_verb(u, v),
667        (Verb::BondRight(u, m), Verb::BondRight(v, n)) => m == n && same_verb(u, v),
668        _ => false,
669    }
670}
671
672/// Optimise a compiled program's sentences: move the values that are only
673/// named for the reader into the sentences that read them, then fuse every
674/// chain that is left.
675pub fn pass(stmts: &mut Vec<Expr>, tol: Tol) {
676    let orig = std::mem::take(stmts);
677    let mut cur = orig.clone();
678    let mut names = 0usize;
679    let mut crossed = false;
680    // A round elides one assignment, so a chain of them — `m =. ...`,
681    // `d =. {x} - m`, `+/ d * d` — takes one round per link.
682    for _ in 0..=orig.len() {
683        match inline_once(&cur, &mut names, tol) {
684            Some(next) => {
685                cur = next;
686                crossed = true;
687            }
688            None => break,
689        }
690    }
691    let mut out: Vec<Expr> = cur.into_iter().map(|e| fuse_expr(e, tol)).collect();
692    if crossed {
693        // What the sentences were, for `unfused` to hold this against.
694        out.insert(0, Expr::Elided { orig, span: Span::new(0, 0) });
695    }
696    *stmts = out;
697}
698
699fn fuse_expr(e: Expr, tol: Tol) -> Expr {
700    if let Some(f) = try_fuse(&e, tol) {
701        return f;
702    }
703    match e {
704        Expr::Assign { name, value, scope, span } => {
705            Expr::Assign { name, value: Box::new(fuse_expr(*value, tol)), scope, span }
706        }
707        Expr::Monad { verb, y, span } => {
708            Expr::Monad { verb, y: Box::new(fuse_expr(*y, tol)), span }
709        }
710        Expr::Dyad { verb, x, y, span } => Expr::Dyad {
711            verb,
712            x: Box::new(fuse_expr(*x, tol)),
713            y: Box::new(fuse_expr(*y, tol)),
714            span,
715        },
716        Expr::PrintPass { value, bare, span } => {
717            Expr::PrintPass { value: Box::new(fuse_expr(*value, tol)), bare, span }
718        }
719        other => other,
720    }
721}
722
723/// The kernel for the chain rooted at `root`, if it carries at least
724/// `least` operations and reads nothing that cannot be replayed.
725fn build<'a>(
726    root: &'a Expr,
727    yields: Yield,
728    least: usize,
729    sub: &mut Option<Inline<'a>>,
730    tol: Tol,
731) -> Option<(FusedKernel, Vec<&'a Expr>)> {
732    if let Some(s) = sub.as_mut() {
733        s.hits = 0;
734    }
735    let plan = plan_of(root, sub.as_ref());
736    let mut lv = Leaves::default();
737    let node = chain(root, &mut lv, sub, plan, Dom::Result);
738    if ops(&node) < least || !lv.inputs.iter().all(|l| replayable(l)) {
739        return None;
740    }
741    let mut code = Vec::new();
742    let lets = lets_of(&node);
743    emit_all(&node, &lets, &mut code);
744    let window = code.iter().find_map(|i| match i {
745        Instr::Window(_, k) => Some(*k),
746        _ => None,
747    });
748    let scans = code.iter().filter(|i| matches!(i, Instr::Scan(_))).count();
749    // A running fold hands its accumulator from one block to the next, so
750    // its blocks run forwards and in order; an absorbed reduction folds
751    // them backwards, which is the insert's own order. A chain that wants
752    // both runs as the passes it was written as.
753    if scans > 0 && matches!(yields, Yield::Reduce(_)) {
754        return None;
755    }
756    let kernel = FusedKernel {
757        slots: slots(&code),
758        code,
759        yields,
760        leaves: lv.order,
761        doms: lv.doms,
762        let_doms: lets.iter().map(|(_, d)| *d).collect(),
763        plan,
764        window,
765        scans,
766        tol,
767    };
768    Some((kernel, lv.inputs))
769}
770
771/// The kernel this node becomes, with the subtree it stands for — the chain
772/// itself where a tally reads only its shape, the whole sentence where a
773/// reduction sits above it.
774///
775/// One elementwise verb on its own already runs as one pass; fusing it
776/// would only add a layer. A reduction to absorb, or a tally that makes the
777/// values unnecessary altogether, makes one verb enough.
778fn kernel_at<'a>(
779    e: &'a Expr,
780    sub: &mut Option<Inline<'a>>,
781    tol: Tol,
782) -> Option<(FusedKernel, Vec<&'a Expr>, &'a Expr)> {
783    if let Expr::Monad { verb, y, .. } = e {
784        if is_tally(verb) && let Some((k, l)) = build(y, Yield::Tally, 1, sub, tol) {
785            return Some((k, l, e));
786        }
787        if let Some(op) = absorbable_reduce(verb)
788            && let Some((k, l)) = build(y, Yield::Reduce(op), 1, sub, tol)
789        {
790            return Some((k, l, e));
791        }
792    }
793    let (k, l) = build(e, Yield::Values, 2, sub, tol)?;
794    Some((k, l, e))
795}
796
797/// The fused node for the chain rooted at `e`, if there is one worth making.
798fn try_fuse(e: &Expr, tol: Tol) -> Option<Expr> {
799    let (kernel, leaves, orig) = kernel_at(e, &mut None, tol)?;
800    let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
801    Some(Expr::Fused {
802        kernel,
803        inputs,
804        orig: Box::new(orig.clone()),
805        span: e.span(),
806    })
807}
808
809/// The chain a fused node came from, with its leaves replaced by the values
810/// already computed for them.
811///
812/// This is what runs when the kernel declines. Rebuilding the tree costs a
813/// handful of small allocations and saves evaluating the leaves a second
814/// time, which for a leaf like `19 }. {close}` is a whole array.
815pub(crate) fn fallback_tree(k: &FusedKernel, orig: &Expr, values: &[Array]) -> Expr {
816    let mut next = 0;
817    let plan = k.plan;
818    let tree = match orig {
819        // An absorbed reduction sits above the chain; only the chain's own
820        // leaves were evaluated.
821        Expr::Monad { verb, y, span } if matches!(k.yields, Yield::Reduce(_)) => Expr::Monad {
822            verb: verb.clone(),
823            y: Box::new(substitute(y, values, k, &mut next, plan, Dom::Result)),
824            span: *span,
825        },
826        // A tally is not applied at all: the chain alone runs, and the
827        // count of what it made is what the node yields.
828        Expr::Monad { verb, y, .. } if k.yields == Yield::Tally && is_tally(verb) => {
829            substitute(y, values, k, &mut next, plan, Dom::Result)
830        }
831        e => substitute(e, values, k, &mut next, plan, Dom::Result),
832    };
833    debug_assert_eq!(next, k.leaves.len(), "the fallback found different leaves");
834    tree
835}
836
837/// What the kernel would have made of the value its chain produced. A tally
838/// skips the chain entirely when it runs, and counts the items of it when
839/// the chain has had to run instead.
840pub(crate) fn fallback_finish(k: &FusedKernel, v: Array) -> Array {
841    match k.yields {
842        Yield::Tally => Array::scalar_i64(v.items() as i64),
843        _ => v,
844    }
845}
846
847/// Walk the chain exactly as [`chain`] walked it, so the leaves take their
848/// values in the order they were numbered in.
849fn substitute(
850    e: &Expr,
851    values: &[Array],
852    k: &FusedKernel,
853    next: &mut usize,
854    plan: Plan,
855    dom: Dom,
856) -> Expr {
857    match e {
858        Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
859            verb: verb.clone(),
860            y: Box::new(substitute(y, values, k, next, plan, dom)),
861            span: *span,
862        },
863        Expr::Monad { verb, y, span } if plan.scan && absorbable_scan(e).is_some() => {
864            Expr::Monad {
865                verb: verb.clone(),
866                y: Box::new(substitute(y, values, k, next, plan, dom)),
867                span: *span,
868            }
869        }
870        Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => {
871            let ry = substitute(y, values, k, next, plan, dom);
872            let rx = substitute(x, values, k, next, plan, dom);
873            Expr::Dyad { verb: verb.clone(), x: Box::new(rx), y: Box::new(ry), span: *span }
874        }
875        Expr::Dyad { verb, x, y, span }
876            if dom == Dom::Result
877                && absorbable_window(e).map(|(_, k)| k) == plan.window
878                && plan.window.is_some() =>
879        {
880            Expr::Dyad {
881                verb: verb.clone(),
882                x: x.clone(),
883                y: Box::new(substitute(y, values, k, next, plan, Dom::Wide)),
884                span: *span,
885            }
886        }
887        leaf => {
888            let v = values[k.leaves[*next]].clone();
889            *next += 1;
890            Expr::Const(v, leaf.span())
891        }
892    }
893}
894
895// ------------------------------------------- across sentence boundaries
896//
897// `d =. {x} - m` and then `+/ d * d` is the same computation as the one
898// sentence that spells it out, but the assignment writes `d` to memory in
899// full and the next sentence reads it back — the traffic the kernel exists
900// to remove. Nothing there needs the array: the name is for the reader.
901//
902// So the pass moves the value into the sentences that read it, and hoists
903// the value's own leaves — the mean's `+/ {x}` — into sentences of their
904// own first, so that copying the chain does not copy the work. What comes
905// out is the two-phase shape a hand-written kernel has: one pass for the
906// reductions the chain reads as scalars, one for the map-reduce over them.
907
908/// Names the pass introduces for the values it hoists. `·` starts no name
909/// either frontend accepts, so these cannot collide with the program's.
910fn hoisted_name(n: &mut usize) -> String {
911    *n += 1;
912    format!("·{}", *n - 1)
913}
914
915/// Elide the first assignment whose value can move into the sentences that
916/// read it, and report the sentences that leaves; None when none can.
917///
918/// The value moves only where the name is pure dataflow:
919///
920/// - the value is replayable and is a chain, so that moving it moves
921///   arithmetic into a kernel rather than moving a whole pass;
922/// - no later sentence assigns the name again, or any name the value reads,
923///   so every copy means what the original meant;
924/// - every use lands inside a kernel, so no copy materialises the value.
925///   A tally counts as landing inside one: it reads the chain's shape.
926///
927/// The assignment's own sentence stays, as the tally of the chain: that
928/// reaches every leaf and every rule the kernel has, so whatever the
929/// assignment would have raised is raised where it was raised before, and
930/// nothing else is computed.
931fn inline_once(stmts: &[Expr], names: &mut usize, tol: Tol) -> Option<Vec<Expr>> {
932    for (i, stmt) in stmts.iter().enumerate() {
933        let Expr::Assign { name, value, span, .. } = stmt else { continue };
934        if !inlinable(stmts, i, name, value, tol) {
935            continue;
936        }
937        if let Some(out) = rewrite(stmts, i, name, value, *span, names, tol) {
938            return Some(out);
939        }
940    }
941    None
942}
943
944fn inlinable(stmts: &[Expr], i: usize, name: &str, value: &Expr, tol: Tol) -> bool {
945    if !replayable(value) || mentions(value, name) {
946        return false;
947    }
948    let mut lv = Leaves::default();
949    if ops(&chain(value, &mut lv, &mut None, plan_of(value, None), Dom::Result)) < 1 {
950        return false;
951    }
952    let mut guarded = vec![name.to_string()];
953    free_names(value, &mut guarded);
954    let later = &stmts[i + 1..];
955    if later.iter().any(|s| assigns_any(s, &guarded)) {
956        return false;
957    }
958    let mut uses = 0;
959    for stmt in later {
960        match uses_land(stmt, name, value, tol) {
961            Some(n) => uses += n,
962            None => return false,
963        }
964    }
965    uses > 0
966}
967
968/// How many uses of `name` this sentence would take into a kernel, or None
969/// when one of them would have to materialise the value instead.
970fn uses_land(e: &Expr, name: &str, def: &Expr, tol: Tol) -> Option<usize> {
971    let mut sub = Some(Inline { name, def, hits: 0 });
972    if let Some((_, leaves, _)) = kernel_at(e, &mut sub, tol) {
973        let mut n = sub.map_or(0, |s| s.hits);
974        for l in leaves {
975            n += uses_land(l, name, def, tol)?;
976        }
977        return Some(n);
978    }
979    match e {
980        Expr::Name(n, _) if n == name => None,
981        Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => Some(0),
982        Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => uses_land(value, name, def, tol),
983        Expr::Monad { y, .. } => uses_land(y, name, def, tol),
984        Expr::Dyad { x, y, .. } => Some(uses_land(x, name, def, tol)? + uses_land(y, name, def, tol)?),
985        Expr::Fused { .. }
986        | Expr::Elided { .. }
987        | Expr::Input { .. }
988        | Expr::Control(..)
989        | Expr::AmendIndex { .. }
990        | Expr::VerbDef { .. }
991        | Expr::ModDef { .. } => None,
992    }
993}
994
995/// The sentences that replace `stmts`, with the assignment at `i` elided.
996fn rewrite(
997    stmts: &[Expr],
998    i: usize,
999    name: &str,
1000    value: &Expr,
1001    span: Span,
1002    names: &mut usize,
1003    tol: Tol,
1004) -> Option<Vec<Expr>> {
1005    let mut lv = Leaves::default();
1006    let plan = plan_of(value, None);
1007    chain(value, &mut lv, &mut None, plan, Dom::Result);
1008    // A leaf that is more than a name or a constant becomes a sentence of
1009    // its own, evaluated once and where it was evaluated before.
1010    let mut hoists = Vec::new();
1011    let mut bound: Vec<Option<String>> = Vec::new();
1012    for l in &lv.inputs {
1013        if matches!(l, Expr::Const(..) | Expr::Param(..) | Expr::Name(..)) {
1014            bound.push(None);
1015            continue;
1016        }
1017        let n = hoisted_name(names);
1018        hoists.push(Expr::Assign {
1019            name: n.clone(),
1020            value: Box::new((*l).clone()),
1021            scope: Scope::Local,
1022            span: l.span(),
1023        });
1024        bound.push(Some(n));
1025    }
1026    let def = with_leaves(value, &lv, &bound, plan, Dom::Result);
1027    let (kernel, leaves) = build(&def, Yield::Tally, 1, &mut None, tol)?;
1028    let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
1029    let guard = Expr::Assign {
1030        name: hoisted_name(names),
1031        value: Box::new(Expr::Fused {
1032            kernel,
1033            inputs,
1034            orig: Box::new(def.clone()),
1035            span,
1036        }),
1037        scope: Scope::Local,
1038        span,
1039    };
1040    let mut out = stmts[..i].to_vec();
1041    out.extend(hoists);
1042    out.push(guard);
1043    out.extend(stmts[i + 1..].iter().map(|s| replace_name(s, name, &def)));
1044    Some(out)
1045}
1046
1047/// The chain with its hoisted leaves replaced by the names they were bound
1048/// to. Walks exactly as [`chain`] walks, so the leaves are the same ones.
1049fn with_leaves(e: &Expr, lv: &Leaves<'_>, bound: &[Option<String>], plan: Plan, dom: Dom) -> Expr {
1050    match e {
1051        Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
1052            verb: verb.clone(),
1053            y: Box::new(with_leaves(y, lv, bound, plan, dom)),
1054            span: *span,
1055        },
1056        Expr::Monad { verb, y, span } if plan.scan && absorbable_scan(e).is_some() => {
1057            Expr::Monad {
1058                verb: verb.clone(),
1059                y: Box::new(with_leaves(y, lv, bound, plan, dom)),
1060                span: *span,
1061            }
1062        }
1063        Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => Expr::Dyad {
1064            verb: verb.clone(),
1065            x: Box::new(with_leaves(x, lv, bound, plan, dom)),
1066            y: Box::new(with_leaves(y, lv, bound, plan, dom)),
1067            span: *span,
1068        },
1069        Expr::Dyad { verb, x, y, span }
1070            if dom == Dom::Result
1071                && absorbable_window(e).map(|(_, k)| k) == plan.window
1072                && plan.window.is_some() =>
1073        {
1074            Expr::Dyad {
1075                verb: verb.clone(),
1076                x: x.clone(),
1077                y: Box::new(with_leaves(y, lv, bound, plan, Dom::Wide)),
1078                span: *span,
1079            }
1080        }
1081        leaf => {
1082            let bind = lv
1083                .inputs
1084                .iter()
1085                .position(|&p| same(p, leaf))
1086                .and_then(|i| bound[i].as_ref());
1087            match bind {
1088                Some(n) => Expr::Name(n.clone(), leaf.span()),
1089                None => leaf.clone(),
1090            }
1091        }
1092    }
1093}
1094
1095fn replace_name(e: &Expr, name: &str, def: &Expr) -> Expr {
1096    match e {
1097        Expr::Name(n, _) if n == name => def.clone(),
1098        Expr::Assign { name: a, value, scope, span } => Expr::Assign {
1099            scope: *scope,
1100            name: a.clone(),
1101            value: Box::new(replace_name(value, name, def)),
1102            span: *span,
1103        },
1104        Expr::PrintPass { value, bare, span } => Expr::PrintPass {
1105            value: Box::new(replace_name(value, name, def)),
1106            bare: *bare,
1107            span: *span,
1108        },
1109        Expr::Monad { verb, y, span } => Expr::Monad {
1110            verb: verb.clone(),
1111            y: Box::new(replace_name(y, name, def)),
1112            span: *span,
1113        },
1114        Expr::Dyad { verb, x, y, span } => Expr::Dyad {
1115            verb: verb.clone(),
1116            x: Box::new(replace_name(x, name, def)),
1117            y: Box::new(replace_name(y, name, def)),
1118            span: *span,
1119        },
1120        other => other.clone(),
1121    }
1122}
1123
1124fn mentions(e: &Expr, name: &str) -> bool {
1125    let mut names = Vec::new();
1126    free_names(e, &mut names);
1127    names.iter().any(|n| n == name)
1128}
1129
1130/// Every name this subtree reads.
1131fn free_names(e: &Expr, out: &mut Vec<String>) {
1132    match e {
1133        Expr::Name(n, _) => out.push(n.clone()),
1134        Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => free_names(value, out),
1135        Expr::Monad { y, .. } => free_names(y, out),
1136        Expr::Dyad { x, y, .. } => {
1137            free_names(x, out);
1138            free_names(y, out);
1139        }
1140        Expr::Fused { inputs, .. } => inputs.iter().for_each(|i| free_names(i, out)),
1141        Expr::Const(..)
1142        | Expr::Param(..)
1143        | Expr::Elided { .. }
1144        | Expr::Input { .. }
1145        | Expr::Control(..)
1146        | Expr::AmendIndex { .. }
1147        | Expr::VerbDef { .. }
1148        | Expr::ModDef { .. } => {}
1149    }
1150}
1151
1152/// Does this sentence assign any of these names, at any depth?
1153fn assigns_any(e: &Expr, names: &[String]) -> bool {
1154    match e {
1155        Expr::Assign { name, value, .. } => {
1156            names.iter().any(|n| n == name) || assigns_any(value, names)
1157        }
1158        Expr::PrintPass { value, .. } => assigns_any(value, names),
1159        Expr::Monad { y, .. } => assigns_any(y, names),
1160        Expr::Dyad { x, y, .. } => assigns_any(x, names) || assigns_any(y, names),
1161        Expr::Fused { inputs, .. } => inputs.iter().any(|i| assigns_any(i, names)),
1162        Expr::Const(..)
1163        | Expr::Param(..)
1164        | Expr::Name(..)
1165        | Expr::Elided { .. }
1166        | Expr::Input { .. }
1167        | Expr::Control(..)
1168        | Expr::AmendIndex { .. }
1169        | Expr::VerbDef { .. }
1170        | Expr::ModDef { .. } => false,
1171    }
1172}
1173
1174/// Does any sentence of this program run a fused kernel?
1175pub fn is_fused(p: &Program) -> bool {
1176    fn any(e: &Expr) -> bool {
1177        match e {
1178            Expr::Fused { .. } => true,
1179            Expr::Const(..)
1180            | Expr::Param(..)
1181            | Expr::Name(..)
1182            | Expr::Elided { .. }
1183            | Expr::Input { .. }
1184            | Expr::Control(..)
1185            | Expr::AmendIndex { .. }
1186            | Expr::VerbDef { .. }
1187            | Expr::ModDef { .. } => false,
1188            Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => any(value),
1189            Expr::Monad { y, .. } => any(y),
1190            Expr::Dyad { x, y, .. } => any(x) || any(y),
1191        }
1192    }
1193    p.stmts.iter().any(any)
1194}
1195
1196/// Did the pass move a named value into the sentences that read it?
1197pub fn is_inlined(p: &Program) -> bool {
1198    matches!(p.stmts.first(), Some(Expr::Elided { .. }))
1199}
1200
1201/// The program as the plain evaluator would run it: the sentences it was
1202/// compiled from, with every fused node replaced by the subtree it came
1203/// from. The two must compute the same thing; tests hold them to it.
1204pub fn unfused(p: &Program) -> Program {
1205    fn strip(e: &Expr) -> Expr {
1206        match e {
1207            Expr::Fused { orig, .. } => strip(orig),
1208            Expr::Assign { name, value, scope, span } => {
1209                Expr::Assign {
1210                    name: name.clone(),
1211                    value: Box::new(strip(value)),
1212                    scope: *scope,
1213                    span: *span,
1214                }
1215            }
1216            Expr::PrintPass { value, bare, span } => {
1217                Expr::PrintPass { value: Box::new(strip(value)), bare: *bare, span: *span }
1218            }
1219            Expr::Monad { verb, y, span } => {
1220                Expr::Monad { verb: verb.clone(), y: Box::new(strip(y)), span: *span }
1221            }
1222            Expr::Dyad { verb, x, y, span } => Expr::Dyad {
1223                verb: verb.clone(),
1224                x: Box::new(strip(x)),
1225                y: Box::new(strip(y)),
1226                span: *span,
1227            },
1228            other => other.clone(),
1229        }
1230    }
1231    let mut out = p.clone();
1232    // A program the pass rewrote across sentences kept the sentences it
1233    // rewrote; those, not the rewriting, are what the evaluator would run.
1234    let stmts = match p.stmts.first() {
1235        Some(Expr::Elided { orig, .. }) => orig,
1236        _ => &p.stmts,
1237    };
1238    out.stmts = stmts.iter().map(strip).collect();
1239    out
1240}
1241
1242// ------------------------------------------------------------- dtype rules
1243
1244/// The dtype the unfused pipeline gives this monad's result. None where it
1245/// depends on the values (`<.` of a float is an integer only if every
1246/// rounded value fits one), which the kernel declines rather than guess.
1247fn monad_type(op: ScalarMonad, a: DType) -> Option<DType> {
1248    use DType::*;
1249    use ScalarMonad::*;
1250    // The kernel computes in one real type; complex values are not one of
1251    // them, so a chain that touches one declines and runs unfused.
1252    if a == Complex {
1253        return None;
1254    }
1255    Some(match op {
1256        Recip | Halve | Exp => F64,
1257        // Identity and magnitude keep a boolean boolean.
1258        Conj | Abs | OneMinus => a,
1259        Neg | Signum | Inc | Dec | Double | Square => match a {
1260            Bool | I64 => I64,
1261            other => other,
1262        },
1263        Floor | Ceil => match a {
1264            Bool | I64 => I64,
1265            _ => return None,
1266        },
1267        _ => return None,
1268    })
1269}
1270
1271/// The dtype the unfused pipeline gives this dyad's result, on the path
1272/// where no integer step overflows (one that does falls back).
1273fn dyad_type(op: ScalarDyad, a: DType, b: DType) -> Option<DType> {
1274    use ScalarDyad::*;
1275    if a == DType::Complex || b == DType::Complex {
1276        return None;
1277    }
1278    match op {
1279        Eq | Ne | Lt | Le | Gt | Ge => Some(DType::Bool),
1280        DivJ => Some(DType::F64),
1281        Add | Sub | Mul | Min | Max | Residue => match DType::promote(a, b)? {
1282            DType::Bool => Some(DType::I64),
1283            DType::Char | DType::Symbol => None,
1284            t => Some(t),
1285        },
1286        _ => None,
1287    }
1288}
1289
1290/// The dtype the unfused pipeline gives a fold over items of this type —
1291/// a moving window's, or a running one's. Booleans fold as the integers
1292/// they count as, which is what the windowed and scanning fast paths do.
1293fn fold_type(op: ScalarDyad, a: DType) -> Option<DType> {
1294    use ScalarDyad::*;
1295    if !matches!(op, Add | Mul | Min | Max) {
1296        return None;
1297    }
1298    match a {
1299        DType::Bool | DType::I64 => Some(DType::I64),
1300        DType::F64 => Some(DType::F64),
1301        _ => None,
1302    }
1303}
1304
1305/// The type the kernel computes in, and the dtype of its mapped result.
1306///
1307/// Every value in the program is computed in one type, so it must be one
1308/// that holds them all: integers when nothing in the chain leaves them,
1309/// floats otherwise. That leaves one case the kernel cannot serve — a chain
1310/// that computes an integer somewhere along a float path, as
1311/// `(x > 0) + (y > 0)` or `({a} + {b}) % 2` do. Its unfused pipeline holds
1312/// those steps in i64, exactly, past where f64 stops being exact, and its
1313/// result may be an integer array. Rather than compute them in the wrong
1314/// type, the kernel declines and the chain runs.
1315///
1316/// A boolean is not such a case: a comparison yields 0 and 1, which f64
1317/// holds exactly, and only the dtype of a result made from one has to be
1318/// narrowed at the end.
1319///
1320/// This is the kernel's main blind spot — a random chain over mixed integer
1321/// and float arguments declines about half the time — and the way out is a
1322/// stack whose entries carry their own type rather than one type per
1323/// kernel. Nothing measured so far needs it.
1324pub(crate) fn working_type(k: &FusedKernel, inputs: &[Array]) -> Option<(DType, DType)> {
1325    let mut stack: Vec<DType> = Vec::with_capacity(k.slots);
1326    let mut lets: Vec<DType> = Vec::new();
1327    let mut float = false;
1328    let mut integer_step = false;
1329    // The exact types and the complex ones have no blockwise kernel: a
1330    // fused chain over them declines and the general path evaluates it.
1331    if inputs.iter().any(|a| a.dtype() == DType::Complex || a.dtype().is_exact()) {
1332        return None;
1333    }
1334    for ins in &k.code {
1335        let t = match ins {
1336            Instr::Load(i) => inputs[*i].dtype(),
1337            Instr::Monad(op) => monad_type(*op, stack.pop()?)?,
1338            Instr::Window(op, _) | Instr::Scan(op) => fold_type(*op, stack.pop()?)?,
1339            Instr::Dyad(op) => {
1340                let b = stack.pop()?;
1341                let a = stack.pop()?;
1342                dyad_type(*op, a, b)?
1343            }
1344            Instr::Store(k) => {
1345                let t = stack.pop()?;
1346                if lets.len() != *k {
1347                    return None;
1348                }
1349                lets.push(t);
1350                continue;
1351            }
1352            // Reading a let is not a step: the value was accounted for
1353            // where it was computed.
1354            Instr::Let(k) => {
1355                let t = *lets.get(*k)?;
1356                float |= t == DType::F64;
1357                stack.push(t);
1358                continue;
1359            }
1360        };
1361        // Only numbers: everything else — characters, boxes — is a type
1362        // the kernel has no arithmetic for and the chain must handle.
1363        if !t.is_numeric() {
1364            return None;
1365        }
1366        float |= t == DType::F64;
1367        // An argument's own values are exact in either type; a step's are
1368        // not, once they are integers wider than f64's 53 bits.
1369        integer_step |= t == DType::I64 && !matches!(ins, Instr::Load(_));
1370        stack.push(t);
1371    }
1372    let root = stack.pop()?;
1373    let working = if float { DType::F64 } else { DType::I64 };
1374    if working == DType::F64 && integer_step {
1375        return None;
1376    }
1377    Some((working, root))
1378}
1379
1380// ------------------------------------------------------------- execution
1381
1382/// Where a value inside a block stands. A repeated scalar stands wherever
1383/// it is read, so it takes the axis of whatever it is combined with.
1384#[derive(Clone, Copy, PartialEq, Eq)]
1385enum On {
1386    Result,
1387    Wide,
1388    Either,
1389}
1390
1391fn combine(a: On, b: On) -> On {
1392    if a == On::Either {
1393        b
1394    } else {
1395        a
1396    }
1397}
1398
1399fn placed(d: Option<Dom>) -> On {
1400    match d {
1401        Some(Dom::Result) => On::Result,
1402        Some(Dom::Wide) => On::Wide,
1403        None => On::Either,
1404    }
1405}
1406
1407/// One block of work: the result items it writes, and the items of the wide
1408/// axis its window steps read to write them.
1409#[derive(Clone, Copy)]
1410struct Extent {
1411    start: usize,
1412    len: usize,
1413    wide_start: usize,
1414    wide_len: usize,
1415}
1416
1417impl Extent {
1418    /// The block that writes result items `start .. start + len`.
1419    ///
1420    /// The window fold cuts the wide axis into runs of `k` counted from the
1421    /// axis's own start and joins one run's suffix to the next run's
1422    /// prefix, so which items a window is folded from, and in what
1423    /// grouping, depend on where the window lies and never on where a block
1424    /// boundary fell. This block therefore reads from the start of the run
1425    /// its first window begins in to the end of the run its last item lies
1426    /// in: the same arithmetic, item for item, as one pass over the whole
1427    /// axis.
1428    fn of(start: usize, len: usize, window: Option<usize>, wide: usize) -> Extent {
1429        let Some(k) = window else {
1430            return Extent { start, len, wide_start: start, wide_len: len };
1431        };
1432        let lo = start - start % k;
1433        let hi = ((start + len + k - 2) / k + 1) * k;
1434        Extent { start, len, wide_start: lo, wide_len: hi.min(wide) - lo }
1435    }
1436
1437    fn len_on(&self, dom: On) -> usize {
1438        match dom {
1439            On::Wide => self.wide_len,
1440            _ => self.len,
1441        }
1442    }
1443}
1444
1445/// One input, in the working type: either the values themselves or one
1446/// value repeated, which is how a rank-0 argument reaches every element.
1447#[derive(Clone, Copy)]
1448struct Loaded<'a, T> {
1449    data: &'a [T],
1450    splat: bool,
1451    /// The axis the chain reads this input on.
1452    on: On,
1453    /// The index `data[0]` stands at on that axis. Zero for an argument's
1454    /// own buffer, and the block's own start for one staged a block at a
1455    /// time.
1456    base: usize,
1457}
1458
1459impl<T> Loaded<'_, T> {
1460    #[inline]
1461    fn block(&self, at: &Extent, dom: On) -> &[T] {
1462        if self.splat {
1463            return &self.data[..at.len_on(dom)];
1464        }
1465        let (start, len) = read_over(self.on, at);
1466        &self.data[start - self.base..start - self.base + len]
1467    }
1468}
1469
1470/// The range of its own axis an input standing on `on` is read over for one
1471/// block.
1472#[inline]
1473fn read_over(on: On, at: &Extent) -> (usize, usize) {
1474    match on {
1475        On::Wide => (at.wide_start, at.wide_len),
1476        _ => (at.start, at.len),
1477    }
1478}
1479
1480/// An argument narrower than the working type, read where it lies.
1481///
1482/// A whole widened copy of such an argument costs two round trips of the
1483/// working set — writing 160 MB of freshly faulted pages, then reading them
1484/// back — where the values themselves are needed one block at a time and a
1485/// block fits in cache. So the promotion happens at the block, into a
1486/// staging buffer each thread reuses; the values are what the widened copy
1487/// would have held, element for element.
1488#[derive(Clone, Copy)]
1489enum Narrow<'a> {
1490    I64(&'a [i64]),
1491    Bool(&'a [u8]),
1492}
1493
1494/// One input as the kernel will read it: its own buffer when that already
1495/// holds the working type, a narrower buffer to be promoted otherwise.
1496enum Source<'a, T> {
1497    Ready(Loaded<'a, T>),
1498    /// The narrow values, and the axis the chain reads them on.
1499    Staged(Narrow<'a>, On),
1500}
1501
1502/// The staging buffer's element type, filled from a narrow argument.
1503trait FromNarrow: Copy {
1504    fn fill(src: Narrow<'_>, at: usize, dst: &mut [Self]);
1505}
1506
1507impl FromNarrow for f64 {
1508    #[inline]
1509    fn fill(src: Narrow<'_>, at: usize, dst: &mut [f64]) {
1510        match src {
1511            Narrow::I64(v) => {
1512                for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1513                    *slot = x as f64;
1514                }
1515            }
1516            Narrow::Bool(v) => {
1517                for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1518                    *slot = x as f64;
1519                }
1520            }
1521        }
1522    }
1523}
1524
1525impl FromNarrow for i64 {
1526    #[inline]
1527    fn fill(src: Narrow<'_>, at: usize, dst: &mut [i64]) {
1528        match src {
1529            // The working type is integer only when no input is a float.
1530            Narrow::I64(v) => dst.copy_from_slice(&v[at..at + dst.len()]),
1531            Narrow::Bool(v) => {
1532                for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1533                    *slot = x as i64;
1534                }
1535            }
1536        }
1537    }
1538}
1539
1540/// The inputs of one run, and how wide a staging block has to be.
1541struct Sources<'a, T> {
1542    of: Vec<Source<'a, T>>,
1543    staged: usize,
1544}
1545
1546impl<T: FromNarrow + Default> Sources<'_, T> {
1547    /// Promote every staged input's block into `stage` and hand the whole
1548    /// input list, block-local, to `f`.
1549    ///
1550    /// `stage` holds one region of `width` per staged input, so the regions
1551    /// are the same from block to block and a thread faults them once.
1552    fn with_block<R>(
1553        &self,
1554        at: &Extent,
1555        stage: &mut [T],
1556        width: usize,
1557        f: impl FnOnce(&[Loaded<'_, T>]) -> R,
1558    ) -> R {
1559        let mut k = 0;
1560        for s in &self.of {
1561            let Source::Staged(src, on) = s else { continue };
1562            let (start, len) = read_over(*on, at);
1563            T::fill(*src, start, &mut stage[k * width..k * width + len]);
1564            k += 1;
1565        }
1566        let mut k = 0;
1567        let loaded: Vec<Loaded<'_, T>> = self
1568            .of
1569            .iter()
1570            .map(|s| match s {
1571                Source::Ready(l) => *l,
1572                Source::Staged(_, on) => {
1573                    let (start, len) = read_over(*on, at);
1574                    let d = &stage[k * width..k * width + len];
1575                    k += 1;
1576                    Loaded { data: d, splat: false, on: *on, base: start }
1577                }
1578            })
1579            .collect();
1580        f(&loaded)
1581    }
1582}
1583
1584/// What a stack entry refers to: an input, or a block buffer and the axis
1585/// the value in it stands on.
1586#[derive(Clone, Copy)]
1587enum Slot {
1588    Input(usize),
1589    Block(usize, On),
1590}
1591
1592/// The buffers one thread reuses from block to block.
1593struct Scratch<T> {
1594    cells: Vec<T>,
1595    /// Elements one block buffer holds: a block's result items, and the
1596    /// halo of the wide axis its window steps read around them.
1597    width: usize,
1598    free: Vec<usize>,
1599    stack: Vec<Slot>,
1600    lets: Vec<usize>,
1601    /// One accumulator per running fold in the code, carried from block to
1602    /// block so that the fold is the one the unfused scan performs.
1603    carry: Vec<Option<T>>,
1604}
1605
1606impl<T: Copy + Default> Scratch<T> {
1607    /// Room for one thread's blocks of `w` result items each.
1608    ///
1609    /// A window step reads the run its first window begins in and the run
1610    /// its last item lies in, so a block of `w` items reads fewer than
1611    /// `w + 3k` items of the wide axis, and every buffer is that wide.
1612    fn new(k: &FusedKernel, w: usize) -> Scratch<T> {
1613        let width = w + 3 * k.window.unwrap_or(0);
1614        Scratch {
1615            cells: vec![T::default(); k.slots * width],
1616            width,
1617            free: Vec::with_capacity(k.slots),
1618            stack: Vec::with_capacity(k.slots),
1619            lets: Vec::new(),
1620            carry: vec![None; k.scans],
1621        }
1622    }
1623}
1624
1625/// The leaf loops one working type runs, one block of one instruction at a
1626/// time. All of a kernel's arithmetic goes through these four.
1627struct Steps<M, D, W, S> {
1628    monad: M,
1629    dyad: D,
1630    window: W,
1631    scan: S,
1632}
1633
1634/// Block buffer `d` for writing, plus read-only access to the others.
1635fn split_slots<'s, T>(
1636    scratch: &'s mut [T],
1637    w: usize,
1638    d: usize,
1639) -> (&'s mut [T], impl Fn(usize) -> &'s [T]) {
1640    let (lo, hi) = scratch.split_at_mut(d * w);
1641    let (dst, hi) = hi.split_at_mut(w);
1642    let lo: &[T] = lo;
1643    let hi: &[T] = hi;
1644    (dst, move |i: usize| {
1645        if i < d {
1646            &lo[i * w..(i + 1) * w]
1647        } else {
1648            &hi[(i - d - 1) * w..(i - d) * w]
1649        }
1650    })
1651}
1652
1653/// Run the kernel over one block.
1654///
1655/// `out`, when given, receives the last instruction's result directly and
1656/// the returned index means nothing; otherwise the result stays in the
1657/// block buffer that index names. None means a step left the working type
1658/// and the caller must fall back.
1659fn exec_block<T, M, D, W, S>(
1660    k: &FusedKernel,
1661    srcs: &[Loaded<'_, T>],
1662    at: &Extent,
1663    sc: &mut Scratch<T>,
1664    out: Option<&mut [T]>,
1665    steps: &Steps<M, D, W, S>,
1666) -> Option<usize>
1667where
1668    T: Copy,
1669    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1670    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1671    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool,
1672    S: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T>,
1673{
1674    let Scratch { cells, width, free, stack, lets, carry } = sc;
1675    let w = *width;
1676    stack.clear();
1677    free.clear();
1678    lets.clear();
1679    let nslots = cells.len() / w;
1680    free.extend((0..nslots).rev());
1681    let place = |s: &Slot| match s {
1682        Slot::Input(j) => srcs[*j].on,
1683        Slot::Block(_, o) => *o,
1684    };
1685    let last = k.code.len() - 1;
1686    let head = if out.is_some() { last } else { k.code.len() };
1687    let mut scanned = 0usize;
1688    for ins in &k.code[..head] {
1689        match ins {
1690            Instr::Load(j) => stack.push(Slot::Input(*j)),
1691            Instr::Monad(op) => {
1692                let a = stack.pop()?;
1693                let dom = place(&a);
1694                let len = at.len_on(dom);
1695                let d = free.pop()?;
1696                let (dst, get) = split_slots(cells, w, d);
1697                let av = match a {
1698                    Slot::Input(j) => srcs[j].block(at, dom),
1699                    Slot::Block(i, _) => &get(i)[..len],
1700                };
1701                if !(steps.monad)(*op, av, &mut dst[..len]) {
1702                    return None;
1703                }
1704                release(free, lets, a);
1705                stack.push(Slot::Block(d, dom));
1706            }
1707            Instr::Scan(op) => {
1708                let a = stack.pop()?;
1709                let dom = place(&a);
1710                let len = at.len_on(dom);
1711                let d = free.pop()?;
1712                let (dst, get) = split_slots(cells, w, d);
1713                let av = match a {
1714                    Slot::Input(j) => srcs[j].block(at, dom),
1715                    Slot::Block(i, _) => &get(i)[..len],
1716                };
1717                carry[scanned] = Some((steps.scan)(*op, av, carry[scanned], &mut dst[..len])?);
1718                scanned += 1;
1719                release(free, lets, a);
1720                stack.push(Slot::Block(d, dom));
1721            }
1722            Instr::Window(op, size) => {
1723                let a = stack.pop()?;
1724                let d = free.pop()?;
1725                let (dst, get) = split_slots(cells, w, d);
1726                let av = match a {
1727                    Slot::Input(j) => srcs[j].block(at, On::Wide),
1728                    Slot::Block(i, _) => &get(i)[..at.wide_len],
1729                };
1730                let first = at.start - at.wide_start;
1731                if !(steps.window)(*op, *size, av, first, &mut dst[..at.len]) {
1732                    return None;
1733                }
1734                release(free, lets, a);
1735                stack.push(Slot::Block(d, On::Result));
1736            }
1737            Instr::Dyad(op) => {
1738                let b = stack.pop()?;
1739                let a = stack.pop()?;
1740                let dom = combine(place(&a), place(&b));
1741                let len = at.len_on(dom);
1742                let d = free.pop()?;
1743                let (dst, get) = split_slots(cells, w, d);
1744                let av = match a {
1745                    Slot::Input(j) => srcs[j].block(at, dom),
1746                    Slot::Block(i, _) => &get(i)[..len],
1747                };
1748                let bv = match b {
1749                    Slot::Input(j) => srcs[j].block(at, dom),
1750                    Slot::Block(i, _) => &get(i)[..len],
1751                };
1752                if !(steps.dyad)(*op, av, bv, &mut dst[..len]) {
1753                    return None;
1754                }
1755                for s in [a, b] {
1756                    release(free, lets, s);
1757                }
1758                stack.push(Slot::Block(d, dom));
1759            }
1760            Instr::Store(j) => {
1761                let Slot::Block(i, _) = stack.pop()? else { return None };
1762                if lets.len() != *j {
1763                    return None;
1764                }
1765                lets.push(i);
1766            }
1767            // A let stands where the pass computed it, which is the one
1768            // axis every repeat it stands for was written on.
1769            Instr::Let(j) => {
1770                stack.push(Slot::Block(*lets.get(*j)?, placed(Some(*k.let_doms.get(*j)?))))
1771            }
1772        }
1773    }
1774    let Some(dst) = out else {
1775        return match stack.pop()? {
1776            Slot::Block(i, _) => Some(i),
1777            // Every kernel ends in an operation, so the result is a buffer.
1778            Slot::Input(_) => None,
1779        };
1780    };
1781    // The last instruction writes the caller's buffer instead of a block.
1782    // The chain's root stands on the result's own axis, whatever its
1783    // operands stand on.
1784    let dst = &mut dst[..at.len];
1785    let view = |s: Slot, dom: On| match s {
1786        Slot::Input(j) => srcs[j].block(at, dom),
1787        Slot::Block(i, o) => &cells[i * w..i * w + at.len_on(o)],
1788    };
1789    let ok = match k.code[last] {
1790        Instr::Monad(op) => {
1791            let a = stack.pop()?;
1792            (steps.monad)(op, view(a, On::Result), dst)
1793        }
1794        Instr::Scan(op) => {
1795            let a = stack.pop()?;
1796            match (steps.scan)(op, view(a, On::Result), carry[scanned], dst) {
1797                Some(c) => {
1798                    carry[scanned] = Some(c);
1799                    true
1800                }
1801                None => false,
1802            }
1803        }
1804        Instr::Window(op, size) => {
1805            let a = stack.pop()?;
1806            let first = at.start - at.wide_start;
1807            (steps.window)(op, size, view(a, On::Wide), first, dst)
1808        }
1809        Instr::Dyad(op) => {
1810            let b = stack.pop()?;
1811            let a = stack.pop()?;
1812            let dom = combine(place(&a), place(&b));
1813            (steps.dyad)(op, view(a, dom), view(b, dom), dst)
1814        }
1815        // A kernel ends in the operation that makes its result.
1816        Instr::Load(_) | Instr::Store(_) | Instr::Let(_) => return None,
1817    };
1818    ok.then_some(usize::MAX)
1819}
1820
1821/// Give a block buffer back, unless a let is holding it for the rest of
1822/// the block.
1823fn release(free: &mut Vec<usize>, lets: &[usize], s: Slot) {
1824    if let Slot::Block(i, _) = s
1825        && !lets.contains(&i)
1826    {
1827        free.push(i);
1828    }
1829}
1830
1831/// The whole mapped result, one block at a time. None on integer overflow.
1832fn map_pass<T, M, D, W, S>(
1833    k: &FusedKernel,
1834    srcs: &Sources<'_, T>,
1835    n: usize,
1836    wide: usize,
1837    steps: &Steps<M, D, W, S>,
1838) -> Option<Vec<T>>
1839where
1840    T: FromNarrow + Default + Send + Sync,
1841    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1842    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1843    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool + Sync + Send,
1844    S: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T> + Sync + Send,
1845{
1846    let run = |start: usize, part: &mut [T]| {
1847        let w = BLOCK.min(part.len()).max(1);
1848        let mut sc = Scratch::new(k, w);
1849        let width = sc.width;
1850        let mut stage = vec![T::default(); srcs.staged * width];
1851        for (b, chunk) in part.chunks_mut(w).enumerate() {
1852            let at = Extent::of(start + b * w, chunk.len(), k.window, wide);
1853            let done = srcs.with_block(&at, &mut stage, width, |loaded| {
1854                exec_block(k, loaded, &at, &mut sc, Some(chunk), steps).is_some()
1855            });
1856            if !done {
1857                return false;
1858            }
1859        }
1860        true
1861    };
1862    if k.scans > 0 {
1863        // A running fold hands its accumulator to the next block, so the
1864        // blocks run in one order on one thread. That is the order the
1865        // unfused scan runs them in, and it rounds where that rounds.
1866        let mut out = vec![T::default(); n];
1867        return run(0, &mut out).then_some(out);
1868    }
1869    let (out, ok) = par::fill(n, run);
1870    ok.then_some(out)
1871}
1872
1873/// Independent accumulators the fold over a block keeps in flight, and the
1874/// block length below which one accumulator is cheaper. The reasoning is
1875/// the one `verb::FOLD_LANES` carries: a single accumulator makes the fold
1876/// a chain of dependent steps, and only an associative step is ever
1877/// absorbed here, so the lanes are a regrouping the float contract already
1878/// allows (§5.9).
1879const FOLD_LANES: usize = 8;
1880const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
1881
1882/// Fold one block of mapped values right to left, in lanes. None when a
1883/// step left the element type.
1884#[inline(always)]
1885fn fold_block_body<T, S>(v: &[T], step: &S) -> Option<T>
1886where
1887    T: Copy,
1888    S: Fn(T, T) -> Option<T>,
1889{
1890    let n = v.len();
1891    if n < MIN_LANE_WORK {
1892        let mut acc = v[n - 1];
1893        for &x in v[..n - 1].iter().rev() {
1894            acc = step(x, acc)?;
1895        }
1896        return Some(acc);
1897    }
1898    let rows = n / FOLD_LANES;
1899    let head = n - rows * FOLD_LANES;
1900    let last = head + (rows - 1) * FOLD_LANES;
1901    let mut acc = [v[last]; FOLD_LANES];
1902    acc.copy_from_slice(&v[last..last + FOLD_LANES]);
1903    for r in (0..rows - 1).rev() {
1904        let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
1905        for (slot, &x) in acc.iter_mut().zip(row) {
1906            *slot = step(x, *slot)?;
1907        }
1908    }
1909    let mut a = acc[FOLD_LANES - 1];
1910    for &x in acc[..FOLD_LANES - 1].iter().rev() {
1911        a = step(x, a)?;
1912    }
1913    for &x in v[..head].iter().rev() {
1914        a = step(x, a)?;
1915    }
1916    Some(a)
1917}
1918
1919multiversioned! {
1920    /// One block's values folded into one, at the CPU's own width.
1921    fn fold_block[T: Copy, S: Fn(T, T) -> Option<T>](
1922        v: &[T],
1923        step: &S,
1924    ) -> Option<T> = fold_block_body;
1925}
1926
1927/// Fold the mapped values of `lo .. hi` right to left, block by block.
1928fn fold_range<T, M, D, W, C, S>(
1929    k: &FusedKernel,
1930    srcs: &Sources<'_, T>,
1931    lo: usize,
1932    hi: usize,
1933    wide: usize,
1934    steps: &Steps<M, D, W, C>,
1935    step: &S,
1936) -> Option<T>
1937where
1938    T: FromNarrow + Default,
1939    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1940    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1941    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool,
1942    C: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T>,
1943    S: Fn(T, T) -> Option<T>,
1944{
1945    let w = BLOCK.min(hi - lo).max(1);
1946    let mut sc = Scratch::new(k, w);
1947    let width = sc.width;
1948    let mut stage = vec![T::default(); srcs.staged * width];
1949    let mut acc: Option<T> = None;
1950    // Blocks run backwards and the accumulator carries across them, so the
1951    // fold is the insert's own right-to-left order over the whole range.
1952    // Nothing a block computes depends on the block before it: a running
1953    // fold, which would, is never absorbed under a reduction.
1954    for b in (0..(hi - lo).div_ceil(w)).rev() {
1955        let start = lo + b * w;
1956        let len = (hi - start).min(w);
1957        let at = Extent::of(start, len, k.window, wide);
1958        let slot = srcs
1959            .with_block(&at, &mut stage, width, |loaded| {
1960                exec_block(k, loaded, &at, &mut sc, None, steps)
1961            })?;
1962        let block = fold_block(&sc.cells[slot * sc.width..slot * sc.width + len], step)?;
1963        acc = Some(match acc {
1964            None => block,
1965            Some(a) => step(block, a)?,
1966        });
1967    }
1968    acc
1969}
1970
1971/// The mapped values folded into one. None on integer overflow.
1972fn reduce_pass<T, M, D, W, C, S>(
1973    k: &FusedKernel,
1974    srcs: &Sources<'_, T>,
1975    n: usize,
1976    wide: usize,
1977    steps: &Steps<M, D, W, C>,
1978    step: S,
1979) -> Option<T>
1980where
1981    T: FromNarrow + Default + Send + Sync,
1982    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1983    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1984    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool + Sync + Send,
1985    C: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T> + Sync + Send,
1986    S: Fn(T, T) -> Option<T> + Sync + Send,
1987{
1988    let chunks = par::chunks(n, n * k.code.len());
1989    if chunks < 2 {
1990        return fold_range(k, srcs, 0, n, wide, steps, &step);
1991    }
1992    let per = n.div_ceil(chunks);
1993    let parts = par::map_indexed(n.div_ceil(per), |c| {
1994        fold_range(k, srcs, c * per, ((c + 1) * per).min(n), wide, steps, &step)
1995    });
1996    // The chunks combine right to left, the order they were folded in. That
1997    // regroups an associative float fold, which is the §5.9 contract; only
1998    // associative operations are absorbed.
1999    let mut it = parts.into_iter().rev();
2000    let mut acc = it.next()??;
2001    for part in it {
2002        acc = step(part?, acc)?;
2003    }
2004    Some(acc)
2005}
2006
2007// ------------------------------------------------------------ the kernels
2008//
2009// Each pass picks its operation before the loop and then runs one plain
2010// loop over slices, which is the shape the compiler vectorises. Nothing in
2011// here is hand-written SIMD, and nothing may become it.
2012//
2013// These four are the whole arithmetic of a kernel, so they are also where
2014// the CPU feature levels are chosen: each is compiled once per level (see
2015// `simd`) and the call dispatches on what the machine runs. One block of
2016// one instruction is thousands of elements, so the dispatch costs nothing
2017// measurable.
2018
2019macro_rules! each {
2020    ($a:expr, $dst:expr, $f:expr) => {{
2021        let f = $f;
2022        for (slot, &x) in $dst.iter_mut().zip($a) {
2023            *slot = f(x);
2024        }
2025        return true;
2026    }};
2027}
2028
2029macro_rules! zip {
2030    ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
2031        let f = $f;
2032        for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
2033            *slot = f(x, y);
2034        }
2035        return true;
2036    }};
2037}
2038
2039#[inline(always)]
2040fn monad_f64_body(op: ScalarMonad, a: &[f64], dst: &mut [f64], tol: Tol) -> bool {
2041    use ScalarMonad::*;
2042    match op {
2043        Conj => each!(a, dst, |x: f64| x),
2044        Neg => each!(a, dst, |x: f64| -x),
2045        Abs => each!(a, dst, f64::abs),
2046        // A magnitude the dialect's tolerance reads as zero has no sign,
2047        // exactly as unfused.
2048        Signum => each!(a, dst, |x: f64| if tol.is_zero(x) {
2049            0.0
2050        } else if x > 0.0 {
2051            1.0
2052        } else if x < 0.0 {
2053            -1.0
2054        } else {
2055            0.0
2056        }),
2057        // `% 0` is infinity, the J rule the unfused monad follows; under
2058        // APL's rules it is a DOMAIN ERROR, which only the unfused monad
2059        // can raise, so a zero here declines the whole kernel.
2060        Recip => {
2061            if !tol.is_j() && a.contains(&0.0) {
2062                return false;
2063            }
2064            each!(a, dst, |x: f64| if x == 0.0 { f64::INFINITY } else { 1.0 / x })
2065        }
2066        // Reached only through an integer chain, where they are the
2067        // identity: rounding a float narrows its dtype, which is declined.
2068        Floor => each!(a, dst, f64::floor),
2069        Ceil => each!(a, dst, f64::ceil),
2070        Inc => each!(a, dst, |x: f64| x + 1.0),
2071        Dec => each!(a, dst, |x: f64| x - 1.0),
2072        Double => each!(a, dst, |x: f64| x + x),
2073        Halve => each!(a, dst, |x: f64| x / 2.0),
2074        Square => each!(a, dst, |x: f64| x * x),
2075        OneMinus => each!(a, dst, |x: f64| 1.0 - x),
2076        Exp => each!(a, dst, f64::exp),
2077        _ => false,
2078    }
2079}
2080
2081#[inline(always)]
2082fn dyad_f64_body(op: ScalarDyad, a: &[f64], b: &[f64], dst: &mut [f64], tol: Tol) -> bool {
2083    use ScalarDyad::*;
2084    match op {
2085        Add => zip!(a, b, dst, |x: f64, y: f64| x + y),
2086        Sub => zip!(a, b, dst, |x: f64, y: f64| x - y),
2087        Mul => zip!(a, b, dst, |x: f64, y: f64| x * y),
2088        Min => zip!(a, b, dst, f64::min),
2089        Max => zip!(a, b, dst, f64::max),
2090        DivJ => zip!(a, b, dst, |x: f64, y: f64| if y == 0.0 {
2091            if x == 0.0 { 0.0 } else { f64::INFINITY.copysign(x) }
2092        } else {
2093            x / y
2094        }),
2095        // The quotient is rounded with the dialect's tolerance, exactly as
2096        // unfused: the fused answer must not differ from the plain one.
2097        Residue => zip!(a, b, dst, |x: f64, y: f64| tol.residue(x, y)),
2098        // A comparison is a number here, as it is in J: the boolean only
2099        // shows in the dtype of a result, which the caller narrows. Floats
2100        // compare with the dialect's tolerance, as they do unfused.
2101        Eq | Ne | Lt | Le | Gt | Ge => {
2102            zip!(a, b, dst, |x: f64, y: f64| tol_cmp(op, x, y, tol) as u8 as f64)
2103        }
2104        _ => false,
2105    }
2106}
2107
2108/// Whether a float block is one the fused kernel may keep.
2109///
2110/// A NaN is where IEEE arithmetic stops agreeing with the languages: J's
2111/// `0 * _` is 0 and its `_ - _` is refused, and both of those are an IEEE
2112/// NaN. Those rules live in the unfused verbs, so a block holding one
2113/// declines and the sentence is redone unfused — what an integer overflow
2114/// already does. An INFINITY needs no such treatment: it is an ordinary
2115/// value in J and, in APL, the operations that would refuse one (`÷` `⍟`
2116/// `!` `⋆` `○`) do not fuse at all, `÷`'s zero being caught in the kernel
2117/// above. The pass vectorises and finds nothing on ordinary data.
2118#[inline(always)]
2119fn no_nan_block(dst: &[f64]) -> bool {
2120    !dst.iter().any(|x| x.is_nan())
2121}
2122
2123/// Whether a device's answer holds no NaN. See [`no_nan_block`].
2124#[inline(always)]
2125fn finite_block(dst: &[f64]) -> bool {
2126    dst.iter().all(|x| x.is_finite())
2127}
2128
2129/// One step of a blockwise float fold, scan or window: a NaN abandons the
2130/// block for the unfused path, as [`no_nan_block`] does. It is how
2131/// `*/ 0 , _` reaches J's zero-factor rule and `+/ _ , __` its refusal.
2132#[inline(always)]
2133fn block_f64(r: f64) -> (f64, bool) {
2134    (r, r.is_nan())
2135}
2136
2137/// Integer passes fold overflow into a flag instead of branching out of the
2138/// loop: the whole evaluation is thrown away and redone unfused either way.
2139macro_rules! each_over {
2140    ($a:expr, $dst:expr, $f:expr) => {{
2141        let f = $f;
2142        let mut over = false;
2143        for (slot, &x) in $dst.iter_mut().zip($a) {
2144            let (v, o) = f(x);
2145            *slot = v;
2146            over |= o;
2147        }
2148        return !over;
2149    }};
2150}
2151
2152macro_rules! zip_over {
2153    ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
2154        let f = $f;
2155        let mut over = false;
2156        for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
2157            let (v, o) = f(x, y);
2158            *slot = v;
2159            over |= o;
2160        }
2161        return !over;
2162    }};
2163}
2164
2165#[inline(always)]
2166fn monad_i64_body(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool {
2167    use ScalarMonad::*;
2168    match op {
2169        Conj | Floor | Ceil => each!(a, dst, |x: i64| x),
2170        Neg => each_over!(a, dst, i64::overflowing_neg),
2171        Abs => each_over!(a, dst, i64::overflowing_abs),
2172        Signum => each!(a, dst, i64::signum),
2173        Inc => each_over!(a, dst, |x: i64| x.overflowing_add(1)),
2174        Dec => each_over!(a, dst, |x: i64| x.overflowing_sub(1)),
2175        Double => each_over!(a, dst, |x: i64| x.overflowing_add(x)),
2176        Square => each_over!(a, dst, |x: i64| x.overflowing_mul(x)),
2177        OneMinus => each_over!(a, dst, |x: i64| 1i64.overflowing_sub(x)),
2178        _ => false,
2179    }
2180}
2181
2182#[inline(always)]
2183fn dyad_i64_body(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool {
2184    use ScalarDyad::*;
2185    match op {
2186        Add => zip_over!(a, b, dst, i64::overflowing_add),
2187        Sub => zip_over!(a, b, dst, i64::overflowing_sub),
2188        Mul => zip_over!(a, b, dst, i64::overflowing_mul),
2189        Min => zip!(a, b, dst, i64::min),
2190        Max => zip!(a, b, dst, i64::max),
2191        Residue => zip!(a, b, dst, |x: i64, y: i64| if x == 0 {
2192            y
2193        } else {
2194            // wrapping_rem: i64::MIN % -1 is mathematically 0.
2195            let mut r = y.wrapping_rem(x);
2196            if r != 0 && (r < 0) != (x < 0) {
2197                r += x;
2198            }
2199            r
2200        }),
2201        Eq => zip!(a, b, dst, |x: i64, y: i64| (x == y) as i64),
2202        Ne => zip!(a, b, dst, |x: i64, y: i64| (x != y) as i64),
2203        Lt => zip!(a, b, dst, |x: i64, y: i64| (x < y) as i64),
2204        Le => zip!(a, b, dst, |x: i64, y: i64| (x <= y) as i64),
2205        Gt => zip!(a, b, dst, |x: i64, y: i64| (x > y) as i64),
2206        Ge => zip!(a, b, dst, |x: i64, y: i64| (x >= y) as i64),
2207        _ => false,
2208    }
2209}
2210
2211multiversioned! {
2212    /// One instruction of a kernel over one block of floats: the monadic
2213    /// operations. False is unreachable — every operation a kernel holds is
2214    /// covered — and exists so the two passes have one signature.
2215    fn monad_f64(
2216        op: ScalarMonad,
2217        a: &[f64],
2218        dst: &mut [f64],
2219        tol: Tol,
2220    ) -> bool = monad_f64_body;
2221}
2222
2223multiversioned! {
2224    /// One instruction of a kernel over one block of floats: the dyadic
2225    /// operations.
2226    fn dyad_f64(
2227        op: ScalarDyad,
2228        a: &[f64],
2229        b: &[f64],
2230        dst: &mut [f64],
2231        tol: Tol,
2232    ) -> bool = dyad_f64_body;
2233}
2234
2235multiversioned! {
2236    /// One instruction of a kernel over one block of integers: the monadic
2237    /// operations. False means the block left i64.
2238    fn monad_i64(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool = monad_i64_body;
2239}
2240
2241multiversioned! {
2242    /// One instruction of a kernel over one block of integers: the dyadic
2243    /// operations. False means the block left i64.
2244    fn dyad_i64(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool = dyad_i64_body;
2245}
2246
2247/// One block's running fold, continued from the accumulator the block
2248/// before it left. None when a step left the element type.
2249///
2250/// The accumulator runs the length of the argument, one step per item, in
2251/// the order the unfused scan takes them: what the fused kernel saves is
2252/// the traffic around the scan, not the scan.
2253#[inline(always)]
2254fn scan_block_body<T, F>(v: &[T], carry: Option<T>, dst: &mut [T], step: &F) -> Option<T>
2255where
2256    T: Copy,
2257    F: Fn(T, T) -> (T, bool),
2258{
2259    let mut over = false;
2260    let (mut acc, from) = match carry {
2261        Some(a) => (a, 0),
2262        None => {
2263            // The first item of a scan is the item itself.
2264            dst[0] = v[0];
2265            (v[0], 1)
2266        }
2267    };
2268    for (slot, &x) in dst.iter_mut().zip(v).skip(from) {
2269        let (r, o) = step(acc, x);
2270        acc = r;
2271        over |= o;
2272        *slot = acc;
2273    }
2274    (!over).then_some(acc)
2275}
2276
2277multiversioned! {
2278    /// One block of a running fold. The steps depend on one another, so
2279    /// what a wider vector reaches here is the loop around them.
2280    fn scan_block[T: Copy, F: Fn(T, T) -> (T, bool)](
2281        v: &[T],
2282        carry: Option<T>,
2283        dst: &mut [T],
2284        step: &F,
2285    ) -> Option<T> = scan_block_body;
2286}
2287
2288/// The windows of a block of floats, folded one per result item. The step
2289/// is chosen before the fold so that the fold itself is one plain loop.
2290fn window_pass_f64(op: ScalarDyad, k: usize, v: &[f64], first: usize, dst: &mut [f64]) -> bool {
2291    use ScalarDyad::*;
2292    match op {
2293        Add => windows_into(v, k, first, dst, &|a: f64, b: f64| block_f64(a + b)),
2294        Mul => windows_into(v, k, first, dst, &|a: f64, b: f64| block_f64(a * b)),
2295        Min => windows_into(v, k, first, dst, &|a: f64, b: f64| (a.min(b), false)),
2296        Max => windows_into(v, k, first, dst, &|a: f64, b: f64| (a.max(b), false)),
2297        _ => false,
2298    }
2299}
2300
2301fn window_pass_i64(op: ScalarDyad, k: usize, v: &[i64], first: usize, dst: &mut [i64]) -> bool {
2302    use ScalarDyad::*;
2303    match op {
2304        Add => windows_into(v, k, first, dst, &i64::overflowing_add),
2305        Mul => windows_into(v, k, first, dst, &i64::overflowing_mul),
2306        Min => windows_into(v, k, first, dst, &|a: i64, b: i64| (a.min(b), false)),
2307        Max => windows_into(v, k, first, dst, &|a: i64, b: i64| (a.max(b), false)),
2308        _ => false,
2309    }
2310}
2311
2312fn scan_pass_f64(op: ScalarDyad, v: &[f64], carry: Option<f64>, dst: &mut [f64]) -> Option<f64> {
2313    use ScalarDyad::*;
2314    match op {
2315        Add => scan_block(v, carry, dst, &|a: f64, b: f64| block_f64(a + b)),
2316        Mul => scan_block(v, carry, dst, &|a: f64, b: f64| block_f64(a * b)),
2317        Min => scan_block(v, carry, dst, &|a: f64, b: f64| (a.min(b), false)),
2318        Max => scan_block(v, carry, dst, &|a: f64, b: f64| (a.max(b), false)),
2319        _ => None,
2320    }
2321}
2322
2323fn scan_pass_i64(op: ScalarDyad, v: &[i64], carry: Option<i64>, dst: &mut [i64]) -> Option<i64> {
2324    use ScalarDyad::*;
2325    match op {
2326        Add => scan_block(v, carry, dst, &i64::overflowing_add),
2327        Mul => scan_block(v, carry, dst, &i64::overflowing_mul),
2328        Min => scan_block(v, carry, dst, &|a: i64, b: i64| (a.min(b), false)),
2329        Max => scan_block(v, carry, dst, &|a: i64, b: i64| (a.max(b), false)),
2330        _ => None,
2331    }
2332}
2333
2334/// One fold step of an absorbed reduction. None on integer overflow.
2335fn step_i64(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
2336    use ScalarDyad::*;
2337    match op {
2338        Add => a.checked_add(b),
2339        Mul => a.checked_mul(b),
2340        Min => Some(a.min(b)),
2341        Max => Some(a.max(b)),
2342        _ => None,
2343    }
2344}
2345
2346/// One fold step of an absorbed float reduction, for a backend that mapped
2347/// the values elsewhere and brings its partials back here to combine.
2348pub(crate) fn step(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
2349    step_f64(op, a, b)
2350}
2351
2352fn step_f64(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
2353    use ScalarDyad::*;
2354    match op {
2355        // None here abandons the fused fold for the plain one, which is
2356        // where the dialect's rule for the value lives.
2357        Add => Some(a + b).filter(|r| !r.is_nan()),
2358        Mul => Some(a * b).filter(|r| !r.is_nan()),
2359        Min => Some(a.min(b)),
2360        Max => Some(a.max(b)),
2361        _ => None,
2362    }
2363}
2364
2365// ------------------------------------------------------------- the driver
2366
2367/// A rank-0 argument as one block of the repeated value, which is how it
2368/// reaches every element without an index test. None for an argument with
2369/// items, which is read where it lies.
2370fn splat_f64(a: &Array, w: usize) -> Option<Vec<f64>> {
2371    if a.rank() != 0 {
2372        return None;
2373    }
2374    let v = match &a.data {
2375        Data::Bool(d) => d[0] as f64,
2376        Data::I64(d) => d[0] as f64,
2377        Data::F64(d) => d[0],
2378        Data::Ext(_)
2379        | Data::Rat(_)
2380        | Data::Complex(_)
2381        | Data::Char(_)
2382        | Data::Symbol(_)
2383        | Data::Box(_) => {
2384            return Some(Vec::new());
2385        }
2386    };
2387    Some(vec![v; w])
2388}
2389
2390fn splat_i64(a: &Array, w: usize) -> Option<Vec<i64>> {
2391    if a.rank() != 0 {
2392        return None;
2393    }
2394    let v = match &a.data {
2395        Data::Bool(d) => d[0] as i64,
2396        Data::I64(d) => d[0],
2397        _ => return Some(Vec::new()),
2398    };
2399    Some(vec![v; w])
2400}
2401
2402/// An argument with items, seen by a float kernel: its own buffer when that
2403/// is already f64, the narrow values otherwise.
2404fn narrow_f64(a: &Array) -> Result<&[f64], Narrow<'_>> {
2405    match &a.data {
2406        Data::I64(d) => Err(Narrow::I64(d)),
2407        Data::Bool(d) => Err(Narrow::Bool(d)),
2408        _ => Ok(a.as_f64_slice().unwrap_or(&[])),
2409    }
2410}
2411
2412fn narrow_i64(a: &Array) -> Result<&[i64], Narrow<'_>> {
2413    match &a.data {
2414        Data::Bool(d) => Err(Narrow::Bool(d)),
2415        _ => Ok(a.as_i64_slice().unwrap_or(&[])),
2416    }
2417}
2418
2419/// The input list one run reads: a repeated scalar from `owned`, an
2420/// argument's own buffer, or a narrow buffer to be promoted block by block.
2421fn sources<'a, T>(
2422    inputs: &'a [Array],
2423    owned: &'a [Option<Vec<T>>],
2424    on: &impl Fn(usize) -> On,
2425    narrow: impl Fn(&'a Array) -> Result<&'a [T], Narrow<'a>>,
2426) -> Sources<'a, T> {
2427    let mut staged = 0;
2428    let of = inputs
2429        .iter()
2430        .zip(owned)
2431        .enumerate()
2432        .map(|(j, (a, o))| match o {
2433            Some(v) => Source::Ready(Loaded { data: v, splat: true, on: on(j), base: 0 }),
2434            None => match narrow(a) {
2435                Ok(d) => Source::Ready(Loaded { data: d, splat: false, on: on(j), base: 0 }),
2436                Err(n) => {
2437                    staged += 1;
2438                    Source::Staged(n, on(j))
2439                }
2440            },
2441        })
2442        .collect();
2443    Sources { of, staged }
2444}
2445
2446/// The shape every element of the result has: identical for all non-scalar
2447/// inputs, since anything else needs the agreement machinery.
2448pub(crate) fn common_shape(inputs: &[Array]) -> Option<Option<Vec<usize>>> {
2449    let mut shape: Option<&Vec<usize>> = None;
2450    for a in inputs {
2451        if a.rank() == 0 {
2452            continue;
2453        }
2454        match shape {
2455            None => shape = Some(&a.shape),
2456            Some(s) if *s == a.shape => {}
2457            Some(_) => return None,
2458        }
2459    }
2460    Some(shape.cloned())
2461}
2462
2463/// The axes a kernel's inputs stand on: the shape of its result, and the
2464/// length of the wide axis its window steps read.
2465///
2466/// This is the whole of the alignment rule, and it is decided by shapes
2467/// alone. Where a chain reads an input is settled when the chain is built —
2468/// everything under a window step is wide — so all that is left at run time
2469/// is that the inputs on one axis agree with each other, and that the two
2470/// axes stand `k - 1` items apart. `19 }. y` beside `20 +/\ y` passes
2471/// because it is 19 items shorter; `18 }. y` beside it does not, and the
2472/// chain runs and raises the length error it was going to raise. Nothing is
2473/// shifted or padded here: an input arrives as the items it holds.
2474struct Axes {
2475    shape: Vec<usize>,
2476    /// Items of the wide axis, when there is a window step to read it.
2477    wide: usize,
2478}
2479
2480fn axes(k: &FusedKernel, inputs: &[Array]) -> Option<Axes> {
2481    let Some(window) = k.window else {
2482        // Every input a scalar: no work worth blocking, and a reduction
2483        // would need the leading axis a scalar has not got.
2484        let shape = common_shape(inputs)??;
2485        // A running fold folds items, and a block of this kernel is
2486        // elements: over anything but a vector the two are not the same
2487        // fold, so a higher-rank argument goes the way it went.
2488        if k.scans > 0 && shape.len() != 1 {
2489            return None;
2490        }
2491        return Some(Axes { shape, wide: 0 });
2492    };
2493    let (mut wide, mut result) = (None, None);
2494    for (a, dom) in inputs.iter().zip(&k.doms) {
2495        // A scalar reaches every item of whatever it is combined with, so
2496        // it stands on either axis and constrains neither.
2497        if a.rank() == 0 {
2498            continue;
2499        }
2500        // A window folds the items of a vector. An input the chain reads on
2501        // both axes cannot be two lengths at once.
2502        let (Some(d), 1) = (dom, a.rank()) else { return None };
2503        let seen = if *d == Dom::Wide { &mut wide } else { &mut result };
2504        match seen {
2505            None => *seen = Some(a.shape[0]),
2506            Some(m) if *m == a.shape[0] => {}
2507            Some(_) => return None,
2508        }
2509    }
2510    let wide = wide?;
2511    if wide < window {
2512        // No window fits: the result has no items, which the chain builds
2513        // out of the verb's own answer for an empty argument.
2514        return None;
2515    }
2516    let count = wide - window + 1;
2517    if result.is_some_and(|m| m != count) {
2518        return None;
2519    }
2520    Some(Axes { shape: vec![count], wide })
2521}
2522
2523/// Run a fused node. None means the kernel declined and the caller must
2524/// evaluate the original subtree, which is always allowed to be slower and
2525/// never allowed to differ.
2526pub(crate) fn run(k: &FusedKernel, inputs: &[Array]) -> Option<Array> {
2527    let reducing = matches!(k.yields, Yield::Reduce(_));
2528    let Axes { shape, wide } = axes(k, inputs)?;
2529    let n: usize = shape.iter().product();
2530    if n == 0 {
2531        return None;
2532    }
2533    if reducing && (shape.len() != 1 || n < 2) {
2534        // A one-item reduction yields the item itself, dtype and all, and a
2535        // higher-rank one folds cells rather than elements.
2536        return None;
2537    }
2538    let (working, root) = working_type(k, inputs)?;
2539    if k.yields == Yield::Tally {
2540        // The shapes have already said how many items the chain produces,
2541        // and the type rules have said it would reach them without an
2542        // error. There is nothing else a tally wants from the values.
2543        return Some(Array::scalar_i64(shape[0] as i64));
2544    }
2545    // A repeated scalar is one block long, and a block reads its window
2546    // halo as well as its own items.
2547    let w = BLOCK.min(n).max(1) + 3 * k.window.unwrap_or(0);
2548    // The kernel's comparisons carry the tolerance the program was compiled
2549    // with, so a fused comparison answers as the unfused one does.
2550    let tol = k.tol;
2551    let on = |j: usize| placed(k.doms[j]);
2552
2553    let data = if working == DType::F64 {
2554        let steps = Steps {
2555            monad: move |op, a: &[f64], dst: &mut [f64]| {
2556                monad_f64(op, a, dst, tol) && no_nan_block(dst)
2557            },
2558            dyad: move |op, a: &[f64], b: &[f64], dst: &mut [f64]| {
2559                dyad_f64(op, a, b, dst, tol) && no_nan_block(dst)
2560            },
2561            window: window_pass_f64,
2562            scan: scan_pass_f64,
2563        };
2564        let owned: Vec<Option<Vec<f64>>> = inputs.iter().map(|a| splat_f64(a, w)).collect();
2565        let srcs = sources(inputs, &owned, &on, narrow_f64);
2566        match k.reduce() {
2567            None => {
2568                let out = map_pass(k, &srcs, n, wide, &steps)?;
2569                float_result(out, root)
2570            }
2571            Some(op) => {
2572                let v = reduce_pass(k, &srcs, n, wide, &steps, |a, b| step_f64(op, a, b))?;
2573                // A comparison at the root maps to exact 0 and 1, which the
2574                // fold keeps exact; the reduction of booleans is integer.
2575                match root {
2576                    DType::F64 => Data::F64(vec![v].into()),
2577                    _ => Data::I64(vec![v as i64].into()),
2578                }
2579            }
2580        }
2581    } else {
2582        let steps = Steps {
2583            monad: monad_i64,
2584            dyad: dyad_i64,
2585            window: window_pass_i64,
2586            scan: scan_pass_i64,
2587        };
2588        let owned: Vec<Option<Vec<i64>>> = inputs.iter().map(|a| splat_i64(a, w)).collect();
2589        let srcs = sources(inputs, &owned, &on, narrow_i64);
2590        match k.reduce() {
2591            None => {
2592                let out = map_pass(k, &srcs, n, wide, &steps)?;
2593                int_result(out, root)
2594            }
2595            Some(op) => {
2596                let v = reduce_pass(k, &srcs, n, wide, &steps, |a, b| step_i64(op, a, b))?;
2597                Data::I64(vec![v].into())
2598            }
2599        }
2600    };
2601    Some(Array::new(if reducing { Vec::new() } else { shape }, data))
2602}
2603
2604/// The mapped block values as the array the unfused chain would build. A
2605/// comparison at the root costs one narrowing pass, since the kernel
2606/// computes 0 and 1 in its working type and a boolean array holds bytes.
2607fn float_result(out: Vec<f64>, root: DType) -> Data {
2608    match root {
2609        DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0.0) as u8).into()),
2610        _ => Data::F64(out.into()),
2611    }
2612}
2613
2614fn int_result(out: Vec<i64>, root: DType) -> Data {
2615    match root {
2616        DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0) as u8).into()),
2617        _ => Data::I64(out.into()),
2618    }
2619}
2620
2621// -------------------------------------------------------- describing one
2622//
2623// Read-only descriptions of a compiled kernel, for `Program::explain`.
2624// Nothing here runs a kernel or changes one; the summary is derived from
2625// the code the pass emitted, and the decline reason re-checks the same
2626// preconditions `run` checks before it starts.
2627
2628/// Why a kernel handed its work back to the chain it came from.
2629#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2630pub enum Decline {
2631    /// Inputs disagree on shape, or every input is a scalar: broadcasting
2632    /// and agreement are the chain's business.
2633    Agreement,
2634    /// Nothing to compute.
2635    Empty,
2636    /// An absorbed reduction wants one axis with at least two items.
2637    ReduceShape,
2638    /// One working type cannot hold every step exactly — a chain that
2639    /// computes integers along a float path, or non-numeric data.
2640    WorkingType,
2641    /// The preconditions held, so a step went out of range mid-block:
2642    /// integer overflow, which the chain redoes in a wider type.
2643    Overflow,
2644    /// A window step wants one vector axis longer than the window, and
2645    /// every other input aligned on the window's last item.
2646    Window,
2647}
2648
2649impl Decline {
2650    pub fn reason(self) -> &'static str {
2651        match self {
2652            Decline::Agreement => "the inputs need agreement or are all scalars",
2653            Decline::Empty => "there is nothing to compute",
2654            Decline::ReduceShape => "the reduction needs one axis of two or more items",
2655            Decline::WorkingType => "no single working type holds every step exactly",
2656            Decline::Overflow => "an integer step left 64-bit range",
2657            Decline::Window => "the window does not fit the axis, or the inputs are not aligned with it",
2658        }
2659    }
2660}
2661
2662/// Why this kernel would decline these inputs, or None if it would run.
2663///
2664/// A read-only mirror of the preconditions at the top of `run`: it looks
2665/// at shapes and dtypes only, never at values, so the one thing it cannot
2666/// see in advance is an overflow — which is what is left when every
2667/// precondition holds.
2668pub fn decline_reason(k: &FusedKernel, inputs: &[Array]) -> Option<Decline> {
2669    let Some(Axes { shape, .. }) = axes(k, inputs) else {
2670        return Some(if k.window.is_some() { Decline::Window } else { Decline::Agreement });
2671    };
2672    let n: usize = shape.iter().product();
2673    if n == 0 {
2674        return Some(Decline::Empty);
2675    }
2676    if matches!(k.yields, Yield::Reduce(_)) && (shape.len() != 1 || n < 2) {
2677        return Some(Decline::ReduceShape);
2678    }
2679    if working_type(k, inputs).is_none() {
2680        return Some(Decline::WorkingType);
2681    }
2682    Some(Decline::Overflow)
2683}
2684
2685/// What a compiled kernel is made of.
2686#[derive(Clone, Debug, PartialEq, Eq)]
2687pub struct Summary {
2688    /// Arithmetic steps: the monads and dyads, not the loads and stores.
2689    pub ops: usize,
2690    /// Those steps in the order the kernel performs them.
2691    pub op_names: Vec<String>,
2692    /// The reduction folded into the same pass, if there is one.
2693    pub reduce: Option<&'static str>,
2694    /// True when the whole chain collapsed to a count of its own items.
2695    pub tally: bool,
2696    /// Values the kernel keeps for a second read within one block.
2697    pub lets: usize,
2698    /// Subtrees the chain reads.
2699    pub inputs: usize,
2700    /// Elements one block buffer holds.
2701    pub block: usize,
2702    /// The window every window step folds, when the kernel has one.
2703    pub window: Option<usize>,
2704    /// Running folds the kernel carries from block to block.
2705    pub scans: usize,
2706}
2707
2708impl std::fmt::Display for Summary {
2709    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2710        write!(f, "{} op{}", self.ops, if self.ops == 1 { "" } else { "s" })?;
2711        if !self.op_names.is_empty() {
2712            write!(f, ": {}", self.op_names.join(" "))?;
2713        }
2714        if let Some(r) = self.reduce {
2715            write!(f, "; {r}/ absorbed")?;
2716        }
2717        if self.tally {
2718            write!(f, "; tally only")?;
2719        }
2720        if self.lets > 0 {
2721            write!(f, "; {} let slot{}", self.lets, if self.lets == 1 { "" } else { "s" })?;
2722        }
2723        if let Some(k) = self.window {
2724            write!(f, "; window {k}")?;
2725        }
2726        if self.scans > 0 {
2727            write!(f, "; {} running fold{}", self.scans, if self.scans == 1 { "" } else { "s" })?;
2728        }
2729        write!(f, "; block {}", self.block)
2730    }
2731}
2732
2733/// Describe a compiled kernel: what it computes, and with what.
2734pub fn summary(k: &FusedKernel) -> Summary {
2735    let mut op_names: Vec<String> = Vec::new();
2736    let mut lets = 0usize;
2737    for ins in &k.code {
2738        match ins {
2739            Instr::Monad(op) => op_names.push(monad_name(*op).to_string()),
2740            Instr::Dyad(op) => op_names.push(dyad_name(*op).to_string()),
2741            Instr::Window(op, k) => op_names.push(format!("{k} {}/\\", dyad_name(*op))),
2742            Instr::Scan(op) => op_names.push(format!("{}/\\", dyad_name(*op))),
2743            Instr::Store(_) => lets += 1,
2744            Instr::Load(_) | Instr::Let(_) => {}
2745        }
2746    }
2747    Summary {
2748        ops: op_names.len(),
2749        op_names,
2750        reduce: k.reduce().map(dyad_name),
2751        tally: k.yields == Yield::Tally,
2752        lets,
2753        inputs: k.leaves.iter().copied().max().map_or(0, |m| m + 1),
2754        block: BLOCK,
2755        window: k.window,
2756        scans: k.scans,
2757    }
2758}
2759
2760/// The names the pass took out of the program: values it moved into the
2761/// kernels that read them, so no sentence computes them as arrays any more.
2762pub fn inlined_names(p: &Program) -> Vec<String> {
2763    let Some(Expr::Elided { orig, .. }) = p.stmts.first() else { return Vec::new() };
2764    let assigned = |stmts: &[Expr]| -> Vec<String> {
2765        stmts
2766            .iter()
2767            .filter_map(|s| match s {
2768                Expr::Assign { name, .. } => Some(name.clone()),
2769                _ => None,
2770            })
2771            .collect()
2772    };
2773    let kept = assigned(&p.stmts);
2774    assigned(orig).into_iter().filter(|n| !kept.contains(n)).collect()
2775}
2776
2777/// J spellings for the elementwise operations a kernel can hold. Only the
2778/// naming lives here; the meanings are [`crate::verb`]'s.
2779fn monad_name(op: ScalarMonad) -> &'static str {
2780    use ScalarMonad::*;
2781    match op {
2782        Conj => "+",
2783        Neg => "-",
2784        Signum => "*",
2785        Recip => "%",
2786        Sqrt => "%:",
2787        Exp => "^",
2788        Abs => "|",
2789        Floor => "<.",
2790        Ceil => ">.",
2791        Not => "-.",
2792        OneMinus => "-.",
2793        Inc => ">:",
2794        Dec => "<:",
2795        Double => "+:",
2796        Halve => "-:",
2797        Square => "*:",
2798        Ln => "^.",
2799        Pi => "o.",
2800        Factorial => "!",
2801        Imaginary => "j.",
2802        Polar => "r.",
2803    }
2804}
2805
2806pub(crate) fn dyad_name(op: ScalarDyad) -> &'static str {
2807    use ScalarDyad::*;
2808    match op {
2809        Add => "+",
2810        Sub => "-",
2811        Mul => "*",
2812        DivJ | DivApl => "%",
2813        Min => "<.",
2814        Max => ">.",
2815        Pow => "^",
2816        Residue => "|",
2817        Eq => "=",
2818        Ne => "~:",
2819        Lt => "<",
2820        Le => "<:",
2821        Gt => ">",
2822        Ge => ">:",
2823        Lcm => "*.",
2824        Gcd => "+.",
2825        Log => "^.",
2826        Root => "%:",
2827        Circle => "o.",
2828        Binomial => "!",
2829        MakeComplex => "j.",
2830        PolarBy => "r.",
2831    }
2832}
2833
2834/// Evaluate a fused node from its already-evaluated inputs, or report that
2835/// the original subtree must run instead.
2836///
2837/// This is the one place a device gets to run libjay's arithmetic. With a
2838/// device attached the kernel is offered to it first; everything it will not
2839/// take comes back here with a reason, and the CPU path runs exactly as it
2840/// runs with no device in sight. The device therefore cannot change a
2841/// result's shape, dtype or error — only where the arithmetic happened.
2842pub(crate) fn eval_on(
2843    device: Option<&crate::device::Device>,
2844    k: &FusedKernel,
2845    inputs: &[Array],
2846) -> (Option<Array>, crate::device::Placement) {
2847    use crate::device::Placement;
2848    let mut placement = Placement::Default;
2849    // A block kernel reads every input and writes every slot at the same
2850    // index, so the order the buffers are laid out in cannot reach the
2851    // result — as long as every non-scalar input is laid out the same way.
2852    // Then the answer is laid out that way too, and no transpose is made.
2853    let materialised: Vec<Array>;
2854    let (inputs, layout) = match kernel_layout(inputs) {
2855        Some(l) => (inputs, l),
2856        None => {
2857            materialised = inputs.iter().map(Array::to_row_major).collect();
2858            (&materialised[..], Layout::RowMajor)
2859        }
2860    };
2861    // The device is offered row-major work only: uploading a matrix that is
2862    // faster to fold where it lies would be the wrong trade anyway.
2863    if layout == Layout::RowMajor && let Some(d) = device.filter(|d| d.is_gpu()) {
2864        match crate::device::try_run(d, k, inputs) {
2865            // A shader computes in IEEE arithmetic and knows none of the
2866            // dialect's rules for what an infinity or a NaN means, so an
2867            // answer holding one is handed back and the sentence is redone
2868            // on the CPU — the same decline the CPU kernels make among
2869            // themselves. Without it the device WOULD change a result.
2870            Ok(a) if finite_answer(&a) => return (Some(a), Placement::Gpu),
2871            Ok(_) => placement = Placement::Cpu(crate::device::Refusal::NonFinite),
2872            Err(why) => placement = Placement::Cpu(why),
2873        }
2874    }
2875    let r = run(k, inputs).map(|a| a.with_layout(layout));
2876    if r.is_none() {
2877        note_fallback();
2878    }
2879    (r, placement)
2880}
2881
2882/// Whether a device's answer holds no infinity and no NaN, and may
2883/// therefore be kept. See [`finite_block`], whose rule this is.
2884fn finite_answer(a: &Array) -> bool {
2885    match &a.data {
2886        crate::array::Data::F64(v) => finite_block(v),
2887        _ => true,
2888    }
2889}
2890
2891/// The layout a fused kernel's answer keeps, or None when its inputs
2892/// disagree and the caller must materialise the rows of each.
2893fn kernel_layout(inputs: &[Array]) -> Option<Layout> {
2894    let mut found: Option<Layout> = None;
2895    for a in inputs {
2896        // A scalar is one value repeated into every block: it has no layout
2897        // to agree or disagree with.
2898        if a.rank() == 0 {
2899            continue;
2900        }
2901        match found {
2902            None => found = Some(a.layout()),
2903            Some(l) if l == a.layout() => {}
2904            Some(_) => return None,
2905        }
2906    }
2907    Some(found.unwrap_or_default())
2908}
2909
2910#[cfg(test)]
2911mod tests {
2912    use super::*;
2913    use crate::frontend::{compile, Dialect, Lang};
2914
2915    fn program(src: &str) -> Program {
2916        compile(Lang::J, src, &Dialect::default()).expect("compile")
2917    }
2918
2919    #[test]
2920    fn a_chain_of_two_scalar_verbs_fuses() {
2921        assert!(is_fused(&program("1 + 2 * {x}")));
2922        assert!(is_fused(&program("+/ {w} * {x}")));
2923        assert!(is_fused(&program("+/ ^ {x}")));
2924    }
2925
2926    #[test]
2927    fn one_verb_on_its_own_is_left_alone() {
2928        assert!(!is_fused(&program("2 * {x}")));
2929        assert!(!is_fused(&program("+/ {x}")));
2930        assert!(!is_fused(&program("{x}")));
2931    }
2932
2933    #[test]
2934    fn a_verb_the_kernel_does_not_cover_breaks_the_chain() {
2935        // `%:` can fail elementwise, so it stays outside; the chain under it
2936        // still fuses.
2937        assert!(!is_fused(&program("%: 2 * {x}")));
2938        assert!(is_fused(&program("%: 1 + 2 * {x}")));
2939    }
2940
2941    #[test]
2942    fn an_effect_in_a_leaf_keeps_the_chain_unfused() {
2943        assert!(!is_fused(&program("1 + 2 * echo {x}")));
2944    }
2945
2946    #[test]
2947    fn the_postfix_program_pushes_the_left_operand_first() {
2948        let p = program("{w} - {x} - 1");
2949        let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2950        assert_eq!(
2951            kernel.code(),
2952            [
2953                Instr::Load(2),
2954                Instr::Load(1),
2955                Instr::Load(0),
2956                Instr::Dyad(ScalarDyad::Sub),
2957                Instr::Dyad(ScalarDyad::Sub),
2958            ]
2959        );
2960        // One buffer holds the inner difference, one takes the outer one.
2961        assert_eq!(kernel.slots, 2);
2962    }
2963
2964    #[test]
2965    fn a_value_the_chain_reads_twice_becomes_a_let() {
2966        // What `d =. {x} + 1` then `+/ d * d` comes to once the name has
2967        // moved into the kernel: the sum is computed once per block.
2968        let p = program("+/ ({x} + 1) * ({x} + 1)");
2969        let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2970        assert_eq!(
2971            kernel.code(),
2972            [
2973                Instr::Load(1),
2974                Instr::Load(0),
2975                Instr::Dyad(ScalarDyad::Add),
2976                Instr::Store(0),
2977                Instr::Let(0),
2978                Instr::Let(0),
2979                Instr::Dyad(ScalarDyad::Mul),
2980            ]
2981        );
2982        // One buffer for the let, one for the product it feeds.
2983        assert_eq!(kernel.slots, 2);
2984    }
2985
2986    #[test]
2987    fn a_named_value_moves_into_the_sentence_that_reads_it() {
2988        let p = program("d =. {x} + 1\n+/ d * d");
2989        assert!(is_inlined(&p));
2990        // Three sentences: what the program was, the check that stands
2991        // where the assignment stood, and the sum, which is now the chain
2992        // of the test above.
2993        assert_eq!(p.stmts.len(), 3);
2994        let Expr::Fused { kernel, .. } = &p.stmts[2] else { panic!("the sum did not fuse") };
2995        assert!(kernel.code().contains(&Instr::Store(0)));
2996        assert_eq!(unfused(&p).stmts.len(), 2);
2997    }
2998}