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