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) | Verb::NWise(u) => u,
222 Verb::Rank(u, r) if r[0] >= 1 => match &**u {
225 Verb::Reduce(inner) | Verb::NWise(inner) => inner,
226 _ => return None,
227 },
228 _ => return None,
229 };
230 let Verb::Prim(p) = &**inner else { return None };
231 let DyadOp::Scalar(op) = p.dyad else { return None };
232 matches!(op, Add | Mul | Min | Max).then_some(op)
233}
234
235fn absorbable_window(e: &Expr) -> Option<(ScalarDyad, usize)> {
242 let (op, x) = match e {
243 Expr::Dyad { verb: Verb::Windowed(u, WindowKind::Prefix), x, .. } => {
244 (absorbable_reduce(u)?, x)
245 }
246 Expr::Dyad { verb: v @ (Verb::NWise(_) | Verb::Rank(..)), x, .. }
249 if is_nwise(v) =>
250 {
251 (absorbable_reduce(v)?, x)
252 }
253 _ => return None,
254 };
255 let Expr::Const(a, _) = &**x else { return None };
256 if a.rank() != 0 {
257 return None;
258 }
259 let k = *a.to_i64_vec()?.first()?;
260 (1..=MAX_WINDOW as i64).contains(&k).then_some((op, k as usize))
264}
265
266fn is_nwise(v: &Verb) -> bool {
269 match v {
270 Verb::NWise(_) => true,
271 Verb::Rank(u, r) => r[0] >= 1 && matches!(&**u, Verb::NWise(_)),
272 _ => false,
273 }
274}
275
276fn absorbable_scan(e: &Expr) -> Option<ScalarDyad> {
280 let Expr::Monad { verb, .. } = e else { return None };
281 let inner = match verb {
284 Verb::Rank(u, r) if r[0] >= 1 => &**u,
285 v => v,
286 };
287 let Verb::Windowed(u, kind) = inner else { return None };
288 if *kind == WindowKind::Suffix {
289 return None;
290 }
291 absorbable_reduce(u)
292}
293
294fn is_tally(v: &Verb) -> bool {
297 matches!(v, Verb::Prim(p) if p.monad == MonadOp::Tally && p.ranks[0] == RANK_INF)
298}
299
300#[derive(Clone, PartialEq)]
304enum Node {
305 Leaf(usize),
307 Monad(ScalarMonad, Box<Node>),
308 Dyad(ScalarDyad, Box<Node>, Box<Node>),
309 Window(ScalarDyad, usize, Box<Node>),
312 Scan(ScalarDyad, Box<Node>),
313}
314
315#[derive(Default)]
324struct Leaves<'a> {
325 inputs: Vec<&'a Expr>,
326 order: Vec<usize>,
328 doms: Vec<Option<Dom>>,
331}
332
333impl<'a> Leaves<'a> {
334 fn push(&mut self, e: &'a Expr, dom: Dom) -> usize {
335 let i = match self.inputs.iter().position(|&p| same(p, e)) {
336 Some(i) => i,
337 None => {
338 self.inputs.push(e);
339 self.doms.push(Some(dom));
340 self.inputs.len() - 1
341 }
342 };
343 if self.doms[i] != Some(dom) {
344 self.doms[i] = None;
345 }
346 self.order.push(i);
347 i
348 }
349}
350
351struct Inline<'a> {
354 name: &'a str,
355 def: &'a Expr,
356 hits: usize,
357}
358
359fn read_through<'a>(e: &Expr, sub: Option<&Inline<'a>>) -> Option<&'a Expr> {
361 match (e, sub) {
362 (Expr::Name(n, _), Some(s)) if n == s.name => Some(s.def),
363 _ => None,
364 }
365}
366
367fn plan_of(e: &Expr, sub: Option<&Inline<'_>>) -> Plan {
375 fn walk(e: &Expr, sub: Option<&Inline<'_>>, inside: bool, ks: &mut Vec<usize>, s: &mut bool) {
376 if let Some(def) = read_through(e, sub) {
377 return walk(def, sub, inside, ks, s);
378 }
379 match e {
380 Expr::Monad { verb, y, .. } if fusable_monad(verb).is_some() => {
381 walk(y, sub, inside, ks, s)
382 }
383 Expr::Dyad { verb, x, y, .. } if fusable_dyad(verb).is_some() => {
384 walk(y, sub, inside, ks, s);
385 walk(x, sub, inside, ks, s);
386 }
387 Expr::Dyad { y, .. } if !inside && absorbable_window(e).is_some() => {
388 ks.push(absorbable_window(e).expect("just matched").1);
389 walk(y, sub, true, ks, s);
390 }
391 Expr::Monad { y, .. } if absorbable_scan(e).is_some() => {
392 *s = true;
393 walk(y, sub, inside, ks, s);
394 }
395 _ => {}
396 }
397 }
398 let (mut ks, mut scan) = (Vec::new(), false);
399 walk(e, sub, false, &mut ks, &mut scan);
400 let window = match ks.split_first() {
401 Some((k, rest)) if rest.iter().all(|r| r == k) => Some(*k),
402 _ => None,
403 };
404 Plan { window, scan: window.is_none() && scan }
405}
406
407fn chain<'a>(
409 e: &'a Expr,
410 lv: &mut Leaves<'a>,
411 sub: &mut Option<Inline<'a>>,
412 plan: Plan,
413 dom: Dom,
414) -> Node {
415 if read_through(e, sub.as_ref()).is_some() {
416 let def = read_through(e, sub.as_ref()).expect("just matched");
417 if let Some(s) = sub.as_mut() {
418 s.hits += 1;
419 }
420 return chain(def, lv, sub, plan, dom);
421 }
422 match e {
423 Expr::Monad { verb, y, .. } => {
424 if let Some(op) = fusable_monad(verb) {
425 return Node::Monad(op, Box::new(chain(y, lv, sub, plan, dom)));
426 }
427 if plan.scan && let Some(op) = absorbable_scan(e) {
428 return Node::Scan(op, Box::new(chain(y, lv, sub, plan, dom)));
429 }
430 Node::Leaf(lv.push(e, dom))
431 }
432 Expr::Dyad { verb, x, y, .. } => {
433 if let Some(op) = fusable_dyad(verb) {
434 let ry = chain(y, lv, sub, plan, dom);
435 let rx = chain(x, lv, sub, plan, dom);
436 return Node::Dyad(op, Box::new(rx), Box::new(ry));
437 }
438 if dom == Dom::Result
441 && let Some((op, k)) = absorbable_window(e)
442 && plan.window == Some(k)
443 {
444 return Node::Window(op, k, Box::new(chain(y, lv, sub, plan, Dom::Wide)));
445 }
446 Node::Leaf(lv.push(e, dom))
447 }
448 _ => Node::Leaf(lv.push(e, dom)),
449 }
450}
451
452fn ops(n: &Node) -> usize {
453 match n {
454 Node::Leaf(_) => 0,
455 Node::Monad(_, y) | Node::Window(_, _, y) | Node::Scan(_, y) => 1 + ops(y),
456 Node::Dyad(_, x, y) => 1 + ops(x) + ops(y),
457 }
458}
459
460fn subtrees<'a>(n: &'a Node, dom: Dom, out: &mut Vec<(&'a Node, Dom)>) {
463 if ops(n) == 0 {
464 return;
465 }
466 out.push((n, dom));
467 match n {
468 Node::Leaf(_) => {}
469 Node::Monad(_, y) | Node::Scan(_, y) => subtrees(y, dom, out),
470 Node::Window(_, _, y) => subtrees(y, Dom::Wide, out),
471 Node::Dyad(_, x, y) => {
472 subtrees(x, dom, out);
473 subtrees(y, dom, out);
474 }
475 }
476}
477
478fn lets_of(n: &Node) -> Vec<(Node, Dom)> {
491 let mut all = Vec::new();
492 subtrees(n, Dom::Result, &mut all);
493 let mut out = Vec::new();
494 fn walk(n: &Node, dom: Dom, all: &[(&Node, Dom)], out: &mut Vec<(Node, Dom)>) {
495 let count = |d: Dom| all.iter().filter(|(m, md)| *m == n && *md == d).count();
496 if ops(n) >= 1 && count(dom) >= 2 && count(other(dom)) == 0 {
497 if !out.iter().any(|(m, _)| m == n) {
498 out.push((n.clone(), dom));
499 }
500 return;
501 }
502 match n {
503 Node::Leaf(_) => {}
504 Node::Monad(_, y) | Node::Scan(_, y) => walk(y, dom, all, out),
505 Node::Window(_, _, y) => walk(y, Dom::Wide, all, out),
506 Node::Dyad(_, x, y) => {
507 walk(x, dom, all, out);
508 walk(y, dom, all, out);
509 }
510 }
511 }
512 walk(n, Dom::Result, &all, &mut out);
513 out
514}
515
516fn other(d: Dom) -> Dom {
517 match d {
518 Dom::Result => Dom::Wide,
519 Dom::Wide => Dom::Result,
520 }
521}
522
523fn emit_all(n: &Node, lets: &[(Node, Dom)], code: &mut Vec<Instr>) {
526 for (k, (l, _)) in lets.iter().enumerate() {
527 emit(l, &lets[..k], code);
530 code.push(Instr::Store(k));
531 }
532 emit(n, lets, code);
533}
534
535fn emit(n: &Node, lets: &[(Node, Dom)], code: &mut Vec<Instr>) {
537 if let Some(k) = lets.iter().position(|(l, _)| l == n) {
538 code.push(Instr::Let(k));
539 return;
540 }
541 match n {
542 Node::Leaf(i) => code.push(Instr::Load(*i)),
543 Node::Monad(op, y) => {
544 emit(y, lets, code);
545 code.push(Instr::Monad(*op));
546 }
547 Node::Window(op, k, y) => {
548 emit(y, lets, code);
549 code.push(Instr::Window(*op, *k));
550 }
551 Node::Scan(op, y) => {
552 emit(y, lets, code);
553 code.push(Instr::Scan(*op));
554 }
555 Node::Dyad(op, x, y) => {
556 emit(x, lets, code);
557 emit(y, lets, code);
558 code.push(Instr::Dyad(*op));
559 }
560 }
561}
562
563fn slots(code: &[Instr]) -> usize {
569 let mut stack: Vec<bool> = Vec::new();
570 let mut live = 0usize;
571 let mut max = 1usize;
572 for ins in code {
573 let operands = match ins {
574 Instr::Load(_) => {
575 stack.push(false);
576 continue;
577 }
578 Instr::Let(_) => {
582 stack.push(false);
583 continue;
584 }
585 Instr::Store(_) => {
586 stack.pop();
587 continue;
588 }
589 Instr::Monad(_) | Instr::Window(..) | Instr::Scan(_) => 1,
590 Instr::Dyad(_) => 2,
591 };
592 max = max.max(live + 1);
593 for _ in 0..operands {
594 if stack.pop().unwrap_or(false) {
595 live -= 1;
596 }
597 }
598 live += 1;
599 stack.push(true);
600 }
601 max
602}
603
604fn replayable(e: &Expr) -> bool {
611 match e {
612 Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => true,
613 Expr::Assign { .. }
614 | Expr::PrintPass { .. }
615 | Expr::Input { .. }
616 | Expr::Elided { .. }
617 | Expr::Control(..)
618 | Expr::AmendIndex { .. }
619 | Expr::VerbDef { .. }
620 | Expr::ModDef { .. } => false,
621 Expr::Monad { verb, y, .. } => verb.is_pure() && replayable(y),
622 Expr::Dyad { verb, x, y, .. } => verb.is_pure() && replayable(x) && replayable(y),
623 Expr::Fused { inputs, .. } => inputs.iter().all(replayable),
624 }
625}
626
627fn same(a: &Expr, b: &Expr) -> bool {
633 match (a, b) {
634 (Expr::Const(p, _), Expr::Const(q, _)) => p == q,
635 (Expr::Param(p, _), Expr::Param(q, _)) => p == q,
636 (Expr::Name(p, _), Expr::Name(q, _)) => p == q,
637 (Expr::Monad { verb: u, y: p, .. }, Expr::Monad { verb: v, y: q, .. }) => {
638 same_verb(u, v) && same(p, q)
639 }
640 (
641 Expr::Dyad { verb: u, x: px, y: py, .. },
642 Expr::Dyad { verb: v, x: qx, y: qy, .. },
643 ) => same_verb(u, v) && same(px, qx) && same(py, qy),
644 _ => false,
645 }
646}
647
648fn same_verb(a: &Verb, b: &Verb) -> bool {
649 match (a, b) {
650 (Verb::Prim(p), Verb::Prim(q)) => p == q,
651 (Verb::Rank(u, r), Verb::Rank(v, s)) => r == s && same_verb(u, v),
652 (Verb::Reduce(u), Verb::Reduce(v))
653 | (Verb::NWise(u), Verb::NWise(v))
654 | (Verb::Commute(u), Verb::Commute(v)) => same_verb(u, v),
655 (Verb::Windowed(u, j), Verb::Windowed(v, k)) => j == k && same_verb(u, v),
656 (Verb::PowerN(u, m), Verb::PowerN(v, n)) => m == n && same_verb(u, v),
657 (Verb::Fork(f, g, h), Verb::Fork(f2, g2, h2)) => {
658 same_verb(f, f2) && same_verb(g, g2) && same_verb(h, h2)
659 }
660 (Verb::NounFork(m, g, h), Verb::NounFork(n, g2, h2)) => {
661 m == n && same_verb(g, g2) && same_verb(h, h2)
662 }
663 (Verb::Hook(g, h), Verb::Hook(g2, h2))
664 | (Verb::Atop(g, h), Verb::Atop(g2, h2))
665 | (Verb::Compose(g, h), Verb::Compose(g2, h2)) => same_verb(g, g2) && same_verb(h, h2),
666 (Verb::BondLeft(m, u), Verb::BondLeft(n, v)) => m == n && same_verb(u, v),
667 (Verb::BondRight(u, m), Verb::BondRight(v, n)) => m == n && same_verb(u, v),
668 _ => false,
669 }
670}
671
672pub fn pass(stmts: &mut Vec<Expr>, tol: Tol) {
676 let orig = std::mem::take(stmts);
677 let mut cur = orig.clone();
678 let mut names = 0usize;
679 let mut crossed = false;
680 for _ in 0..=orig.len() {
683 match inline_once(&cur, &mut names, tol) {
684 Some(next) => {
685 cur = next;
686 crossed = true;
687 }
688 None => break,
689 }
690 }
691 let mut out: Vec<Expr> = cur.into_iter().map(|e| fuse_expr(e, tol)).collect();
692 if crossed {
693 out.insert(0, Expr::Elided { orig, span: Span::new(0, 0) });
695 }
696 *stmts = out;
697}
698
699fn fuse_expr(e: Expr, tol: Tol) -> Expr {
700 if let Some(f) = try_fuse(&e, tol) {
701 return f;
702 }
703 match e {
704 Expr::Assign { name, value, scope, span } => {
705 Expr::Assign { name, value: Box::new(fuse_expr(*value, tol)), scope, span }
706 }
707 Expr::Monad { verb, y, span } => {
708 Expr::Monad { verb, y: Box::new(fuse_expr(*y, tol)), span }
709 }
710 Expr::Dyad { verb, x, y, span } => Expr::Dyad {
711 verb,
712 x: Box::new(fuse_expr(*x, tol)),
713 y: Box::new(fuse_expr(*y, tol)),
714 span,
715 },
716 Expr::PrintPass { value, bare, span } => {
717 Expr::PrintPass { value: Box::new(fuse_expr(*value, tol)), bare, span }
718 }
719 other => other,
720 }
721}
722
723fn build<'a>(
726 root: &'a Expr,
727 yields: Yield,
728 least: usize,
729 sub: &mut Option<Inline<'a>>,
730 tol: Tol,
731) -> Option<(FusedKernel, Vec<&'a Expr>)> {
732 if let Some(s) = sub.as_mut() {
733 s.hits = 0;
734 }
735 let plan = plan_of(root, sub.as_ref());
736 let mut lv = Leaves::default();
737 let node = chain(root, &mut lv, sub, plan, Dom::Result);
738 if ops(&node) < least || !lv.inputs.iter().all(|l| replayable(l)) {
739 return None;
740 }
741 let mut code = Vec::new();
742 let lets = lets_of(&node);
743 emit_all(&node, &lets, &mut code);
744 let window = code.iter().find_map(|i| match i {
745 Instr::Window(_, k) => Some(*k),
746 _ => None,
747 });
748 let scans = code.iter().filter(|i| matches!(i, Instr::Scan(_))).count();
749 if scans > 0 && matches!(yields, Yield::Reduce(_)) {
754 return None;
755 }
756 let kernel = FusedKernel {
757 slots: slots(&code),
758 code,
759 yields,
760 leaves: lv.order,
761 doms: lv.doms,
762 let_doms: lets.iter().map(|(_, d)| *d).collect(),
763 plan,
764 window,
765 scans,
766 tol,
767 };
768 Some((kernel, lv.inputs))
769}
770
771fn kernel_at<'a>(
779 e: &'a Expr,
780 sub: &mut Option<Inline<'a>>,
781 tol: Tol,
782) -> Option<(FusedKernel, Vec<&'a Expr>, &'a Expr)> {
783 if let Expr::Monad { verb, y, .. } = e {
784 if is_tally(verb) && let Some((k, l)) = build(y, Yield::Tally, 1, sub, tol) {
785 return Some((k, l, e));
786 }
787 if let Some(op) = absorbable_reduce(verb)
788 && let Some((k, l)) = build(y, Yield::Reduce(op), 1, sub, tol)
789 {
790 return Some((k, l, e));
791 }
792 }
793 let (k, l) = build(e, Yield::Values, 2, sub, tol)?;
794 Some((k, l, e))
795}
796
797fn try_fuse(e: &Expr, tol: Tol) -> Option<Expr> {
799 let (kernel, leaves, orig) = kernel_at(e, &mut None, tol)?;
800 let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
801 Some(Expr::Fused {
802 kernel,
803 inputs,
804 orig: Box::new(orig.clone()),
805 span: e.span(),
806 })
807}
808
809pub(crate) fn fallback_tree(k: &FusedKernel, orig: &Expr, values: &[Array]) -> Expr {
816 let mut next = 0;
817 let plan = k.plan;
818 let tree = match orig {
819 Expr::Monad { verb, y, span } if matches!(k.yields, Yield::Reduce(_)) => Expr::Monad {
822 verb: verb.clone(),
823 y: Box::new(substitute(y, values, k, &mut next, plan, Dom::Result)),
824 span: *span,
825 },
826 Expr::Monad { verb, y, .. } if k.yields == Yield::Tally && is_tally(verb) => {
829 substitute(y, values, k, &mut next, plan, Dom::Result)
830 }
831 e => substitute(e, values, k, &mut next, plan, Dom::Result),
832 };
833 debug_assert_eq!(next, k.leaves.len(), "the fallback found different leaves");
834 tree
835}
836
837pub(crate) fn fallback_finish(k: &FusedKernel, v: Array) -> Array {
841 match k.yields {
842 Yield::Tally => Array::scalar_i64(v.items() as i64),
843 _ => v,
844 }
845}
846
847fn substitute(
850 e: &Expr,
851 values: &[Array],
852 k: &FusedKernel,
853 next: &mut usize,
854 plan: Plan,
855 dom: Dom,
856) -> Expr {
857 match e {
858 Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
859 verb: verb.clone(),
860 y: Box::new(substitute(y, values, k, next, plan, dom)),
861 span: *span,
862 },
863 Expr::Monad { verb, y, span } if plan.scan && absorbable_scan(e).is_some() => {
864 Expr::Monad {
865 verb: verb.clone(),
866 y: Box::new(substitute(y, values, k, next, plan, dom)),
867 span: *span,
868 }
869 }
870 Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => {
871 let ry = substitute(y, values, k, next, plan, dom);
872 let rx = substitute(x, values, k, next, plan, dom);
873 Expr::Dyad { verb: verb.clone(), x: Box::new(rx), y: Box::new(ry), span: *span }
874 }
875 Expr::Dyad { verb, x, y, span }
876 if dom == Dom::Result
877 && absorbable_window(e).map(|(_, k)| k) == plan.window
878 && plan.window.is_some() =>
879 {
880 Expr::Dyad {
881 verb: verb.clone(),
882 x: x.clone(),
883 y: Box::new(substitute(y, values, k, next, plan, Dom::Wide)),
884 span: *span,
885 }
886 }
887 leaf => {
888 let v = values[k.leaves[*next]].clone();
889 *next += 1;
890 Expr::Const(v, leaf.span())
891 }
892 }
893}
894
895fn hoisted_name(n: &mut usize) -> String {
911 *n += 1;
912 format!("·{}", *n - 1)
913}
914
915fn inline_once(stmts: &[Expr], names: &mut usize, tol: Tol) -> Option<Vec<Expr>> {
932 for (i, stmt) in stmts.iter().enumerate() {
933 let Expr::Assign { name, value, span, .. } = stmt else { continue };
934 if !inlinable(stmts, i, name, value, tol) {
935 continue;
936 }
937 if let Some(out) = rewrite(stmts, i, name, value, *span, names, tol) {
938 return Some(out);
939 }
940 }
941 None
942}
943
944fn inlinable(stmts: &[Expr], i: usize, name: &str, value: &Expr, tol: Tol) -> bool {
945 if !replayable(value) || mentions(value, name) {
946 return false;
947 }
948 let mut lv = Leaves::default();
949 if ops(&chain(value, &mut lv, &mut None, plan_of(value, None), Dom::Result)) < 1 {
950 return false;
951 }
952 let mut guarded = vec![name.to_string()];
953 free_names(value, &mut guarded);
954 let later = &stmts[i + 1..];
955 if later.iter().any(|s| assigns_any(s, &guarded)) {
956 return false;
957 }
958 let mut uses = 0;
959 for stmt in later {
960 match uses_land(stmt, name, value, tol) {
961 Some(n) => uses += n,
962 None => return false,
963 }
964 }
965 uses > 0
966}
967
968fn uses_land(e: &Expr, name: &str, def: &Expr, tol: Tol) -> Option<usize> {
971 let mut sub = Some(Inline { name, def, hits: 0 });
972 if let Some((_, leaves, _)) = kernel_at(e, &mut sub, tol) {
973 let mut n = sub.map_or(0, |s| s.hits);
974 for l in leaves {
975 n += uses_land(l, name, def, tol)?;
976 }
977 return Some(n);
978 }
979 match e {
980 Expr::Name(n, _) if n == name => None,
981 Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => Some(0),
982 Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => uses_land(value, name, def, tol),
983 Expr::Monad { y, .. } => uses_land(y, name, def, tol),
984 Expr::Dyad { x, y, .. } => Some(uses_land(x, name, def, tol)? + uses_land(y, name, def, tol)?),
985 Expr::Fused { .. }
986 | Expr::Elided { .. }
987 | Expr::Input { .. }
988 | Expr::Control(..)
989 | Expr::AmendIndex { .. }
990 | Expr::VerbDef { .. }
991 | Expr::ModDef { .. } => None,
992 }
993}
994
995fn rewrite(
997 stmts: &[Expr],
998 i: usize,
999 name: &str,
1000 value: &Expr,
1001 span: Span,
1002 names: &mut usize,
1003 tol: Tol,
1004) -> Option<Vec<Expr>> {
1005 let mut lv = Leaves::default();
1006 let plan = plan_of(value, None);
1007 chain(value, &mut lv, &mut None, plan, Dom::Result);
1008 let mut hoists = Vec::new();
1011 let mut bound: Vec<Option<String>> = Vec::new();
1012 for l in &lv.inputs {
1013 if matches!(l, Expr::Const(..) | Expr::Param(..) | Expr::Name(..)) {
1014 bound.push(None);
1015 continue;
1016 }
1017 let n = hoisted_name(names);
1018 hoists.push(Expr::Assign {
1019 name: n.clone(),
1020 value: Box::new((*l).clone()),
1021 scope: Scope::Local,
1022 span: l.span(),
1023 });
1024 bound.push(Some(n));
1025 }
1026 let def = with_leaves(value, &lv, &bound, plan, Dom::Result);
1027 let (kernel, leaves) = build(&def, Yield::Tally, 1, &mut None, tol)?;
1028 let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
1029 let guard = Expr::Assign {
1030 name: hoisted_name(names),
1031 value: Box::new(Expr::Fused {
1032 kernel,
1033 inputs,
1034 orig: Box::new(def.clone()),
1035 span,
1036 }),
1037 scope: Scope::Local,
1038 span,
1039 };
1040 let mut out = stmts[..i].to_vec();
1041 out.extend(hoists);
1042 out.push(guard);
1043 out.extend(stmts[i + 1..].iter().map(|s| replace_name(s, name, &def)));
1044 Some(out)
1045}
1046
1047fn with_leaves(e: &Expr, lv: &Leaves<'_>, bound: &[Option<String>], plan: Plan, dom: Dom) -> Expr {
1050 match e {
1051 Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
1052 verb: verb.clone(),
1053 y: Box::new(with_leaves(y, lv, bound, plan, dom)),
1054 span: *span,
1055 },
1056 Expr::Monad { verb, y, span } if plan.scan && absorbable_scan(e).is_some() => {
1057 Expr::Monad {
1058 verb: verb.clone(),
1059 y: Box::new(with_leaves(y, lv, bound, plan, dom)),
1060 span: *span,
1061 }
1062 }
1063 Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => Expr::Dyad {
1064 verb: verb.clone(),
1065 x: Box::new(with_leaves(x, lv, bound, plan, dom)),
1066 y: Box::new(with_leaves(y, lv, bound, plan, dom)),
1067 span: *span,
1068 },
1069 Expr::Dyad { verb, x, y, span }
1070 if dom == Dom::Result
1071 && absorbable_window(e).map(|(_, k)| k) == plan.window
1072 && plan.window.is_some() =>
1073 {
1074 Expr::Dyad {
1075 verb: verb.clone(),
1076 x: x.clone(),
1077 y: Box::new(with_leaves(y, lv, bound, plan, Dom::Wide)),
1078 span: *span,
1079 }
1080 }
1081 leaf => {
1082 let bind = lv
1083 .inputs
1084 .iter()
1085 .position(|&p| same(p, leaf))
1086 .and_then(|i| bound[i].as_ref());
1087 match bind {
1088 Some(n) => Expr::Name(n.clone(), leaf.span()),
1089 None => leaf.clone(),
1090 }
1091 }
1092 }
1093}
1094
1095fn replace_name(e: &Expr, name: &str, def: &Expr) -> Expr {
1096 match e {
1097 Expr::Name(n, _) if n == name => def.clone(),
1098 Expr::Assign { name: a, value, scope, span } => Expr::Assign {
1099 scope: *scope,
1100 name: a.clone(),
1101 value: Box::new(replace_name(value, name, def)),
1102 span: *span,
1103 },
1104 Expr::PrintPass { value, bare, span } => Expr::PrintPass {
1105 value: Box::new(replace_name(value, name, def)),
1106 bare: *bare,
1107 span: *span,
1108 },
1109 Expr::Monad { verb, y, span } => Expr::Monad {
1110 verb: verb.clone(),
1111 y: Box::new(replace_name(y, name, def)),
1112 span: *span,
1113 },
1114 Expr::Dyad { verb, x, y, span } => Expr::Dyad {
1115 verb: verb.clone(),
1116 x: Box::new(replace_name(x, name, def)),
1117 y: Box::new(replace_name(y, name, def)),
1118 span: *span,
1119 },
1120 other => other.clone(),
1121 }
1122}
1123
1124fn mentions(e: &Expr, name: &str) -> bool {
1125 let mut names = Vec::new();
1126 free_names(e, &mut names);
1127 names.iter().any(|n| n == name)
1128}
1129
1130fn free_names(e: &Expr, out: &mut Vec<String>) {
1132 match e {
1133 Expr::Name(n, _) => out.push(n.clone()),
1134 Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => free_names(value, out),
1135 Expr::Monad { y, .. } => free_names(y, out),
1136 Expr::Dyad { x, y, .. } => {
1137 free_names(x, out);
1138 free_names(y, out);
1139 }
1140 Expr::Fused { inputs, .. } => inputs.iter().for_each(|i| free_names(i, out)),
1141 Expr::Const(..)
1142 | Expr::Param(..)
1143 | Expr::Elided { .. }
1144 | Expr::Input { .. }
1145 | Expr::Control(..)
1146 | Expr::AmendIndex { .. }
1147 | Expr::VerbDef { .. }
1148 | Expr::ModDef { .. } => {}
1149 }
1150}
1151
1152fn assigns_any(e: &Expr, names: &[String]) -> bool {
1154 match e {
1155 Expr::Assign { name, value, .. } => {
1156 names.iter().any(|n| n == name) || assigns_any(value, names)
1157 }
1158 Expr::PrintPass { value, .. } => assigns_any(value, names),
1159 Expr::Monad { y, .. } => assigns_any(y, names),
1160 Expr::Dyad { x, y, .. } => assigns_any(x, names) || assigns_any(y, names),
1161 Expr::Fused { inputs, .. } => inputs.iter().any(|i| assigns_any(i, names)),
1162 Expr::Const(..)
1163 | Expr::Param(..)
1164 | Expr::Name(..)
1165 | Expr::Elided { .. }
1166 | Expr::Input { .. }
1167 | Expr::Control(..)
1168 | Expr::AmendIndex { .. }
1169 | Expr::VerbDef { .. }
1170 | Expr::ModDef { .. } => false,
1171 }
1172}
1173
1174pub fn is_fused(p: &Program) -> bool {
1176 fn any(e: &Expr) -> bool {
1177 match e {
1178 Expr::Fused { .. } => true,
1179 Expr::Const(..)
1180 | Expr::Param(..)
1181 | Expr::Name(..)
1182 | Expr::Elided { .. }
1183 | Expr::Input { .. }
1184 | Expr::Control(..)
1185 | Expr::AmendIndex { .. }
1186 | Expr::VerbDef { .. }
1187 | Expr::ModDef { .. } => false,
1188 Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => any(value),
1189 Expr::Monad { y, .. } => any(y),
1190 Expr::Dyad { x, y, .. } => any(x) || any(y),
1191 }
1192 }
1193 p.stmts.iter().any(any)
1194}
1195
1196pub fn is_inlined(p: &Program) -> bool {
1198 matches!(p.stmts.first(), Some(Expr::Elided { .. }))
1199}
1200
1201pub fn unfused(p: &Program) -> Program {
1205 fn strip(e: &Expr) -> Expr {
1206 match e {
1207 Expr::Fused { orig, .. } => strip(orig),
1208 Expr::Assign { name, value, scope, span } => {
1209 Expr::Assign {
1210 name: name.clone(),
1211 value: Box::new(strip(value)),
1212 scope: *scope,
1213 span: *span,
1214 }
1215 }
1216 Expr::PrintPass { value, bare, span } => {
1217 Expr::PrintPass { value: Box::new(strip(value)), bare: *bare, span: *span }
1218 }
1219 Expr::Monad { verb, y, span } => {
1220 Expr::Monad { verb: verb.clone(), y: Box::new(strip(y)), span: *span }
1221 }
1222 Expr::Dyad { verb, x, y, span } => Expr::Dyad {
1223 verb: verb.clone(),
1224 x: Box::new(strip(x)),
1225 y: Box::new(strip(y)),
1226 span: *span,
1227 },
1228 other => other.clone(),
1229 }
1230 }
1231 let mut out = p.clone();
1232 let stmts = match p.stmts.first() {
1235 Some(Expr::Elided { orig, .. }) => orig,
1236 _ => &p.stmts,
1237 };
1238 out.stmts = stmts.iter().map(strip).collect();
1239 out
1240}
1241
1242fn monad_type(op: ScalarMonad, a: DType) -> Option<DType> {
1248 use DType::*;
1249 use ScalarMonad::*;
1250 if a == Complex {
1253 return None;
1254 }
1255 Some(match op {
1256 Recip | Halve | Exp => F64,
1257 Conj | Abs | OneMinus => a,
1259 Neg | Signum | Inc | Dec | Double | Square => match a {
1260 Bool | I64 => I64,
1261 other => other,
1262 },
1263 Floor | Ceil => match a {
1264 Bool | I64 => I64,
1265 _ => return None,
1266 },
1267 _ => return None,
1268 })
1269}
1270
1271fn dyad_type(op: ScalarDyad, a: DType, b: DType) -> Option<DType> {
1274 use ScalarDyad::*;
1275 if a == DType::Complex || b == DType::Complex {
1276 return None;
1277 }
1278 match op {
1279 Eq | Ne | Lt | Le | Gt | Ge => Some(DType::Bool),
1280 DivJ => Some(DType::F64),
1281 Add | Sub | Mul | Min | Max | Residue => match DType::promote(a, b)? {
1282 DType::Bool => Some(DType::I64),
1283 DType::Char | DType::Symbol => None,
1284 t => Some(t),
1285 },
1286 _ => None,
1287 }
1288}
1289
1290fn fold_type(op: ScalarDyad, a: DType) -> Option<DType> {
1294 use ScalarDyad::*;
1295 if !matches!(op, Add | Mul | Min | Max) {
1296 return None;
1297 }
1298 match a {
1299 DType::Bool | DType::I64 => Some(DType::I64),
1300 DType::F64 => Some(DType::F64),
1301 _ => None,
1302 }
1303}
1304
1305pub(crate) fn working_type(k: &FusedKernel, inputs: &[Array]) -> Option<(DType, DType)> {
1325 let mut stack: Vec<DType> = Vec::with_capacity(k.slots);
1326 let mut lets: Vec<DType> = Vec::new();
1327 let mut float = false;
1328 let mut integer_step = false;
1329 if inputs.iter().any(|a| a.dtype() == DType::Complex || a.dtype().is_exact()) {
1332 return None;
1333 }
1334 for ins in &k.code {
1335 let t = match ins {
1336 Instr::Load(i) => inputs[*i].dtype(),
1337 Instr::Monad(op) => monad_type(*op, stack.pop()?)?,
1338 Instr::Window(op, _) | Instr::Scan(op) => fold_type(*op, stack.pop()?)?,
1339 Instr::Dyad(op) => {
1340 let b = stack.pop()?;
1341 let a = stack.pop()?;
1342 dyad_type(*op, a, b)?
1343 }
1344 Instr::Store(k) => {
1345 let t = stack.pop()?;
1346 if lets.len() != *k {
1347 return None;
1348 }
1349 lets.push(t);
1350 continue;
1351 }
1352 Instr::Let(k) => {
1355 let t = *lets.get(*k)?;
1356 float |= t == DType::F64;
1357 stack.push(t);
1358 continue;
1359 }
1360 };
1361 if !t.is_numeric() {
1364 return None;
1365 }
1366 float |= t == DType::F64;
1367 integer_step |= t == DType::I64 && !matches!(ins, Instr::Load(_));
1370 stack.push(t);
1371 }
1372 let root = stack.pop()?;
1373 let working = if float { DType::F64 } else { DType::I64 };
1374 if working == DType::F64 && integer_step {
1375 return None;
1376 }
1377 Some((working, root))
1378}
1379
1380#[derive(Clone, Copy, PartialEq, Eq)]
1385enum On {
1386 Result,
1387 Wide,
1388 Either,
1389}
1390
1391fn combine(a: On, b: On) -> On {
1392 if a == On::Either {
1393 b
1394 } else {
1395 a
1396 }
1397}
1398
1399fn placed(d: Option<Dom>) -> On {
1400 match d {
1401 Some(Dom::Result) => On::Result,
1402 Some(Dom::Wide) => On::Wide,
1403 None => On::Either,
1404 }
1405}
1406
1407#[derive(Clone, Copy)]
1410struct Extent {
1411 start: usize,
1412 len: usize,
1413 wide_start: usize,
1414 wide_len: usize,
1415}
1416
1417impl Extent {
1418 fn of(start: usize, len: usize, window: Option<usize>, wide: usize) -> Extent {
1429 let Some(k) = window else {
1430 return Extent { start, len, wide_start: start, wide_len: len };
1431 };
1432 let lo = start - start % k;
1433 let hi = ((start + len + k - 2) / k + 1) * k;
1434 Extent { start, len, wide_start: lo, wide_len: hi.min(wide) - lo }
1435 }
1436
1437 fn len_on(&self, dom: On) -> usize {
1438 match dom {
1439 On::Wide => self.wide_len,
1440 _ => self.len,
1441 }
1442 }
1443}
1444
1445#[derive(Clone, Copy)]
1448struct Loaded<'a, T> {
1449 data: &'a [T],
1450 splat: bool,
1451 on: On,
1453 base: usize,
1457}
1458
1459impl<T> Loaded<'_, T> {
1460 #[inline]
1461 fn block(&self, at: &Extent, dom: On) -> &[T] {
1462 if self.splat {
1463 return &self.data[..at.len_on(dom)];
1464 }
1465 let (start, len) = read_over(self.on, at);
1466 &self.data[start - self.base..start - self.base + len]
1467 }
1468}
1469
1470#[inline]
1473fn read_over(on: On, at: &Extent) -> (usize, usize) {
1474 match on {
1475 On::Wide => (at.wide_start, at.wide_len),
1476 _ => (at.start, at.len),
1477 }
1478}
1479
1480#[derive(Clone, Copy)]
1489enum Narrow<'a> {
1490 I64(&'a [i64]),
1491 Bool(&'a [u8]),
1492}
1493
1494enum Source<'a, T> {
1497 Ready(Loaded<'a, T>),
1498 Staged(Narrow<'a>, On),
1500}
1501
1502trait FromNarrow: Copy {
1504 fn fill(src: Narrow<'_>, at: usize, dst: &mut [Self]);
1505}
1506
1507impl FromNarrow for f64 {
1508 #[inline]
1509 fn fill(src: Narrow<'_>, at: usize, dst: &mut [f64]) {
1510 match src {
1511 Narrow::I64(v) => {
1512 for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1513 *slot = x as f64;
1514 }
1515 }
1516 Narrow::Bool(v) => {
1517 for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1518 *slot = x as f64;
1519 }
1520 }
1521 }
1522 }
1523}
1524
1525impl FromNarrow for i64 {
1526 #[inline]
1527 fn fill(src: Narrow<'_>, at: usize, dst: &mut [i64]) {
1528 match src {
1529 Narrow::I64(v) => dst.copy_from_slice(&v[at..at + dst.len()]),
1531 Narrow::Bool(v) => {
1532 for (slot, &x) in dst.iter_mut().zip(&v[at..]) {
1533 *slot = x as i64;
1534 }
1535 }
1536 }
1537 }
1538}
1539
1540struct Sources<'a, T> {
1542 of: Vec<Source<'a, T>>,
1543 staged: usize,
1544}
1545
1546impl<T: FromNarrow + Default> Sources<'_, T> {
1547 fn with_block<R>(
1553 &self,
1554 at: &Extent,
1555 stage: &mut [T],
1556 width: usize,
1557 f: impl FnOnce(&[Loaded<'_, T>]) -> R,
1558 ) -> R {
1559 let mut k = 0;
1560 for s in &self.of {
1561 let Source::Staged(src, on) = s else { continue };
1562 let (start, len) = read_over(*on, at);
1563 T::fill(*src, start, &mut stage[k * width..k * width + len]);
1564 k += 1;
1565 }
1566 let mut k = 0;
1567 let loaded: Vec<Loaded<'_, T>> = self
1568 .of
1569 .iter()
1570 .map(|s| match s {
1571 Source::Ready(l) => *l,
1572 Source::Staged(_, on) => {
1573 let (start, len) = read_over(*on, at);
1574 let d = &stage[k * width..k * width + len];
1575 k += 1;
1576 Loaded { data: d, splat: false, on: *on, base: start }
1577 }
1578 })
1579 .collect();
1580 f(&loaded)
1581 }
1582}
1583
1584#[derive(Clone, Copy)]
1587enum Slot {
1588 Input(usize),
1589 Block(usize, On),
1590}
1591
1592struct Scratch<T> {
1594 cells: Vec<T>,
1595 width: usize,
1598 free: Vec<usize>,
1599 stack: Vec<Slot>,
1600 lets: Vec<usize>,
1601 carry: Vec<Option<T>>,
1604}
1605
1606impl<T: Copy + Default> Scratch<T> {
1607 fn new(k: &FusedKernel, w: usize) -> Scratch<T> {
1613 let width = w + 3 * k.window.unwrap_or(0);
1614 Scratch {
1615 cells: vec![T::default(); k.slots * width],
1616 width,
1617 free: Vec::with_capacity(k.slots),
1618 stack: Vec::with_capacity(k.slots),
1619 lets: Vec::new(),
1620 carry: vec![None; k.scans],
1621 }
1622 }
1623}
1624
1625struct Steps<M, D, W, S> {
1628 monad: M,
1629 dyad: D,
1630 window: W,
1631 scan: S,
1632}
1633
1634fn split_slots<'s, T>(
1636 scratch: &'s mut [T],
1637 w: usize,
1638 d: usize,
1639) -> (&'s mut [T], impl Fn(usize) -> &'s [T]) {
1640 let (lo, hi) = scratch.split_at_mut(d * w);
1641 let (dst, hi) = hi.split_at_mut(w);
1642 let lo: &[T] = lo;
1643 let hi: &[T] = hi;
1644 (dst, move |i: usize| {
1645 if i < d {
1646 &lo[i * w..(i + 1) * w]
1647 } else {
1648 &hi[(i - d - 1) * w..(i - d) * w]
1649 }
1650 })
1651}
1652
1653fn exec_block<T, M, D, W, S>(
1660 k: &FusedKernel,
1661 srcs: &[Loaded<'_, T>],
1662 at: &Extent,
1663 sc: &mut Scratch<T>,
1664 out: Option<&mut [T]>,
1665 steps: &Steps<M, D, W, S>,
1666) -> Option<usize>
1667where
1668 T: Copy,
1669 M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1670 D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1671 W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool,
1672 S: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T>,
1673{
1674 let Scratch { cells, width, free, stack, lets, carry } = sc;
1675 let w = *width;
1676 stack.clear();
1677 free.clear();
1678 lets.clear();
1679 let nslots = cells.len() / w;
1680 free.extend((0..nslots).rev());
1681 let place = |s: &Slot| match s {
1682 Slot::Input(j) => srcs[*j].on,
1683 Slot::Block(_, o) => *o,
1684 };
1685 let last = k.code.len() - 1;
1686 let head = if out.is_some() { last } else { k.code.len() };
1687 let mut scanned = 0usize;
1688 for ins in &k.code[..head] {
1689 match ins {
1690 Instr::Load(j) => stack.push(Slot::Input(*j)),
1691 Instr::Monad(op) => {
1692 let a = stack.pop()?;
1693 let dom = place(&a);
1694 let len = at.len_on(dom);
1695 let d = free.pop()?;
1696 let (dst, get) = split_slots(cells, w, d);
1697 let av = match a {
1698 Slot::Input(j) => srcs[j].block(at, dom),
1699 Slot::Block(i, _) => &get(i)[..len],
1700 };
1701 if !(steps.monad)(*op, av, &mut dst[..len]) {
1702 return None;
1703 }
1704 release(free, lets, a);
1705 stack.push(Slot::Block(d, dom));
1706 }
1707 Instr::Scan(op) => {
1708 let a = stack.pop()?;
1709 let dom = place(&a);
1710 let len = at.len_on(dom);
1711 let d = free.pop()?;
1712 let (dst, get) = split_slots(cells, w, d);
1713 let av = match a {
1714 Slot::Input(j) => srcs[j].block(at, dom),
1715 Slot::Block(i, _) => &get(i)[..len],
1716 };
1717 carry[scanned] = Some((steps.scan)(*op, av, carry[scanned], &mut dst[..len])?);
1718 scanned += 1;
1719 release(free, lets, a);
1720 stack.push(Slot::Block(d, dom));
1721 }
1722 Instr::Window(op, size) => {
1723 let a = stack.pop()?;
1724 let d = free.pop()?;
1725 let (dst, get) = split_slots(cells, w, d);
1726 let av = match a {
1727 Slot::Input(j) => srcs[j].block(at, On::Wide),
1728 Slot::Block(i, _) => &get(i)[..at.wide_len],
1729 };
1730 let first = at.start - at.wide_start;
1731 if !(steps.window)(*op, *size, av, first, &mut dst[..at.len]) {
1732 return None;
1733 }
1734 release(free, lets, a);
1735 stack.push(Slot::Block(d, On::Result));
1736 }
1737 Instr::Dyad(op) => {
1738 let b = stack.pop()?;
1739 let a = stack.pop()?;
1740 let dom = combine(place(&a), place(&b));
1741 let len = at.len_on(dom);
1742 let d = free.pop()?;
1743 let (dst, get) = split_slots(cells, w, d);
1744 let av = match a {
1745 Slot::Input(j) => srcs[j].block(at, dom),
1746 Slot::Block(i, _) => &get(i)[..len],
1747 };
1748 let bv = match b {
1749 Slot::Input(j) => srcs[j].block(at, dom),
1750 Slot::Block(i, _) => &get(i)[..len],
1751 };
1752 if !(steps.dyad)(*op, av, bv, &mut dst[..len]) {
1753 return None;
1754 }
1755 for s in [a, b] {
1756 release(free, lets, s);
1757 }
1758 stack.push(Slot::Block(d, dom));
1759 }
1760 Instr::Store(j) => {
1761 let Slot::Block(i, _) = stack.pop()? else { return None };
1762 if lets.len() != *j {
1763 return None;
1764 }
1765 lets.push(i);
1766 }
1767 Instr::Let(j) => {
1770 stack.push(Slot::Block(*lets.get(*j)?, placed(Some(*k.let_doms.get(*j)?))))
1771 }
1772 }
1773 }
1774 let Some(dst) = out else {
1775 return match stack.pop()? {
1776 Slot::Block(i, _) => Some(i),
1777 Slot::Input(_) => None,
1779 };
1780 };
1781 let dst = &mut dst[..at.len];
1785 let view = |s: Slot, dom: On| match s {
1786 Slot::Input(j) => srcs[j].block(at, dom),
1787 Slot::Block(i, o) => &cells[i * w..i * w + at.len_on(o)],
1788 };
1789 let ok = match k.code[last] {
1790 Instr::Monad(op) => {
1791 let a = stack.pop()?;
1792 (steps.monad)(op, view(a, On::Result), dst)
1793 }
1794 Instr::Scan(op) => {
1795 let a = stack.pop()?;
1796 match (steps.scan)(op, view(a, On::Result), carry[scanned], dst) {
1797 Some(c) => {
1798 carry[scanned] = Some(c);
1799 true
1800 }
1801 None => false,
1802 }
1803 }
1804 Instr::Window(op, size) => {
1805 let a = stack.pop()?;
1806 let first = at.start - at.wide_start;
1807 (steps.window)(op, size, view(a, On::Wide), first, dst)
1808 }
1809 Instr::Dyad(op) => {
1810 let b = stack.pop()?;
1811 let a = stack.pop()?;
1812 let dom = combine(place(&a), place(&b));
1813 (steps.dyad)(op, view(a, dom), view(b, dom), dst)
1814 }
1815 Instr::Load(_) | Instr::Store(_) | Instr::Let(_) => return None,
1817 };
1818 ok.then_some(usize::MAX)
1819}
1820
1821fn release(free: &mut Vec<usize>, lets: &[usize], s: Slot) {
1824 if let Slot::Block(i, _) = s
1825 && !lets.contains(&i)
1826 {
1827 free.push(i);
1828 }
1829}
1830
1831fn map_pass<T, M, D, W, S>(
1833 k: &FusedKernel,
1834 srcs: &Sources<'_, T>,
1835 n: usize,
1836 wide: usize,
1837 steps: &Steps<M, D, W, S>,
1838) -> Option<Vec<T>>
1839where
1840 T: FromNarrow + Default + Send + Sync,
1841 M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1842 D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1843 W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool + Sync + Send,
1844 S: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T> + Sync + Send,
1845{
1846 let run = |start: usize, part: &mut [T]| {
1847 let w = BLOCK.min(part.len()).max(1);
1848 let mut sc = Scratch::new(k, w);
1849 let width = sc.width;
1850 let mut stage = vec![T::default(); srcs.staged * width];
1851 for (b, chunk) in part.chunks_mut(w).enumerate() {
1852 let at = Extent::of(start + b * w, chunk.len(), k.window, wide);
1853 let done = srcs.with_block(&at, &mut stage, width, |loaded| {
1854 exec_block(k, loaded, &at, &mut sc, Some(chunk), steps).is_some()
1855 });
1856 if !done {
1857 return false;
1858 }
1859 }
1860 true
1861 };
1862 if k.scans > 0 {
1863 let mut out = vec![T::default(); n];
1867 return run(0, &mut out).then_some(out);
1868 }
1869 let (out, ok) = par::fill(n, run);
1870 ok.then_some(out)
1871}
1872
1873const FOLD_LANES: usize = 8;
1880const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
1881
1882#[inline(always)]
1885fn fold_block_body<T, S>(v: &[T], step: &S) -> Option<T>
1886where
1887 T: Copy,
1888 S: Fn(T, T) -> Option<T>,
1889{
1890 let n = v.len();
1891 if n < MIN_LANE_WORK {
1892 let mut acc = v[n - 1];
1893 for &x in v[..n - 1].iter().rev() {
1894 acc = step(x, acc)?;
1895 }
1896 return Some(acc);
1897 }
1898 let rows = n / FOLD_LANES;
1899 let head = n - rows * FOLD_LANES;
1900 let last = head + (rows - 1) * FOLD_LANES;
1901 let mut acc = [v[last]; FOLD_LANES];
1902 acc.copy_from_slice(&v[last..last + FOLD_LANES]);
1903 for r in (0..rows - 1).rev() {
1904 let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
1905 for (slot, &x) in acc.iter_mut().zip(row) {
1906 *slot = step(x, *slot)?;
1907 }
1908 }
1909 let mut a = acc[FOLD_LANES - 1];
1910 for &x in acc[..FOLD_LANES - 1].iter().rev() {
1911 a = step(x, a)?;
1912 }
1913 for &x in v[..head].iter().rev() {
1914 a = step(x, a)?;
1915 }
1916 Some(a)
1917}
1918
1919multiversioned! {
1920 fn fold_block[T: Copy, S: Fn(T, T) -> Option<T>](
1922 v: &[T],
1923 step: &S,
1924 ) -> Option<T> = fold_block_body;
1925}
1926
1927fn fold_range<T, M, D, W, C, S>(
1929 k: &FusedKernel,
1930 srcs: &Sources<'_, T>,
1931 lo: usize,
1932 hi: usize,
1933 wide: usize,
1934 steps: &Steps<M, D, W, C>,
1935 step: &S,
1936) -> Option<T>
1937where
1938 T: FromNarrow + Default,
1939 M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
1940 D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
1941 W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool,
1942 C: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T>,
1943 S: Fn(T, T) -> Option<T>,
1944{
1945 let w = BLOCK.min(hi - lo).max(1);
1946 let mut sc = Scratch::new(k, w);
1947 let width = sc.width;
1948 let mut stage = vec![T::default(); srcs.staged * width];
1949 let mut acc: Option<T> = None;
1950 for b in (0..(hi - lo).div_ceil(w)).rev() {
1955 let start = lo + b * w;
1956 let len = (hi - start).min(w);
1957 let at = Extent::of(start, len, k.window, wide);
1958 let slot = srcs
1959 .with_block(&at, &mut stage, width, |loaded| {
1960 exec_block(k, loaded, &at, &mut sc, None, steps)
1961 })?;
1962 let block = fold_block(&sc.cells[slot * sc.width..slot * sc.width + len], step)?;
1963 acc = Some(match acc {
1964 None => block,
1965 Some(a) => step(block, a)?,
1966 });
1967 }
1968 acc
1969}
1970
1971fn reduce_pass<T, M, D, W, C, S>(
1973 k: &FusedKernel,
1974 srcs: &Sources<'_, T>,
1975 n: usize,
1976 wide: usize,
1977 steps: &Steps<M, D, W, C>,
1978 step: S,
1979) -> Option<T>
1980where
1981 T: FromNarrow + Default + Send + Sync,
1982 M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
1983 D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
1984 W: Fn(ScalarDyad, usize, &[T], usize, &mut [T]) -> bool + Sync + Send,
1985 C: Fn(ScalarDyad, &[T], Option<T>, &mut [T]) -> Option<T> + Sync + Send,
1986 S: Fn(T, T) -> Option<T> + Sync + Send,
1987{
1988 let chunks = par::chunks(n, n * k.code.len());
1989 if chunks < 2 {
1990 return fold_range(k, srcs, 0, n, wide, steps, &step);
1991 }
1992 let per = n.div_ceil(chunks);
1993 let parts = par::map_indexed(n.div_ceil(per), |c| {
1994 fold_range(k, srcs, c * per, ((c + 1) * per).min(n), wide, steps, &step)
1995 });
1996 let mut it = parts.into_iter().rev();
2000 let mut acc = it.next()??;
2001 for part in it {
2002 acc = step(part?, acc)?;
2003 }
2004 Some(acc)
2005}
2006
2007macro_rules! each {
2020 ($a:expr, $dst:expr, $f:expr) => {{
2021 let f = $f;
2022 for (slot, &x) in $dst.iter_mut().zip($a) {
2023 *slot = f(x);
2024 }
2025 return true;
2026 }};
2027}
2028
2029macro_rules! zip {
2030 ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
2031 let f = $f;
2032 for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
2033 *slot = f(x, y);
2034 }
2035 return true;
2036 }};
2037}
2038
2039#[inline(always)]
2040fn monad_f64_body(op: ScalarMonad, a: &[f64], dst: &mut [f64], tol: Tol) -> bool {
2041 use ScalarMonad::*;
2042 match op {
2043 Conj => each!(a, dst, |x: f64| x),
2044 Neg => each!(a, dst, |x: f64| -x),
2045 Abs => each!(a, dst, f64::abs),
2046 Signum => each!(a, dst, |x: f64| if tol.is_zero(x) {
2049 0.0
2050 } else if x > 0.0 {
2051 1.0
2052 } else if x < 0.0 {
2053 -1.0
2054 } else {
2055 0.0
2056 }),
2057 Recip => {
2061 if !tol.is_j() && a.contains(&0.0) {
2062 return false;
2063 }
2064 each!(a, dst, |x: f64| if x == 0.0 { f64::INFINITY } else { 1.0 / x })
2065 }
2066 Floor => each!(a, dst, f64::floor),
2069 Ceil => each!(a, dst, f64::ceil),
2070 Inc => each!(a, dst, |x: f64| x + 1.0),
2071 Dec => each!(a, dst, |x: f64| x - 1.0),
2072 Double => each!(a, dst, |x: f64| x + x),
2073 Halve => each!(a, dst, |x: f64| x / 2.0),
2074 Square => each!(a, dst, |x: f64| x * x),
2075 OneMinus => each!(a, dst, |x: f64| 1.0 - x),
2076 Exp => each!(a, dst, f64::exp),
2077 _ => false,
2078 }
2079}
2080
2081#[inline(always)]
2082fn dyad_f64_body(op: ScalarDyad, a: &[f64], b: &[f64], dst: &mut [f64], tol: Tol) -> bool {
2083 use ScalarDyad::*;
2084 match op {
2085 Add => zip!(a, b, dst, |x: f64, y: f64| x + y),
2086 Sub => zip!(a, b, dst, |x: f64, y: f64| x - y),
2087 Mul => zip!(a, b, dst, |x: f64, y: f64| x * y),
2088 Min => zip!(a, b, dst, f64::min),
2089 Max => zip!(a, b, dst, f64::max),
2090 DivJ => zip!(a, b, dst, |x: f64, y: f64| if y == 0.0 {
2091 if x == 0.0 { 0.0 } else { f64::INFINITY.copysign(x) }
2092 } else {
2093 x / y
2094 }),
2095 Residue => zip!(a, b, dst, |x: f64, y: f64| tol.residue(x, y)),
2098 Eq | Ne | Lt | Le | Gt | Ge => {
2102 zip!(a, b, dst, |x: f64, y: f64| tol_cmp(op, x, y, tol) as u8 as f64)
2103 }
2104 _ => false,
2105 }
2106}
2107
2108#[inline(always)]
2119fn no_nan_block(dst: &[f64]) -> bool {
2120 !dst.iter().any(|x| x.is_nan())
2121}
2122
2123#[inline(always)]
2125fn finite_block(dst: &[f64]) -> bool {
2126 dst.iter().all(|x| x.is_finite())
2127}
2128
2129#[inline(always)]
2133fn block_f64(r: f64) -> (f64, bool) {
2134 (r, r.is_nan())
2135}
2136
2137macro_rules! each_over {
2140 ($a:expr, $dst:expr, $f:expr) => {{
2141 let f = $f;
2142 let mut over = false;
2143 for (slot, &x) in $dst.iter_mut().zip($a) {
2144 let (v, o) = f(x);
2145 *slot = v;
2146 over |= o;
2147 }
2148 return !over;
2149 }};
2150}
2151
2152macro_rules! zip_over {
2153 ($a:expr, $b:expr, $dst:expr, $f:expr) => {{
2154 let f = $f;
2155 let mut over = false;
2156 for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
2157 let (v, o) = f(x, y);
2158 *slot = v;
2159 over |= o;
2160 }
2161 return !over;
2162 }};
2163}
2164
2165#[inline(always)]
2166fn monad_i64_body(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool {
2167 use ScalarMonad::*;
2168 match op {
2169 Conj | Floor | Ceil => each!(a, dst, |x: i64| x),
2170 Neg => each_over!(a, dst, i64::overflowing_neg),
2171 Abs => each_over!(a, dst, i64::overflowing_abs),
2172 Signum => each!(a, dst, i64::signum),
2173 Inc => each_over!(a, dst, |x: i64| x.overflowing_add(1)),
2174 Dec => each_over!(a, dst, |x: i64| x.overflowing_sub(1)),
2175 Double => each_over!(a, dst, |x: i64| x.overflowing_add(x)),
2176 Square => each_over!(a, dst, |x: i64| x.overflowing_mul(x)),
2177 OneMinus => each_over!(a, dst, |x: i64| 1i64.overflowing_sub(x)),
2178 _ => false,
2179 }
2180}
2181
2182#[inline(always)]
2183fn dyad_i64_body(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool {
2184 use ScalarDyad::*;
2185 match op {
2186 Add => zip_over!(a, b, dst, i64::overflowing_add),
2187 Sub => zip_over!(a, b, dst, i64::overflowing_sub),
2188 Mul => zip_over!(a, b, dst, i64::overflowing_mul),
2189 Min => zip!(a, b, dst, i64::min),
2190 Max => zip!(a, b, dst, i64::max),
2191 Residue => zip!(a, b, dst, |x: i64, y: i64| if x == 0 {
2192 y
2193 } else {
2194 let mut r = y.wrapping_rem(x);
2196 if r != 0 && (r < 0) != (x < 0) {
2197 r += x;
2198 }
2199 r
2200 }),
2201 Eq => zip!(a, b, dst, |x: i64, y: i64| (x == y) as i64),
2202 Ne => zip!(a, b, dst, |x: i64, y: i64| (x != y) as i64),
2203 Lt => zip!(a, b, dst, |x: i64, y: i64| (x < y) as i64),
2204 Le => zip!(a, b, dst, |x: i64, y: i64| (x <= y) as i64),
2205 Gt => zip!(a, b, dst, |x: i64, y: i64| (x > y) as i64),
2206 Ge => zip!(a, b, dst, |x: i64, y: i64| (x >= y) as i64),
2207 _ => false,
2208 }
2209}
2210
2211multiversioned! {
2212 fn monad_f64(
2216 op: ScalarMonad,
2217 a: &[f64],
2218 dst: &mut [f64],
2219 tol: Tol,
2220 ) -> bool = monad_f64_body;
2221}
2222
2223multiversioned! {
2224 fn dyad_f64(
2227 op: ScalarDyad,
2228 a: &[f64],
2229 b: &[f64],
2230 dst: &mut [f64],
2231 tol: Tol,
2232 ) -> bool = dyad_f64_body;
2233}
2234
2235multiversioned! {
2236 fn monad_i64(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool = monad_i64_body;
2239}
2240
2241multiversioned! {
2242 fn dyad_i64(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool = dyad_i64_body;
2245}
2246
2247#[inline(always)]
2254fn scan_block_body<T, F>(v: &[T], carry: Option<T>, dst: &mut [T], step: &F) -> Option<T>
2255where
2256 T: Copy,
2257 F: Fn(T, T) -> (T, bool),
2258{
2259 let mut over = false;
2260 let (mut acc, from) = match carry {
2261 Some(a) => (a, 0),
2262 None => {
2263 dst[0] = v[0];
2265 (v[0], 1)
2266 }
2267 };
2268 for (slot, &x) in dst.iter_mut().zip(v).skip(from) {
2269 let (r, o) = step(acc, x);
2270 acc = r;
2271 over |= o;
2272 *slot = acc;
2273 }
2274 (!over).then_some(acc)
2275}
2276
2277multiversioned! {
2278 fn scan_block[T: Copy, F: Fn(T, T) -> (T, bool)](
2281 v: &[T],
2282 carry: Option<T>,
2283 dst: &mut [T],
2284 step: &F,
2285 ) -> Option<T> = scan_block_body;
2286}
2287
2288fn window_pass_f64(op: ScalarDyad, k: usize, v: &[f64], first: usize, dst: &mut [f64]) -> bool {
2291 use ScalarDyad::*;
2292 match op {
2293 Add => windows_into(v, k, first, dst, &|a: f64, b: f64| block_f64(a + b)),
2294 Mul => windows_into(v, k, first, dst, &|a: f64, b: f64| block_f64(a * b)),
2295 Min => windows_into(v, k, first, dst, &|a: f64, b: f64| (a.min(b), false)),
2296 Max => windows_into(v, k, first, dst, &|a: f64, b: f64| (a.max(b), false)),
2297 _ => false,
2298 }
2299}
2300
2301fn window_pass_i64(op: ScalarDyad, k: usize, v: &[i64], first: usize, dst: &mut [i64]) -> bool {
2302 use ScalarDyad::*;
2303 match op {
2304 Add => windows_into(v, k, first, dst, &i64::overflowing_add),
2305 Mul => windows_into(v, k, first, dst, &i64::overflowing_mul),
2306 Min => windows_into(v, k, first, dst, &|a: i64, b: i64| (a.min(b), false)),
2307 Max => windows_into(v, k, first, dst, &|a: i64, b: i64| (a.max(b), false)),
2308 _ => false,
2309 }
2310}
2311
2312fn scan_pass_f64(op: ScalarDyad, v: &[f64], carry: Option<f64>, dst: &mut [f64]) -> Option<f64> {
2313 use ScalarDyad::*;
2314 match op {
2315 Add => scan_block(v, carry, dst, &|a: f64, b: f64| block_f64(a + b)),
2316 Mul => scan_block(v, carry, dst, &|a: f64, b: f64| block_f64(a * b)),
2317 Min => scan_block(v, carry, dst, &|a: f64, b: f64| (a.min(b), false)),
2318 Max => scan_block(v, carry, dst, &|a: f64, b: f64| (a.max(b), false)),
2319 _ => None,
2320 }
2321}
2322
2323fn scan_pass_i64(op: ScalarDyad, v: &[i64], carry: Option<i64>, dst: &mut [i64]) -> Option<i64> {
2324 use ScalarDyad::*;
2325 match op {
2326 Add => scan_block(v, carry, dst, &i64::overflowing_add),
2327 Mul => scan_block(v, carry, dst, &i64::overflowing_mul),
2328 Min => scan_block(v, carry, dst, &|a: i64, b: i64| (a.min(b), false)),
2329 Max => scan_block(v, carry, dst, &|a: i64, b: i64| (a.max(b), false)),
2330 _ => None,
2331 }
2332}
2333
2334fn step_i64(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
2336 use ScalarDyad::*;
2337 match op {
2338 Add => a.checked_add(b),
2339 Mul => a.checked_mul(b),
2340 Min => Some(a.min(b)),
2341 Max => Some(a.max(b)),
2342 _ => None,
2343 }
2344}
2345
2346pub(crate) fn step(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
2349 step_f64(op, a, b)
2350}
2351
2352fn step_f64(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
2353 use ScalarDyad::*;
2354 match op {
2355 Add => Some(a + b).filter(|r| !r.is_nan()),
2358 Mul => Some(a * b).filter(|r| !r.is_nan()),
2359 Min => Some(a.min(b)),
2360 Max => Some(a.max(b)),
2361 _ => None,
2362 }
2363}
2364
2365fn splat_f64(a: &Array, w: usize) -> Option<Vec<f64>> {
2371 if a.rank() != 0 {
2372 return None;
2373 }
2374 let v = match &a.data {
2375 Data::Bool(d) => d[0] as f64,
2376 Data::I64(d) => d[0] as f64,
2377 Data::F64(d) => d[0],
2378 Data::Ext(_)
2379 | Data::Rat(_)
2380 | Data::Complex(_)
2381 | Data::Char(_)
2382 | Data::Symbol(_)
2383 | Data::Box(_) => {
2384 return Some(Vec::new());
2385 }
2386 };
2387 Some(vec![v; w])
2388}
2389
2390fn splat_i64(a: &Array, w: usize) -> Option<Vec<i64>> {
2391 if a.rank() != 0 {
2392 return None;
2393 }
2394 let v = match &a.data {
2395 Data::Bool(d) => d[0] as i64,
2396 Data::I64(d) => d[0],
2397 _ => return Some(Vec::new()),
2398 };
2399 Some(vec![v; w])
2400}
2401
2402fn narrow_f64(a: &Array) -> Result<&[f64], Narrow<'_>> {
2405 match &a.data {
2406 Data::I64(d) => Err(Narrow::I64(d)),
2407 Data::Bool(d) => Err(Narrow::Bool(d)),
2408 _ => Ok(a.as_f64_slice().unwrap_or(&[])),
2409 }
2410}
2411
2412fn narrow_i64(a: &Array) -> Result<&[i64], Narrow<'_>> {
2413 match &a.data {
2414 Data::Bool(d) => Err(Narrow::Bool(d)),
2415 _ => Ok(a.as_i64_slice().unwrap_or(&[])),
2416 }
2417}
2418
2419fn sources<'a, T>(
2422 inputs: &'a [Array],
2423 owned: &'a [Option<Vec<T>>],
2424 on: &impl Fn(usize) -> On,
2425 narrow: impl Fn(&'a Array) -> Result<&'a [T], Narrow<'a>>,
2426) -> Sources<'a, T> {
2427 let mut staged = 0;
2428 let of = inputs
2429 .iter()
2430 .zip(owned)
2431 .enumerate()
2432 .map(|(j, (a, o))| match o {
2433 Some(v) => Source::Ready(Loaded { data: v, splat: true, on: on(j), base: 0 }),
2434 None => match narrow(a) {
2435 Ok(d) => Source::Ready(Loaded { data: d, splat: false, on: on(j), base: 0 }),
2436 Err(n) => {
2437 staged += 1;
2438 Source::Staged(n, on(j))
2439 }
2440 },
2441 })
2442 .collect();
2443 Sources { of, staged }
2444}
2445
2446pub(crate) fn common_shape(inputs: &[Array]) -> Option<Option<Vec<usize>>> {
2449 let mut shape: Option<&Vec<usize>> = None;
2450 for a in inputs {
2451 if a.rank() == 0 {
2452 continue;
2453 }
2454 match shape {
2455 None => shape = Some(&a.shape),
2456 Some(s) if *s == a.shape => {}
2457 Some(_) => return None,
2458 }
2459 }
2460 Some(shape.cloned())
2461}
2462
2463struct Axes {
2475 shape: Vec<usize>,
2476 wide: usize,
2478}
2479
2480fn axes(k: &FusedKernel, inputs: &[Array]) -> Option<Axes> {
2481 let Some(window) = k.window else {
2482 let shape = common_shape(inputs)??;
2485 if k.scans > 0 && shape.len() != 1 {
2489 return None;
2490 }
2491 return Some(Axes { shape, wide: 0 });
2492 };
2493 let (mut wide, mut result) = (None, None);
2494 for (a, dom) in inputs.iter().zip(&k.doms) {
2495 if a.rank() == 0 {
2498 continue;
2499 }
2500 let (Some(d), 1) = (dom, a.rank()) else { return None };
2503 let seen = if *d == Dom::Wide { &mut wide } else { &mut result };
2504 match seen {
2505 None => *seen = Some(a.shape[0]),
2506 Some(m) if *m == a.shape[0] => {}
2507 Some(_) => return None,
2508 }
2509 }
2510 let wide = wide?;
2511 if wide < window {
2512 return None;
2515 }
2516 let count = wide - window + 1;
2517 if result.is_some_and(|m| m != count) {
2518 return None;
2519 }
2520 Some(Axes { shape: vec![count], wide })
2521}
2522
2523pub(crate) fn run(k: &FusedKernel, inputs: &[Array]) -> Option<Array> {
2527 let reducing = matches!(k.yields, Yield::Reduce(_));
2528 let Axes { shape, wide } = axes(k, inputs)?;
2529 let n: usize = shape.iter().product();
2530 if n == 0 {
2531 return None;
2532 }
2533 if reducing && (shape.len() != 1 || n < 2) {
2534 return None;
2537 }
2538 let (working, root) = working_type(k, inputs)?;
2539 if k.yields == Yield::Tally {
2540 return Some(Array::scalar_i64(shape[0] as i64));
2544 }
2545 let w = BLOCK.min(n).max(1) + 3 * k.window.unwrap_or(0);
2548 let tol = k.tol;
2551 let on = |j: usize| placed(k.doms[j]);
2552
2553 let data = if working == DType::F64 {
2554 let steps = Steps {
2555 monad: move |op, a: &[f64], dst: &mut [f64]| {
2556 monad_f64(op, a, dst, tol) && no_nan_block(dst)
2557 },
2558 dyad: move |op, a: &[f64], b: &[f64], dst: &mut [f64]| {
2559 dyad_f64(op, a, b, dst, tol) && no_nan_block(dst)
2560 },
2561 window: window_pass_f64,
2562 scan: scan_pass_f64,
2563 };
2564 let owned: Vec<Option<Vec<f64>>> = inputs.iter().map(|a| splat_f64(a, w)).collect();
2565 let srcs = sources(inputs, &owned, &on, narrow_f64);
2566 match k.reduce() {
2567 None => {
2568 let out = map_pass(k, &srcs, n, wide, &steps)?;
2569 float_result(out, root)
2570 }
2571 Some(op) => {
2572 let v = reduce_pass(k, &srcs, n, wide, &steps, |a, b| step_f64(op, a, b))?;
2573 match root {
2576 DType::F64 => Data::F64(vec![v].into()),
2577 _ => Data::I64(vec![v as i64].into()),
2578 }
2579 }
2580 }
2581 } else {
2582 let steps = Steps {
2583 monad: monad_i64,
2584 dyad: dyad_i64,
2585 window: window_pass_i64,
2586 scan: scan_pass_i64,
2587 };
2588 let owned: Vec<Option<Vec<i64>>> = inputs.iter().map(|a| splat_i64(a, w)).collect();
2589 let srcs = sources(inputs, &owned, &on, narrow_i64);
2590 match k.reduce() {
2591 None => {
2592 let out = map_pass(k, &srcs, n, wide, &steps)?;
2593 int_result(out, root)
2594 }
2595 Some(op) => {
2596 let v = reduce_pass(k, &srcs, n, wide, &steps, |a, b| step_i64(op, a, b))?;
2597 Data::I64(vec![v].into())
2598 }
2599 }
2600 };
2601 Some(Array::new(if reducing { Vec::new() } else { shape }, data))
2602}
2603
2604fn float_result(out: Vec<f64>, root: DType) -> Data {
2608 match root {
2609 DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0.0) as u8).into()),
2610 _ => Data::F64(out.into()),
2611 }
2612}
2613
2614fn int_result(out: Vec<i64>, root: DType) -> Data {
2615 match root {
2616 DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0) as u8).into()),
2617 _ => Data::I64(out.into()),
2618 }
2619}
2620
2621#[derive(Clone, Copy, Debug, PartialEq, Eq)]
2630pub enum Decline {
2631 Agreement,
2634 Empty,
2636 ReduceShape,
2638 WorkingType,
2641 Overflow,
2644 Window,
2647}
2648
2649impl Decline {
2650 pub fn reason(self) -> &'static str {
2651 match self {
2652 Decline::Agreement => "the inputs need agreement or are all scalars",
2653 Decline::Empty => "there is nothing to compute",
2654 Decline::ReduceShape => "the reduction needs one axis of two or more items",
2655 Decline::WorkingType => "no single working type holds every step exactly",
2656 Decline::Overflow => "an integer step left 64-bit range",
2657 Decline::Window => "the window does not fit the axis, or the inputs are not aligned with it",
2658 }
2659 }
2660}
2661
2662pub fn decline_reason(k: &FusedKernel, inputs: &[Array]) -> Option<Decline> {
2669 let Some(Axes { shape, .. }) = axes(k, inputs) else {
2670 return Some(if k.window.is_some() { Decline::Window } else { Decline::Agreement });
2671 };
2672 let n: usize = shape.iter().product();
2673 if n == 0 {
2674 return Some(Decline::Empty);
2675 }
2676 if matches!(k.yields, Yield::Reduce(_)) && (shape.len() != 1 || n < 2) {
2677 return Some(Decline::ReduceShape);
2678 }
2679 if working_type(k, inputs).is_none() {
2680 return Some(Decline::WorkingType);
2681 }
2682 Some(Decline::Overflow)
2683}
2684
2685#[derive(Clone, Debug, PartialEq, Eq)]
2687pub struct Summary {
2688 pub ops: usize,
2690 pub op_names: Vec<String>,
2692 pub reduce: Option<&'static str>,
2694 pub tally: bool,
2696 pub lets: usize,
2698 pub inputs: usize,
2700 pub block: usize,
2702 pub window: Option<usize>,
2704 pub scans: usize,
2706}
2707
2708impl std::fmt::Display for Summary {
2709 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2710 write!(f, "{} op{}", self.ops, if self.ops == 1 { "" } else { "s" })?;
2711 if !self.op_names.is_empty() {
2712 write!(f, ": {}", self.op_names.join(" "))?;
2713 }
2714 if let Some(r) = self.reduce {
2715 write!(f, "; {r}/ absorbed")?;
2716 }
2717 if self.tally {
2718 write!(f, "; tally only")?;
2719 }
2720 if self.lets > 0 {
2721 write!(f, "; {} let slot{}", self.lets, if self.lets == 1 { "" } else { "s" })?;
2722 }
2723 if let Some(k) = self.window {
2724 write!(f, "; window {k}")?;
2725 }
2726 if self.scans > 0 {
2727 write!(f, "; {} running fold{}", self.scans, if self.scans == 1 { "" } else { "s" })?;
2728 }
2729 write!(f, "; block {}", self.block)
2730 }
2731}
2732
2733pub fn summary(k: &FusedKernel) -> Summary {
2735 let mut op_names: Vec<String> = Vec::new();
2736 let mut lets = 0usize;
2737 for ins in &k.code {
2738 match ins {
2739 Instr::Monad(op) => op_names.push(monad_name(*op).to_string()),
2740 Instr::Dyad(op) => op_names.push(dyad_name(*op).to_string()),
2741 Instr::Window(op, k) => op_names.push(format!("{k} {}/\\", dyad_name(*op))),
2742 Instr::Scan(op) => op_names.push(format!("{}/\\", dyad_name(*op))),
2743 Instr::Store(_) => lets += 1,
2744 Instr::Load(_) | Instr::Let(_) => {}
2745 }
2746 }
2747 Summary {
2748 ops: op_names.len(),
2749 op_names,
2750 reduce: k.reduce().map(dyad_name),
2751 tally: k.yields == Yield::Tally,
2752 lets,
2753 inputs: k.leaves.iter().copied().max().map_or(0, |m| m + 1),
2754 block: BLOCK,
2755 window: k.window,
2756 scans: k.scans,
2757 }
2758}
2759
2760pub fn inlined_names(p: &Program) -> Vec<String> {
2763 let Some(Expr::Elided { orig, .. }) = p.stmts.first() else { return Vec::new() };
2764 let assigned = |stmts: &[Expr]| -> Vec<String> {
2765 stmts
2766 .iter()
2767 .filter_map(|s| match s {
2768 Expr::Assign { name, .. } => Some(name.clone()),
2769 _ => None,
2770 })
2771 .collect()
2772 };
2773 let kept = assigned(&p.stmts);
2774 assigned(orig).into_iter().filter(|n| !kept.contains(n)).collect()
2775}
2776
2777fn monad_name(op: ScalarMonad) -> &'static str {
2780 use ScalarMonad::*;
2781 match op {
2782 Conj => "+",
2783 Neg => "-",
2784 Signum => "*",
2785 Recip => "%",
2786 Sqrt => "%:",
2787 Exp => "^",
2788 Abs => "|",
2789 Floor => "<.",
2790 Ceil => ">.",
2791 Not => "-.",
2792 OneMinus => "-.",
2793 Inc => ">:",
2794 Dec => "<:",
2795 Double => "+:",
2796 Halve => "-:",
2797 Square => "*:",
2798 Ln => "^.",
2799 Pi => "o.",
2800 Factorial => "!",
2801 Imaginary => "j.",
2802 Polar => "r.",
2803 }
2804}
2805
2806pub(crate) fn dyad_name(op: ScalarDyad) -> &'static str {
2807 use ScalarDyad::*;
2808 match op {
2809 Add => "+",
2810 Sub => "-",
2811 Mul => "*",
2812 DivJ | DivApl => "%",
2813 Min => "<.",
2814 Max => ">.",
2815 Pow => "^",
2816 Residue => "|",
2817 Eq => "=",
2818 Ne => "~:",
2819 Lt => "<",
2820 Le => "<:",
2821 Gt => ">",
2822 Ge => ">:",
2823 Lcm => "*.",
2824 Gcd => "+.",
2825 Log => "^.",
2826 Root => "%:",
2827 Circle => "o.",
2828 Binomial => "!",
2829 MakeComplex => "j.",
2830 PolarBy => "r.",
2831 }
2832}
2833
2834pub(crate) fn eval_on(
2843 device: Option<&crate::device::Device>,
2844 k: &FusedKernel,
2845 inputs: &[Array],
2846) -> (Option<Array>, crate::device::Placement) {
2847 use crate::device::Placement;
2848 let mut placement = Placement::Default;
2849 let materialised: Vec<Array>;
2854 let (inputs, layout) = match kernel_layout(inputs) {
2855 Some(l) => (inputs, l),
2856 None => {
2857 materialised = inputs.iter().map(Array::to_row_major).collect();
2858 (&materialised[..], Layout::RowMajor)
2859 }
2860 };
2861 if layout == Layout::RowMajor && let Some(d) = device.filter(|d| d.is_gpu()) {
2864 match crate::device::try_run(d, k, inputs) {
2865 Ok(a) if finite_answer(&a) => return (Some(a), Placement::Gpu),
2871 Ok(_) => placement = Placement::Cpu(crate::device::Refusal::NonFinite),
2872 Err(why) => placement = Placement::Cpu(why),
2873 }
2874 }
2875 let r = run(k, inputs).map(|a| a.with_layout(layout));
2876 if r.is_none() {
2877 note_fallback();
2878 }
2879 (r, placement)
2880}
2881
2882fn finite_answer(a: &Array) -> bool {
2885 match &a.data {
2886 crate::array::Data::F64(v) => finite_block(v),
2887 _ => true,
2888 }
2889}
2890
2891fn kernel_layout(inputs: &[Array]) -> Option<Layout> {
2894 let mut found: Option<Layout> = None;
2895 for a in inputs {
2896 if a.rank() == 0 {
2899 continue;
2900 }
2901 match found {
2902 None => found = Some(a.layout()),
2903 Some(l) if l == a.layout() => {}
2904 Some(_) => return None,
2905 }
2906 }
2907 Some(found.unwrap_or_default())
2908}
2909
2910#[cfg(test)]
2911mod tests {
2912 use super::*;
2913 use crate::frontend::{compile, Dialect, Lang};
2914
2915 fn program(src: &str) -> Program {
2916 compile(Lang::J, src, &Dialect::default()).expect("compile")
2917 }
2918
2919 #[test]
2920 fn a_chain_of_two_scalar_verbs_fuses() {
2921 assert!(is_fused(&program("1 + 2 * {x}")));
2922 assert!(is_fused(&program("+/ {w} * {x}")));
2923 assert!(is_fused(&program("+/ ^ {x}")));
2924 }
2925
2926 #[test]
2927 fn one_verb_on_its_own_is_left_alone() {
2928 assert!(!is_fused(&program("2 * {x}")));
2929 assert!(!is_fused(&program("+/ {x}")));
2930 assert!(!is_fused(&program("{x}")));
2931 }
2932
2933 #[test]
2934 fn a_verb_the_kernel_does_not_cover_breaks_the_chain() {
2935 assert!(!is_fused(&program("%: 2 * {x}")));
2938 assert!(is_fused(&program("%: 1 + 2 * {x}")));
2939 }
2940
2941 #[test]
2942 fn an_effect_in_a_leaf_keeps_the_chain_unfused() {
2943 assert!(!is_fused(&program("1 + 2 * echo {x}")));
2944 }
2945
2946 #[test]
2947 fn the_postfix_program_pushes_the_left_operand_first() {
2948 let p = program("{w} - {x} - 1");
2949 let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2950 assert_eq!(
2951 kernel.code(),
2952 [
2953 Instr::Load(2),
2954 Instr::Load(1),
2955 Instr::Load(0),
2956 Instr::Dyad(ScalarDyad::Sub),
2957 Instr::Dyad(ScalarDyad::Sub),
2958 ]
2959 );
2960 assert_eq!(kernel.slots, 2);
2962 }
2963
2964 #[test]
2965 fn a_value_the_chain_reads_twice_becomes_a_let() {
2966 let p = program("+/ ({x} + 1) * ({x} + 1)");
2969 let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
2970 assert_eq!(
2971 kernel.code(),
2972 [
2973 Instr::Load(1),
2974 Instr::Load(0),
2975 Instr::Dyad(ScalarDyad::Add),
2976 Instr::Store(0),
2977 Instr::Let(0),
2978 Instr::Let(0),
2979 Instr::Dyad(ScalarDyad::Mul),
2980 ]
2981 );
2982 assert_eq!(kernel.slots, 2);
2984 }
2985
2986 #[test]
2987 fn a_named_value_moves_into_the_sentence_that_reads_it() {
2988 let p = program("d =. {x} + 1\n+/ d * d");
2989 assert!(is_inlined(&p));
2990 assert_eq!(p.stmts.len(), 3);
2994 let Expr::Fused { kernel, .. } = &p.stmts[2] else { panic!("the sum did not fuse") };
2995 assert!(kernel.code().contains(&Instr::Store(0)));
2996 assert_eq!(unfused(&p).stmts.len(), 2);
2997 }
2998}