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 | DType::Symbol => 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.
1424#[derive(Clone, Copy)]
1425struct Loaded<'a, T> {
1426    data: &'a [T],
1427    splat: bool,
1428    /// The axis the chain reads this input on.
1429    on: On,
1430    /// The index `data[0]` stands at on that axis. Zero for an argument's
1431    /// own buffer, and the block's own start for one staged a block at a
1432    /// time.
1433    base: usize,
1434}
1435
1436impl<T> Loaded<'_, T> {
1437    #[inline]
1438    fn block(&self, at: &Extent, dom: On) -> &[T] {
1439        if self.splat {
1440            return &self.data[..at.len_on(dom)];
1441        }
1442        let (start, len) = read_over(self.on, at);
1443        &self.data[start - self.base..start - self.base + len]
1444    }
1445}
1446
1447/// The range of its own axis an input standing on `on` is read over for one
1448/// block.
1449#[inline]
1450fn read_over(on: On, at: &Extent) -> (usize, usize) {
1451    match on {
1452        On::Wide => (at.wide_start, at.wide_len),
1453        _ => (at.start, at.len),
1454    }
1455}
1456
1457/// An argument narrower than the working type, read where it lies.
1458///
1459/// A whole widened copy of such an argument costs two round trips of the
1460/// working set — writing 160 MB of freshly faulted pages, then reading them
1461/// back — where the values themselves are needed one block at a time and a
1462/// block fits in cache. So the promotion happens at the block, into a
1463/// staging buffer each thread reuses; the values are what the widened copy
1464/// would have held, element for element.
1465#[derive(Clone, Copy)]
1466enum Narrow<'a> {
1467    I64(&'a [i64]),
1468    Bool(&'a [u8]),
1469}
1470
1471/// One input as the kernel will read it: its own buffer when that already
1472/// holds the working type, a narrower buffer to be promoted otherwise.
1473enum Source<'a, T> {
1474    Ready(Loaded<'a, T>),
1475    /// The narrow values, and the axis the chain reads them on.
1476    Staged(Narrow<'a>, On),
1477}
1478
1479/// The staging buffer's element type, filled from a narrow argument.
1480trait FromNarrow: Copy {
1481    fn fill(src: Narrow<'_>, at: usize, dst: &mut [Self]);
1482}
1483
1484impl FromNarrow for f64 {
1485    #[inline]
1486    fn fill(src: Narrow<'_>, at: usize, dst: &mut [f64]) {
1487        match src {
1488            Narrow::I64(v) => {
1489                for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1490                    *slot = x as f64;
1491                }
1492            }
1493            Narrow::Bool(v) => {
1494                for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1495                    *slot = x as f64;
1496                }
1497            }
1498        }
1499    }
1500}
1501
1502impl FromNarrow for i64 {
1503    #[inline]
1504    fn fill(src: Narrow<'_>, at: usize, dst: &mut [i64]) {
1505        match src {
1506            // The working type is integer only when no input is a float.
1507            Narrow::I64(v) => dst.copy_from_slice(&v[at..at + dst.len()]),
1508            Narrow::Bool(v) => {
1509                for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1510                    *slot = x as i64;
1511                }
1512            }
1513        }
1514    }
1515}
1516
1517/// The inputs of one run, and how wide a staging block has to be.
1518struct Sources<'a, T> {
1519    of: Vec<Source<'a, T>>,
1520    staged: usize,
1521}
1522
1523impl<T: FromNarrow + Default> Sources<'_, T> {
1524    /// Promote every staged input's block into `stage` and hand the whole
1525    /// input list, block-local, to `f`.
1526    ///
1527    /// `stage` holds one region of `width` per staged input, so the regions
1528    /// are the same from block to block and a thread faults them once.
1529    fn with_block<R>(
1530        &self,
1531        at: &Extent,
1532        stage: &mut [T],
1533        width: usize,
1534        f: impl FnOnce(&[Loaded<'_, T>]) -> R,
1535    ) -> R {
1536        let mut k = 0;
1537        for s in &self.of {
1538            let Source::Staged(src, on) = s else { continue };
1539            let (start, len) = read_over(*on, at);
1540            T::fill(*src, start, &mut stage[k * width..k * width + len]);
1541            k += 1;
1542        }
1543        let mut k = 0;
1544        let loaded: Vec<Loaded<'_, T>> = self
1545            .of
1546            .iter()
1547            .map(|s| match s {
1548                Source::Ready(l) => *l,
1549                Source::Staged(_, on) => {
1550                    let (start, len) = read_over(*on, at);
1551                    let d = &stage[k * width..k * width + len];
1552                    k += 1;
1553                    Loaded { data: d, splat: false, on: *on, base: start }
1554                }
1555            })
1556            .collect();
1557        f(&loaded)
1558    }
1559}
1560
1561/// What a stack entry refers to: an input, or a block buffer and the axis
1562/// the value in it stands on.
1563#[derive(Clone, Copy)]
1564enum Slot {
1565    Input(usize),
1566    Block(usize, On),
1567}
1568
1569/// The buffers one thread reuses from block to block.
1570struct Scratch<T> {
1571    cells: Vec<T>,
1572    /// Elements one block buffer holds: a block's result items, and the
1573    /// halo of the wide axis its window steps read around them.
1574    width: usize,
1575    free: Vec<usize>,
1576    stack: Vec<Slot>,
1577    lets: Vec<usize>,
1578    /// One accumulator per running fold in the code, carried from block to
1579    /// block so that the fold is the one the unfused scan performs.
1580    carry: Vec<Option<T>>,
1581}
1582
1583impl<T: Copy + Default> Scratch<T> {
1584    /// Room for one thread's blocks of `w` result items each.
1585    ///
1586    /// A window step reads the run its first window begins in and the run
1587    /// its last item lies in, so a block of `w` items reads fewer than
1588    /// `w + 3k` items of the wide axis, and every buffer is that wide.
1589    fn new(k: &FusedKernel, w: usize) -> Scratch<T> {
1590        let width = w + 3 * k.window.unwrap_or(0);
1591        Scratch {
1592            cells: vec![T::default(); k.slots * width],
1593            width,
1594            free: Vec::with_capacity(k.slots),
1595            stack: Vec::with_capacity(k.slots),
1596            lets: Vec::new(),
1597            carry: vec![None; k.scans],
1598        }
1599    }
1600}
1601
1602/// The leaf loops one working type runs, one block of one instruction at a
1603/// time. All of a kernel's arithmetic goes through these four.
1604struct Steps<M, D, W, S> {
1605    monad: M,
1606    dyad: D,
1607    window: W,
1608    scan: S,
1609}
1610
1611/// Block buffer `d` for writing, plus read-only access to the others.
1612fn split_slots<'s, T>(
1613    scratch: &'s mut [T],
1614    w: usize,
1615    d: usize,
1616) -> (&'s mut [T], impl Fn(usize) -> &'s [T]) {
1617    let (lo, hi) = scratch.split_at_mut(d * w);
1618    let (dst, hi) = hi.split_at_mut(w);
1619    let lo: &[T] = lo;
1620    let hi: &[T] = hi;
1621    (dst, move |i: usize| {
1622        if i < d {
1623            &lo[i * w..(i + 1) * w]
1624        } else {
1625            &hi[(i - d - 1) * w..(i - d) * w]
1626        }
1627    })
1628}
1629
1630/// Run the kernel over one block.
1631///
1632/// `out`, when given, receives the last instruction's result directly and
1633/// the returned index means nothing; otherwise the result stays in the
1634/// block buffer that index names. None means a step left the working type
1635/// and the caller must fall back.
1636fn exec_block<T, M, D, W, S>(
1637    k: &FusedKernel,
1638    srcs: &[Loaded<'_, T>],
1639    at: &Extent,
1640    sc: &mut Scratch<T>,
1641    out: Option<&mut [T]>,
1642    steps: &Steps<M, D, W, S>,
1643) -> Option<usize>
1644where
1645    T: Copy,
1646    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1647    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1648    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool,
1649    S: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T>,
1650{
1651    let Scratch { cells, width, free, stack, lets, carry } = sc;
1652    let w = *width;
1653    stack.clear();
1654    free.clear();
1655    lets.clear();
1656    let nslots = cells.len() / w;
1657    free.extend((0..nslots).rev());
1658    let place = |s: &Slot| match s {
1659        Slot::Input(j) => srcs[*j].on,
1660        Slot::Block(_, o) => *o,
1661    };
1662    let last = k.code.len() - 1;
1663    let head = if out.is_some() { last } else { k.code.len() };
1664    let mut scanned = 0usize;
1665    for ins in &k.code[..head] {
1666        match ins {
1667            Instr::Load(j) => stack.push(Slot::Input(*j)),
1668            Instr::Monad(op) => {
1669                let a = stack.pop()?;
1670                let dom = place(&a);
1671                let len = at.len_on(dom);
1672                let d = free.pop()?;
1673                let (dst, get) = split_slots(cells, w, d);
1674                let av = match a {
1675                    Slot::Input(j) => srcs[j].block(at, dom),
1676                    Slot::Block(i, _) => &get(i)[..len],
1677                };
1678                if !(steps.monad)(*op, av, &mut dst[..len]) {
1679                    return None;
1680                }
1681                release(free, lets, a);
1682                stack.push(Slot::Block(d, dom));
1683            }
1684            Instr::Scan(op) => {
1685                let a = stack.pop()?;
1686                let dom = place(&a);
1687                let len = at.len_on(dom);
1688                let d = free.pop()?;
1689                let (dst, get) = split_slots(cells, w, d);
1690                let av = match a {
1691                    Slot::Input(j) => srcs[j].block(at, dom),
1692                    Slot::Block(i, _) => &get(i)[..len],
1693                };
1694                carry[scanned] = Some((steps.scan)(*op, av, carry[scanned], &mut dst[..len])?);
1695                scanned += 1;
1696                release(free, lets, a);
1697                stack.push(Slot::Block(d, dom));
1698            }
1699            Instr::Window(op, size) => {
1700                let a = stack.pop()?;
1701                let d = free.pop()?;
1702                let (dst, get) = split_slots(cells, w, d);
1703                let av = match a {
1704                    Slot::Input(j) => srcs[j].block(at, On::Wide),
1705                    Slot::Block(i, _) => &get(i)[..at.wide_len],
1706                };
1707                let first = at.start - at.wide_start;
1708                if !(steps.window)(*op, *size, av, first, &mut dst[..at.len]) {
1709                    return None;
1710                }
1711                release(free, lets, a);
1712                stack.push(Slot::Block(d, On::Result));
1713            }
1714            Instr::Dyad(op) => {
1715                let b = stack.pop()?;
1716                let a = stack.pop()?;
1717                let dom = combine(place(&a), place(&b));
1718                let len = at.len_on(dom);
1719                let d = free.pop()?;
1720                let (dst, get) = split_slots(cells, w, d);
1721                let av = match a {
1722                    Slot::Input(j) => srcs[j].block(at, dom),
1723                    Slot::Block(i, _) => &get(i)[..len],
1724                };
1725                let bv = match b {
1726                    Slot::Input(j) => srcs[j].block(at, dom),
1727                    Slot::Block(i, _) => &get(i)[..len],
1728                };
1729                if !(steps.dyad)(*op, av, bv, &mut dst[..len]) {
1730                    return None;
1731                }
1732                for s in [a, b] {
1733                    release(free, lets, s);
1734                }
1735                stack.push(Slot::Block(d, dom));
1736            }
1737            Instr::Store(j) => {
1738                let Slot::Block(i, _) = stack.pop()? else { return None };
1739                if lets.len() != *j {
1740                    return None;
1741                }
1742                lets.push(i);
1743            }
1744            // A let stands where the pass computed it, which is the one
1745            // axis every repeat it stands for was written on.
1746            Instr::Let(j) => {
1747                stack.push(Slot::Block(*lets.get(*j)?, placed(Some(*k.let_doms.get(*j)?))))
1748            }
1749        }
1750    }
1751    let Some(dst) = out else {
1752        return match stack.pop()? {
1753            Slot::Block(i, _) => Some(i),
1754            // Every kernel ends in an operation, so the result is a buffer.
1755            Slot::Input(_) => None,
1756        };
1757    };
1758    // The last instruction writes the caller's buffer instead of a block.
1759    // The chain's root stands on the result's own axis, whatever its
1760    // operands stand on.
1761    let dst = &mut dst[..at.len];
1762    let view = |s: Slot, dom: On| match s {
1763        Slot::Input(j) => srcs[j].block(at, dom),
1764        Slot::Block(i, o) => &cells[i * w..i * w + at.len_on(o)],
1765    };
1766    let ok = match k.code[last] {
1767        Instr::Monad(op) => {
1768            let a = stack.pop()?;
1769            (steps.monad)(op, view(a, On::Result), dst)
1770        }
1771        Instr::Scan(op) => {
1772            let a = stack.pop()?;
1773            match (steps.scan)(op, view(a, On::Result), carry[scanned], dst) {
1774                Some(c) => {
1775                    carry[scanned] = Some(c);
1776                    true
1777                }
1778                None => false,
1779            }
1780        }
1781        Instr::Window(op, size) => {
1782            let a = stack.pop()?;
1783            let first = at.start - at.wide_start;
1784            (steps.window)(op, size, view(a, On::Wide), first, dst)
1785        }
1786        Instr::Dyad(op) => {
1787            let b = stack.pop()?;
1788            let a = stack.pop()?;
1789            let dom = combine(place(&a), place(&b));
1790            (steps.dyad)(op, view(a, dom), view(b, dom), dst)
1791        }
1792        // A kernel ends in the operation that makes its result.
1793        Instr::Load(_) | Instr::Store(_) | Instr::Let(_) => return None,
1794    };
1795    ok.then_some(usize::MAX)
1796}
1797
1798/// Give a block buffer back, unless a let is holding it for the rest of
1799/// the block.
1800fn release(free: &mut Vec<usize>, lets: &[usize], s: Slot) {
1801    if let Slot::Block(i, _) = s
1802        && !lets.contains(&i)
1803    {
1804        free.push(i);
1805    }
1806}
1807
1808/// The whole mapped result, one block at a time. None on integer overflow.
1809fn map_pass<T, M, D, W, S>(
1810    k: &FusedKernel,
1811    srcs: &Sources<'_, T>,
1812    n: usize,
1813    wide: usize,
1814    steps: &Steps<M, D, W, S>,
1815) -> Option<Vec<T>>
1816where
1817    T: FromNarrow + Default + Send + Sync,
1818    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1819    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1820    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool + Sync + Send,
1821    S: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T> + Sync + Send,
1822{
1823    let run = |start: usize, part: &mut [T]| {
1824        let w = BLOCK.min(part.len()).max(1);
1825        let mut sc = Scratch::new(k, w);
1826        let width = sc.width;
1827        let mut stage = vec![T::default(); srcs.staged * width];
1828        for (b, chunk) in part.chunks_mut(w).enumerate() {
1829            let at = Extent::of(start + b * w, chunk.len(), k.window, wide);
1830            let done = srcs.with_block(&at, &mut stage, width, |loaded| {
1831                exec_block(k, loaded, &at, &mut sc, Some(chunk), steps).is_some()
1832            });
1833            if !done {
1834                return false;
1835            }
1836        }
1837        true
1838    };
1839    if k.scans > 0 {
1840        // A running fold hands its accumulator to the next block, so the
1841        // blocks run in one order on one thread. That is the order the
1842        // unfused scan runs them in, and it rounds where that rounds.
1843        let mut out = vec![T::default(); n];
1844        return run(0, &mut out).then_some(out);
1845    }
1846    let (out, ok) = par::fill(n, run);
1847    ok.then_some(out)
1848}
1849
1850/// Independent accumulators the fold over a block keeps in flight, and the
1851/// block length below which one accumulator is cheaper. The reasoning is
1852/// the one `verb::FOLD_LANES` carries: a single accumulator makes the fold
1853/// a chain of dependent steps, and only an associative step is ever
1854/// absorbed here, so the lanes are a regrouping the float contract already
1855/// allows (§5.9).
1856const FOLD_LANES: usize = 8;
1857const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
1858
1859/// Fold one block of mapped values right to left, in lanes. None when a
1860/// step left the element type.
1861#[inline(always)]
1862fn fold_block_body<T, S>(v: &[T], step: &S) -> Option<T>
1863where
1864    T: Copy,
1865    S: Fn(T, T) -> Option<T>,
1866{
1867    let n = v.len();
1868    if n < MIN_LANE_WORK {
1869        let mut acc = v[n - 1];
1870        for &x in v[..n - 1].iter().rev() {
1871            acc = step(x, acc)?;
1872        }
1873        return Some(acc);
1874    }
1875    let rows = n / FOLD_LANES;
1876    let head = n - rows * FOLD_LANES;
1877    let last = head + (rows - 1) * FOLD_LANES;
1878    let mut acc = [v[last]; FOLD_LANES];
1879    acc.copy_from_slice(&v[last..last + FOLD_LANES]);
1880    for r in (0..rows - 1).rev() {
1881        let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
1882        for (slot, &x) in acc.iter_mut().zip(row) {
1883            *slot = step(x, *slot)?;
1884        }
1885    }
1886    let mut a = acc[FOLD_LANES - 1];
1887    for &x in acc[..FOLD_LANES - 1].iter().rev() {
1888        a = step(x, a)?;
1889    }
1890    for &x in v[..head].iter().rev() {
1891        a = step(x, a)?;
1892    }
1893    Some(a)
1894}
1895
1896multiversioned! {
1897    /// One block's values folded into one, at the CPU's own width.
1898    fn fold_block[T: Copy, S: Fn(T, T) -> Option<T>](
1899        v: &[T],
1900        step: &S,
1901    ) -> Option<T> = fold_block_body;
1902}
1903
1904/// Fold the mapped values of `lo .. hi` right to left, block by block.
1905fn fold_range<T, M, D, W, C, S>(
1906    k: &FusedKernel,
1907    srcs: &Sources<'_, T>,
1908    lo: usize,
1909    hi: usize,
1910    wide: usize,
1911    steps: &Steps<M, D, W, C>,
1912    step: &S,
1913) -> Option<T>
1914where
1915    T: FromNarrow + Default,
1916    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1917    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1918    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool,
1919    C: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T>,
1920    S: Fn(T, T) -> Option<T>,
1921{
1922    let w = BLOCK.min(hi - lo).max(1);
1923    let mut sc = Scratch::new(k, w);
1924    let width = sc.width;
1925    let mut stage = vec![T::default(); srcs.staged * width];
1926    let mut acc: Option<T> = None;
1927    // Blocks run backwards and the accumulator carries across them, so the
1928    // fold is the insert's own right-to-left order over the whole range.
1929    // Nothing a block computes depends on the block before it: a running
1930    // fold, which would, is never absorbed under a reduction.
1931    for b in (0..(hi - lo).div_ceil(w)).rev() {
1932        let start = lo + b * w;
1933        let len = (hi - start).min(w);
1934        let at = Extent::of(start, len, k.window, wide);
1935        let slot = srcs
1936            .with_block(&at, &mut stage, width, |loaded| {
1937                exec_block(k, loaded, &at, &mut sc, None, steps)
1938            })?;
1939        let block = fold_block(&sc.cells[slot * sc.width..slot * sc.width + len], step)?;
1940        acc = Some(match acc {
1941            None => block,
1942            Some(a) => step(block, a)?,
1943        });
1944    }
1945    acc
1946}
1947
1948/// The mapped values folded into one. None on integer overflow.
1949fn reduce_pass<T, M, D, W, C, S>(
1950    k: &FusedKernel,
1951    srcs: &Sources<'_, T>,
1952    n: usize,
1953    wide: usize,
1954    steps: &Steps<M, D, W, C>,
1955    step: S,
1956) -> Option<T>
1957where
1958    T: FromNarrow + Default + Send + Sync,
1959    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1960    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1961    W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool + Sync + Send,
1962    C: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T> + Sync + Send,
1963    S: Fn(T, T) -> Option<T> + Sync + Send,
1964{
1965    let chunks = par::chunks(n, n * k.code.len());
1966    if chunks < 2 {
1967        return fold_range(k, srcs, 0, n, wide, steps, &step);
1968    }
1969    let per = n.div_ceil(chunks);
1970    let parts = par::map_indexed(n.div_ceil(per), |c| {
1971        fold_range(k, srcs, c * per, ((c + 1) * per).min(n), wide, steps, &step)
1972    });
1973    // The chunks combine right to left, the order they were folded in. That
1974    // regroups an associative float fold, which is the §5.9 contract; only
1975    // associative operations are absorbed.
1976    let mut it = parts.into_iter().rev();
1977    let mut acc = it.next()??;
1978    for part in it {
1979        acc = step(part?, acc)?;
1980    }
1981    Some(acc)
1982}
1983
1984// ------------------------------------------------------------ the kernels
1985//
1986// Each pass picks its operation before the loop and then runs one plain
1987// loop over slices, which is the shape the compiler vectorises. Nothing in
1988// here is hand-written SIMD, and nothing may become it.
1989//
1990// These four are the whole arithmetic of a kernel, so they are also where
1991// the CPU feature levels are chosen: each is compiled once per level (see
1992// `simd`) and the call dispatches on what the machine runs. One block of
1993// one instruction is thousands of elements, so the dispatch costs nothing
1994// measurable.
1995
1996macro_rules! each {
1997    ($a:expr, $dst:expr, $f:expr) => {{
1998        let f = $f;
1999        for (slot, &x) in $dst.iter_mut().zip($a) {
2000            *slot = f(x);
2001        }
2002        return true;
2003    }};
2004}
2005
2006macro_rules! zip {
2007    ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
2008        let f = $f;
2009        for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
2010            *slot = f(x, y);
2011        }
2012        return true;
2013    }};
2014}
2015
2016#[inline(always)]
2017fn monad_f64_body(op: ScalarMonad, a: &[f64], dst: &mut [f64], tol: Tol) -> bool {
2018    use ScalarMonad::*;
2019    match op {
2020        Conj => each!(a, dst, |x: f64| x),
2021        Neg => each!(a, dst, |x: f64| -x),
2022        Abs => each!(a, dst, f64::abs),
2023        // A magnitude the dialect's tolerance reads as zero has no sign,
2024        // exactly as unfused.
2025        Signum => each!(a, dst, |x: f64| if tol.is_zero(x) {
2026            0.0
2027        } else if x > 0.0 {
2028            1.0
2029        } else if x < 0.0 {
2030            -1.0
2031        } else {
2032            0.0
2033        }),
2034        // `% 0` is infinity, the J rule the unfused monad follows.
2035        Recip => each!(a, dst, |x: f64| if x == 0.0 { f64::INFINITY } else { 1.0 / x }),
2036        // Reached only through an integer chain, where they are the
2037        // identity: rounding a float narrows its dtype, which is declined.
2038        Floor => each!(a, dst, f64::floor),
2039        Ceil => each!(a, dst, f64::ceil),
2040        Inc => each!(a, dst, |x: f64| x + 1.0),
2041        Dec => each!(a, dst, |x: f64| x - 1.0),
2042        Double => each!(a, dst, |x: f64| x + x),
2043        Halve => each!(a, dst, |x: f64| x / 2.0),
2044        Square => each!(a, dst, |x: f64| x * x),
2045        OneMinus => each!(a, dst, |x: f64| 1.0 - x),
2046        Exp => each!(a, dst, f64::exp),
2047        _ => false,
2048    }
2049}
2050
2051#[inline(always)]
2052fn dyad_f64_body(op: ScalarDyad, a: &[f64], b: &[f64], dst: &mut [f64], tol: Tol) -> bool {
2053    use ScalarDyad::*;
2054    match op {
2055        Add => zip!(a, b, dst, |x: f64, y: f64| x + y),
2056        Sub => zip!(a, b, dst, |x: f64, y: f64| x - y),
2057        Mul => zip!(a, b, dst, |x: f64, y: f64| x * y),
2058        Min => zip!(a, b, dst, f64::min),
2059        Max => zip!(a, b, dst, f64::max),
2060        DivJ => zip!(a, b, dst, |x: f64, y: f64| if y == 0.0 {
2061            if x == 0.0 { 0.0 } else { f64::INFINITY.copysign(x) }
2062        } else {
2063            x / y
2064        }),
2065        // An infinite modulus leaves a value of its own sign alone and
2066        // sends the other one to that infinity, exactly as unfused.
2067        Residue => zip!(a, b, dst, |x: f64, y: f64| if x.is_infinite() {
2068            if y == 0.0 || (y > 0.0) == (x > 0.0) { y } else { x }
2069        } else if x == 0.0 {
2070            y
2071        } else {
2072            y - x * (y / x).floor()
2073        }),
2074        // A comparison is a number here, as it is in J: the boolean only
2075        // shows in the dtype of a result, which the caller narrows. Floats
2076        // compare with the dialect's tolerance, as they do unfused.
2077        Eq | Ne | Lt | Le | Gt | Ge => {
2078            zip!(a, b, dst, |x: f64, y: f64| tol_cmp(op, x, y, tol) as u8 as f64)
2079        }
2080        _ => false,
2081    }
2082}
2083
2084/// Integer passes fold overflow into a flag instead of branching out of the
2085/// loop: the whole evaluation is thrown away and redone unfused either way.
2086macro_rules! each_over {
2087    ($a:expr, $dst:expr, $f:expr) => {{
2088        let f = $f;
2089        let mut over = false;
2090        for (slot, &x) in $dst.iter_mut().zip($a) {
2091            let (v, o) = f(x);
2092            *slot = v;
2093            over |= o;
2094        }
2095        return !over;
2096    }};
2097}
2098
2099macro_rules! zip_over {
2100    ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
2101        let f = $f;
2102        let mut over = false;
2103        for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
2104            let (v, o) = f(x, y);
2105            *slot = v;
2106            over |= o;
2107        }
2108        return !over;
2109    }};
2110}
2111
2112#[inline(always)]
2113fn monad_i64_body(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool {
2114    use ScalarMonad::*;
2115    match op {
2116        Conj | Floor | Ceil => each!(a, dst, |x: i64| x),
2117        Neg => each_over!(a, dst, i64::overflowing_neg),
2118        Abs => each_over!(a, dst, i64::overflowing_abs),
2119        Signum => each!(a, dst, i64::signum),
2120        Inc => each_over!(a, dst, |x: i64| x.overflowing_add(1)),
2121        Dec => each_over!(a, dst, |x: i64| x.overflowing_sub(1)),
2122        Double => each_over!(a, dst, |x: i64| x.overflowing_add(x)),
2123        Square => each_over!(a, dst, |x: i64| x.overflowing_mul(x)),
2124        OneMinus => each_over!(a, dst, |x: i64| 1i64.overflowing_sub(x)),
2125        _ => false,
2126    }
2127}
2128
2129#[inline(always)]
2130fn dyad_i64_body(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool {
2131    use ScalarDyad::*;
2132    match op {
2133        Add => zip_over!(a, b, dst, i64::overflowing_add),
2134        Sub => zip_over!(a, b, dst, i64::overflowing_sub),
2135        Mul => zip_over!(a, b, dst, i64::overflowing_mul),
2136        Min => zip!(a, b, dst, i64::min),
2137        Max => zip!(a, b, dst, i64::max),
2138        Residue => zip!(a, b, dst, |x: i64, y: i64| if x == 0 {
2139            y
2140        } else {
2141            // wrapping_rem: i64::MIN % -1 is mathematically 0.
2142            let mut r = y.wrapping_rem(x);
2143            if r != 0 && (r < 0) != (x < 0) {
2144                r += x;
2145            }
2146            r
2147        }),
2148        Eq => zip!(a, b, dst, |x: i64, y: i64| (x == y) as i64),
2149        Ne => zip!(a, b, dst, |x: i64, y: i64| (x != y) as i64),
2150        Lt => zip!(a, b, dst, |x: i64, y: i64| (x < y) as i64),
2151        Le => zip!(a, b, dst, |x: i64, y: i64| (x <= y) as i64),
2152        Gt => zip!(a, b, dst, |x: i64, y: i64| (x > y) as i64),
2153        Ge => zip!(a, b, dst, |x: i64, y: i64| (x >= y) as i64),
2154        _ => false,
2155    }
2156}
2157
2158multiversioned! {
2159    /// One instruction of a kernel over one block of floats: the monadic
2160    /// operations. False is unreachable — every operation a kernel holds is
2161    /// covered — and exists so the two passes have one signature.
2162    fn monad_f64(
2163        op: ScalarMonad,
2164        a: &[f64],
2165        dst: &mut [f64],
2166        tol: Tol,
2167    ) -> bool = monad_f64_body;
2168}
2169
2170multiversioned! {
2171    /// One instruction of a kernel over one block of floats: the dyadic
2172    /// operations.
2173    fn dyad_f64(
2174        op: ScalarDyad,
2175        a: &[f64],
2176        b: &[f64],
2177        dst: &mut [f64],
2178        tol: Tol,
2179    ) -> bool = dyad_f64_body;
2180}
2181
2182multiversioned! {
2183    /// One instruction of a kernel over one block of integers: the monadic
2184    /// operations. False means the block left i64.
2185    fn monad_i64(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool = monad_i64_body;
2186}
2187
2188multiversioned! {
2189    /// One instruction of a kernel over one block of integers: the dyadic
2190    /// operations. False means the block left i64.
2191    fn dyad_i64(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool = dyad_i64_body;
2192}
2193
2194/// One block's running fold, continued from the accumulator the block
2195/// before it left. None when a step left the element type.
2196///
2197/// The accumulator runs the length of the argument, one step per item, in
2198/// the order the unfused scan takes them: what the fused kernel saves is
2199/// the traffic around the scan, not the scan.
2200#[inline(always)]
2201fn scan_block_body<T, F>(v: &[T], carry: Option<T>, dst: &mut [T], step: &F) -> Option<T>
2202where
2203    T: Copy,
2204    F: Fn(T, T) -> (T, bool),
2205{
2206    let mut over = false;
2207    let (mut acc, from) = match carry {
2208        Some(a) => (a, 0),
2209        None => {
2210            // The first item of a scan is the item itself.
2211            dst[0] = v[0];
2212            (v[0], 1)
2213        }
2214    };
2215    for (slot, &x) in dst.iter_mut().zip(v).skip(from) {
2216        let (r, o) = step(acc, x);
2217        acc = r;
2218        over |= o;
2219        *slot = acc;
2220    }
2221    (!over).then_some(acc)
2222}
2223
2224multiversioned! {
2225    /// One block of a running fold. The steps depend on one another, so
2226    /// what a wider vector reaches here is the loop around them.
2227    fn scan_block[T: Copy, F: Fn(T, T) -> (T, bool)](
2228        v: &[T],
2229        carry: Option<T>,
2230        dst: &mut [T],
2231        step: &F,
2232    ) -> Option<T> = scan_block_body;
2233}
2234
2235/// The windows of a block of floats, folded one per result item. The step
2236/// is chosen before the fold so that the fold itself is one plain loop.
2237fn window_pass_f64(op: ScalarDyad, k: usize, v: &[f64], first: usize, dst: &mut [f64]) -> bool {
2238    use ScalarDyad::*;
2239    match op {
2240        Add => windows_into(v, k, first, dst, &|a: f64, b: f64| (a + b, false)),
2241        Mul => windows_into(v, k, first, dst, &|a: f64, b: f64| (a * b, false)),
2242        Min => windows_into(v, k, first, dst, &|a: f64, b: f64| (a.min(b), false)),
2243        Max => windows_into(v, k, first, dst, &|a: f64, b: f64| (a.max(b), false)),
2244        _ => false,
2245    }
2246}
2247
2248fn window_pass_i64(op: ScalarDyad, k: usize, v: &[i64], first: usize, dst: &mut [i64]) -> bool {
2249    use ScalarDyad::*;
2250    match op {
2251        Add => windows_into(v, k, first, dst, &i64::overflowing_add),
2252        Mul => windows_into(v, k, first, dst, &i64::overflowing_mul),
2253        Min => windows_into(v, k, first, dst, &|a: i64, b: i64| (a.min(b), false)),
2254        Max => windows_into(v, k, first, dst, &|a: i64, b: i64| (a.max(b), false)),
2255        _ => false,
2256    }
2257}
2258
2259fn scan_pass_f64(op: ScalarDyad, v: &[f64], carry: Option<f64>, dst: &mut [f64]) -> Option<f64> {
2260    use ScalarDyad::*;
2261    match op {
2262        Add => scan_block(v, carry, dst, &|a: f64, b: f64| (a + b, false)),
2263        Mul => scan_block(v, carry, dst, &|a: f64, b: f64| (a * b, false)),
2264        Min => scan_block(v, carry, dst, &|a: f64, b: f64| (a.min(b), false)),
2265        Max => scan_block(v, carry, dst, &|a: f64, b: f64| (a.max(b), false)),
2266        _ => None,
2267    }
2268}
2269
2270fn scan_pass_i64(op: ScalarDyad, v: &[i64], carry: Option<i64>, dst: &mut [i64]) -> Option<i64> {
2271    use ScalarDyad::*;
2272    match op {
2273        Add => scan_block(v, carry, dst, &i64::overflowing_add),
2274        Mul => scan_block(v, carry, dst, &i64::overflowing_mul),
2275        Min => scan_block(v, carry, dst, &|a: i64, b: i64| (a.min(b), false)),
2276        Max => scan_block(v, carry, dst, &|a: i64, b: i64| (a.max(b), false)),
2277        _ => None,
2278    }
2279}
2280
2281/// One fold step of an absorbed reduction. None on integer overflow.
2282fn step_i64(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
2283    use ScalarDyad::*;
2284    match op {
2285        Add => a.checked_add(b),
2286        Mul => a.checked_mul(b),
2287        Min => Some(a.min(b)),
2288        Max => Some(a.max(b)),
2289        _ => None,
2290    }
2291}
2292
2293/// One fold step of an absorbed float reduction, for a backend that mapped
2294/// the values elsewhere and brings its partials back here to combine.
2295pub(crate) fn step(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
2296    step_f64(op, a, b)
2297}
2298
2299fn step_f64(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
2300    use ScalarDyad::*;
2301    match op {
2302        Add => Some(a + b),
2303        Mul => Some(a * b),
2304        Min => Some(a.min(b)),
2305        Max => Some(a.max(b)),
2306        _ => None,
2307    }
2308}
2309
2310// ------------------------------------------------------------- the driver
2311
2312/// A rank-0 argument as one block of the repeated value, which is how it
2313/// reaches every element without an index test. None for an argument with
2314/// items, which is read where it lies.
2315fn splat_f64(a: &Array, w: usize) -> Option<Vec<f64>> {
2316    if a.rank() != 0 {
2317        return None;
2318    }
2319    let v = match &a.data {
2320        Data::Bool(d) => d[0] as f64,
2321        Data::I64(d) => d[0] as f64,
2322        Data::F64(d) => d[0],
2323        Data::Ext(_)
2324        | Data::Rat(_)
2325        | Data::Complex(_)
2326        | Data::Char(_)
2327        | Data::Symbol(_)
2328        | Data::Box(_) => {
2329            return Some(Vec::new());
2330        }
2331    };
2332    Some(vec![v; w])
2333}
2334
2335fn splat_i64(a: &Array, w: usize) -> Option<Vec<i64>> {
2336    if a.rank() != 0 {
2337        return None;
2338    }
2339    let v = match &a.data {
2340        Data::Bool(d) => d[0] as i64,
2341        Data::I64(d) => d[0],
2342        _ => return Some(Vec::new()),
2343    };
2344    Some(vec![v; w])
2345}
2346
2347/// An argument with items, seen by a float kernel: its own buffer when that
2348/// is already f64, the narrow values otherwise.
2349fn narrow_f64(a: &Array) -> Result<&[f64], Narrow<'_>> {
2350    match &a.data {
2351        Data::I64(d) => Err(Narrow::I64(d)),
2352        Data::Bool(d) => Err(Narrow::Bool(d)),
2353        _ => Ok(a.as_f64_slice().unwrap_or(&[])),
2354    }
2355}
2356
2357fn narrow_i64(a: &Array) -> Result<&[i64], Narrow<'_>> {
2358    match &a.data {
2359        Data::Bool(d) => Err(Narrow::Bool(d)),
2360        _ => Ok(a.as_i64_slice().unwrap_or(&[])),
2361    }
2362}
2363
2364/// The input list one run reads: a repeated scalar from `owned`, an
2365/// argument's own buffer, or a narrow buffer to be promoted block by block.
2366fn sources<'a, T>(
2367    inputs: &'a [Array],
2368    owned: &'a [Option<Vec<T>>],
2369    on: &impl Fn(usize) -> On,
2370    narrow: impl Fn(&'a Array) -> Result<&'a [T], Narrow<'a>>,
2371) -> Sources<'a, T> {
2372    let mut staged = 0;
2373    let of = inputs
2374        .iter()
2375        .zip(owned)
2376        .enumerate()
2377        .map(|(j, (a, o))| match o {
2378            Some(v) => Source::Ready(Loaded { data: v, splat: true, on: on(j), base: 0 }),
2379            None => match narrow(a) {
2380                Ok(d) => Source::Ready(Loaded { data: d, splat: false, on: on(j), base: 0 }),
2381                Err(n) => {
2382                    staged += 1;
2383                    Source::Staged(n, on(j))
2384                }
2385            },
2386        })
2387        .collect();
2388    Sources { of, staged }
2389}
2390
2391/// The shape every element of the result has: identical for all non-scalar
2392/// inputs, since anything else needs the agreement machinery.
2393pub(crate) fn common_shape(inputs: &[Array]) -> Option<Option<Vec<usize>>> {
2394    let mut shape: Option<&Vec<usize>> = None;
2395    for a in inputs {
2396        if a.rank() == 0 {
2397            continue;
2398        }
2399        match shape {
2400            None => shape = Some(&a.shape),
2401            Some(s) if *s == a.shape => {}
2402            Some(_) => return None,
2403        }
2404    }
2405    Some(shape.cloned())
2406}
2407
2408/// The axes a kernel's inputs stand on: the shape of its result, and the
2409/// length of the wide axis its window steps read.
2410///
2411/// This is the whole of the alignment rule, and it is decided by shapes
2412/// alone. Where a chain reads an input is settled when the chain is built —
2413/// everything under a window step is wide — so all that is left at run time
2414/// is that the inputs on one axis agree with each other, and that the two
2415/// axes stand `k - 1` items apart. `19 }. y` beside `20 +/\ y` passes
2416/// because it is 19 items shorter; `18 }. y` beside it does not, and the
2417/// chain runs and raises the length error it was going to raise. Nothing is
2418/// shifted or padded here: an input arrives as the items it holds.
2419struct Axes {
2420    shape: Vec<usize>,
2421    /// Items of the wide axis, when there is a window step to read it.
2422    wide: usize,
2423}
2424
2425fn axes(k: &FusedKernel, inputs: &[Array]) -> Option<Axes> {
2426    let Some(window) = k.window else {
2427        // Every input a scalar: no work worth blocking, and a reduction
2428        // would need the leading axis a scalar has not got.
2429        let shape = common_shape(inputs)??;
2430        // A running fold folds items, and a block of this kernel is
2431        // elements: over anything but a vector the two are not the same
2432        // fold, so a higher-rank argument goes the way it went.
2433        if k.scans > 0 && shape.len() != 1 {
2434            return None;
2435        }
2436        return Some(Axes { shape, wide: 0 });
2437    };
2438    let (mut wide, mut result) = (None, None);
2439    for (a, dom) in inputs.iter().zip(&k.doms) {
2440        // A scalar reaches every item of whatever it is combined with, so
2441        // it stands on either axis and constrains neither.
2442        if a.rank() == 0 {
2443            continue;
2444        }
2445        // A window folds the items of a vector. An input the chain reads on
2446        // both axes cannot be two lengths at once.
2447        let (Some(d), 1) = (dom, a.rank()) else { return None };
2448        let seen = if *d == Dom::Wide { &mut wide } else { &mut result };
2449        match seen {
2450            None => *seen = Some(a.shape[0]),
2451            Some(m) if *m == a.shape[0] => {}
2452            Some(_) => return None,
2453        }
2454    }
2455    let wide = wide?;
2456    if wide < window {
2457        // No window fits: the result has no items, which the chain builds
2458        // out of the verb's own answer for an empty argument.
2459        return None;
2460    }
2461    let count = wide - window + 1;
2462    if result.is_some_and(|m| m != count) {
2463        return None;
2464    }
2465    Some(Axes { shape: vec![count], wide })
2466}
2467
2468/// Run a fused node. None means the kernel declined and the caller must
2469/// evaluate the original subtree, which is always allowed to be slower and
2470/// never allowed to differ.
2471pub(crate) fn run(k: &FusedKernel, inputs: &[Array]) -> Option<Array> {
2472    let reducing = matches!(k.yields, Yield::Reduce(_));
2473    let Axes { shape, wide } = axes(k, inputs)?;
2474    let n: usize = shape.iter().product();
2475    if n == 0 {
2476        return None;
2477    }
2478    if reducing && (shape.len() != 1 || n < 2) {
2479        // A one-item reduction yields the item itself, dtype and all, and a
2480        // higher-rank one folds cells rather than elements.
2481        return None;
2482    }
2483    let (working, root) = working_type(k, inputs)?;
2484    if k.yields == Yield::Tally {
2485        // The shapes have already said how many items the chain produces,
2486        // and the type rules have said it would reach them without an
2487        // error. There is nothing else a tally wants from the values.
2488        return Some(Array::scalar_i64(shape[0] as i64));
2489    }
2490    // A repeated scalar is one block long, and a block reads its window
2491    // halo as well as its own items.
2492    let w = BLOCK.min(n).max(1) + 3 * k.window.unwrap_or(0);
2493    // The kernel's comparisons carry the tolerance the program was compiled
2494    // with, so a fused comparison answers as the unfused one does.
2495    let tol = k.tol;
2496    let on = |j: usize| placed(k.doms[j]);
2497
2498    let data = if working == DType::F64 {
2499        let steps = Steps {
2500            monad: move |op, a: &[f64], dst: &mut [f64]| monad_f64(op, a, dst, tol),
2501            dyad: move |op, a: &[f64], b: &[f64], dst: &mut [f64]| dyad_f64(op, a, b, dst, tol),
2502            window: window_pass_f64,
2503            scan: scan_pass_f64,
2504        };
2505        let owned: Vec<Option<Vec<f64>>> = inputs.iter().map(|a| splat_f64(a, w)).collect();
2506        let srcs = sources(inputs, &owned, &on, narrow_f64);
2507        match k.reduce() {
2508            None => {
2509                let out = map_pass(k, &srcs, n, wide, &steps)?;
2510                float_result(out, root)
2511            }
2512            Some(op) => {
2513                let v = reduce_pass(k, &srcs, n, wide, &steps, |a, b| step_f64(op, a, b))?;
2514                // A comparison at the root maps to exact 0 and 1, which the
2515                // fold keeps exact; the reduction of booleans is integer.
2516                match root {
2517                    DType::F64 => Data::F64(vec![v].into()),
2518                    _ => Data::I64(vec![v as i64].into()),
2519                }
2520            }
2521        }
2522    } else {
2523        let steps = Steps {
2524            monad: monad_i64,
2525            dyad: dyad_i64,
2526            window: window_pass_i64,
2527            scan: scan_pass_i64,
2528        };
2529        let owned: Vec<Option<Vec<i64>>> = inputs.iter().map(|a| splat_i64(a, w)).collect();
2530        let srcs = sources(inputs, &owned, &on, narrow_i64);
2531        match k.reduce() {
2532            None => {
2533                let out = map_pass(k, &srcs, n, wide, &steps)?;
2534                int_result(out, root)
2535            }
2536            Some(op) => {
2537                let v = reduce_pass(k, &srcs, n, wide, &steps, |a, b| step_i64(op, a, b))?;
2538                Data::I64(vec![v].into())
2539            }
2540        }
2541    };
2542    Some(Array::new(if reducing { Vec::new() } else { shape }, data))
2543}
2544
2545/// The mapped block values as the array the unfused chain would build. A
2546/// comparison at the root costs one narrowing pass, since the kernel
2547/// computes 0 and 1 in its working type and a boolean array holds bytes.
2548fn float_result(out: Vec<f64>, root: DType) -> Data {
2549    match root {
2550        DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0.0) as u8).into()),
2551        _ => Data::F64(out.into()),
2552    }
2553}
2554
2555fn int_result(out: Vec<i64>, root: DType) -> Data {
2556    match root {
2557        DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0) as u8).into()),
2558        _ => Data::I64(out.into()),
2559    }
2560}
2561
2562// -------------------------------------------------------- describing one
2563//
2564// Read-only descriptions of a compiled kernel, for `Program::explain`.
2565// Nothing here runs a kernel or changes one; the summary is derived from
2566// the code the pass emitted, and the decline reason re-checks the same
2567// preconditions `run` checks before it starts.
2568
2569/// Why a kernel handed its work back to the chain it came from.
2570#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2571pub enum Decline {
2572    /// Inputs disagree on shape, or every input is a scalar: broadcasting
2573    /// and agreement are the chain's business.
2574    Agreement,
2575    /// Nothing to compute.
2576    Empty,
2577    /// An absorbed reduction wants one axis with at least two items.
2578    ReduceShape,
2579    /// One working type cannot hold every step exactly — a chain that
2580    /// computes integers along a float path, or non-numeric data.
2581    WorkingType,
2582    /// The preconditions held, so a step went out of range mid-block:
2583    /// integer overflow, which the chain redoes in a wider type.
2584    Overflow,
2585    /// A window step wants one vector axis longer than the window, and
2586    /// every other input aligned on the window's last item.
2587    Window,
2588}
2589
2590impl Decline {
2591    pub fn reason(self) -> &'static str {
2592        match self {
2593            Decline::Agreement => "the inputs need agreement or are all scalars",
2594            Decline::Empty => "there is nothing to compute",
2595            Decline::ReduceShape => "the reduction needs one axis of two or more items",
2596            Decline::WorkingType => "no single working type holds every step exactly",
2597            Decline::Overflow => "an integer step left 64-bit range",
2598            Decline::Window => "the window does not fit the axis, or the inputs are not aligned with it",
2599        }
2600    }
2601}
2602
2603/// Why this kernel would decline these inputs, or None if it would run.
2604///
2605/// A read-only mirror of the preconditions at the top of `run`: it looks
2606/// at shapes and dtypes only, never at values, so the one thing it cannot
2607/// see in advance is an overflow — which is what is left when every
2608/// precondition holds.
2609pub fn decline_reason(k: &FusedKernel, inputs: &[Array]) -> Option<Decline> {
2610    let Some(Axes { shape, .. }) = axes(k, inputs) else {
2611        return Some(if k.window.is_some() { Decline::Window } else { Decline::Agreement });
2612    };
2613    let n: usize = shape.iter().product();
2614    if n == 0 {
2615        return Some(Decline::Empty);
2616    }
2617    if matches!(k.yields, Yield::Reduce(_)) && (shape.len() != 1 || n < 2) {
2618        return Some(Decline::ReduceShape);
2619    }
2620    if working_type(k, inputs).is_none() {
2621        return Some(Decline::WorkingType);
2622    }
2623    Some(Decline::Overflow)
2624}
2625
2626/// What a compiled kernel is made of.
2627#[derive(Clone, Debug, PartialEq, Eq)]
2628pub struct Summary {
2629    /// Arithmetic steps: the monads and dyads, not the loads and stores.
2630    pub ops: usize,
2631    /// Those steps in the order the kernel performs them.
2632    pub op_names: Vec<String>,
2633    /// The reduction folded into the same pass, if there is one.
2634    pub reduce: Option<&'static str>,
2635    /// True when the whole chain collapsed to a count of its own items.
2636    pub tally: bool,
2637    /// Values the kernel keeps for a second read within one block.
2638    pub lets: usize,
2639    /// Subtrees the chain reads.
2640    pub inputs: usize,
2641    /// Elements one block buffer holds.
2642    pub block: usize,
2643    /// The window every window step folds, when the kernel has one.
2644    pub window: Option<usize>,
2645    /// Running folds the kernel carries from block to block.
2646    pub scans: usize,
2647}
2648
2649impl std::fmt::Display for Summary {
2650    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2651        write!(f, "{} op{}", self.ops, if self.ops == 1 { "" } else { "s" })?;
2652        if !self.op_names.is_empty() {
2653            write!(f, ": {}", self.op_names.join(" "))?;
2654        }
2655        if let Some(r) = self.reduce {
2656            write!(f, "; {r}/ absorbed")?;
2657        }
2658        if self.tally {
2659            write!(f, "; tally only")?;
2660        }
2661        if self.lets > 0 {
2662            write!(f, "; {} let slot{}", self.lets, if self.lets == 1 { "" } else { "s" })?;
2663        }
2664        if let Some(k) = self.window {
2665            write!(f, "; window {k}")?;
2666        }
2667        if self.scans > 0 {
2668            write!(f, "; {} running fold{}", self.scans, if self.scans == 1 { "" } else { "s" })?;
2669        }
2670        write!(f, "; block {}", self.block)
2671    }
2672}
2673
2674/// Describe a compiled kernel: what it computes, and with what.
2675pub fn summary(k: &FusedKernel) -> Summary {
2676    let mut op_names: Vec<String> = Vec::new();
2677    let mut lets = 0usize;
2678    for ins in &k.code {
2679        match ins {
2680            Instr::Monad(op) => op_names.push(monad_name(*op).to_string()),
2681            Instr::Dyad(op) => op_names.push(dyad_name(*op).to_string()),
2682            Instr::Window(op, k) => op_names.push(format!("{k} {}/\\", dyad_name(*op))),
2683            Instr::Scan(op) => op_names.push(format!("{}/\\", dyad_name(*op))),
2684            Instr::Store(_) => lets += 1,
2685            Instr::Load(_) | Instr::Let(_) => {}
2686        }
2687    }
2688    Summary {
2689        ops: op_names.len(),
2690        op_names,
2691        reduce: k.reduce().map(dyad_name),
2692        tally: k.yields == Yield::Tally,
2693        lets,
2694        inputs: k.leaves.iter().copied().max().map_or(0, |m| m + 1),
2695        block: BLOCK,
2696        window: k.window,
2697        scans: k.scans,
2698    }
2699}
2700
2701/// The names the pass took out of the program: values it moved into the
2702/// kernels that read them, so no sentence computes them as arrays any more.
2703pub fn inlined_names(p: &Program) -> Vec<String> {
2704    let Some(Expr::Elided { orig, .. }) = p.stmts.first() else { return Vec::new() };
2705    let assigned = |stmts: &[Expr]| -> Vec<String> {
2706        stmts
2707            .iter()
2708            .filter_map(|s| match s {
2709                Expr::Assign { name, .. } => Some(name.clone()),
2710                _ => None,
2711            })
2712            .collect()
2713    };
2714    let kept = assigned(&p.stmts);
2715    assigned(orig).into_iter().filter(|n| !kept.contains(n)).collect()
2716}
2717
2718/// J spellings for the elementwise operations a kernel can hold. Only the
2719/// naming lives here; the meanings are [`crate::verb`]'s.
2720fn monad_name(op: ScalarMonad) -> &'static str {
2721    use ScalarMonad::*;
2722    match op {
2723        Conj => "+",
2724        Neg => "-",
2725        Signum => "*",
2726        Recip => "%",
2727        Sqrt => "%:",
2728        Exp => "^",
2729        Abs => "|",
2730        Floor => "<.",
2731        Ceil => ">.",
2732        Not => "-.",
2733        OneMinus => "-.",
2734        Inc => ">:",
2735        Dec => "<:",
2736        Double => "+:",
2737        Halve => "-:",
2738        Square => "*:",
2739        Ln => "^.",
2740        Pi => "o.",
2741        Factorial => "!",
2742        Imaginary => "j.",
2743        Polar => "r.",
2744    }
2745}
2746
2747fn dyad_name(op: ScalarDyad) -> &'static str {
2748    use ScalarDyad::*;
2749    match op {
2750        Add => "+",
2751        Sub => "-",
2752        Mul => "*",
2753        DivJ | DivApl => "%",
2754        Min => "<.",
2755        Max => ">.",
2756        Pow => "^",
2757        Residue => "|",
2758        Eq => "=",
2759        Ne => "~:",
2760        Lt => "<",
2761        Le => "<:",
2762        Gt => ">",
2763        Ge => ">:",
2764        Lcm => "*.",
2765        Gcd => "+.",
2766        Log => "^.",
2767        Root => "%:",
2768        Circle => "o.",
2769        Binomial => "!",
2770        MakeComplex => "j.",
2771        PolarBy => "r.",
2772    }
2773}
2774
2775/// Evaluate a fused node from its already-evaluated inputs, or report that
2776/// the original subtree must run instead.
2777///
2778/// This is the one place a device gets to run libjay's arithmetic. With a
2779/// device attached the kernel is offered to it first; everything it will not
2780/// take comes back here with a reason, and the CPU path runs exactly as it
2781/// runs with no device in sight. The device therefore cannot change a
2782/// result's shape, dtype or error — only where the arithmetic happened.
2783pub(crate) fn eval_on(
2784    device: Option<&crate::device::Device>,
2785    k: &FusedKernel,
2786    inputs: &[Array],
2787) -> (Option<Array>, crate::device::Placement) {
2788    use crate::device::Placement;
2789    let mut placement = Placement::Default;
2790    // A block kernel reads every input and writes every slot at the same
2791    // index, so the order the buffers are laid out in cannot reach the
2792    // result — as long as every non-scalar input is laid out the same way.
2793    // Then the answer is laid out that way too, and no transpose is made.
2794    let materialised: Vec<Array>;
2795    let (inputs, layout) = match kernel_layout(inputs) {
2796        Some(l) => (inputs, l),
2797        None => {
2798            materialised = inputs.iter().map(Array::to_row_major).collect();
2799            (&materialised[..], Layout::RowMajor)
2800        }
2801    };
2802    // The device is offered row-major work only: uploading a matrix that is
2803    // faster to fold where it lies would be the wrong trade anyway.
2804    if layout == Layout::RowMajor && let Some(d) = device.filter(|d| d.is_gpu()) {
2805        match crate::device::try_run(d, k, inputs) {
2806            Ok(a) => return (Some(a), Placement::Gpu),
2807            Err(why) => placement = Placement::Cpu(why),
2808        }
2809    }
2810    let r = run(k, inputs).map(|a| a.with_layout(layout));
2811    if r.is_none() {
2812        note_fallback();
2813    }
2814    (r, placement)
2815}
2816
2817/// The layout a fused kernel's answer keeps, or None when its inputs
2818/// disagree and the caller must materialise the rows of each.
2819fn kernel_layout(inputs: &[Array]) -> Option<Layout> {
2820    let mut found: Option<Layout> = None;
2821    for a in inputs {
2822        // A scalar is one value repeated into every block: it has no layout
2823        // to agree or disagree with.
2824        if a.rank() == 0 {
2825            continue;
2826        }
2827        match found {
2828            None => found = Some(a.layout()),
2829            Some(l) if l == a.layout() => {}
2830            Some(_) => return None,
2831        }
2832    }
2833    Some(found.unwrap_or_default())
2834}
2835
2836#[cfg(test)]
2837mod tests {
2838    use super::*;
2839    use crate::frontend::{compile, Dialect, Lang};
2840
2841    fn program(src: &str) -> Program {
2842        compile(Lang::J, src, &Dialect::default()).expect("compile")
2843    }
2844
2845    #[test]
2846    fn a_chain_of_two_scalar_verbs_fuses() {
2847        assert!(is_fused(&program("1 + 2 * {x}")));
2848        assert!(is_fused(&program("+/ {w} * {x}")));
2849        assert!(is_fused(&program("+/ ^ {x}")));
2850    }
2851
2852    #[test]
2853    fn one_verb_on_its_own_is_left_alone() {
2854        assert!(!is_fused(&program("2 * {x}")));
2855        assert!(!is_fused(&program("+/ {x}")));
2856        assert!(!is_fused(&program("{x}")));
2857    }
2858
2859    #[test]
2860    fn a_verb_the_kernel_does_not_cover_breaks_the_chain() {
2861        // `%:` can fail elementwise, so it stays outside; the chain under it
2862        // still fuses.
2863        assert!(!is_fused(&program("%: 2 * {x}")));
2864        assert!(is_fused(&program("%: 1 + 2 * {x}")));
2865    }
2866
2867    #[test]
2868    fn an_effect_in_a_leaf_keeps_the_chain_unfused() {
2869        assert!(!is_fused(&program("1 + 2 * echo {x}")));
2870    }
2871
2872    #[test]
2873    fn the_postfix_program_pushes_the_left_operand_first() {
2874        let p = program("{w} - {x} - 1");
2875        let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2876        assert_eq!(
2877            kernel.code(),
2878            [
2879                Instr::Load(2),
2880                Instr::Load(1),
2881                Instr::Load(0),
2882                Instr::Dyad(ScalarDyad::Sub),
2883                Instr::Dyad(ScalarDyad::Sub),
2884            ]
2885        );
2886        // One buffer holds the inner difference, one takes the outer one.
2887        assert_eq!(kernel.slots, 2);
2888    }
2889
2890    #[test]
2891    fn a_value_the_chain_reads_twice_becomes_a_let() {
2892        // What `d =. {x} + 1` then `+/ d * d` comes to once the name has
2893        // moved into the kernel: the sum is computed once per block.
2894        let p = program("+/ ({x} + 1) * ({x} + 1)");
2895        let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2896        assert_eq!(
2897            kernel.code(),
2898            [
2899                Instr::Load(1),
2900                Instr::Load(0),
2901                Instr::Dyad(ScalarDyad::Add),
2902                Instr::Store(0),
2903                Instr::Let(0),
2904                Instr::Let(0),
2905                Instr::Dyad(ScalarDyad::Mul),
2906            ]
2907        );
2908        // One buffer for the let, one for the product it feeds.
2909        assert_eq!(kernel.slots, 2);
2910    }
2911
2912    #[test]
2913    fn a_named_value_moves_into_the_sentence_that_reads_it() {
2914        let p = program("d =. {x} + 1\n+/ d * d");
2915        assert!(is_inlined(&p));
2916        // Three sentences: what the program was, the check that stands
2917        // where the assignment stood, and the sum, which is now the chain
2918        // of the test above.
2919        assert_eq!(p.stmts.len(), 3);
2920        let Expr::Fused { kernel, .. } = &p.stmts[2] else { panic!("the sum did not fuse") };
2921        assert!(kernel.code().contains(&Instr::Store(0)));
2922        assert_eq!(unfused(&p).stmts.len(), 2);
2923    }
2924}