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};
25use crate::dtype::DType;
26use crate::error::Span;
27use crate::ir::{Expr, Program, Scope};
28use crate::par;
29use crate::simd::multiversioned;
30use crate::verb::{tol_cmp, DyadOp, MonadOp, ScalarDyad, ScalarMonad, Tol, Verb, RANK_INF};
31
32/// Elements a block buffer holds.
33///
34/// The working set is `slots` buffers of this size — two or three for the
35/// benchmark kernels — so 8,192 f64 is 128 to 192 KB and stays inside a
36/// 256 KB L2. The value is not delicate: measured at 2,048 / 4,096 / 8,192 /
37/// 16,384 / 32,768 on `+/ w * x` and `+/ ^ x` over 20M rows, the whole range
38/// lands within a few per cent of the best, because what the kernel is
39/// really bounded by is streaming the leaves in from memory once.
40pub const BLOCK: usize = 8_192;
41
42/// One step of a kernel: postfix, so operands are already on the stack.
43#[derive(Clone, Copy, Debug, PartialEq, Eq)]
44pub enum Instr {
45    /// Push input `k`.
46    Load(usize),
47    /// Replace the top of the stack.
48    Monad(ScalarMonad),
49    /// Replace the top two, left below right.
50    Dyad(ScalarDyad),
51    /// Keep the top of the stack as let `k`, a value the rest of the
52    /// program reads more than once. It holds its block buffer until the
53    /// block is finished; nothing pops it.
54    Store(usize),
55    /// Push let `k` again.
56    Let(usize),
57}
58
59/// What one evaluation of a kernel produces.
60#[derive(Clone, Copy, Debug, PartialEq, Eq)]
61pub enum Yield {
62    /// The mapped values, as an array of the chain's own shape.
63    Values,
64    /// The mapped values folded into one by an absorbed reduction.
65    Reduce(ScalarDyad),
66    /// How many items the mapped values would have — `#` over a chain. The
67    /// shapes answer that before any arithmetic runs, so none runs.
68    Tally,
69}
70
71/// A fused elementwise chain and what is made of its values.
72#[derive(Clone, Debug)]
73pub struct FusedKernel {
74    code: Vec<Instr>,
75    /// Block buffers one evaluation needs at once.
76    slots: usize,
77    yields: Yield,
78    /// The input each leaf of the chain reads, in the order the chain
79    /// reaches them. Two leaves that are the same subtree share one input,
80    /// so this is not the identity, and the fallback needs it to give every
81    /// leaf back the value it was given.
82    leaves: Vec<usize>,
83    /// The dialect's comparison tolerance, so that a comparison inside the
84    /// kernel answers exactly as the same comparison outside it does.
85    tol: Tol,
86}
87
88impl FusedKernel {
89    pub fn code(&self) -> &[Instr] {
90        &self.code
91    }
92
93    pub fn yields(&self) -> Yield {
94        self.yields
95    }
96
97    pub fn reduce(&self) -> Option<ScalarDyad> {
98        match self.yields {
99            Yield::Reduce(op) => Some(op),
100            _ => None,
101        }
102    }
103
104    /// The comparison tolerance the program was compiled with. A backend
105    /// that generates its own code for this kernel needs it, so that a
106    /// comparison answers there as it answers everywhere else.
107    pub fn tol(&self) -> Tol {
108        self.tol
109    }
110}
111
112/// How often a fused node has handed its work back to the original subtree.
113/// A counter rather than a log: the fallback is correct, only slower, and
114/// what a caller wants to know is whether it is happening at all.
115static FALLBACKS: AtomicU64 = AtomicU64::new(0);
116
117/// Number of fallbacks since the process started.
118pub fn fallback_count() -> u64 {
119    FALLBACKS.load(Ordering::Relaxed)
120}
121
122fn note_fallback() {
123    FALLBACKS.fetch_add(1, Ordering::Relaxed);
124}
125
126// ------------------------------------------------------------- the op set
127//
128// A verb may join a kernel only if it cannot fail on numeric data: the
129// kernel reports no errors of its own, so anything that could raise one
130// (APL's `÷` by zero, `%:` and `^.` of a negative, `^`'s zero to a negative
131// power, APL's `~` off 0/1) stays outside and breaks the chain there.
132
133/// The elementwise monad this verb performs, if the kernel covers it.
134fn fusable_monad(v: &Verb) -> Option<ScalarMonad> {
135    use ScalarMonad::*;
136    let Verb::Prim(p) = v else { return None };
137    let MonadOp::Scalar(op) = p.monad else { return None };
138    matches!(
139        op,
140        Conj | Neg | Abs | Signum | Recip | Floor | Ceil | Inc | Dec | Double | Halve | Square
141            | OneMinus | Exp
142    )
143    .then_some(op)
144}
145
146/// The elementwise dyad this verb performs, if the kernel covers it.
147fn fusable_dyad(v: &Verb) -> Option<ScalarDyad> {
148    use ScalarDyad::*;
149    let Verb::Prim(p) = v else { return None };
150    let DyadOp::Scalar(op) = p.dyad else { return None };
151    matches!(op, Add | Sub | Mul | DivJ | Min | Max | Residue | Eq | Ne | Lt | Le | Gt | Ge)
152        .then_some(op)
153}
154
155/// The reduction this verb performs over the leading axis, if the kernel can
156/// absorb it: an associative arithmetic primitive, applied at full rank.
157/// APL's `+/` is the same thing under a rank wrapper.
158fn absorbable_reduce(v: &Verb) -> Option<ScalarDyad> {
159    use ScalarDyad::*;
160    let inner = match v {
161        Verb::Reduce(u) => u,
162        // The wrapper applies the reduction to cells of rank >= 1; over the
163        // rank-1 argument this kernel insists on, that is the whole array.
164        Verb::Rank(u, r) if r[0] >= 1 => match &**u {
165            Verb::Reduce(inner) => inner,
166            _ => return None,
167        },
168        _ => return None,
169    };
170    let Verb::Prim(p) = &**inner else { return None };
171    let DyadOp::Scalar(op) = p.dyad else { return None };
172    matches!(op, Add | Mul | Min | Max).then_some(op)
173}
174
175/// Is this the tally, applied to the array as a whole? `#"1` and its like
176/// count the items of cells instead, which is not what the shape says.
177fn is_tally(v: &Verb) -> bool {
178    matches!(v, Verb::Prim(p) if p.monad == MonadOp::Tally && p.ranks[0] == RANK_INF)
179}
180
181// ------------------------------------------------------------- the pass
182
183/// The chain as a tree, before it becomes postfix code.
184#[derive(Clone, PartialEq)]
185enum Node {
186    /// A subtree the kernel does not cover: an input, with its index.
187    Leaf(usize),
188    Monad(ScalarMonad, Box<Node>),
189    Dyad(ScalarDyad, Box<Node>, Box<Node>),
190}
191
192/// The subtrees a chain reads.
193///
194/// Inputs are numbered in the order the evaluator would reach them — a
195/// dyad's right argument first — so that a fused node evaluates its leaves
196/// exactly when and where the unfused tree does. Two leaves that are the
197/// same subtree take the same input: nothing inside a chain can assign, so
198/// the second writing of `+/ {x}` reads what the first one read, and
199/// evaluating it once is what the sentence means either way.
200#[derive(Default)]
201struct Leaves<'a> {
202    inputs: Vec<&'a Expr>,
203    /// The input each leaf position reads, in chain order.
204    order: Vec<usize>,
205}
206
207impl<'a> Leaves<'a> {
208    fn push(&mut self, e: &'a Expr) -> usize {
209        let i = match self.inputs.iter().position(|&p| same(p, e)) {
210            Some(i) => i,
211            None => {
212                self.inputs.push(e);
213                self.inputs.len() - 1
214            }
215        };
216        self.order.push(i);
217        i
218    }
219}
220
221/// A name the chain reads through to the value assigned to it, as inlining
222/// that assignment would; `hits` counts the uses it absorbed.
223struct Inline<'a> {
224    name: &'a str,
225    def: &'a Expr,
226    hits: usize,
227}
228
229/// Build the chain rooted at `e`, collecting the subtrees that feed it.
230fn chain<'a>(e: &'a Expr, lv: &mut Leaves<'a>, sub: &mut Option<Inline<'a>>) -> Node {
231    let read_through = match (e, sub.as_ref()) {
232        (Expr::Name(n, _), Some(s)) if n == s.name => Some(s.def),
233        _ => None,
234    };
235    if let Some(def) = read_through {
236        if let Some(s) = sub.as_mut() {
237            s.hits += 1;
238        }
239        return chain(def, lv, sub);
240    }
241    match e {
242        Expr::Monad { verb, y, .. } => match fusable_monad(verb) {
243            Some(op) => Node::Monad(op, Box::new(chain(y, lv, sub))),
244            None => Node::Leaf(lv.push(e)),
245        },
246        Expr::Dyad { verb, x, y, .. } => match fusable_dyad(verb) {
247            Some(op) => {
248                let ry = chain(y, lv, sub);
249                let rx = chain(x, lv, sub);
250                Node::Dyad(op, Box::new(rx), Box::new(ry))
251            }
252            None => Node::Leaf(lv.push(e)),
253        },
254        _ => Node::Leaf(lv.push(e)),
255    }
256}
257
258fn ops(n: &Node) -> usize {
259    match n {
260        Node::Leaf(_) => 0,
261        Node::Monad(_, y) => 1 + ops(y),
262        Node::Dyad(_, x, y) => 1 + ops(x) + ops(y),
263    }
264}
265
266/// Every subtree of the chain that computes something.
267fn subtrees<'a>(n: &'a Node, out: &mut Vec<&'a Node>) {
268    if ops(n) == 0 {
269        return;
270    }
271    out.push(n);
272    match n {
273        Node::Leaf(_) => {}
274        Node::Monad(_, y) => subtrees(y, out),
275        Node::Dyad(_, x, y) => {
276            subtrees(x, out);
277            subtrees(y, out);
278        }
279    }
280}
281
282/// The values the chain computes more than once, largest first.
283///
284/// `+/ d * d` over an inlined `d` writes the same arithmetic twice, and a
285/// block-at-a-time kernel can do what the assignment did: compute it once
286/// and read it twice. Each of these becomes a let — a block buffer of its
287/// own, held for the length of the block. Only maximal repeats are taken,
288/// so a repeat inside a let is part of that let rather than one more.
289fn lets_of(n: &Node) -> Vec<Node> {
290    let mut all = Vec::new();
291    subtrees(n, &mut all);
292    let mut out = Vec::new();
293    fn walk(n: &Node, all: &[&Node], out: &mut Vec<Node>) {
294        if ops(n) >= 1 && all.iter().filter(|m| **m == n).count() >= 2 {
295            if !out.contains(n) {
296                out.push(n.clone());
297            }
298            return;
299        }
300        match n {
301            Node::Leaf(_) => {}
302            Node::Monad(_, y) => walk(y, all, out),
303            Node::Dyad(_, x, y) => {
304                walk(x, all, out);
305                walk(y, all, out);
306            }
307        }
308    }
309    walk(n, &all, &mut out);
310    out
311}
312
313/// Postfix code for the chain: the lets first, each into a slot of its own,
314/// then the chain that reads them.
315fn emit_all(n: &Node, lets: &[Node], code: &mut Vec<Instr>) {
316    for (k, l) in lets.iter().enumerate() {
317        // A let is emitted from the lets before it, so it cannot read
318        // itself; maximal repeats never nest, so there is nothing else.
319        emit(l, &lets[..k], code);
320        code.push(Instr::Store(k));
321    }
322    emit(n, lets, code);
323}
324
325/// Postfix code for the chain: a dyad's left operand is pushed first.
326fn emit(n: &Node, lets: &[Node], code: &mut Vec<Instr>) {
327    if let Some(k) = lets.iter().position(|l| l == n) {
328        code.push(Instr::Let(k));
329        return;
330    }
331    match n {
332        Node::Leaf(i) => code.push(Instr::Load(*i)),
333        Node::Monad(op, y) => {
334            emit(y, lets, code);
335            code.push(Instr::Monad(*op));
336        }
337        Node::Dyad(op, x, y) => {
338            emit(x, lets, code);
339            emit(y, lets, code);
340            code.push(Instr::Dyad(*op));
341        }
342    }
343}
344
345/// Block buffers the postfix program needs at once.
346///
347/// Only a computed value holds one — an input is read where it lies — and
348/// the buffer being written is allocated before the operands are released,
349/// so the peak is the live count at some operation plus one.
350fn slots(code: &[Instr]) -> usize {
351    let mut stack: Vec<bool> = Vec::new();
352    let mut live = 0usize;
353    let mut max = 1usize;
354    for ins in code {
355        let operands = match ins {
356            Instr::Load(_) => {
357                stack.push(false);
358                continue;
359            }
360            // A let holds its buffer for the whole block: it is never
361            // released, so the count it added when it was computed stands
362            // and reading it takes nothing.
363            Instr::Let(_) => {
364                stack.push(false);
365                continue;
366            }
367            Instr::Store(_) => {
368                stack.pop();
369                continue;
370            }
371            Instr::Monad(_) => 1,
372            Instr::Dyad(_) => 2,
373        };
374        max = max.max(live + 1);
375        for _ in 0..operands {
376            if stack.pop().unwrap_or(false) {
377                live -= 1;
378            }
379        }
380        live += 1;
381        stack.push(true);
382    }
383    max
384}
385
386/// Is this subtree free of effects?
387///
388/// A fused node evaluates its leaves in the order the unfused tree would,
389/// so an effect in one would still happen exactly once — but a node that
390/// can fall back is easier to be sure of when nothing inside it can act on
391/// the world, and a chain with `echo` in it is not the kind worth fusing.
392fn replayable(e: &Expr) -> bool {
393    match e {
394        Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => true,
395        Expr::Assign { .. }
396        | Expr::PrintPass { .. }
397        | Expr::Elided { .. }
398        | Expr::Control(..)
399        | Expr::AmendIndex { .. }
400        | Expr::VerbDef { .. } => false,
401        Expr::Monad { verb, y, .. } => verb.is_pure() && replayable(y),
402        Expr::Dyad { verb, x, y, .. } => verb.is_pure() && replayable(x) && replayable(y),
403        Expr::Fused { inputs, .. } => inputs.iter().all(replayable),
404    }
405}
406
407/// Are these the same computation? Two writings of one subexpression differ
408/// in their spans, which are positions in the source and mean nothing to
409/// the value, so spans are not compared. Assignments, output and fused
410/// nodes are never the same as anything: only leaves of a chain reach here,
411/// and a chain holds none of those.
412fn same(a: &Expr, b: &Expr) -> bool {
413    match (a, b) {
414        (Expr::Const(p, _), Expr::Const(q, _)) => p == q,
415        (Expr::Param(p, _), Expr::Param(q, _)) => p == q,
416        (Expr::Name(p, _), Expr::Name(q, _)) => p == q,
417        (Expr::Monad { verb: u, y: p, .. }, Expr::Monad { verb: v, y: q, .. }) => {
418            same_verb(u, v) && same(p, q)
419        }
420        (
421            Expr::Dyad { verb: u, x: px, y: py, .. },
422            Expr::Dyad { verb: v, x: qx, y: qy, .. },
423        ) => same_verb(u, v) && same(px, qx) && same(py, qy),
424        _ => false,
425    }
426}
427
428fn same_verb(a: &Verb, b: &Verb) -> bool {
429    match (a, b) {
430        (Verb::Prim(p), Verb::Prim(q)) => p == q,
431        (Verb::Rank(u, r), Verb::Rank(v, s)) => r == s && same_verb(u, v),
432        (Verb::Reduce(u), Verb::Reduce(v)) | (Verb::Commute(u), Verb::Commute(v)) => {
433            same_verb(u, v)
434        }
435        (Verb::Windowed(u, j), Verb::Windowed(v, k)) => j == k && same_verb(u, v),
436        (Verb::PowerN(u, m), Verb::PowerN(v, n)) => m == n && same_verb(u, v),
437        (Verb::Fork(f, g, h), Verb::Fork(f2, g2, h2)) => {
438            same_verb(f, f2) && same_verb(g, g2) && same_verb(h, h2)
439        }
440        (Verb::NounFork(m, g, h), Verb::NounFork(n, g2, h2)) => {
441            m == n && same_verb(g, g2) && same_verb(h, h2)
442        }
443        (Verb::Hook(g, h), Verb::Hook(g2, h2))
444        | (Verb::Atop(g, h), Verb::Atop(g2, h2))
445        | (Verb::Compose(g, h), Verb::Compose(g2, h2)) => same_verb(g, g2) && same_verb(h, h2),
446        (Verb::BondLeft(m, u), Verb::BondLeft(n, v)) => m == n && same_verb(u, v),
447        (Verb::BondRight(u, m), Verb::BondRight(v, n)) => m == n && same_verb(u, v),
448        _ => false,
449    }
450}
451
452/// Optimise a compiled program's sentences: move the values that are only
453/// named for the reader into the sentences that read them, then fuse every
454/// chain that is left.
455pub fn pass(stmts: &mut Vec<Expr>, tol: Tol) {
456    let orig = std::mem::take(stmts);
457    let mut cur = orig.clone();
458    let mut names = 0usize;
459    let mut crossed = false;
460    // A round elides one assignment, so a chain of them — `m =. ...`,
461    // `d =. {x} - m`, `+/ d * d` — takes one round per link.
462    for _ in 0..=orig.len() {
463        match inline_once(&cur, &mut names, tol) {
464            Some(next) => {
465                cur = next;
466                crossed = true;
467            }
468            None => break,
469        }
470    }
471    let mut out: Vec<Expr> = cur.into_iter().map(|e| fuse_expr(e, tol)).collect();
472    if crossed {
473        // What the sentences were, for `unfused` to hold this against.
474        out.insert(0, Expr::Elided { orig, span: Span::new(0, 0) });
475    }
476    *stmts = out;
477}
478
479fn fuse_expr(e: Expr, tol: Tol) -> Expr {
480    if let Some(f) = try_fuse(&e, tol) {
481        return f;
482    }
483    match e {
484        Expr::Assign { name, value, scope, span } => {
485            Expr::Assign { name, value: Box::new(fuse_expr(*value, tol)), scope, span }
486        }
487        Expr::Monad { verb, y, span } => {
488            Expr::Monad { verb, y: Box::new(fuse_expr(*y, tol)), span }
489        }
490        Expr::Dyad { verb, x, y, span } => Expr::Dyad {
491            verb,
492            x: Box::new(fuse_expr(*x, tol)),
493            y: Box::new(fuse_expr(*y, tol)),
494            span,
495        },
496        Expr::PrintPass { value, span } => {
497            Expr::PrintPass { value: Box::new(fuse_expr(*value, tol)), span }
498        }
499        other => other,
500    }
501}
502
503/// The kernel for the chain rooted at `root`, if it carries at least
504/// `least` operations and reads nothing that cannot be replayed.
505fn build<'a>(
506    root: &'a Expr,
507    yields: Yield,
508    least: usize,
509    sub: &mut Option<Inline<'a>>,
510    tol: Tol,
511) -> Option<(FusedKernel, Vec<&'a Expr>)> {
512    if let Some(s) = sub.as_mut() {
513        s.hits = 0;
514    }
515    let mut lv = Leaves::default();
516    let node = chain(root, &mut lv, sub);
517    if ops(&node) < least || !lv.inputs.iter().all(|l| replayable(l)) {
518        return None;
519    }
520    let mut code = Vec::new();
521    emit_all(&node, &lets_of(&node), &mut code);
522    let kernel = FusedKernel { slots: slots(&code), code, yields, leaves: lv.order, tol };
523    Some((kernel, lv.inputs))
524}
525
526/// The kernel this node becomes, with the subtree it stands for — the chain
527/// itself where a tally reads only its shape, the whole sentence where a
528/// reduction sits above it.
529///
530/// One elementwise verb on its own already runs as one pass; fusing it
531/// would only add a layer. A reduction to absorb, or a tally that makes the
532/// values unnecessary altogether, makes one verb enough.
533fn kernel_at<'a>(
534    e: &'a Expr,
535    sub: &mut Option<Inline<'a>>,
536    tol: Tol,
537) -> Option<(FusedKernel, Vec<&'a Expr>, &'a Expr)> {
538    if let Expr::Monad { verb, y, .. } = e {
539        if is_tally(verb) {
540            if let Some((k, l)) = build(y, Yield::Tally, 1, sub, tol) {
541                return Some((k, l, e));
542            }
543        }
544        if let Some(op) = absorbable_reduce(verb) {
545            if let Some((k, l)) = build(y, Yield::Reduce(op), 1, sub, tol) {
546                return Some((k, l, e));
547            }
548        }
549    }
550    let (k, l) = build(e, Yield::Values, 2, sub, tol)?;
551    Some((k, l, e))
552}
553
554/// The fused node for the chain rooted at `e`, if there is one worth making.
555fn try_fuse(e: &Expr, tol: Tol) -> Option<Expr> {
556    let (kernel, leaves, orig) = kernel_at(e, &mut None, tol)?;
557    let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
558    Some(Expr::Fused {
559        kernel,
560        inputs,
561        orig: Box::new(orig.clone()),
562        span: e.span(),
563    })
564}
565
566/// The chain a fused node came from, with its leaves replaced by the values
567/// already computed for them.
568///
569/// This is what runs when the kernel declines. Rebuilding the tree costs a
570/// handful of small allocations and saves evaluating the leaves a second
571/// time, which for a leaf like `19 }. {close}` is a whole array.
572pub(crate) fn fallback_tree(k: &FusedKernel, orig: &Expr, values: &[Array]) -> Expr {
573    let mut next = 0;
574    let tree = match orig {
575        // An absorbed reduction sits above the chain; only the chain's own
576        // leaves were evaluated.
577        Expr::Monad { verb, y, span } if matches!(k.yields, Yield::Reduce(_)) => Expr::Monad {
578            verb: verb.clone(),
579            y: Box::new(substitute(y, values, k, &mut next)),
580            span: *span,
581        },
582        // A tally is not applied at all: the chain alone runs, and the
583        // count of what it made is what the node yields.
584        Expr::Monad { verb, y, .. } if k.yields == Yield::Tally && is_tally(verb) => {
585            substitute(y, values, k, &mut next)
586        }
587        e => substitute(e, values, k, &mut next),
588    };
589    debug_assert_eq!(next, k.leaves.len(), "the fallback found different leaves");
590    tree
591}
592
593/// What the kernel would have made of the value its chain produced. A tally
594/// skips the chain entirely when it runs, and counts the items of it when
595/// the chain has had to run instead.
596pub(crate) fn fallback_finish(k: &FusedKernel, v: Array) -> Array {
597    match k.yields {
598        Yield::Tally => Array::scalar_i64(v.items() as i64),
599        _ => v,
600    }
601}
602
603/// Walk the chain exactly as [`chain`] walked it, so the leaves take their
604/// values in the order they were numbered in.
605fn substitute(e: &Expr, values: &[Array], k: &FusedKernel, next: &mut usize) -> Expr {
606    match e {
607        Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
608            verb: verb.clone(),
609            y: Box::new(substitute(y, values, k, next)),
610            span: *span,
611        },
612        Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => {
613            let ry = substitute(y, values, k, next);
614            let rx = substitute(x, values, k, next);
615            Expr::Dyad { verb: verb.clone(), x: Box::new(rx), y: Box::new(ry), span: *span }
616        }
617        leaf => {
618            let v = values[k.leaves[*next]].clone();
619            *next += 1;
620            Expr::Const(v, leaf.span())
621        }
622    }
623}
624
625// ------------------------------------------- across sentence boundaries
626//
627// `d =. {x} - m` and then `+/ d * d` is the same computation as the one
628// sentence that spells it out, but the assignment writes `d` to memory in
629// full and the next sentence reads it back — the traffic the kernel exists
630// to remove. Nothing there needs the array: the name is for the reader.
631//
632// So the pass moves the value into the sentences that read it, and hoists
633// the value's own leaves — the mean's `+/ {x}` — into sentences of their
634// own first, so that copying the chain does not copy the work. What comes
635// out is the two-phase shape a hand-written kernel has: one pass for the
636// reductions the chain reads as scalars, one for the map-reduce over them.
637
638/// Names the pass introduces for the values it hoists. `·` starts no name
639/// either frontend accepts, so these cannot collide with the program's.
640fn hoisted_name(n: &mut usize) -> String {
641    *n += 1;
642    format!("·{}", *n - 1)
643}
644
645/// Elide the first assignment whose value can move into the sentences that
646/// read it, and report the sentences that leaves; None when none can.
647///
648/// The value moves only where the name is pure dataflow:
649///
650/// - the value is replayable and is a chain, so that moving it moves
651///   arithmetic into a kernel rather than moving a whole pass;
652/// - no later sentence assigns the name again, or any name the value reads,
653///   so every copy means what the original meant;
654/// - every use lands inside a kernel, so no copy materialises the value.
655///   A tally counts as landing inside one: it reads the chain's shape.
656///
657/// The assignment's own sentence stays, as the tally of the chain: that
658/// reaches every leaf and every rule the kernel has, so whatever the
659/// assignment would have raised is raised where it was raised before, and
660/// nothing else is computed.
661fn inline_once(stmts: &[Expr], names: &mut usize, tol: Tol) -> Option<Vec<Expr>> {
662    for (i, stmt) in stmts.iter().enumerate() {
663        let Expr::Assign { name, value, span, .. } = stmt else { continue };
664        if !inlinable(stmts, i, name, value, tol) {
665            continue;
666        }
667        if let Some(out) = rewrite(stmts, i, name, value, *span, names, tol) {
668            return Some(out);
669        }
670    }
671    None
672}
673
674fn inlinable(stmts: &[Expr], i: usize, name: &str, value: &Expr, tol: Tol) -> bool {
675    if !replayable(value) || mentions(value, name) {
676        return false;
677    }
678    let mut lv = Leaves::default();
679    if ops(&chain(value, &mut lv, &mut None)) < 1 {
680        return false;
681    }
682    let mut guarded = vec![name.to_string()];
683    free_names(value, &mut guarded);
684    let later = &stmts[i + 1..];
685    if later.iter().any(|s| assigns_any(s, &guarded)) {
686        return false;
687    }
688    let mut uses = 0;
689    for stmt in later {
690        match uses_land(stmt, name, value, tol) {
691            Some(n) => uses += n,
692            None => return false,
693        }
694    }
695    uses > 0
696}
697
698/// How many uses of `name` this sentence would take into a kernel, or None
699/// when one of them would have to materialise the value instead.
700fn uses_land(e: &Expr, name: &str, def: &Expr, tol: Tol) -> Option<usize> {
701    let mut sub = Some(Inline { name, def, hits: 0 });
702    if let Some((_, leaves, _)) = kernel_at(e, &mut sub, tol) {
703        let mut n = sub.map_or(0, |s| s.hits);
704        for l in leaves {
705            n += uses_land(l, name, def, tol)?;
706        }
707        return Some(n);
708    }
709    match e {
710        Expr::Name(n, _) if n == name => None,
711        Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => Some(0),
712        Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => uses_land(value, name, def, tol),
713        Expr::Monad { y, .. } => uses_land(y, name, def, tol),
714        Expr::Dyad { x, y, .. } => Some(uses_land(x, name, def, tol)? + uses_land(y, name, def, tol)?),
715        Expr::Fused { .. }
716        | Expr::Elided { .. }
717        | Expr::Control(..)
718        | Expr::AmendIndex { .. }
719        | Expr::VerbDef { .. } => None,
720    }
721}
722
723/// The sentences that replace `stmts`, with the assignment at `i` elided.
724fn rewrite(
725    stmts: &[Expr],
726    i: usize,
727    name: &str,
728    value: &Expr,
729    span: Span,
730    names: &mut usize,
731    tol: Tol,
732) -> Option<Vec<Expr>> {
733    let mut lv = Leaves::default();
734    chain(value, &mut lv, &mut None);
735    // A leaf that is more than a name or a constant becomes a sentence of
736    // its own, evaluated once and where it was evaluated before.
737    let mut hoists = Vec::new();
738    let mut bound: Vec<Option<String>> = Vec::new();
739    for l in &lv.inputs {
740        if matches!(l, Expr::Const(..) | Expr::Param(..) | Expr::Name(..)) {
741            bound.push(None);
742            continue;
743        }
744        let n = hoisted_name(names);
745        hoists.push(Expr::Assign {
746            name: n.clone(),
747            value: Box::new((*l).clone()),
748            scope: Scope::Local,
749            span: l.span(),
750        });
751        bound.push(Some(n));
752    }
753    let def = with_leaves(value, &lv, &bound);
754    let (kernel, leaves) = build(&def, Yield::Tally, 1, &mut None, tol)?;
755    let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
756    let guard = Expr::Assign {
757        name: hoisted_name(names),
758        value: Box::new(Expr::Fused {
759            kernel,
760            inputs,
761            orig: Box::new(def.clone()),
762            span,
763        }),
764        scope: Scope::Local,
765        span,
766    };
767    let mut out = stmts[..i].to_vec();
768    out.extend(hoists);
769    out.push(guard);
770    out.extend(stmts[i + 1..].iter().map(|s| replace_name(s, name, &def)));
771    Some(out)
772}
773
774/// The chain with its hoisted leaves replaced by the names they were bound
775/// to. Walks exactly as [`chain`] walks, so the leaves are the same ones.
776fn with_leaves(e: &Expr, lv: &Leaves<'_>, bound: &[Option<String>]) -> Expr {
777    match e {
778        Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
779            verb: verb.clone(),
780            y: Box::new(with_leaves(y, lv, bound)),
781            span: *span,
782        },
783        Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => Expr::Dyad {
784            verb: verb.clone(),
785            x: Box::new(with_leaves(x, lv, bound)),
786            y: Box::new(with_leaves(y, lv, bound)),
787            span: *span,
788        },
789        leaf => {
790            let bind = lv
791                .inputs
792                .iter()
793                .position(|&p| same(p, leaf))
794                .and_then(|i| bound[i].as_ref());
795            match bind {
796                Some(n) => Expr::Name(n.clone(), leaf.span()),
797                None => leaf.clone(),
798            }
799        }
800    }
801}
802
803fn replace_name(e: &Expr, name: &str, def: &Expr) -> Expr {
804    match e {
805        Expr::Name(n, _) if n == name => def.clone(),
806        Expr::Assign { name: a, value, scope, span } => Expr::Assign {
807            scope: *scope,
808            name: a.clone(),
809            value: Box::new(replace_name(value, name, def)),
810            span: *span,
811        },
812        Expr::PrintPass { value, span } => Expr::PrintPass {
813            value: Box::new(replace_name(value, name, def)),
814            span: *span,
815        },
816        Expr::Monad { verb, y, span } => Expr::Monad {
817            verb: verb.clone(),
818            y: Box::new(replace_name(y, name, def)),
819            span: *span,
820        },
821        Expr::Dyad { verb, x, y, span } => Expr::Dyad {
822            verb: verb.clone(),
823            x: Box::new(replace_name(x, name, def)),
824            y: Box::new(replace_name(y, name, def)),
825            span: *span,
826        },
827        other => other.clone(),
828    }
829}
830
831fn mentions(e: &Expr, name: &str) -> bool {
832    let mut names = Vec::new();
833    free_names(e, &mut names);
834    names.iter().any(|n| n == name)
835}
836
837/// Every name this subtree reads.
838fn free_names(e: &Expr, out: &mut Vec<String>) {
839    match e {
840        Expr::Name(n, _) => out.push(n.clone()),
841        Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => free_names(value, out),
842        Expr::Monad { y, .. } => free_names(y, out),
843        Expr::Dyad { x, y, .. } => {
844            free_names(x, out);
845            free_names(y, out);
846        }
847        Expr::Fused { inputs, .. } => inputs.iter().for_each(|i| free_names(i, out)),
848        Expr::Const(..)
849        | Expr::Param(..)
850        | Expr::Elided { .. }
851        | Expr::Control(..)
852        | Expr::AmendIndex { .. }
853        | Expr::VerbDef { .. } => {}
854    }
855}
856
857/// Does this sentence assign any of these names, at any depth?
858fn assigns_any(e: &Expr, names: &[String]) -> bool {
859    match e {
860        Expr::Assign { name, value, .. } => {
861            names.iter().any(|n| n == name) || assigns_any(value, names)
862        }
863        Expr::PrintPass { value, .. } => assigns_any(value, names),
864        Expr::Monad { y, .. } => assigns_any(y, names),
865        Expr::Dyad { x, y, .. } => assigns_any(x, names) || assigns_any(y, names),
866        Expr::Fused { inputs, .. } => inputs.iter().any(|i| assigns_any(i, names)),
867        Expr::Const(..)
868        | Expr::Param(..)
869        | Expr::Name(..)
870        | Expr::Elided { .. }
871        | Expr::Control(..)
872        | Expr::AmendIndex { .. }
873        | Expr::VerbDef { .. } => false,
874    }
875}
876
877/// Does any sentence of this program run a fused kernel?
878pub fn is_fused(p: &Program) -> bool {
879    fn any(e: &Expr) -> bool {
880        match e {
881            Expr::Fused { .. } => true,
882            Expr::Const(..)
883            | Expr::Param(..)
884            | Expr::Name(..)
885            | Expr::Elided { .. }
886            | Expr::Control(..)
887            | Expr::AmendIndex { .. }
888            | Expr::VerbDef { .. } => false,
889            Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => any(value),
890            Expr::Monad { y, .. } => any(y),
891            Expr::Dyad { x, y, .. } => any(x) || any(y),
892        }
893    }
894    p.stmts.iter().any(any)
895}
896
897/// Did the pass move a named value into the sentences that read it?
898pub fn is_inlined(p: &Program) -> bool {
899    matches!(p.stmts.first(), Some(Expr::Elided { .. }))
900}
901
902/// The program as the plain evaluator would run it: the sentences it was
903/// compiled from, with every fused node replaced by the subtree it came
904/// from. The two must compute the same thing; tests hold them to it.
905pub fn unfused(p: &Program) -> Program {
906    fn strip(e: &Expr) -> Expr {
907        match e {
908            Expr::Fused { orig, .. } => strip(orig),
909            Expr::Assign { name, value, scope, span } => {
910                Expr::Assign {
911                    name: name.clone(),
912                    value: Box::new(strip(value)),
913                    scope: *scope,
914                    span: *span,
915                }
916            }
917            Expr::PrintPass { value, span } => {
918                Expr::PrintPass { value: Box::new(strip(value)), span: *span }
919            }
920            Expr::Monad { verb, y, span } => {
921                Expr::Monad { verb: verb.clone(), y: Box::new(strip(y)), span: *span }
922            }
923            Expr::Dyad { verb, x, y, span } => Expr::Dyad {
924                verb: verb.clone(),
925                x: Box::new(strip(x)),
926                y: Box::new(strip(y)),
927                span: *span,
928            },
929            other => other.clone(),
930        }
931    }
932    let mut out = p.clone();
933    // A program the pass rewrote across sentences kept the sentences it
934    // rewrote; those, not the rewriting, are what the evaluator would run.
935    let stmts = match p.stmts.first() {
936        Some(Expr::Elided { orig, .. }) => orig,
937        _ => &p.stmts,
938    };
939    out.stmts = stmts.iter().map(strip).collect();
940    out
941}
942
943// ------------------------------------------------------------- dtype rules
944
945/// The dtype the unfused pipeline gives this monad's result. None where it
946/// depends on the values (`<.` of a float is an integer only if every
947/// rounded value fits one), which the kernel declines rather than guess.
948fn monad_type(op: ScalarMonad, a: DType) -> Option<DType> {
949    use DType::*;
950    use ScalarMonad::*;
951    // The kernel computes in one real type; complex values are not one of
952    // them, so a chain that touches one declines and runs unfused.
953    if a == Complex {
954        return None;
955    }
956    Some(match op {
957        Recip | Halve | Exp => F64,
958        // Identity and magnitude keep a boolean boolean.
959        Conj | Abs | OneMinus => a,
960        Neg | Signum | Inc | Dec | Double | Square => match a {
961            Bool | I64 => I64,
962            other => other,
963        },
964        Floor | Ceil => match a {
965            Bool | I64 => I64,
966            _ => return None,
967        },
968        _ => return None,
969    })
970}
971
972/// The dtype the unfused pipeline gives this dyad's result, on the path
973/// where no integer step overflows (one that does falls back).
974fn dyad_type(op: ScalarDyad, a: DType, b: DType) -> Option<DType> {
975    use ScalarDyad::*;
976    if a == DType::Complex || b == DType::Complex {
977        return None;
978    }
979    match op {
980        Eq | Ne | Lt | Le | Gt | Ge => Some(DType::Bool),
981        DivJ => Some(DType::F64),
982        Add | Sub | Mul | Min | Max | Residue => match DType::promote(a, b)? {
983            DType::Bool => Some(DType::I64),
984            DType::Char => None,
985            t => Some(t),
986        },
987        _ => None,
988    }
989}
990
991/// The type the kernel computes in, and the dtype of its mapped result.
992///
993/// Every value in the program is computed in one type, so it must be one
994/// that holds them all: integers when nothing in the chain leaves them,
995/// floats otherwise. That leaves one case the kernel cannot serve — a chain
996/// that computes an integer somewhere along a float path, as
997/// `(x > 0) + (y > 0)` or `({a} + {b}) % 2` do. Its unfused pipeline holds
998/// those steps in i64, exactly, past where f64 stops being exact, and its
999/// result may be an integer array. Rather than compute them in the wrong
1000/// type, the kernel declines and the chain runs.
1001///
1002/// A boolean is not such a case: a comparison yields 0 and 1, which f64
1003/// holds exactly, and only the dtype of a result made from one has to be
1004/// narrowed at the end.
1005///
1006/// This is the kernel's main blind spot — a random chain over mixed integer
1007/// and float arguments declines about half the time — and the way out is a
1008/// stack whose entries carry their own type rather than one type per
1009/// kernel. Nothing measured so far needs it.
1010pub(crate) fn working_type(k: &FusedKernel, inputs: &[Array]) -> Option<(DType, DType)> {
1011    let mut stack: Vec<DType> = Vec::with_capacity(k.slots);
1012    let mut lets: Vec<DType> = Vec::new();
1013    let mut float = false;
1014    let mut integer_step = false;
1015    // The exact types and the complex ones have no blockwise kernel: a
1016    // fused chain over them declines and the general path evaluates it.
1017    if inputs.iter().any(|a| a.dtype() == DType::Complex || a.dtype().is_exact()) {
1018        return None;
1019    }
1020    for ins in &k.code {
1021        let t = match ins {
1022            Instr::Load(i) => inputs[*i].dtype(),
1023            Instr::Monad(op) => monad_type(*op, stack.pop()?)?,
1024            Instr::Dyad(op) => {
1025                let b = stack.pop()?;
1026                let a = stack.pop()?;
1027                dyad_type(*op, a, b)?
1028            }
1029            Instr::Store(k) => {
1030                let t = stack.pop()?;
1031                if lets.len() != *k {
1032                    return None;
1033                }
1034                lets.push(t);
1035                continue;
1036            }
1037            // Reading a let is not a step: the value was accounted for
1038            // where it was computed.
1039            Instr::Let(k) => {
1040                let t = *lets.get(*k)?;
1041                float |= t == DType::F64;
1042                stack.push(t);
1043                continue;
1044            }
1045        };
1046        // Only numbers: everything else — characters, boxes — is a type
1047        // the kernel has no arithmetic for and the chain must handle.
1048        if !t.is_numeric() {
1049            return None;
1050        }
1051        float |= t == DType::F64;
1052        // An argument's own values are exact in either type; a step's are
1053        // not, once they are integers wider than f64's 53 bits.
1054        integer_step |= t == DType::I64 && !matches!(ins, Instr::Load(_));
1055        stack.push(t);
1056    }
1057    let root = stack.pop()?;
1058    let working = if float { DType::F64 } else { DType::I64 };
1059    if working == DType::F64 && integer_step {
1060        return None;
1061    }
1062    Some((working, root))
1063}
1064
1065// ------------------------------------------------------------- execution
1066
1067/// One input, in the working type: either the values themselves or one
1068/// value repeated, which is how a rank-0 argument reaches every element.
1069struct Loaded<'a, T> {
1070    data: &'a [T],
1071    splat: bool,
1072}
1073
1074impl<T> Loaded<'_, T> {
1075    #[inline]
1076    fn block(&self, start: usize, len: usize) -> &[T] {
1077        if self.splat {
1078            &self.data[..len]
1079        } else {
1080            &self.data[start..start + len]
1081        }
1082    }
1083}
1084
1085/// What a stack entry refers to: an input, or a block buffer.
1086#[derive(Clone, Copy)]
1087enum Slot {
1088    Input(usize),
1089    Block(usize),
1090}
1091
1092/// Block buffer `d` for writing, plus read-only access to the others.
1093fn split_slots<'s, T>(
1094    scratch: &'s mut [T],
1095    w: usize,
1096    d: usize,
1097) -> (&'s mut [T], impl Fn(usize) -> &'s [T]) {
1098    let (lo, hi) = scratch.split_at_mut(d * w);
1099    let (dst, hi) = hi.split_at_mut(w);
1100    let lo: &[T] = lo;
1101    let hi: &[T] = hi;
1102    (dst, move |i: usize| {
1103        if i < d {
1104            &lo[i * w..(i + 1) * w]
1105        } else {
1106            &hi[(i - d - 1) * w..(i - d) * w]
1107        }
1108    })
1109}
1110
1111/// Run the kernel's map over elements `start .. start + len`.
1112///
1113/// `out`, when given, receives the last instruction's result directly and
1114/// the returned index means nothing; otherwise the result stays in the
1115/// block buffer that index names. None means an integer step left i64 and
1116/// the caller must fall back.
1117#[allow(clippy::too_many_arguments)]
1118fn exec_block<T, M, D>(
1119    code: &[Instr],
1120    srcs: &[Loaded<'_, T>],
1121    start: usize,
1122    len: usize,
1123    scratch: &mut [T],
1124    w: usize,
1125    free: &mut Vec<usize>,
1126    stack: &mut Vec<Slot>,
1127    lets: &mut Vec<usize>,
1128    out: Option<&mut [T]>,
1129    mon: &M,
1130    dya: &D,
1131) -> Option<usize>
1132where
1133    T: Copy,
1134    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1135    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1136{
1137    stack.clear();
1138    free.clear();
1139    lets.clear();
1140    let nslots = scratch.len() / w;
1141    free.extend((0..nslots).rev());
1142    let last = code.len() - 1;
1143    let head = if out.is_some() { last } else { code.len() };
1144    for ins in &code[..head] {
1145        match ins {
1146            Instr::Load(k) => stack.push(Slot::Input(*k)),
1147            Instr::Monad(op) => {
1148                let a = stack.pop()?;
1149                let d = free.pop()?;
1150                let (dst, get) = split_slots(scratch, w, d);
1151                let av = match a {
1152                    Slot::Input(k) => srcs[k].block(start, len),
1153                    Slot::Block(i) => &get(i)[..len],
1154                };
1155                if !mon(*op, av, &mut dst[..len]) {
1156                    return None;
1157                }
1158                release(free, lets, a);
1159                stack.push(Slot::Block(d));
1160            }
1161            Instr::Dyad(op) => {
1162                let b = stack.pop()?;
1163                let a = stack.pop()?;
1164                let d = free.pop()?;
1165                let (dst, get) = split_slots(scratch, w, d);
1166                let av = match a {
1167                    Slot::Input(k) => srcs[k].block(start, len),
1168                    Slot::Block(i) => &get(i)[..len],
1169                };
1170                let bv = match b {
1171                    Slot::Input(k) => srcs[k].block(start, len),
1172                    Slot::Block(i) => &get(i)[..len],
1173                };
1174                if !dya(*op, av, bv, &mut dst[..len]) {
1175                    return None;
1176                }
1177                for s in [a, b] {
1178                    release(free, lets, s);
1179                }
1180                stack.push(Slot::Block(d));
1181            }
1182            Instr::Store(k) => {
1183                let Slot::Block(i) = stack.pop()? else { return None };
1184                if lets.len() != *k {
1185                    return None;
1186                }
1187                lets.push(i);
1188            }
1189            Instr::Let(k) => stack.push(Slot::Block(*lets.get(*k)?)),
1190        }
1191    }
1192    let Some(dst) = out else {
1193        return match stack.pop()? {
1194            Slot::Block(i) => Some(i),
1195            // Every kernel ends in an operation, so the result is a buffer.
1196            Slot::Input(_) => None,
1197        };
1198    };
1199    // The last instruction writes the caller's buffer instead of a block.
1200    let dst = &mut dst[..len];
1201    let view = |s: Slot| match s {
1202        Slot::Input(k) => srcs[k].block(start, len),
1203        Slot::Block(i) => &scratch[i * w..i * w + len],
1204    };
1205    let ok = match code[last] {
1206        Instr::Monad(op) => {
1207            let a = view(stack.pop()?);
1208            mon(op, a, dst)
1209        }
1210        Instr::Dyad(op) => {
1211            let b = stack.pop()?;
1212            let a = stack.pop()?;
1213            dya(op, view(a), view(b), dst)
1214        }
1215        // A kernel ends in the operation that makes its result.
1216        Instr::Load(_) | Instr::Store(_) | Instr::Let(_) => return None,
1217    };
1218    ok.then_some(usize::MAX)
1219}
1220
1221/// Give a block buffer back, unless a let is holding it for the rest of
1222/// the block.
1223fn release(free: &mut Vec<usize>, lets: &[usize], s: Slot) {
1224    if let Slot::Block(i) = s {
1225        if !lets.contains(&i) {
1226            free.push(i);
1227        }
1228    }
1229}
1230
1231/// The whole mapped result, one block at a time. None on integer overflow.
1232fn map_pass<T, M, D>(
1233    k: &FusedKernel,
1234    srcs: &[Loaded<'_, T>],
1235    n: usize,
1236    mon: M,
1237    dya: D,
1238) -> Option<Vec<T>>
1239where
1240    T: Copy + Default + Send + Sync,
1241    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1242    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1243{
1244    let (out, ok) = par::fill(n, |start, part: &mut [T]| {
1245        let w = BLOCK.min(part.len()).max(1);
1246        let mut scratch = vec![T::default(); k.slots * w];
1247        let mut free = Vec::with_capacity(k.slots);
1248        let mut stack = Vec::with_capacity(k.slots);
1249        let mut lets = Vec::new();
1250        for (b, chunk) in part.chunks_mut(w).enumerate() {
1251            let len = chunk.len();
1252            let ok = exec_block(
1253                &k.code,
1254                srcs,
1255                start + b * w,
1256                len,
1257                &mut scratch,
1258                w,
1259                &mut free,
1260                &mut stack,
1261                &mut lets,
1262                Some(chunk),
1263                &mon,
1264                &dya,
1265            );
1266            if ok.is_none() {
1267                return false;
1268            }
1269        }
1270        true
1271    });
1272    ok.then_some(out)
1273}
1274
1275/// Independent accumulators the fold over a block keeps in flight, and the
1276/// block length below which one accumulator is cheaper. The reasoning is
1277/// the one `verb::FOLD_LANES` carries: a single accumulator makes the fold
1278/// a chain of dependent steps, and only an associative step is ever
1279/// absorbed here, so the lanes are a regrouping the float contract already
1280/// allows (§5.9).
1281const FOLD_LANES: usize = 8;
1282const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
1283
1284/// Fold one block of mapped values right to left, in lanes. None when a
1285/// step left the element type.
1286#[inline(always)]
1287fn fold_block_body<T, S>(v: &[T], step: &S) -> Option<T>
1288where
1289    T: Copy,
1290    S: Fn(T, T) -> Option<T>,
1291{
1292    let n = v.len();
1293    if n < MIN_LANE_WORK {
1294        let mut acc = v[n - 1];
1295        for &x in v[..n - 1].iter().rev() {
1296            acc = step(x, acc)?;
1297        }
1298        return Some(acc);
1299    }
1300    let rows = n / FOLD_LANES;
1301    let head = n - rows * FOLD_LANES;
1302    let last = head + (rows - 1) * FOLD_LANES;
1303    let mut acc = [v[last]; FOLD_LANES];
1304    acc.copy_from_slice(&v[last..last + FOLD_LANES]);
1305    for r in (0..rows - 1).rev() {
1306        let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
1307        for (slot, &x) in acc.iter_mut().zip(row) {
1308            *slot = step(x, *slot)?;
1309        }
1310    }
1311    let mut a = acc[FOLD_LANES - 1];
1312    for &x in acc[..FOLD_LANES - 1].iter().rev() {
1313        a = step(x, a)?;
1314    }
1315    for &x in v[..head].iter().rev() {
1316        a = step(x, a)?;
1317    }
1318    Some(a)
1319}
1320
1321multiversioned! {
1322    /// One block's values folded into one, at the CPU's own width.
1323    fn fold_block[T: Copy, S: Fn(T, T) -> Option<T>](
1324        v: &[T],
1325        step: &S,
1326    ) -> Option<T> = fold_block_body;
1327}
1328
1329/// Fold the mapped values of `lo .. hi` right to left, block by block.
1330#[allow(clippy::too_many_arguments)]
1331fn fold_range<T, M, D, S>(
1332    k: &FusedKernel,
1333    srcs: &[Loaded<'_, T>],
1334    lo: usize,
1335    hi: usize,
1336    mon: &M,
1337    dya: &D,
1338    step: &S,
1339) -> Option<T>
1340where
1341    T: Copy + Default,
1342    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1343    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1344    S: Fn(T, T) -> Option<T>,
1345{
1346    let w = BLOCK.min(hi - lo).max(1);
1347    let mut scratch = vec![T::default(); k.slots * w];
1348    let mut free = Vec::with_capacity(k.slots);
1349    let mut stack = Vec::with_capacity(k.slots);
1350    let mut lets = Vec::new();
1351    let mut acc: Option<T> = None;
1352    // Blocks run backwards and the accumulator carries across them, so the
1353    // fold is the insert's own right-to-left order over the whole range.
1354    for b in (0..(hi - lo).div_ceil(w)).rev() {
1355        let start = lo + b * w;
1356        let len = (hi - start).min(w);
1357        let slot = exec_block(
1358            &k.code, srcs, start, len, &mut scratch, w, &mut free, &mut stack, &mut lets, None,
1359            mon, dya,
1360        )?;
1361        let block = fold_block(&scratch[slot * w..slot * w + len], step)?;
1362        acc = Some(match acc {
1363            None => block,
1364            Some(a) => step(block, a)?,
1365        });
1366    }
1367    acc
1368}
1369
1370/// The mapped values folded into one. None on integer overflow.
1371fn reduce_pass<T, M, D, S>(
1372    k: &FusedKernel,
1373    srcs: &[Loaded<'_, T>],
1374    n: usize,
1375    mon: M,
1376    dya: D,
1377    step: S,
1378) -> Option<T>
1379where
1380    T: Copy + Default + Send + Sync,
1381    M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1382    D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1383    S: Fn(T, T) -> Option<T> + Sync + Send,
1384{
1385    let chunks = par::chunks(n, n * k.code.len());
1386    if chunks < 2 {
1387        return fold_range(k, srcs, 0, n, &mon, &dya, &step);
1388    }
1389    let per = n.div_ceil(chunks);
1390    let parts = par::map_indexed(n.div_ceil(per), |c| {
1391        fold_range(k, srcs, c * per, ((c + 1) * per).min(n), &mon, &dya, &step)
1392    });
1393    // The chunks combine right to left, the order they were folded in. That
1394    // regroups an associative float fold, which is the §5.9 contract; only
1395    // associative operations are absorbed.
1396    let mut it = parts.into_iter().rev();
1397    let mut acc = it.next()??;
1398    for part in it {
1399        acc = step(part?, acc)?;
1400    }
1401    Some(acc)
1402}
1403
1404// ------------------------------------------------------------ the kernels
1405//
1406// Each pass picks its operation before the loop and then runs one plain
1407// loop over slices, which is the shape the compiler vectorises. Nothing in
1408// here is hand-written SIMD, and nothing may become it.
1409//
1410// These four are the whole arithmetic of a kernel, so they are also where
1411// the CPU feature levels are chosen: each is compiled once per level (see
1412// `simd`) and the call dispatches on what the machine runs. One block of
1413// one instruction is thousands of elements, so the dispatch costs nothing
1414// measurable.
1415
1416macro_rules! each {
1417    ($a:expr, $dst:expr, $f:expr) => {{
1418        let f = $f;
1419        for (slot, &x) in $dst.iter_mut().zip($a) {
1420            *slot = f(x);
1421        }
1422        return true;
1423    }};
1424}
1425
1426macro_rules! zip {
1427    ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
1428        let f = $f;
1429        for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
1430            *slot = f(x, y);
1431        }
1432        return true;
1433    }};
1434}
1435
1436#[inline(always)]
1437fn monad_f64_body(op: ScalarMonad, a: &[f64], dst: &mut [f64], tol: Tol) -> bool {
1438    use ScalarMonad::*;
1439    match op {
1440        Conj => each!(a, dst, |x: f64| x),
1441        Neg => each!(a, dst, |x: f64| -x),
1442        Abs => each!(a, dst, f64::abs),
1443        // A magnitude the dialect's tolerance reads as zero has no sign,
1444        // exactly as unfused.
1445        Signum => each!(a, dst, |x: f64| if tol.is_zero(x) {
1446            0.0
1447        } else if x > 0.0 {
1448            1.0
1449        } else if x < 0.0 {
1450            -1.0
1451        } else {
1452            0.0
1453        }),
1454        // `% 0` is infinity, the J rule the unfused monad follows.
1455        Recip => each!(a, dst, |x: f64| if x == 0.0 { f64::INFINITY } else { 1.0 / x }),
1456        // Reached only through an integer chain, where they are the
1457        // identity: rounding a float narrows its dtype, which is declined.
1458        Floor => each!(a, dst, f64::floor),
1459        Ceil => each!(a, dst, f64::ceil),
1460        Inc => each!(a, dst, |x: f64| x + 1.0),
1461        Dec => each!(a, dst, |x: f64| x - 1.0),
1462        Double => each!(a, dst, |x: f64| x + x),
1463        Halve => each!(a, dst, |x: f64| x / 2.0),
1464        Square => each!(a, dst, |x: f64| x * x),
1465        OneMinus => each!(a, dst, |x: f64| 1.0 - x),
1466        Exp => each!(a, dst, f64::exp),
1467        _ => false,
1468    }
1469}
1470
1471#[inline(always)]
1472fn dyad_f64_body(op: ScalarDyad, a: &[f64], b: &[f64], dst: &mut [f64], tol: Tol) -> bool {
1473    use ScalarDyad::*;
1474    match op {
1475        Add => zip!(a, b, dst, |x: f64, y: f64| x + y),
1476        Sub => zip!(a, b, dst, |x: f64, y: f64| x - y),
1477        Mul => zip!(a, b, dst, |x: f64, y: f64| x * y),
1478        Min => zip!(a, b, dst, f64::min),
1479        Max => zip!(a, b, dst, f64::max),
1480        DivJ => zip!(a, b, dst, |x: f64, y: f64| if y == 0.0 {
1481            if x == 0.0 { 0.0 } else { f64::INFINITY.copysign(x) }
1482        } else {
1483            x / y
1484        }),
1485        // An infinite modulus leaves a value of its own sign alone and
1486        // sends the other one to that infinity, exactly as unfused.
1487        Residue => zip!(a, b, dst, |x: f64, y: f64| if x.is_infinite() {
1488            if y == 0.0 || (y > 0.0) == (x > 0.0) { y } else { x }
1489        } else if x == 0.0 {
1490            y
1491        } else {
1492            y - x * (y / x).floor()
1493        }),
1494        // A comparison is a number here, as it is in J: the boolean only
1495        // shows in the dtype of a result, which the caller narrows. Floats
1496        // compare with the dialect's tolerance, as they do unfused.
1497        Eq | Ne | Lt | Le | Gt | Ge => {
1498            zip!(a, b, dst, |x: f64, y: f64| tol_cmp(op, x, y, tol) as u8 as f64)
1499        }
1500        _ => false,
1501    }
1502}
1503
1504/// Integer passes fold overflow into a flag instead of branching out of the
1505/// loop: the whole evaluation is thrown away and redone unfused either way.
1506macro_rules! each_over {
1507    ($a:expr, $dst:expr, $f:expr) => {{
1508        let f = $f;
1509        let mut over = false;
1510        for (slot, &x) in $dst.iter_mut().zip($a) {
1511            let (v, o) = f(x);
1512            *slot = v;
1513            over |= o;
1514        }
1515        return !over;
1516    }};
1517}
1518
1519macro_rules! zip_over {
1520    ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
1521        let f = $f;
1522        let mut over = false;
1523        for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
1524            let (v, o) = f(x, y);
1525            *slot = v;
1526            over |= o;
1527        }
1528        return !over;
1529    }};
1530}
1531
1532#[inline(always)]
1533fn monad_i64_body(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool {
1534    use ScalarMonad::*;
1535    match op {
1536        Conj | Floor | Ceil => each!(a, dst, |x: i64| x),
1537        Neg => each_over!(a, dst, i64::overflowing_neg),
1538        Abs => each_over!(a, dst, i64::overflowing_abs),
1539        Signum => each!(a, dst, i64::signum),
1540        Inc => each_over!(a, dst, |x: i64| x.overflowing_add(1)),
1541        Dec => each_over!(a, dst, |x: i64| x.overflowing_sub(1)),
1542        Double => each_over!(a, dst, |x: i64| x.overflowing_add(x)),
1543        Square => each_over!(a, dst, |x: i64| x.overflowing_mul(x)),
1544        OneMinus => each_over!(a, dst, |x: i64| 1i64.overflowing_sub(x)),
1545        _ => false,
1546    }
1547}
1548
1549#[inline(always)]
1550fn dyad_i64_body(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool {
1551    use ScalarDyad::*;
1552    match op {
1553        Add => zip_over!(a, b, dst, i64::overflowing_add),
1554        Sub => zip_over!(a, b, dst, i64::overflowing_sub),
1555        Mul => zip_over!(a, b, dst, i64::overflowing_mul),
1556        Min => zip!(a, b, dst, i64::min),
1557        Max => zip!(a, b, dst, i64::max),
1558        Residue => zip!(a, b, dst, |x: i64, y: i64| if x == 0 {
1559            y
1560        } else {
1561            // wrapping_rem: i64::MIN % -1 is mathematically 0.
1562            let mut r = y.wrapping_rem(x);
1563            if r != 0 && (r < 0) != (x < 0) {
1564                r += x;
1565            }
1566            r
1567        }),
1568        Eq => zip!(a, b, dst, |x: i64, y: i64| (x == y) as i64),
1569        Ne => zip!(a, b, dst, |x: i64, y: i64| (x != y) as i64),
1570        Lt => zip!(a, b, dst, |x: i64, y: i64| (x < y) as i64),
1571        Le => zip!(a, b, dst, |x: i64, y: i64| (x <= y) as i64),
1572        Gt => zip!(a, b, dst, |x: i64, y: i64| (x > y) as i64),
1573        Ge => zip!(a, b, dst, |x: i64, y: i64| (x >= y) as i64),
1574        _ => false,
1575    }
1576}
1577
1578multiversioned! {
1579    /// One instruction of a kernel over one block of floats: the monadic
1580    /// operations. False is unreachable — every operation a kernel holds is
1581    /// covered — and exists so the two passes have one signature.
1582    fn monad_f64(
1583        op: ScalarMonad,
1584        a: &[f64],
1585        dst: &mut [f64],
1586        tol: Tol,
1587    ) -> bool = monad_f64_body;
1588}
1589
1590multiversioned! {
1591    /// One instruction of a kernel over one block of floats: the dyadic
1592    /// operations.
1593    fn dyad_f64(
1594        op: ScalarDyad,
1595        a: &[f64],
1596        b: &[f64],
1597        dst: &mut [f64],
1598        tol: Tol,
1599    ) -> bool = dyad_f64_body;
1600}
1601
1602multiversioned! {
1603    /// One instruction of a kernel over one block of integers: the monadic
1604    /// operations. False means the block left i64.
1605    fn monad_i64(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool = monad_i64_body;
1606}
1607
1608multiversioned! {
1609    /// One instruction of a kernel over one block of integers: the dyadic
1610    /// operations. False means the block left i64.
1611    fn dyad_i64(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool = dyad_i64_body;
1612}
1613
1614/// One fold step of an absorbed reduction. None on integer overflow.
1615fn step_i64(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
1616    use ScalarDyad::*;
1617    match op {
1618        Add => a.checked_add(b),
1619        Mul => a.checked_mul(b),
1620        Min => Some(a.min(b)),
1621        Max => Some(a.max(b)),
1622        _ => None,
1623    }
1624}
1625
1626/// One fold step of an absorbed float reduction, for a backend that mapped
1627/// the values elsewhere and brings its partials back here to combine.
1628pub(crate) fn step(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
1629    step_f64(op, a, b)
1630}
1631
1632fn step_f64(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
1633    use ScalarDyad::*;
1634    match op {
1635        Add => Some(a + b),
1636        Mul => Some(a * b),
1637        Min => Some(a.min(b)),
1638        Max => Some(a.max(b)),
1639        _ => None,
1640    }
1641}
1642
1643// ------------------------------------------------------------- the driver
1644
1645/// Elements of `a` as the working type, or None when the array's own buffer
1646/// already is that. A rank-0 argument becomes one block of the repeated
1647/// value, which is how it reaches every element without an index test.
1648fn to_f64(a: &Array, w: usize) -> Option<Vec<f64>> {
1649    if a.rank() == 0 {
1650        let v = match &a.data {
1651            Data::Bool(d) => d[0] as f64,
1652            Data::I64(d) => d[0] as f64,
1653            Data::F64(d) => d[0],
1654            Data::Ext(_) | Data::Rat(_) | Data::Complex(_) | Data::Char(_) | Data::Box(_) => {
1655                return Some(Vec::new());
1656            }
1657        };
1658        return Some(vec![v; w]);
1659    }
1660    match &a.data {
1661        Data::F64(_) => None,
1662        Data::I64(d) => Some(par::map(d, |&x| x as f64)),
1663        Data::Bool(d) => Some(par::map(d, |&x| x as f64)),
1664        Data::Ext(_) | Data::Rat(_) | Data::Complex(_) | Data::Char(_) | Data::Box(_) => {
1665            Some(Vec::new())
1666        }
1667    }
1668}
1669
1670fn to_i64(a: &Array, w: usize) -> Option<Vec<i64>> {
1671    if a.rank() == 0 {
1672        let v = match &a.data {
1673            Data::Bool(d) => d[0] as i64,
1674            Data::I64(d) => d[0],
1675            _ => return Some(Vec::new()),
1676        };
1677        return Some(vec![v; w]);
1678    }
1679    match &a.data {
1680        Data::I64(_) => None,
1681        Data::Bool(d) => Some(par::map(d, |&x| x as i64)),
1682        // The working type is integer only when no input is a float.
1683        _ => Some(Vec::new()),
1684    }
1685}
1686
1687/// The shape every element of the result has: identical for all non-scalar
1688/// inputs, since anything else needs the agreement machinery.
1689pub(crate) fn common_shape(inputs: &[Array]) -> Option<Option<Vec<usize>>> {
1690    let mut shape: Option<&Vec<usize>> = None;
1691    for a in inputs {
1692        if a.rank() == 0 {
1693            continue;
1694        }
1695        match shape {
1696            None => shape = Some(&a.shape),
1697            Some(s) if *s == a.shape => {}
1698            Some(_) => return None,
1699        }
1700    }
1701    Some(shape.cloned())
1702}
1703
1704/// Run a fused node. None means the kernel declined and the caller must
1705/// evaluate the original subtree, which is always allowed to be slower and
1706/// never allowed to differ.
1707pub(crate) fn run(k: &FusedKernel, inputs: &[Array]) -> Option<Array> {
1708    let reducing = matches!(k.yields, Yield::Reduce(_));
1709    // Every input a scalar: no work worth blocking, and a reduction would
1710    // need the leading axis a scalar has not got.
1711    let shape = common_shape(inputs)??;
1712    let n: usize = shape.iter().product();
1713    if n == 0 {
1714        return None;
1715    }
1716    if reducing && (shape.len() != 1 || n < 2) {
1717        // A one-item reduction yields the item itself, dtype and all, and a
1718        // higher-rank one folds cells rather than elements.
1719        return None;
1720    }
1721    let (working, root) = working_type(k, inputs)?;
1722    if k.yields == Yield::Tally {
1723        // The shapes have already said how many items the chain produces,
1724        // and the type rules have said it would reach them without an
1725        // error. There is nothing else a tally wants from the values.
1726        return Some(Array::scalar_i64(shape[0] as i64));
1727    }
1728    let w = BLOCK.min(n).max(1);
1729    // The kernel's comparisons carry the tolerance the program was compiled
1730    // with, so a fused comparison answers as the unfused one does.
1731    let tol = k.tol;
1732    let cmp_f64 = move |op, a: &[f64], b: &[f64], dst: &mut [f64]| dyad_f64(op, a, b, dst, tol);
1733    let sign_f64 = move |op, a: &[f64], dst: &mut [f64]| monad_f64(op, a, dst, tol);
1734
1735    let data = if working == DType::F64 {
1736        let owned: Vec<Option<Vec<f64>>> = inputs.iter().map(|a| to_f64(a, w)).collect();
1737        let srcs: Vec<Loaded<f64>> = inputs
1738            .iter()
1739            .zip(&owned)
1740            .map(|(a, o)| match o {
1741                Some(v) => Loaded { data: v, splat: a.rank() == 0 },
1742                None => Loaded { data: a.as_f64_slice().unwrap_or(&[]), splat: false },
1743            })
1744            .collect();
1745        match k.reduce() {
1746            None => {
1747                let out = map_pass(k, &srcs, n, sign_f64, cmp_f64)?;
1748                float_result(out, root)
1749            }
1750            Some(op) => {
1751                let v = reduce_pass(k, &srcs, n, sign_f64, cmp_f64, |a, b| step_f64(op, a, b))?;
1752                // A comparison at the root maps to exact 0 and 1, which the
1753                // fold keeps exact; the reduction of booleans is integer.
1754                match root {
1755                    DType::F64 => Data::F64(vec![v].into()),
1756                    _ => Data::I64(vec![v as i64].into()),
1757                }
1758            }
1759        }
1760    } else {
1761        let owned: Vec<Option<Vec<i64>>> = inputs.iter().map(|a| to_i64(a, w)).collect();
1762        let srcs: Vec<Loaded<i64>> = inputs
1763            .iter()
1764            .zip(&owned)
1765            .map(|(a, o)| match o {
1766                Some(v) => Loaded { data: v, splat: a.rank() == 0 },
1767                None => Loaded { data: a.as_i64_slice().unwrap_or(&[]), splat: false },
1768            })
1769            .collect();
1770        match k.reduce() {
1771            None => {
1772                let out = map_pass(k, &srcs, n, monad_i64, dyad_i64)?;
1773                int_result(out, root)
1774            }
1775            Some(op) => {
1776                let v =
1777                    reduce_pass(k, &srcs, n, monad_i64, dyad_i64, |a, b| step_i64(op, a, b))?;
1778                Data::I64(vec![v].into())
1779            }
1780        }
1781    };
1782    Some(Array::new(if reducing { Vec::new() } else { shape }, data))
1783}
1784
1785/// The mapped block values as the array the unfused chain would build. A
1786/// comparison at the root costs one narrowing pass, since the kernel
1787/// computes 0 and 1 in its working type and a boolean array holds bytes.
1788fn float_result(out: Vec<f64>, root: DType) -> Data {
1789    match root {
1790        DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0.0) as u8).into()),
1791        _ => Data::F64(out.into()),
1792    }
1793}
1794
1795fn int_result(out: Vec<i64>, root: DType) -> Data {
1796    match root {
1797        DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0) as u8).into()),
1798        _ => Data::I64(out.into()),
1799    }
1800}
1801
1802// -------------------------------------------------------- describing one
1803//
1804// Read-only descriptions of a compiled kernel, for `Program::explain`.
1805// Nothing here runs a kernel or changes one; the summary is derived from
1806// the code the pass emitted, and the decline reason re-checks the same
1807// preconditions `run` checks before it starts.
1808
1809/// Why a kernel handed its work back to the chain it came from.
1810#[derive(Clone, Copy, Debug, PartialEq, Eq)]
1811pub enum Decline {
1812    /// Inputs disagree on shape, or every input is a scalar: broadcasting
1813    /// and agreement are the chain's business.
1814    Agreement,
1815    /// Nothing to compute.
1816    Empty,
1817    /// An absorbed reduction wants one axis with at least two items.
1818    ReduceShape,
1819    /// One working type cannot hold every step exactly — a chain that
1820    /// computes integers along a float path, or non-numeric data.
1821    WorkingType,
1822    /// The preconditions held, so a step went out of range mid-block:
1823    /// integer overflow, which the chain redoes in a wider type.
1824    Overflow,
1825}
1826
1827impl Decline {
1828    pub fn reason(self) -> &'static str {
1829        match self {
1830            Decline::Agreement => "the inputs need agreement or are all scalars",
1831            Decline::Empty => "there is nothing to compute",
1832            Decline::ReduceShape => "the reduction needs one axis of two or more items",
1833            Decline::WorkingType => "no single working type holds every step exactly",
1834            Decline::Overflow => "an integer step left 64-bit range",
1835        }
1836    }
1837}
1838
1839/// Why this kernel would decline these inputs, or None if it would run.
1840///
1841/// A read-only mirror of the preconditions at the top of `run`: it looks
1842/// at shapes and dtypes only, never at values, so the one thing it cannot
1843/// see in advance is an overflow — which is what is left when every
1844/// precondition holds.
1845pub fn decline_reason(k: &FusedKernel, inputs: &[Array]) -> Option<Decline> {
1846    let Some(Some(shape)) = common_shape(inputs) else {
1847        return Some(Decline::Agreement);
1848    };
1849    let n: usize = shape.iter().product();
1850    if n == 0 {
1851        return Some(Decline::Empty);
1852    }
1853    if matches!(k.yields, Yield::Reduce(_)) && (shape.len() != 1 || n < 2) {
1854        return Some(Decline::ReduceShape);
1855    }
1856    if working_type(k, inputs).is_none() {
1857        return Some(Decline::WorkingType);
1858    }
1859    Some(Decline::Overflow)
1860}
1861
1862/// What a compiled kernel is made of.
1863#[derive(Clone, Debug, PartialEq, Eq)]
1864pub struct Summary {
1865    /// Arithmetic steps: the monads and dyads, not the loads and stores.
1866    pub ops: usize,
1867    /// Those steps in the order the kernel performs them.
1868    pub op_names: Vec<&'static str>,
1869    /// The reduction folded into the same pass, if there is one.
1870    pub reduce: Option<&'static str>,
1871    /// True when the whole chain collapsed to a count of its own items.
1872    pub tally: bool,
1873    /// Values the kernel keeps for a second read within one block.
1874    pub lets: usize,
1875    /// Subtrees the chain reads.
1876    pub inputs: usize,
1877    /// Elements one block buffer holds.
1878    pub block: usize,
1879}
1880
1881impl std::fmt::Display for Summary {
1882    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
1883        write!(f, "{} op{}", self.ops, if self.ops == 1 { "" } else { "s" })?;
1884        if !self.op_names.is_empty() {
1885            write!(f, ": {}", self.op_names.join(" "))?;
1886        }
1887        if let Some(r) = self.reduce {
1888            write!(f, "; {r}/ absorbed")?;
1889        }
1890        if self.tally {
1891            write!(f, "; tally only")?;
1892        }
1893        if self.lets > 0 {
1894            write!(f, "; {} let slot{}", self.lets, if self.lets == 1 { "" } else { "s" })?;
1895        }
1896        write!(f, "; block {}", self.block)
1897    }
1898}
1899
1900/// Describe a compiled kernel: what it computes, and with what.
1901pub fn summary(k: &FusedKernel) -> Summary {
1902    let mut op_names = Vec::new();
1903    let mut lets = 0usize;
1904    for ins in &k.code {
1905        match ins {
1906            Instr::Monad(op) => op_names.push(monad_name(*op)),
1907            Instr::Dyad(op) => op_names.push(dyad_name(*op)),
1908            Instr::Store(_) => lets += 1,
1909            Instr::Load(_) | Instr::Let(_) => {}
1910        }
1911    }
1912    Summary {
1913        ops: op_names.len(),
1914        op_names,
1915        reduce: k.reduce().map(dyad_name),
1916        tally: k.yields == Yield::Tally,
1917        lets,
1918        inputs: k.leaves.iter().copied().max().map_or(0, |m| m + 1),
1919        block: BLOCK,
1920    }
1921}
1922
1923/// The names the pass took out of the program: values it moved into the
1924/// kernels that read them, so no sentence computes them as arrays any more.
1925pub fn inlined_names(p: &Program) -> Vec<String> {
1926    let Some(Expr::Elided { orig, .. }) = p.stmts.first() else { return Vec::new() };
1927    let assigned = |stmts: &[Expr]| -> Vec<String> {
1928        stmts
1929            .iter()
1930            .filter_map(|s| match s {
1931                Expr::Assign { name, .. } => Some(name.clone()),
1932                _ => None,
1933            })
1934            .collect()
1935    };
1936    let kept = assigned(&p.stmts);
1937    assigned(orig).into_iter().filter(|n| !kept.contains(n)).collect()
1938}
1939
1940/// J spellings for the elementwise operations a kernel can hold. Only the
1941/// naming lives here; the meanings are [`crate::verb`]'s.
1942fn monad_name(op: ScalarMonad) -> &'static str {
1943    use ScalarMonad::*;
1944    match op {
1945        Conj => "+",
1946        Neg => "-",
1947        Signum => "*",
1948        Recip => "%",
1949        Sqrt => "%:",
1950        Exp => "^",
1951        Abs => "|",
1952        Floor => "<.",
1953        Ceil => ">.",
1954        Not => "-.",
1955        OneMinus => "-.",
1956        Inc => ">:",
1957        Dec => "<:",
1958        Double => "+:",
1959        Halve => "-:",
1960        Square => "*:",
1961        Ln => "^.",
1962        Pi => "o.",
1963        Factorial => "!",
1964        Imaginary => "j.",
1965        Polar => "r.",
1966    }
1967}
1968
1969fn dyad_name(op: ScalarDyad) -> &'static str {
1970    use ScalarDyad::*;
1971    match op {
1972        Add => "+",
1973        Sub => "-",
1974        Mul => "*",
1975        DivJ | DivApl => "%",
1976        Min => "<.",
1977        Max => ">.",
1978        Pow => "^",
1979        Residue => "|",
1980        Eq => "=",
1981        Ne => "~:",
1982        Lt => "<",
1983        Le => "<:",
1984        Gt => ">",
1985        Ge => ">:",
1986        Lcm => "*.",
1987        Gcd => "+.",
1988        Log => "^.",
1989        Root => "%:",
1990        Circle => "o.",
1991        Binomial => "!",
1992        MakeComplex => "j.",
1993        PolarBy => "r.",
1994    }
1995}
1996
1997/// Evaluate a fused node from its already-evaluated inputs, or report that
1998/// the original subtree must run instead.
1999///
2000/// This is the one place a device gets to run libjay's arithmetic. With a
2001/// device attached the kernel is offered to it first; everything it will not
2002/// take comes back here with a reason, and the CPU path runs exactly as it
2003/// runs with no device in sight. The device therefore cannot change a
2004/// result's shape, dtype or error — only where the arithmetic happened.
2005pub(crate) fn eval_on(
2006    device: Option<&crate::device::Device>,
2007    k: &FusedKernel,
2008    inputs: &[Array],
2009) -> (Option<Array>, crate::device::Placement) {
2010    use crate::device::Placement;
2011    let mut placement = Placement::Default;
2012    if let Some(d) = device.filter(|d| d.is_gpu()) {
2013        match crate::device::try_run(d, k, inputs) {
2014            Ok(a) => return (Some(a), Placement::Gpu),
2015            Err(why) => placement = Placement::Cpu(why),
2016        }
2017    }
2018    let r = run(k, inputs);
2019    if r.is_none() {
2020        note_fallback();
2021    }
2022    (r, placement)
2023}
2024
2025#[cfg(test)]
2026mod tests {
2027    use super::*;
2028    use crate::frontend::{compile, Dialect, Lang};
2029
2030    fn program(src: &str) -> Program {
2031        compile(Lang::J, src, &Dialect::default()).expect("compile")
2032    }
2033
2034    #[test]
2035    fn a_chain_of_two_scalar_verbs_fuses() {
2036        assert!(is_fused(&program("1 + 2 * {x}")));
2037        assert!(is_fused(&program("+/ {w} * {x}")));
2038        assert!(is_fused(&program("+/ ^ {x}")));
2039    }
2040
2041    #[test]
2042    fn one_verb_on_its_own_is_left_alone() {
2043        assert!(!is_fused(&program("2 * {x}")));
2044        assert!(!is_fused(&program("+/ {x}")));
2045        assert!(!is_fused(&program("{x}")));
2046    }
2047
2048    #[test]
2049    fn a_verb_the_kernel_does_not_cover_breaks_the_chain() {
2050        // `%:` can fail elementwise, so it stays outside; the chain under it
2051        // still fuses.
2052        assert!(!is_fused(&program("%: 2 * {x}")));
2053        assert!(is_fused(&program("%: 1 + 2 * {x}")));
2054    }
2055
2056    #[test]
2057    fn an_effect_in_a_leaf_keeps_the_chain_unfused() {
2058        assert!(!is_fused(&program("1 + 2 * echo {x}")));
2059    }
2060
2061    #[test]
2062    fn the_postfix_program_pushes_the_left_operand_first() {
2063        let p = program("{w} - {x} - 1");
2064        let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2065        assert_eq!(
2066            kernel.code(),
2067            [
2068                Instr::Load(2),
2069                Instr::Load(1),
2070                Instr::Load(0),
2071                Instr::Dyad(ScalarDyad::Sub),
2072                Instr::Dyad(ScalarDyad::Sub),
2073            ]
2074        );
2075        // One buffer holds the inner difference, one takes the outer one.
2076        assert_eq!(kernel.slots, 2);
2077    }
2078
2079    #[test]
2080    fn a_value_the_chain_reads_twice_becomes_a_let() {
2081        // What `d =. {x} + 1` then `+/ d * d` comes to once the name has
2082        // moved into the kernel: the sum is computed once per block.
2083        let p = program("+/ ({x} + 1) * ({x} + 1)");
2084        let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2085        assert_eq!(
2086            kernel.code(),
2087            [
2088                Instr::Load(1),
2089                Instr::Load(0),
2090                Instr::Dyad(ScalarDyad::Add),
2091                Instr::Store(0),
2092                Instr::Let(0),
2093                Instr::Let(0),
2094                Instr::Dyad(ScalarDyad::Mul),
2095            ]
2096        );
2097        // One buffer for the let, one for the product it feeds.
2098        assert_eq!(kernel.slots, 2);
2099    }
2100
2101    #[test]
2102    fn a_named_value_moves_into_the_sentence_that_reads_it() {
2103        let p = program("d =. {x} + 1\n+/ d * d");
2104        assert!(is_inlined(&p));
2105        // Three sentences: what the program was, the check that stands
2106        // where the assignment stood, and the sum, which is now the chain
2107        // of the test above.
2108        assert_eq!(p.stmts.len(), 3);
2109        let Expr::Fused { kernel, .. } = &p.stmts[2] else { panic!("the sum did not fuse") };
2110        assert!(kernel.code().contains(&Instr::Store(0)));
2111        assert_eq!(unfused(&p).stmts.len(), 2);
2112    }
2113}