use std::sync::atomic::{AtomicU64, Ordering};
use crate::array::{Array, Data};
use crate::dtype::DType;
use crate::error::Span;
use crate::ir::{Expr, Program, Scope};
use crate::par;
use crate::simd::multiversioned;
use crate::verb::{tol_cmp, DyadOp, MonadOp, ScalarDyad, ScalarMonad, Tol, Verb, RANK_INF};
pub const BLOCK: usize = 8_192;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Instr {
Load(usize),
Monad(ScalarMonad),
Dyad(ScalarDyad),
Store(usize),
Let(usize),
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Yield {
Values,
Reduce(ScalarDyad),
Tally,
}
#[derive(Clone, Debug)]
pub struct FusedKernel {
code: Vec<Instr>,
slots: usize,
yields: Yield,
leaves: Vec<usize>,
tol: Tol,
}
impl FusedKernel {
pub fn code(&self) -> &[Instr] {
&self.code
}
pub fn yields(&self) -> Yield {
self.yields
}
pub fn reduce(&self) -> Option<ScalarDyad> {
match self.yields {
Yield::Reduce(op) => Some(op),
_ => None,
}
}
pub fn tol(&self) -> Tol {
self.tol
}
}
static FALLBACKS: AtomicU64 = AtomicU64::new(0);
pub fn fallback_count() -> u64 {
FALLBACKS.load(Ordering::Relaxed)
}
fn note_fallback() {
FALLBACKS.fetch_add(1, Ordering::Relaxed);
}
fn fusable_monad(v: &Verb) -> Option<ScalarMonad> {
use ScalarMonad::*;
let Verb::Prim(p) = v else { return None };
let MonadOp::Scalar(op) = p.monad else { return None };
matches!(
op,
Conj | Neg | Abs | Signum | Recip | Floor | Ceil | Inc | Dec | Double | Halve | Square
| OneMinus | Exp
)
.then_some(op)
}
fn fusable_dyad(v: &Verb) -> Option<ScalarDyad> {
use ScalarDyad::*;
let Verb::Prim(p) = v else { return None };
let DyadOp::Scalar(op) = p.dyad else { return None };
matches!(op, Add | Sub | Mul | DivJ | Min | Max | Residue | Eq | Ne | Lt | Le | Gt | Ge)
.then_some(op)
}
fn absorbable_reduce(v: &Verb) -> Option<ScalarDyad> {
use ScalarDyad::*;
let inner = match v {
Verb::Reduce(u) => u,
Verb::Rank(u, r) if r[0] >= 1 => match &**u {
Verb::Reduce(inner) => inner,
_ => return None,
},
_ => return None,
};
let Verb::Prim(p) = &**inner else { return None };
let DyadOp::Scalar(op) = p.dyad else { return None };
matches!(op, Add | Mul | Min | Max).then_some(op)
}
fn is_tally(v: &Verb) -> bool {
matches!(v, Verb::Prim(p) if p.monad == MonadOp::Tally && p.ranks[0] == RANK_INF)
}
#[derive(Clone, PartialEq)]
enum Node {
Leaf(usize),
Monad(ScalarMonad, Box<Node>),
Dyad(ScalarDyad, Box<Node>, Box<Node>),
}
#[derive(Default)]
struct Leaves<'a> {
inputs: Vec<&'a Expr>,
order: Vec<usize>,
}
impl<'a> Leaves<'a> {
fn push(&mut self, e: &'a Expr) -> usize {
let i = match self.inputs.iter().position(|&p| same(p, e)) {
Some(i) => i,
None => {
self.inputs.push(e);
self.inputs.len() - 1
}
};
self.order.push(i);
i
}
}
struct Inline<'a> {
name: &'a str,
def: &'a Expr,
hits: usize,
}
fn chain<'a>(e: &'a Expr, lv: &mut Leaves<'a>, sub: &mut Option<Inline<'a>>) -> Node {
let read_through = match (e, sub.as_ref()) {
(Expr::Name(n, _), Some(s)) if n == s.name => Some(s.def),
_ => None,
};
if let Some(def) = read_through {
if let Some(s) = sub.as_mut() {
s.hits += 1;
}
return chain(def, lv, sub);
}
match e {
Expr::Monad { verb, y, .. } => match fusable_monad(verb) {
Some(op) => Node::Monad(op, Box::new(chain(y, lv, sub))),
None => Node::Leaf(lv.push(e)),
},
Expr::Dyad { verb, x, y, .. } => match fusable_dyad(verb) {
Some(op) => {
let ry = chain(y, lv, sub);
let rx = chain(x, lv, sub);
Node::Dyad(op, Box::new(rx), Box::new(ry))
}
None => Node::Leaf(lv.push(e)),
},
_ => Node::Leaf(lv.push(e)),
}
}
fn ops(n: &Node) -> usize {
match n {
Node::Leaf(_) => 0,
Node::Monad(_, y) => 1 + ops(y),
Node::Dyad(_, x, y) => 1 + ops(x) + ops(y),
}
}
fn subtrees<'a>(n: &'a Node, out: &mut Vec<&'a Node>) {
if ops(n) == 0 {
return;
}
out.push(n);
match n {
Node::Leaf(_) => {}
Node::Monad(_, y) => subtrees(y, out),
Node::Dyad(_, x, y) => {
subtrees(x, out);
subtrees(y, out);
}
}
}
fn lets_of(n: &Node) -> Vec<Node> {
let mut all = Vec::new();
subtrees(n, &mut all);
let mut out = Vec::new();
fn walk(n: &Node, all: &[&Node], out: &mut Vec<Node>) {
if ops(n) >= 1 && all.iter().filter(|m| **m == n).count() >= 2 {
if !out.contains(n) {
out.push(n.clone());
}
return;
}
match n {
Node::Leaf(_) => {}
Node::Monad(_, y) => walk(y, all, out),
Node::Dyad(_, x, y) => {
walk(x, all, out);
walk(y, all, out);
}
}
}
walk(n, &all, &mut out);
out
}
fn emit_all(n: &Node, lets: &[Node], code: &mut Vec<Instr>) {
for (k, l) in lets.iter().enumerate() {
emit(l, &lets[..k], code);
code.push(Instr::Store(k));
}
emit(n, lets, code);
}
fn emit(n: &Node, lets: &[Node], code: &mut Vec<Instr>) {
if let Some(k) = lets.iter().position(|l| l == n) {
code.push(Instr::Let(k));
return;
}
match n {
Node::Leaf(i) => code.push(Instr::Load(*i)),
Node::Monad(op, y) => {
emit(y, lets, code);
code.push(Instr::Monad(*op));
}
Node::Dyad(op, x, y) => {
emit(x, lets, code);
emit(y, lets, code);
code.push(Instr::Dyad(*op));
}
}
}
fn slots(code: &[Instr]) -> usize {
let mut stack: Vec<bool> = Vec::new();
let mut live = 0usize;
let mut max = 1usize;
for ins in code {
let operands = match ins {
Instr::Load(_) => {
stack.push(false);
continue;
}
Instr::Let(_) => {
stack.push(false);
continue;
}
Instr::Store(_) => {
stack.pop();
continue;
}
Instr::Monad(_) => 1,
Instr::Dyad(_) => 2,
};
max = max.max(live + 1);
for _ in 0..operands {
if stack.pop().unwrap_or(false) {
live -= 1;
}
}
live += 1;
stack.push(true);
}
max
}
fn replayable(e: &Expr) -> bool {
match e {
Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => true,
Expr::Assign { .. }
| Expr::PrintPass { .. }
| Expr::Elided { .. }
| Expr::Control(..)
| Expr::AmendIndex { .. }
| Expr::VerbDef { .. } => false,
Expr::Monad { verb, y, .. } => verb.is_pure() && replayable(y),
Expr::Dyad { verb, x, y, .. } => verb.is_pure() && replayable(x) && replayable(y),
Expr::Fused { inputs, .. } => inputs.iter().all(replayable),
}
}
fn same(a: &Expr, b: &Expr) -> bool {
match (a, b) {
(Expr::Const(p, _), Expr::Const(q, _)) => p == q,
(Expr::Param(p, _), Expr::Param(q, _)) => p == q,
(Expr::Name(p, _), Expr::Name(q, _)) => p == q,
(Expr::Monad { verb: u, y: p, .. }, Expr::Monad { verb: v, y: q, .. }) => {
same_verb(u, v) && same(p, q)
}
(
Expr::Dyad { verb: u, x: px, y: py, .. },
Expr::Dyad { verb: v, x: qx, y: qy, .. },
) => same_verb(u, v) && same(px, qx) && same(py, qy),
_ => false,
}
}
fn same_verb(a: &Verb, b: &Verb) -> bool {
match (a, b) {
(Verb::Prim(p), Verb::Prim(q)) => p == q,
(Verb::Rank(u, r), Verb::Rank(v, s)) => r == s && same_verb(u, v),
(Verb::Reduce(u), Verb::Reduce(v)) | (Verb::Commute(u), Verb::Commute(v)) => {
same_verb(u, v)
}
(Verb::Windowed(u, j), Verb::Windowed(v, k)) => j == k && same_verb(u, v),
(Verb::PowerN(u, m), Verb::PowerN(v, n)) => m == n && same_verb(u, v),
(Verb::Fork(f, g, h), Verb::Fork(f2, g2, h2)) => {
same_verb(f, f2) && same_verb(g, g2) && same_verb(h, h2)
}
(Verb::NounFork(m, g, h), Verb::NounFork(n, g2, h2)) => {
m == n && same_verb(g, g2) && same_verb(h, h2)
}
(Verb::Hook(g, h), Verb::Hook(g2, h2))
| (Verb::Atop(g, h), Verb::Atop(g2, h2))
| (Verb::Compose(g, h), Verb::Compose(g2, h2)) => same_verb(g, g2) && same_verb(h, h2),
(Verb::BondLeft(m, u), Verb::BondLeft(n, v)) => m == n && same_verb(u, v),
(Verb::BondRight(u, m), Verb::BondRight(v, n)) => m == n && same_verb(u, v),
_ => false,
}
}
pub fn pass(stmts: &mut Vec<Expr>, tol: Tol) {
let orig = std::mem::take(stmts);
let mut cur = orig.clone();
let mut names = 0usize;
let mut crossed = false;
for _ in 0..=orig.len() {
match inline_once(&cur, &mut names, tol) {
Some(next) => {
cur = next;
crossed = true;
}
None => break,
}
}
let mut out: Vec<Expr> = cur.into_iter().map(|e| fuse_expr(e, tol)).collect();
if crossed {
out.insert(0, Expr::Elided { orig, span: Span::new(0, 0) });
}
*stmts = out;
}
fn fuse_expr(e: Expr, tol: Tol) -> Expr {
if let Some(f) = try_fuse(&e, tol) {
return f;
}
match e {
Expr::Assign { name, value, scope, span } => {
Expr::Assign { name, value: Box::new(fuse_expr(*value, tol)), scope, span }
}
Expr::Monad { verb, y, span } => {
Expr::Monad { verb, y: Box::new(fuse_expr(*y, tol)), span }
}
Expr::Dyad { verb, x, y, span } => Expr::Dyad {
verb,
x: Box::new(fuse_expr(*x, tol)),
y: Box::new(fuse_expr(*y, tol)),
span,
},
Expr::PrintPass { value, span } => {
Expr::PrintPass { value: Box::new(fuse_expr(*value, tol)), span }
}
other => other,
}
}
fn build<'a>(
root: &'a Expr,
yields: Yield,
least: usize,
sub: &mut Option<Inline<'a>>,
tol: Tol,
) -> Option<(FusedKernel, Vec<&'a Expr>)> {
if let Some(s) = sub.as_mut() {
s.hits = 0;
}
let mut lv = Leaves::default();
let node = chain(root, &mut lv, sub);
if ops(&node) < least || !lv.inputs.iter().all(|l| replayable(l)) {
return None;
}
let mut code = Vec::new();
emit_all(&node, &lets_of(&node), &mut code);
let kernel = FusedKernel { slots: slots(&code), code, yields, leaves: lv.order, tol };
Some((kernel, lv.inputs))
}
fn kernel_at<'a>(
e: &'a Expr,
sub: &mut Option<Inline<'a>>,
tol: Tol,
) -> Option<(FusedKernel, Vec<&'a Expr>, &'a Expr)> {
if let Expr::Monad { verb, y, .. } = e {
if is_tally(verb) {
if let Some((k, l)) = build(y, Yield::Tally, 1, sub, tol) {
return Some((k, l, e));
}
}
if let Some(op) = absorbable_reduce(verb) {
if let Some((k, l)) = build(y, Yield::Reduce(op), 1, sub, tol) {
return Some((k, l, e));
}
}
}
let (k, l) = build(e, Yield::Values, 2, sub, tol)?;
Some((k, l, e))
}
fn try_fuse(e: &Expr, tol: Tol) -> Option<Expr> {
let (kernel, leaves, orig) = kernel_at(e, &mut None, tol)?;
let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
Some(Expr::Fused {
kernel,
inputs,
orig: Box::new(orig.clone()),
span: e.span(),
})
}
pub(crate) fn fallback_tree(k: &FusedKernel, orig: &Expr, values: &[Array]) -> Expr {
let mut next = 0;
let tree = match orig {
Expr::Monad { verb, y, span } if matches!(k.yields, Yield::Reduce(_)) => Expr::Monad {
verb: verb.clone(),
y: Box::new(substitute(y, values, k, &mut next)),
span: *span,
},
Expr::Monad { verb, y, .. } if k.yields == Yield::Tally && is_tally(verb) => {
substitute(y, values, k, &mut next)
}
e => substitute(e, values, k, &mut next),
};
debug_assert_eq!(next, k.leaves.len(), "the fallback found different leaves");
tree
}
pub(crate) fn fallback_finish(k: &FusedKernel, v: Array) -> Array {
match k.yields {
Yield::Tally => Array::scalar_i64(v.items() as i64),
_ => v,
}
}
fn substitute(e: &Expr, values: &[Array], k: &FusedKernel, next: &mut usize) -> Expr {
match e {
Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
verb: verb.clone(),
y: Box::new(substitute(y, values, k, next)),
span: *span,
},
Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => {
let ry = substitute(y, values, k, next);
let rx = substitute(x, values, k, next);
Expr::Dyad { verb: verb.clone(), x: Box::new(rx), y: Box::new(ry), span: *span }
}
leaf => {
let v = values[k.leaves[*next]].clone();
*next += 1;
Expr::Const(v, leaf.span())
}
}
}
fn hoisted_name(n: &mut usize) -> String {
*n += 1;
format!("·{}", *n - 1)
}
fn inline_once(stmts: &[Expr], names: &mut usize, tol: Tol) -> Option<Vec<Expr>> {
for (i, stmt) in stmts.iter().enumerate() {
let Expr::Assign { name, value, span, .. } = stmt else { continue };
if !inlinable(stmts, i, name, value, tol) {
continue;
}
if let Some(out) = rewrite(stmts, i, name, value, *span, names, tol) {
return Some(out);
}
}
None
}
fn inlinable(stmts: &[Expr], i: usize, name: &str, value: &Expr, tol: Tol) -> bool {
if !replayable(value) || mentions(value, name) {
return false;
}
let mut lv = Leaves::default();
if ops(&chain(value, &mut lv, &mut None)) < 1 {
return false;
}
let mut guarded = vec![name.to_string()];
free_names(value, &mut guarded);
let later = &stmts[i + 1..];
if later.iter().any(|s| assigns_any(s, &guarded)) {
return false;
}
let mut uses = 0;
for stmt in later {
match uses_land(stmt, name, value, tol) {
Some(n) => uses += n,
None => return false,
}
}
uses > 0
}
fn uses_land(e: &Expr, name: &str, def: &Expr, tol: Tol) -> Option<usize> {
let mut sub = Some(Inline { name, def, hits: 0 });
if let Some((_, leaves, _)) = kernel_at(e, &mut sub, tol) {
let mut n = sub.map_or(0, |s| s.hits);
for l in leaves {
n += uses_land(l, name, def, tol)?;
}
return Some(n);
}
match e {
Expr::Name(n, _) if n == name => None,
Expr::Const(..) | Expr::Param(..) | Expr::Name(..) => Some(0),
Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => uses_land(value, name, def, tol),
Expr::Monad { y, .. } => uses_land(y, name, def, tol),
Expr::Dyad { x, y, .. } => Some(uses_land(x, name, def, tol)? + uses_land(y, name, def, tol)?),
Expr::Fused { .. }
| Expr::Elided { .. }
| Expr::Control(..)
| Expr::AmendIndex { .. }
| Expr::VerbDef { .. } => None,
}
}
fn rewrite(
stmts: &[Expr],
i: usize,
name: &str,
value: &Expr,
span: Span,
names: &mut usize,
tol: Tol,
) -> Option<Vec<Expr>> {
let mut lv = Leaves::default();
chain(value, &mut lv, &mut None);
let mut hoists = Vec::new();
let mut bound: Vec<Option<String>> = Vec::new();
for l in &lv.inputs {
if matches!(l, Expr::Const(..) | Expr::Param(..) | Expr::Name(..)) {
bound.push(None);
continue;
}
let n = hoisted_name(names);
hoists.push(Expr::Assign {
name: n.clone(),
value: Box::new((*l).clone()),
scope: Scope::Local,
span: l.span(),
});
bound.push(Some(n));
}
let def = with_leaves(value, &lv, &bound);
let (kernel, leaves) = build(&def, Yield::Tally, 1, &mut None, tol)?;
let inputs = leaves.into_iter().map(|l| fuse_expr(l.clone(), tol)).collect();
let guard = Expr::Assign {
name: hoisted_name(names),
value: Box::new(Expr::Fused {
kernel,
inputs,
orig: Box::new(def.clone()),
span,
}),
scope: Scope::Local,
span,
};
let mut out = stmts[..i].to_vec();
out.extend(hoists);
out.push(guard);
out.extend(stmts[i + 1..].iter().map(|s| replace_name(s, name, &def)));
Some(out)
}
fn with_leaves(e: &Expr, lv: &Leaves<'_>, bound: &[Option<String>]) -> Expr {
match e {
Expr::Monad { verb, y, span } if fusable_monad(verb).is_some() => Expr::Monad {
verb: verb.clone(),
y: Box::new(with_leaves(y, lv, bound)),
span: *span,
},
Expr::Dyad { verb, x, y, span } if fusable_dyad(verb).is_some() => Expr::Dyad {
verb: verb.clone(),
x: Box::new(with_leaves(x, lv, bound)),
y: Box::new(with_leaves(y, lv, bound)),
span: *span,
},
leaf => {
let bind = lv
.inputs
.iter()
.position(|&p| same(p, leaf))
.and_then(|i| bound[i].as_ref());
match bind {
Some(n) => Expr::Name(n.clone(), leaf.span()),
None => leaf.clone(),
}
}
}
}
fn replace_name(e: &Expr, name: &str, def: &Expr) -> Expr {
match e {
Expr::Name(n, _) if n == name => def.clone(),
Expr::Assign { name: a, value, scope, span } => Expr::Assign {
scope: *scope,
name: a.clone(),
value: Box::new(replace_name(value, name, def)),
span: *span,
},
Expr::PrintPass { value, span } => Expr::PrintPass {
value: Box::new(replace_name(value, name, def)),
span: *span,
},
Expr::Monad { verb, y, span } => Expr::Monad {
verb: verb.clone(),
y: Box::new(replace_name(y, name, def)),
span: *span,
},
Expr::Dyad { verb, x, y, span } => Expr::Dyad {
verb: verb.clone(),
x: Box::new(replace_name(x, name, def)),
y: Box::new(replace_name(y, name, def)),
span: *span,
},
other => other.clone(),
}
}
fn mentions(e: &Expr, name: &str) -> bool {
let mut names = Vec::new();
free_names(e, &mut names);
names.iter().any(|n| n == name)
}
fn free_names(e: &Expr, out: &mut Vec<String>) {
match e {
Expr::Name(n, _) => out.push(n.clone()),
Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => free_names(value, out),
Expr::Monad { y, .. } => free_names(y, out),
Expr::Dyad { x, y, .. } => {
free_names(x, out);
free_names(y, out);
}
Expr::Fused { inputs, .. } => inputs.iter().for_each(|i| free_names(i, out)),
Expr::Const(..)
| Expr::Param(..)
| Expr::Elided { .. }
| Expr::Control(..)
| Expr::AmendIndex { .. }
| Expr::VerbDef { .. } => {}
}
}
fn assigns_any(e: &Expr, names: &[String]) -> bool {
match e {
Expr::Assign { name, value, .. } => {
names.iter().any(|n| n == name) || assigns_any(value, names)
}
Expr::PrintPass { value, .. } => assigns_any(value, names),
Expr::Monad { y, .. } => assigns_any(y, names),
Expr::Dyad { x, y, .. } => assigns_any(x, names) || assigns_any(y, names),
Expr::Fused { inputs, .. } => inputs.iter().any(|i| assigns_any(i, names)),
Expr::Const(..)
| Expr::Param(..)
| Expr::Name(..)
| Expr::Elided { .. }
| Expr::Control(..)
| Expr::AmendIndex { .. }
| Expr::VerbDef { .. } => false,
}
}
pub fn is_fused(p: &Program) -> bool {
fn any(e: &Expr) -> bool {
match e {
Expr::Fused { .. } => true,
Expr::Const(..)
| Expr::Param(..)
| Expr::Name(..)
| Expr::Elided { .. }
| Expr::Control(..)
| Expr::AmendIndex { .. }
| Expr::VerbDef { .. } => false,
Expr::Assign { value, .. } | Expr::PrintPass { value, .. } => any(value),
Expr::Monad { y, .. } => any(y),
Expr::Dyad { x, y, .. } => any(x) || any(y),
}
}
p.stmts.iter().any(any)
}
pub fn is_inlined(p: &Program) -> bool {
matches!(p.stmts.first(), Some(Expr::Elided { .. }))
}
pub fn unfused(p: &Program) -> Program {
fn strip(e: &Expr) -> Expr {
match e {
Expr::Fused { orig, .. } => strip(orig),
Expr::Assign { name, value, scope, span } => {
Expr::Assign {
name: name.clone(),
value: Box::new(strip(value)),
scope: *scope,
span: *span,
}
}
Expr::PrintPass { value, span } => {
Expr::PrintPass { value: Box::new(strip(value)), span: *span }
}
Expr::Monad { verb, y, span } => {
Expr::Monad { verb: verb.clone(), y: Box::new(strip(y)), span: *span }
}
Expr::Dyad { verb, x, y, span } => Expr::Dyad {
verb: verb.clone(),
x: Box::new(strip(x)),
y: Box::new(strip(y)),
span: *span,
},
other => other.clone(),
}
}
let mut out = p.clone();
let stmts = match p.stmts.first() {
Some(Expr::Elided { orig, .. }) => orig,
_ => &p.stmts,
};
out.stmts = stmts.iter().map(strip).collect();
out
}
fn monad_type(op: ScalarMonad, a: DType) -> Option<DType> {
use DType::*;
use ScalarMonad::*;
if a == Complex {
return None;
}
Some(match op {
Recip | Halve | Exp => F64,
Conj | Abs | OneMinus => a,
Neg | Signum | Inc | Dec | Double | Square => match a {
Bool | I64 => I64,
other => other,
},
Floor | Ceil => match a {
Bool | I64 => I64,
_ => return None,
},
_ => return None,
})
}
fn dyad_type(op: ScalarDyad, a: DType, b: DType) -> Option<DType> {
use ScalarDyad::*;
if a == DType::Complex || b == DType::Complex {
return None;
}
match op {
Eq | Ne | Lt | Le | Gt | Ge => Some(DType::Bool),
DivJ => Some(DType::F64),
Add | Sub | Mul | Min | Max | Residue => match DType::promote(a, b)? {
DType::Bool => Some(DType::I64),
DType::Char => None,
t => Some(t),
},
_ => None,
}
}
pub(crate) fn working_type(k: &FusedKernel, inputs: &[Array]) -> Option<(DType, DType)> {
let mut stack: Vec<DType> = Vec::with_capacity(k.slots);
let mut lets: Vec<DType> = Vec::new();
let mut float = false;
let mut integer_step = false;
if inputs.iter().any(|a| a.dtype() == DType::Complex || a.dtype().is_exact()) {
return None;
}
for ins in &k.code {
let t = match ins {
Instr::Load(i) => inputs[*i].dtype(),
Instr::Monad(op) => monad_type(*op, stack.pop()?)?,
Instr::Dyad(op) => {
let b = stack.pop()?;
let a = stack.pop()?;
dyad_type(*op, a, b)?
}
Instr::Store(k) => {
let t = stack.pop()?;
if lets.len() != *k {
return None;
}
lets.push(t);
continue;
}
Instr::Let(k) => {
let t = *lets.get(*k)?;
float |= t == DType::F64;
stack.push(t);
continue;
}
};
if !t.is_numeric() {
return None;
}
float |= t == DType::F64;
integer_step |= t == DType::I64 && !matches!(ins, Instr::Load(_));
stack.push(t);
}
let root = stack.pop()?;
let working = if float { DType::F64 } else { DType::I64 };
if working == DType::F64 && integer_step {
return None;
}
Some((working, root))
}
struct Loaded<'a, T> {
data: &'a [T],
splat: bool,
}
impl<T> Loaded<'_, T> {
#[inline]
fn block(&self, start: usize, len: usize) -> &[T] {
if self.splat {
&self.data[..len]
} else {
&self.data[start..start + len]
}
}
}
#[derive(Clone, Copy)]
enum Slot {
Input(usize),
Block(usize),
}
fn split_slots<'s, T>(
scratch: &'s mut [T],
w: usize,
d: usize,
) -> (&'s mut [T], impl Fn(usize) -> &'s [T]) {
let (lo, hi) = scratch.split_at_mut(d * w);
let (dst, hi) = hi.split_at_mut(w);
let lo: &[T] = lo;
let hi: &[T] = hi;
(dst, move |i: usize| {
if i < d {
&lo[i * w..(i + 1) * w]
} else {
&hi[(i - d - 1) * w..(i - d) * w]
}
})
}
#[allow(clippy::too_many_arguments)]
fn exec_block<T, M, D>(
code: &[Instr],
srcs: &[Loaded<'_, T>],
start: usize,
len: usize,
scratch: &mut [T],
w: usize,
free: &mut Vec<usize>,
stack: &mut Vec<Slot>,
lets: &mut Vec<usize>,
out: Option<&mut [T]>,
mon: &M,
dya: &D,
) -> Option<usize>
where
T: Copy,
M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
{
stack.clear();
free.clear();
lets.clear();
let nslots = scratch.len() / w;
free.extend((0..nslots).rev());
let last = code.len() - 1;
let head = if out.is_some() { last } else { code.len() };
for ins in &code[..head] {
match ins {
Instr::Load(k) => stack.push(Slot::Input(*k)),
Instr::Monad(op) => {
let a = stack.pop()?;
let d = free.pop()?;
let (dst, get) = split_slots(scratch, w, d);
let av = match a {
Slot::Input(k) => srcs[k].block(start, len),
Slot::Block(i) => &get(i)[..len],
};
if !mon(*op, av, &mut dst[..len]) {
return None;
}
release(free, lets, a);
stack.push(Slot::Block(d));
}
Instr::Dyad(op) => {
let b = stack.pop()?;
let a = stack.pop()?;
let d = free.pop()?;
let (dst, get) = split_slots(scratch, w, d);
let av = match a {
Slot::Input(k) => srcs[k].block(start, len),
Slot::Block(i) => &get(i)[..len],
};
let bv = match b {
Slot::Input(k) => srcs[k].block(start, len),
Slot::Block(i) => &get(i)[..len],
};
if !dya(*op, av, bv, &mut dst[..len]) {
return None;
}
for s in [a, b] {
release(free, lets, s);
}
stack.push(Slot::Block(d));
}
Instr::Store(k) => {
let Slot::Block(i) = stack.pop()? else { return None };
if lets.len() != *k {
return None;
}
lets.push(i);
}
Instr::Let(k) => stack.push(Slot::Block(*lets.get(*k)?)),
}
}
let Some(dst) = out else {
return match stack.pop()? {
Slot::Block(i) => Some(i),
Slot::Input(_) => None,
};
};
let dst = &mut dst[..len];
let view = |s: Slot| match s {
Slot::Input(k) => srcs[k].block(start, len),
Slot::Block(i) => &scratch[i * w..i * w + len],
};
let ok = match code[last] {
Instr::Monad(op) => {
let a = view(stack.pop()?);
mon(op, a, dst)
}
Instr::Dyad(op) => {
let b = stack.pop()?;
let a = stack.pop()?;
dya(op, view(a), view(b), dst)
}
Instr::Load(_) | Instr::Store(_) | Instr::Let(_) => return None,
};
ok.then_some(usize::MAX)
}
fn release(free: &mut Vec<usize>, lets: &[usize], s: Slot) {
if let Slot::Block(i) = s {
if !lets.contains(&i) {
free.push(i);
}
}
}
fn map_pass<T, M, D>(
k: &FusedKernel,
srcs: &[Loaded<'_, T>],
n: usize,
mon: M,
dya: D,
) -> Option<Vec<T>>
where
T: Copy + Default + Send + Sync,
M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
{
let (out, ok) = par::fill(n, |start, part: &mut [T]| {
let w = BLOCK.min(part.len()).max(1);
let mut scratch = vec![T::default(); k.slots * w];
let mut free = Vec::with_capacity(k.slots);
let mut stack = Vec::with_capacity(k.slots);
let mut lets = Vec::new();
for (b, chunk) in part.chunks_mut(w).enumerate() {
let len = chunk.len();
let ok = exec_block(
&k.code,
srcs,
start + b * w,
len,
&mut scratch,
w,
&mut free,
&mut stack,
&mut lets,
Some(chunk),
&mon,
&dya,
);
if ok.is_none() {
return false;
}
}
true
});
ok.then_some(out)
}
const FOLD_LANES: usize = 8;
const MIN_LANE_WORK: usize = 8 * FOLD_LANES;
#[inline(always)]
fn fold_block_body<T, S>(v: &[T], step: &S) -> Option<T>
where
T: Copy,
S: Fn(T, T) -> Option<T>,
{
let n = v.len();
if n < MIN_LANE_WORK {
let mut acc = v[n - 1];
for &x in v[..n - 1].iter().rev() {
acc = step(x, acc)?;
}
return Some(acc);
}
let rows = n / FOLD_LANES;
let head = n - rows * FOLD_LANES;
let last = head + (rows - 1) * FOLD_LANES;
let mut acc = [v[last]; FOLD_LANES];
acc.copy_from_slice(&v[last..last + FOLD_LANES]);
for r in (0..rows - 1).rev() {
let row = &v[head + r * FOLD_LANES..head + (r + 1) * FOLD_LANES];
for (slot, &x) in acc.iter_mut().zip(row) {
*slot = step(x, *slot)?;
}
}
let mut a = acc[FOLD_LANES - 1];
for &x in acc[..FOLD_LANES - 1].iter().rev() {
a = step(x, a)?;
}
for &x in v[..head].iter().rev() {
a = step(x, a)?;
}
Some(a)
}
multiversioned! {
fn fold_block[T: Copy, S: Fn(T, T) -> Option<T>](
v: &[T],
step: &S,
) -> Option<T> = fold_block_body;
}
#[allow(clippy::too_many_arguments)]
fn fold_range<T, M, D, S>(
k: &FusedKernel,
srcs: &[Loaded<'_, T>],
lo: usize,
hi: usize,
mon: &M,
dya: &D,
step: &S,
) -> Option<T>
where
T: Copy + Default,
M: Fn(ScalarMonad, &[T], &mut [T]) -> bool,
D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool,
S: Fn(T, T) -> Option<T>,
{
let w = BLOCK.min(hi - lo).max(1);
let mut scratch = vec![T::default(); k.slots * w];
let mut free = Vec::with_capacity(k.slots);
let mut stack = Vec::with_capacity(k.slots);
let mut lets = Vec::new();
let mut acc: Option<T> = None;
for b in (0..(hi - lo).div_ceil(w)).rev() {
let start = lo + b * w;
let len = (hi - start).min(w);
let slot = exec_block(
&k.code, srcs, start, len, &mut scratch, w, &mut free, &mut stack, &mut lets, None,
mon, dya,
)?;
let block = fold_block(&scratch[slot * w..slot * w + len], step)?;
acc = Some(match acc {
None => block,
Some(a) => step(block, a)?,
});
}
acc
}
fn reduce_pass<T, M, D, S>(
k: &FusedKernel,
srcs: &[Loaded<'_, T>],
n: usize,
mon: M,
dya: D,
step: S,
) -> Option<T>
where
T: Copy + Default + Send + Sync,
M: Fn(ScalarMonad, &[T], &mut [T]) -> bool + Sync + Send,
D: Fn(ScalarDyad, &[T], &[T], &mut [T]) -> bool + Sync + Send,
S: Fn(T, T) -> Option<T> + Sync + Send,
{
let chunks = par::chunks(n, n * k.code.len());
if chunks < 2 {
return fold_range(k, srcs, 0, n, &mon, &dya, &step);
}
let per = n.div_ceil(chunks);
let parts = par::map_indexed(n.div_ceil(per), |c| {
fold_range(k, srcs, c * per, ((c + 1) * per).min(n), &mon, &dya, &step)
});
let mut it = parts.into_iter().rev();
let mut acc = it.next()??;
for part in it {
acc = step(part?, acc)?;
}
Some(acc)
}
macro_rules! each {
($a:expr, $dst:expr, $f:expr) => {{
let f = $f;
for (slot, &x) in $dst.iter_mut().zip($a) {
*slot = f(x);
}
return true;
}};
}
macro_rules! zip {
($a:expr, $b:expr, $dst:expr, $f:expr) => {{
let f = $f;
for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
*slot = f(x, y);
}
return true;
}};
}
#[inline(always)]
fn monad_f64_body(op: ScalarMonad, a: &[f64], dst: &mut [f64], tol: Tol) -> bool {
use ScalarMonad::*;
match op {
Conj => each!(a, dst, |x: f64| x),
Neg => each!(a, dst, |x: f64| -x),
Abs => each!(a, dst, f64::abs),
Signum => each!(a, dst, |x: f64| if tol.is_zero(x) {
0.0
} else if x > 0.0 {
1.0
} else if x < 0.0 {
-1.0
} else {
0.0
}),
Recip => each!(a, dst, |x: f64| if x == 0.0 { f64::INFINITY } else { 1.0 / x }),
Floor => each!(a, dst, f64::floor),
Ceil => each!(a, dst, f64::ceil),
Inc => each!(a, dst, |x: f64| x + 1.0),
Dec => each!(a, dst, |x: f64| x - 1.0),
Double => each!(a, dst, |x: f64| x + x),
Halve => each!(a, dst, |x: f64| x / 2.0),
Square => each!(a, dst, |x: f64| x * x),
OneMinus => each!(a, dst, |x: f64| 1.0 - x),
Exp => each!(a, dst, f64::exp),
_ => false,
}
}
#[inline(always)]
fn dyad_f64_body(op: ScalarDyad, a: &[f64], b: &[f64], dst: &mut [f64], tol: Tol) -> bool {
use ScalarDyad::*;
match op {
Add => zip!(a, b, dst, |x: f64, y: f64| x + y),
Sub => zip!(a, b, dst, |x: f64, y: f64| x - y),
Mul => zip!(a, b, dst, |x: f64, y: f64| x * y),
Min => zip!(a, b, dst, f64::min),
Max => zip!(a, b, dst, f64::max),
DivJ => zip!(a, b, dst, |x: f64, y: f64| if y == 0.0 {
if x == 0.0 { 0.0 } else { f64::INFINITY.copysign(x) }
} else {
x / y
}),
Residue => zip!(a, b, dst, |x: f64, y: f64| if x.is_infinite() {
if y == 0.0 || (y > 0.0) == (x > 0.0) { y } else { x }
} else if x == 0.0 {
y
} else {
y - x * (y / x).floor()
}),
Eq | Ne | Lt | Le | Gt | Ge => {
zip!(a, b, dst, |x: f64, y: f64| tol_cmp(op, x, y, tol) as u8 as f64)
}
_ => false,
}
}
macro_rules! each_over {
($a:expr, $dst:expr, $f:expr) => {{
let f = $f;
let mut over = false;
for (slot, &x) in $dst.iter_mut().zip($a) {
let (v, o) = f(x);
*slot = v;
over |= o;
}
return !over;
}};
}
macro_rules! zip_over {
($a:expr, $b:expr, $dst:expr, $f:expr) => {{
let f = $f;
let mut over = false;
for ((slot, &x), &y) in $dst.iter_mut().zip($a).zip($b) {
let (v, o) = f(x, y);
*slot = v;
over |= o;
}
return !over;
}};
}
#[inline(always)]
fn monad_i64_body(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool {
use ScalarMonad::*;
match op {
Conj | Floor | Ceil => each!(a, dst, |x: i64| x),
Neg => each_over!(a, dst, i64::overflowing_neg),
Abs => each_over!(a, dst, i64::overflowing_abs),
Signum => each!(a, dst, i64::signum),
Inc => each_over!(a, dst, |x: i64| x.overflowing_add(1)),
Dec => each_over!(a, dst, |x: i64| x.overflowing_sub(1)),
Double => each_over!(a, dst, |x: i64| x.overflowing_add(x)),
Square => each_over!(a, dst, |x: i64| x.overflowing_mul(x)),
OneMinus => each_over!(a, dst, |x: i64| 1i64.overflowing_sub(x)),
_ => false,
}
}
#[inline(always)]
fn dyad_i64_body(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool {
use ScalarDyad::*;
match op {
Add => zip_over!(a, b, dst, i64::overflowing_add),
Sub => zip_over!(a, b, dst, i64::overflowing_sub),
Mul => zip_over!(a, b, dst, i64::overflowing_mul),
Min => zip!(a, b, dst, i64::min),
Max => zip!(a, b, dst, i64::max),
Residue => zip!(a, b, dst, |x: i64, y: i64| if x == 0 {
y
} else {
let mut r = y.wrapping_rem(x);
if r != 0 && (r < 0) != (x < 0) {
r += x;
}
r
}),
Eq => zip!(a, b, dst, |x: i64, y: i64| (x == y) as i64),
Ne => zip!(a, b, dst, |x: i64, y: i64| (x != y) as i64),
Lt => zip!(a, b, dst, |x: i64, y: i64| (x < y) as i64),
Le => zip!(a, b, dst, |x: i64, y: i64| (x <= y) as i64),
Gt => zip!(a, b, dst, |x: i64, y: i64| (x > y) as i64),
Ge => zip!(a, b, dst, |x: i64, y: i64| (x >= y) as i64),
_ => false,
}
}
multiversioned! {
fn monad_f64(
op: ScalarMonad,
a: &[f64],
dst: &mut [f64],
tol: Tol,
) -> bool = monad_f64_body;
}
multiversioned! {
fn dyad_f64(
op: ScalarDyad,
a: &[f64],
b: &[f64],
dst: &mut [f64],
tol: Tol,
) -> bool = dyad_f64_body;
}
multiversioned! {
fn monad_i64(op: ScalarMonad, a: &[i64], dst: &mut [i64]) -> bool = monad_i64_body;
}
multiversioned! {
fn dyad_i64(op: ScalarDyad, a: &[i64], b: &[i64], dst: &mut [i64]) -> bool = dyad_i64_body;
}
fn step_i64(op: ScalarDyad, a: i64, b: i64) -> Option<i64> {
use ScalarDyad::*;
match op {
Add => a.checked_add(b),
Mul => a.checked_mul(b),
Min => Some(a.min(b)),
Max => Some(a.max(b)),
_ => None,
}
}
pub(crate) fn step(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
step_f64(op, a, b)
}
fn step_f64(op: ScalarDyad, a: f64, b: f64) -> Option<f64> {
use ScalarDyad::*;
match op {
Add => Some(a + b),
Mul => Some(a * b),
Min => Some(a.min(b)),
Max => Some(a.max(b)),
_ => None,
}
}
fn to_f64(a: &Array, w: usize) -> Option<Vec<f64>> {
if a.rank() == 0 {
let v = match &a.data {
Data::Bool(d) => d[0] as f64,
Data::I64(d) => d[0] as f64,
Data::F64(d) => d[0],
Data::Ext(_) | Data::Rat(_) | Data::Complex(_) | Data::Char(_) | Data::Box(_) => {
return Some(Vec::new());
}
};
return Some(vec![v; w]);
}
match &a.data {
Data::F64(_) => None,
Data::I64(d) => Some(par::map(d, |&x| x as f64)),
Data::Bool(d) => Some(par::map(d, |&x| x as f64)),
Data::Ext(_) | Data::Rat(_) | Data::Complex(_) | Data::Char(_) | Data::Box(_) => {
Some(Vec::new())
}
}
}
fn to_i64(a: &Array, w: usize) -> Option<Vec<i64>> {
if a.rank() == 0 {
let v = match &a.data {
Data::Bool(d) => d[0] as i64,
Data::I64(d) => d[0],
_ => return Some(Vec::new()),
};
return Some(vec![v; w]);
}
match &a.data {
Data::I64(_) => None,
Data::Bool(d) => Some(par::map(d, |&x| x as i64)),
_ => Some(Vec::new()),
}
}
pub(crate) fn common_shape(inputs: &[Array]) -> Option<Option<Vec<usize>>> {
let mut shape: Option<&Vec<usize>> = None;
for a in inputs {
if a.rank() == 0 {
continue;
}
match shape {
None => shape = Some(&a.shape),
Some(s) if *s == a.shape => {}
Some(_) => return None,
}
}
Some(shape.cloned())
}
pub(crate) fn run(k: &FusedKernel, inputs: &[Array]) -> Option<Array> {
let reducing = matches!(k.yields, Yield::Reduce(_));
let shape = common_shape(inputs)??;
let n: usize = shape.iter().product();
if n == 0 {
return None;
}
if reducing && (shape.len() != 1 || n < 2) {
return None;
}
let (working, root) = working_type(k, inputs)?;
if k.yields == Yield::Tally {
return Some(Array::scalar_i64(shape[0] as i64));
}
let w = BLOCK.min(n).max(1);
let tol = k.tol;
let cmp_f64 = move |op, a: &[f64], b: &[f64], dst: &mut [f64]| dyad_f64(op, a, b, dst, tol);
let sign_f64 = move |op, a: &[f64], dst: &mut [f64]| monad_f64(op, a, dst, tol);
let data = if working == DType::F64 {
let owned: Vec<Option<Vec<f64>>> = inputs.iter().map(|a| to_f64(a, w)).collect();
let srcs: Vec<Loaded<f64>> = inputs
.iter()
.zip(&owned)
.map(|(a, o)| match o {
Some(v) => Loaded { data: v, splat: a.rank() == 0 },
None => Loaded { data: a.as_f64_slice().unwrap_or(&[]), splat: false },
})
.collect();
match k.reduce() {
None => {
let out = map_pass(k, &srcs, n, sign_f64, cmp_f64)?;
float_result(out, root)
}
Some(op) => {
let v = reduce_pass(k, &srcs, n, sign_f64, cmp_f64, |a, b| step_f64(op, a, b))?;
match root {
DType::F64 => Data::F64(vec![v].into()),
_ => Data::I64(vec![v as i64].into()),
}
}
}
} else {
let owned: Vec<Option<Vec<i64>>> = inputs.iter().map(|a| to_i64(a, w)).collect();
let srcs: Vec<Loaded<i64>> = inputs
.iter()
.zip(&owned)
.map(|(a, o)| match o {
Some(v) => Loaded { data: v, splat: a.rank() == 0 },
None => Loaded { data: a.as_i64_slice().unwrap_or(&[]), splat: false },
})
.collect();
match k.reduce() {
None => {
let out = map_pass(k, &srcs, n, monad_i64, dyad_i64)?;
int_result(out, root)
}
Some(op) => {
let v =
reduce_pass(k, &srcs, n, monad_i64, dyad_i64, |a, b| step_i64(op, a, b))?;
Data::I64(vec![v].into())
}
}
};
Some(Array::new(if reducing { Vec::new() } else { shape }, data))
}
fn float_result(out: Vec<f64>, root: DType) -> Data {
match root {
DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0.0) as u8).into()),
_ => Data::F64(out.into()),
}
}
fn int_result(out: Vec<i64>, root: DType) -> Data {
match root {
DType::Bool => Data::Bool(par::map(&out, |&v| (v != 0) as u8).into()),
_ => Data::I64(out.into()),
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Decline {
Agreement,
Empty,
ReduceShape,
WorkingType,
Overflow,
}
impl Decline {
pub fn reason(self) -> &'static str {
match self {
Decline::Agreement => "the inputs need agreement or are all scalars",
Decline::Empty => "there is nothing to compute",
Decline::ReduceShape => "the reduction needs one axis of two or more items",
Decline::WorkingType => "no single working type holds every step exactly",
Decline::Overflow => "an integer step left 64-bit range",
}
}
}
pub fn decline_reason(k: &FusedKernel, inputs: &[Array]) -> Option<Decline> {
let Some(Some(shape)) = common_shape(inputs) else {
return Some(Decline::Agreement);
};
let n: usize = shape.iter().product();
if n == 0 {
return Some(Decline::Empty);
}
if matches!(k.yields, Yield::Reduce(_)) && (shape.len() != 1 || n < 2) {
return Some(Decline::ReduceShape);
}
if working_type(k, inputs).is_none() {
return Some(Decline::WorkingType);
}
Some(Decline::Overflow)
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct Summary {
pub ops: usize,
pub op_names: Vec<&'static str>,
pub reduce: Option<&'static str>,
pub tally: bool,
pub lets: usize,
pub inputs: usize,
pub block: usize,
}
impl std::fmt::Display for Summary {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(f, "{} op{}", self.ops, if self.ops == 1 { "" } else { "s" })?;
if !self.op_names.is_empty() {
write!(f, ": {}", self.op_names.join(" "))?;
}
if let Some(r) = self.reduce {
write!(f, "; {r}/ absorbed")?;
}
if self.tally {
write!(f, "; tally only")?;
}
if self.lets > 0 {
write!(f, "; {} let slot{}", self.lets, if self.lets == 1 { "" } else { "s" })?;
}
write!(f, "; block {}", self.block)
}
}
pub fn summary(k: &FusedKernel) -> Summary {
let mut op_names = Vec::new();
let mut lets = 0usize;
for ins in &k.code {
match ins {
Instr::Monad(op) => op_names.push(monad_name(*op)),
Instr::Dyad(op) => op_names.push(dyad_name(*op)),
Instr::Store(_) => lets += 1,
Instr::Load(_) | Instr::Let(_) => {}
}
}
Summary {
ops: op_names.len(),
op_names,
reduce: k.reduce().map(dyad_name),
tally: k.yields == Yield::Tally,
lets,
inputs: k.leaves.iter().copied().max().map_or(0, |m| m + 1),
block: BLOCK,
}
}
pub fn inlined_names(p: &Program) -> Vec<String> {
let Some(Expr::Elided { orig, .. }) = p.stmts.first() else { return Vec::new() };
let assigned = |stmts: &[Expr]| -> Vec<String> {
stmts
.iter()
.filter_map(|s| match s {
Expr::Assign { name, .. } => Some(name.clone()),
_ => None,
})
.collect()
};
let kept = assigned(&p.stmts);
assigned(orig).into_iter().filter(|n| !kept.contains(n)).collect()
}
fn monad_name(op: ScalarMonad) -> &'static str {
use ScalarMonad::*;
match op {
Conj => "+",
Neg => "-",
Signum => "*",
Recip => "%",
Sqrt => "%:",
Exp => "^",
Abs => "|",
Floor => "<.",
Ceil => ">.",
Not => "-.",
OneMinus => "-.",
Inc => ">:",
Dec => "<:",
Double => "+:",
Halve => "-:",
Square => "*:",
Ln => "^.",
Pi => "o.",
Factorial => "!",
Imaginary => "j.",
Polar => "r.",
}
}
fn dyad_name(op: ScalarDyad) -> &'static str {
use ScalarDyad::*;
match op {
Add => "+",
Sub => "-",
Mul => "*",
DivJ | DivApl => "%",
Min => "<.",
Max => ">.",
Pow => "^",
Residue => "|",
Eq => "=",
Ne => "~:",
Lt => "<",
Le => "<:",
Gt => ">",
Ge => ">:",
Lcm => "*.",
Gcd => "+.",
Log => "^.",
Root => "%:",
Circle => "o.",
Binomial => "!",
MakeComplex => "j.",
PolarBy => "r.",
}
}
pub(crate) fn eval_on(
device: Option<&crate::device::Device>,
k: &FusedKernel,
inputs: &[Array],
) -> (Option<Array>, crate::device::Placement) {
use crate::device::Placement;
let mut placement = Placement::Default;
if let Some(d) = device.filter(|d| d.is_gpu()) {
match crate::device::try_run(d, k, inputs) {
Ok(a) => return (Some(a), Placement::Gpu),
Err(why) => placement = Placement::Cpu(why),
}
}
let r = run(k, inputs);
if r.is_none() {
note_fallback();
}
(r, placement)
}
#[cfg(test)]
mod tests {
use super::*;
use crate::frontend::{compile, Dialect, Lang};
fn program(src: &str) -> Program {
compile(Lang::J, src, &Dialect::default()).expect("compile")
}
#[test]
fn a_chain_of_two_scalar_verbs_fuses() {
assert!(is_fused(&program("1 + 2 * {x}")));
assert!(is_fused(&program("+/ {w} * {x}")));
assert!(is_fused(&program("+/ ^ {x}")));
}
#[test]
fn one_verb_on_its_own_is_left_alone() {
assert!(!is_fused(&program("2 * {x}")));
assert!(!is_fused(&program("+/ {x}")));
assert!(!is_fused(&program("{x}")));
}
#[test]
fn a_verb_the_kernel_does_not_cover_breaks_the_chain() {
assert!(!is_fused(&program("%: 2 * {x}")));
assert!(is_fused(&program("%: 1 + 2 * {x}")));
}
#[test]
fn an_effect_in_a_leaf_keeps_the_chain_unfused() {
assert!(!is_fused(&program("1 + 2 * echo {x}")));
}
#[test]
fn the_postfix_program_pushes_the_left_operand_first() {
let p = program("{w} - {x} - 1");
let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
assert_eq!(
kernel.code(),
[
Instr::Load(2),
Instr::Load(1),
Instr::Load(0),
Instr::Dyad(ScalarDyad::Sub),
Instr::Dyad(ScalarDyad::Sub),
]
);
assert_eq!(kernel.slots, 2);
}
#[test]
fn a_value_the_chain_reads_twice_becomes_a_let() {
let p = program("+/ ({x} + 1) * ({x} + 1)");
let Expr::Fused { kernel, .. } = &p.stmts[0] else { panic!("not fused") };
assert_eq!(
kernel.code(),
[
Instr::Load(1),
Instr::Load(0),
Instr::Dyad(ScalarDyad::Add),
Instr::Store(0),
Instr::Let(0),
Instr::Let(0),
Instr::Dyad(ScalarDyad::Mul),
]
);
assert_eq!(kernel.slots, 2);
}
#[test]
fn a_named_value_moves_into_the_sentence_that_reads_it() {
let p = program("d =. {x} + 1\n+/ d * d");
assert!(is_inlined(&p));
assert_eq!(p.stmts.len(), 3);
let Expr::Fused { kernel, .. } = &p.stmts[2] else { panic!("the sum did not fuse") };
assert!(kernel.code().contains(&Instr::Store(0)));
assert_eq!(unfused(&p).stmts.len(), 2);
}
}