1use 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
35pub const BLOCK: usize = 8_192;
44
45pub const MAX_WINDOW: usize = 1_024;
53
54#[derive(Clone, Copy, Debug, PartialEq, Eq)]
56pub enum Instr {
57 Load(usize),
59 Monad(ScalarMonad),
61 Dyad(ScalarDyad),
63 Store(usize),
67 Let(usize),
69 Window(ScalarDyad, usize),
73 Scan(ScalarDyad),
76}
77
78#[derive(Clone, Copy, Debug, PartialEq, Eq)]
85enum Dom {
86 Result,
87 Wide,
88}
89
90#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
99struct Plan {
100 window: Option<usize>,
101 scan: bool,
102}
103
104#[derive(Clone, Copy, Debug, PartialEq, Eq)]
106pub enum Yield {
107 Values,
109 Reduce(ScalarDyad),
111 Tally,
114}
115
116#[derive(Clone, Debug)]
118pub struct FusedKernel {
119 code: Vec<Instr>,
120 slots: usize,
122 yields: Yield,
123 leaves: Vec<usize>,
128 doms: Vec<Option<Dom>>,
131 let_doms: Vec<Dom>,
134 plan: Plan,
137 window: Option<usize>,
139 scans: usize,
142 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 pub fn tol(&self) -> Tol {
167 self.tol
168 }
169}
170
171static FALLBACKS: AtomicU64 = AtomicU64::new(0);
175
176pub fn fallback_count() -> u64 {
178 FALLBACKS.load(Ordering::Relaxed)
179}
180
181fn note_fallback() {
182 FALLBACKS.fetch_add(1, Ordering::Relaxed);
183}
184
185fn 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
205fn 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
214fn absorbable_reduce(v: &Verb) -> Option<ScalarDyad> {
218 use ScalarDyad::*;
219 let inner = match v {
220 Verb::Reduce(u) => u,
221 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
234fn 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 (1..=MAX_WINDOW as i64).contains(&k).then_some((op, k as usize))
251}
252
253fn absorbable_scan(e: &Expr) -> Option<ScalarDyad> {
257 let Expr::Monad { verb, .. } = e else { return None };
258 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
271fn is_tally(v: &Verb) -> bool {
274 matches!(v, Verb::Prim(p) if p.monad == MonadOp::Tally && p.ranks[0] == RANK_INF)
275}
276
277#[derive(Clone, PartialEq)]
281enum Node {
282 Leaf(usize),
284 Monad(ScalarMonad, Box<Node>),
285 Dyad(ScalarDyad, Box<Node>, Box<Node>),
286 Window(ScalarDyad, usize, Box<Node>),
289 Scan(ScalarDyad, Box<Node>),
290}
291
292#[derive(Default)]
301struct Leaves<'a> {
302 inputs: Vec<&'a Expr>,
303 order: Vec<usize>,
305 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
328struct Inline<'a> {
331 name: &'a str,
332 def: &'a Expr,
333 hits: usize,
334}
335
336fn 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
344fn 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
384fn 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 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
437fn 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
455fn 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
500fn emit_all(n: &Node, lets: &[(Node, Dom)], code: &mut Vec<Instr>) {
503 for (k, (l, _)) in lets.iter().enumerate() {
504 emit(l, &lets[..k], code);
507 code.push(Instr::Store(k));
508 }
509 emit(n, lets, code);
510}
511
512fn 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
540fn 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 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
581fn 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
604fn 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
649pub 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 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 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
700fn 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 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
748fn 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
774fn 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
786pub(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 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 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
814pub(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
824fn 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
872fn hoisted_name(n: &mut usize) -> String {
888 *n += 1;
889 format!("·{}", *n - 1)
890}
891
892fn 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
945fn 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
972fn 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 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
1024fn 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
1107fn 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
1129fn 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
1151pub 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
1173pub fn is_inlined(p: &Program) -> bool {
1175 matches!(p.stmts.first(), Some(Expr::Elided { .. }))
1176}
1177
1178pub 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 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
1219fn monad_type(op: ScalarMonad, a: DType) -> Option<DType> {
1225 use DType::*;
1226 use ScalarMonad::*;
1227 if a == Complex {
1230 return None;
1231 }
1232 Some(match op {
1233 Recip | Halve | Exp => F64,
1234 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
1248fn 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
1267fn 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
1282pub(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 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 Instr::Let(k) => {
1332 let t = *lets.get(*k)?;
1333 float |= t == DType::F64;
1334 stack.push(t);
1335 continue;
1336 }
1337 };
1338 if !t.is_numeric() {
1341 return None;
1342 }
1343 float |= t == DType::F64;
1344 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#[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#[derive(Clone, Copy)]
1387struct Extent {
1388 start: usize,
1389 len: usize,
1390 wide_start: usize,
1391 wide_len: usize,
1392}
1393
1394impl Extent {
1395 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#[derive(Clone, Copy)]
1425struct Loaded<'a, T> {
1426 data: &'a [T],
1427 splat: bool,
1428 on: On,
1430 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#[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#[derive(Clone, Copy)]
1466enum Narrow<'a> {
1467 I64(&'a [i64]),
1468 Bool(&'a [u8]),
1469}
1470
1471enum Source<'a, T> {
1474 Ready(Loaded<'a, T>),
1475 Staged(Narrow<'a>, On),
1477}
1478
1479trait 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 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
1517struct Sources<'a, T> {
1519 of: Vec<Source<'a, T>>,
1520 staged: usize,
1521}
1522
1523impl<T: FromNarrow + Default> Sources<'_, T> {
1524 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#[derive(Clone, Copy)]
1564enum Slot {
1565 Input(usize),
1566 Block(usize, On),
1567}
1568
1569struct Scratch<T> {
1571 cells: Vec<T>,
1572 width: usize,
1575 free: Vec<usize>,
1576 stack: Vec<Slot>,
1577 lets: Vec<usize>,
1578 carry: Vec<Option<T>>,
1581}
1582
1583impl<T: Copy + Default> Scratch<T> {
1584 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
1602struct Steps<M, D, W, S> {
1605 monad: M,
1606 dyad: D,
1607 window: W,
1608 scan: S,
1609}
1610
1611fn 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
1630fn 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 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 Slot::Input(_) => None,
1756 };
1757 };
1758 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 Instr::Load(_) | Instr::Store(_) | Instr::Let(_) => return None,
1794 };
1795 ok.then_some(usize::MAX)
1796}
1797
1798fn 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
1808fn 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 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
1850const FOLD_LANES: usize = 8;
1857const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
1858
1859#[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 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
1904fn 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 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
1948fn 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 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
1984macro_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 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 Recip => each!(a, dst, |x: f64| if x == 0.0 { f64::INFINITY } else { 1.0 / x }),
2036 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 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 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
2084macro_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 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 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 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 fn monad_i64(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool = monad_i64_body;
2186}
2187
2188multiversioned! {
2189 fn dyad_i64(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool = dyad_i64_body;
2192}
2193
2194#[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 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 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
2235fn 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
2281fn 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
2293pub(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
2310fn 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
2347fn 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
2364fn 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
2391pub(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
2408struct Axes {
2420 shape: Vec<usize>,
2421 wide: usize,
2423}
2424
2425fn axes(k: &FusedKernel, inputs: &[Array]) -> Option<Axes> {
2426 let Some(window) = k.window else {
2427 let shape = common_shape(inputs)??;
2430 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 if a.rank() == 0 {
2443 continue;
2444 }
2445 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 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
2468pub(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 return None;
2482 }
2483 let (working, root) = working_type(k, inputs)?;
2484 if k.yields == Yield::Tally {
2485 return Some(Array::scalar_i64(shape[0] as i64));
2489 }
2490 let w = BLOCK.min(n).max(1) + 3 * k.window.unwrap_or(0);
2493 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 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
2545fn 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#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2571pub enum Decline {
2572 Agreement,
2575 Empty,
2577 ReduceShape,
2579 WorkingType,
2582 Overflow,
2585 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
2603pub 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#[derive(Clone, Debug, PartialEq, Eq)]
2628pub struct Summary {
2629 pub ops: usize,
2631 pub op_names: Vec<String>,
2633 pub reduce: Option<&'static str>,
2635 pub tally: bool,
2637 pub lets: usize,
2639 pub inputs: usize,
2641 pub block: usize,
2643 pub window: Option<usize>,
2645 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
2674pub 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
2701pub 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
2718fn 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
2775pub(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 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 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
2817fn kernel_layout(inputs: &[Array]) -> Option<Layout> {
2820 let mut found: Option<Layout> = None;
2821 for a in inputs {
2822 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 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 assert_eq!(kernel.slots, 2);
2888 }
2889
2890 #[test]
2891 fn a_value_the_chain_reads_twice_becomes_a_let() {
2892 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 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 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}