#[cfg(test)]
mod array_tests;
#[cfg(test)]
mod data_gather_tests;
mod exec;
#[cfg(test)]
mod forcing_tests;
mod fuse;
mod geom;
mod ir;
mod layout;
mod lower;
#[cfg(test)]
mod lowering_limit_tests;
mod prune;
#[cfg(test)]
mod refexec;
mod reroll;
mod scan_shift;
mod sink;
#[cfg(test)]
mod tests;
#[cfg(test)]
mod tuple_tests;
#[cfg(feature = "xla")]
pub mod xla_emit;
#[cfg(test)]
pub(crate) use exec::section_primes;
pub(in crate::simulate_array) use exec::{
CallInputs, ForcingView, TapeCtx, run_tape_call, run_tape_inputs,
};
pub(crate) use ir::*;
#[cfg(not(target_arch = "wasm32"))]
pub(crate) use lower::StandaloneError;
use lower::build_tape_program;
pub(crate) use lower::{TileAxisDom, TileLowering, TileObserved, TilePiece};
#[cfg(not(target_arch = "wasm32"))]
pub(in crate::simulate_array) use lower::{collect_variable_names, eval_standalone_rule};
use super::ArrayCompiled;
use std::cell::RefCell;
use std::collections::HashSet;
use std::fmt;
use std::rc::Rc;
#[derive(Debug, Clone, Default)]
pub struct TapeBuildReport {
pub n_rules: usize,
pub n_taped: usize,
pub fallbacks: Vec<(String, String)>,
pub n_instr_const: usize,
pub n_instr_segment: usize,
pub n_instr_continuous: usize,
pub opcode_counts: Vec<(String, usize)>,
pub n_slots: usize,
pub n_storages: usize,
pub n_gather_plans: usize,
pub n_exports: usize,
pub n_dy_writes: usize,
pub slab_bytes: usize,
pub slab_bytes_const: usize,
pub slab_bytes_segment: usize,
pub slab_bytes_recycled: usize,
pub vn_scope_hits: usize,
pub vn_hoist_hits: usize,
pub fuse: FuseStats,
}
impl fmt::Display for TapeBuildReport {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
writeln!(f, "tape build report")?;
writeln!(
f,
" rules: {} taped / {} declined (of {})",
self.n_taped,
self.fallbacks.len(),
self.n_rules
)?;
for (name, reason) in &self.fallbacks {
writeln!(f, " declined {name}: {reason}")?;
}
if !self.fallbacks.is_empty() {
return writeln!(f, " no program: a declined rule leaves no tape");
}
writeln!(
f,
" instructions: {} CONST + {} SEGMENT + {} CONTINUOUS = {}",
self.n_instr_const,
self.n_instr_segment,
self.n_instr_continuous,
self.n_instr_const + self.n_instr_segment + self.n_instr_continuous
)?;
write!(f, " opcodes:")?;
for (op, n) in &self.opcode_counts {
write!(f, " {op}={n}")?;
}
writeln!(f)?;
writeln!(
f,
" slots: {} ({} storages, {} gather plans, {} dy writes, {} exports)",
self.n_slots, self.n_storages, self.n_gather_plans, self.n_dy_writes, self.n_exports
)?;
writeln!(
f,
" slab: {:.2} MiB total ({:.2} MiB CONST + {:.2} MiB SEGMENT dedicated, {:.2} MiB recycled)",
self.slab_bytes as f64 / (1024.0 * 1024.0),
self.slab_bytes_const as f64 / (1024.0 * 1024.0),
self.slab_bytes_segment as f64 / (1024.0 * 1024.0),
self.slab_bytes_recycled as f64 / (1024.0 * 1024.0),
)?;
writeln!(
f,
" value numbering: {} scope hits, {} hoist hits",
self.vn_scope_hits, self.vn_hoist_hits
)?;
if self.fuse.n_groups > 0 {
writeln!(
f,
" fusion: {} -> {} instructions; {} groups absorbing {} \
({} gathers folded, {} kept, {} reductions folded); group sizes \
[2-3]={} [4-7]={} [8-15]={} [16-31]={} [32-63]={} [64+]={}",
self.fuse.instrs_before,
self.fuse.instrs_after,
self.fuse.n_groups,
self.fuse.n_member_instrs,
self.fuse.n_gathers_folded,
self.fuse.n_gathers_kept,
self.fuse.n_reduces_folded,
self.fuse.group_size_hist[0],
self.fuse.group_size_hist[1],
self.fuse.group_size_hist[2],
self.fuse.group_size_hist[3],
self.fuse.group_size_hist[4],
self.fuse.group_size_hist[5],
)?;
}
Ok(())
}
}
#[derive(Clone, Debug)]
pub(crate) struct TapeRuleRecord {
pub name: String,
pub kind: &'static str,
pub cadence: &'static str,
pub fallback_reason: Option<String>,
}
pub(crate) fn make_report(prog: &TapeProgram, vn_hits: (usize, usize)) -> TapeBuildReport {
let mut opcode_counts: std::collections::HashMap<&'static str, usize> =
std::collections::HashMap::new();
for i in &prog.instrs {
*opcode_counts.entry(i.opcode()).or_insert(0) += 1;
}
let mut opcode_counts: Vec<(String, usize)> = opcode_counts
.into_iter()
.map(|(k, v)| (k.to_string(), v))
.collect();
opcode_counts.sort_by(|a, b| b.1.cmp(&a.1).then(a.0.cmp(&b.0)));
let n_cont = prog.instrs.len() - (prog.n_const + prog.n_segment) as usize;
TapeBuildReport {
n_rules: prog.rules.len(),
n_taped: prog.rules.len(),
fallbacks: Vec::new(),
n_instr_const: prog.n_const as usize,
n_instr_segment: prog.n_segment as usize,
n_instr_continuous: n_cont,
opcode_counts,
n_slots: prog.slots.len(),
n_storages: prog.slab.storages.len(),
n_gather_plans: prog.plans.len(),
n_exports: prog.exports.len(),
n_dy_writes: prog.dy_writes.len(),
slab_bytes: prog.slab.total_elems * 8,
slab_bytes_const: prog.slab.const_elems * 8,
slab_bytes_segment: prog.slab.segment_elems * 8,
slab_bytes_recycled: prog.slab.recycled_elems * 8,
vn_scope_hits: vn_hits.0,
vn_hoist_hits: vn_hits.1,
fuse: prog.fuse_stats.clone(),
}
}
pub(crate) fn make_declined_report(
declined: &DeclinedTape,
vn_hits: (usize, usize),
) -> TapeBuildReport {
let fallbacks = declined.declined_rules();
TapeBuildReport {
n_rules: declined.rules.len(),
n_taped: declined.rules.len() - fallbacks.len(),
fallbacks,
vn_scope_hits: vn_hits.0,
vn_hoist_hits: vn_hits.1,
..TapeBuildReport::default()
}
}
impl DeclinedTape {
pub(crate) fn declined_rules(&self) -> Vec<(String, String)> {
self.rules
.iter()
.zip(&self.reasons)
.filter_map(|(r, why)| why.as_ref().map(|why| (r.name.clone(), why.clone())))
.collect()
}
}
impl RuleInfo {
pub(crate) fn kind_name(&self) -> &'static str {
match self.kind {
RuleKind::Observed(_) => "observed",
RuleKind::Rhs(_) => "state derivative",
}
}
pub(crate) fn cadence_name(&self) -> &'static str {
match self.cadence {
Cadence::Const => "const",
Cadence::Segment => "segment",
Cadence::Continuous => "continuous",
}
}
}
#[cfg(test)]
pub(crate) fn taped(
(prog, report): (Result<TapeProgram, DeclinedTape>, TapeBuildReport),
) -> (TapeProgram, TapeBuildReport) {
match prog {
Ok(prog) => (prog, report),
Err(_) => panic!("the tape declined a rule:\n{report}"),
}
}
pub(crate) type TapeBuild = Result<std::sync::Arc<TapeProgram>, Rc<DeclinedTape>>;
pub(crate) struct TapeCache {
entries: RefCell<Vec<CachedTape>>,
}
struct CachedTape {
discrete_forcing: Vec<String>,
forcing_buffer: Vec<(String, Vec<usize>)>,
precision: crate::precision::Precision,
serves_passes: bool,
prog: TapeBuild,
report: Rc<TapeBuildReport>,
}
impl TapeCache {
pub(crate) fn new() -> Self {
TapeCache {
entries: RefCell::new(Vec::new()),
}
}
}
impl ArrayCompiled {
pub(crate) fn tape(
&self,
discrete_forcing: &HashSet<String>,
) -> (TapeBuild, Rc<TapeBuildReport>) {
let mut forcing: Vec<String> = discrete_forcing.iter().cloned().collect();
forcing.sort();
let forcing_buffer = self.forcing_buffer_signature();
let precision = crate::precision::active();
let serves_passes = self.tape_serves_passes();
if let Some(e) = self.tape_cache.entries.borrow().iter().find(|e| {
e.discrete_forcing == forcing
&& e.forcing_buffer == forcing_buffer
&& e.precision == precision
&& e.serves_passes == serves_passes
}) {
return (e.prog.clone(), Rc::clone(&e.report));
}
let (prog, report) = self.build_tape(discrete_forcing);
let (prog, report) = (
prog.map(std::sync::Arc::new).map_err(Rc::new),
Rc::new(report),
);
self.tape_cache.entries.borrow_mut().push(CachedTape {
discrete_forcing: forcing,
forcing_buffer,
precision,
serves_passes,
prog: prog.clone(),
report: Rc::clone(&report),
});
(prog, report)
}
pub(crate) fn strict_tape(
&self,
discrete_forcing: &HashSet<String>,
) -> Result<std::sync::Arc<TapeProgram>, crate::simulate::SimulateError> {
let declined = match self.tape(discrete_forcing).0 {
Ok(prog) => return Ok(prog),
Err(declined) => declined,
};
let (k, reason) = declined
.reasons
.iter()
.enumerate()
.find_map(|(k, why)| why.clone().map(|why| (k, why)))
.expect("a declined tape names at least one rule");
let rule = &declined.rules[k];
Err(crate::simulate::SimulateError::Compile(
crate::compile_error::CompileError::CompilerRefusedRule {
compiler: if self.is_xla() { "xla" } else { "native" },
kind: rule.kind_name(),
rule: rule.name.clone(),
tier: rule.cadence_name(),
reason,
},
))
}
fn forcing_buffer_signature(&self) -> Vec<(String, Vec<usize>)> {
let mut sig: Vec<(String, Vec<usize>)> = self
.forcing
.borrow()
.iter()
.map(|(name, a)| (name.clone(), a.shape().to_vec()))
.collect();
sig.sort();
sig
}
fn forcing_inputs<'a>(&self, discrete: &'a HashSet<String>) -> lower::ForcingInputs<'a> {
let mut shapes: rustc_hash::FxHashMap<String, super::DimU> = self
.forcing
.borrow()
.iter()
.map(|(name, a)| (name.clone(), a.shape().iter().copied().collect()))
.collect();
let mut unsized_ = HashSet::new();
for (name, decl) in &self.forcing_decls {
if shapes.contains_key(name) {
continue;
}
match decl {
Some(shape) => {
shapes.insert(name.clone(), shape.iter().copied().collect());
}
None => {
unsized_.insert(name.clone());
}
}
}
lower::ForcingInputs {
shapes,
unsized_,
discrete,
}
}
pub(crate) fn build_tape(
&self,
discrete_forcing: &HashSet<String>,
) -> (Result<TapeProgram, DeclinedTape>, TapeBuildReport) {
self.build_tape_opts(discrete_forcing, Some(fuse::SuperopCfg::from_env()))
}
pub(crate) fn build_tape_opts(
&self,
discrete_forcing: &HashSet<String>,
fuse: Option<fuse::SuperopCfg>,
) -> (Result<TapeProgram, DeclinedTape>, TapeBuildReport) {
self.build_tape_layout(discrete_forcing, fuse, true)
}
pub(crate) fn build_tape_layout(
&self,
discrete_forcing: &HashSet<String>,
fuse: Option<fuse::SuperopCfg>,
align_layout: bool,
) -> (Result<TapeProgram, DeclinedTape>, TapeBuildReport) {
let const_names = self.classify_static_observeds(discrete_forcing);
let seg_names = self.classify_segment_invariant_observeds(discrete_forcing, true);
let forcing = self.forcing_inputs(discrete_forcing);
let (prog, vn_hits) = build_tape_program(
self,
&const_names,
&seg_names,
&forcing,
fuse,
align_layout,
None,
);
let report = match &prog {
Ok(prog) => make_report(prog, vn_hits),
Err(declined) => make_declined_report(declined, vn_hits),
};
(prog, report)
}
pub(crate) fn build_tile_tape(&self, tile: &TileLowering) -> Result<TapeProgram, DeclinedTape> {
let discrete: HashSet<String> = self.forcing_decls.keys().cloned().collect();
let const_names = self.classify_static_observeds(&discrete);
let seg_names = self.classify_segment_invariant_observeds(&discrete, true);
let forcing = self.forcing_inputs(&discrete);
build_tape_program(
self,
&const_names,
&seg_names,
&forcing,
Some(fuse::SuperopCfg::from_env()),
true,
Some(tile),
)
.0
}
pub fn debug_build_tape_report(&self) -> TapeBuildReport {
(*self.tape(&HashSet::new()).1).clone()
}
#[doc(hidden)]
pub fn debug_tape_listing(&self) -> String {
use std::fmt::Write;
let (prog, report) = self.tape(&HashSet::new());
let Ok(prog) = prog else {
return report.to_string();
};
let mut out = String::new();
for (pc, ins) in prog.instrs.iter().enumerate() {
let boxed = ins
.out()
.map(|o| {
let d = &prog.slots[o as usize];
format!(" -> s{o}{:?}", &d.shape[..])
})
.unwrap_or_default();
let _ = writeln!(out, "{pc:4} {:?} {ins:?}{boxed}", prog.section_of(pc));
}
for (i, fs) in prog.fused.iter().enumerate() {
let _ = writeln!(
out,
"fused {i}: box {:?}, {} micro-ops, {} inputs ({} shifted), {} scalars, {} outputs, {} runs ({} schedule nodes), reduce {:?}",
&fs.shape[..],
fs.micro.len(),
fs.inputs.len(),
fs.inputs.iter().filter(|x| x.shifted_ix.is_some()).count(),
fs.scalars.len(),
fs.outputs.len(),
fs.schedule.n_runs,
fs.schedule.nodes.len(),
fs.reduce
);
for (k, inp) in fs.inputs.iter().enumerate() {
let _ = writeln!(out, " in {k}: {inp:?}");
}
for op in &fs.micro {
let _ = writeln!(out, " {op:?}");
}
}
for (i, ls) in prog.lanes.iter().enumerate() {
let _ = writeln!(
out,
"lanes {i}: {} lanes, {} micro-ops, {} inputs, {} scalars, {} registers, {} writes",
ls.lanes,
ls.micro.len(),
ls.inputs.len(),
ls.scalars.len(),
ls.n_regs,
ls.writes.len()
);
for (k, inp) in ls.inputs.iter().enumerate() {
let _ = writeln!(out, " in {k}: {:?} {:?}", inp.kind, inp.ix);
}
for op in &ls.micro {
let _ = writeln!(out, " {op:?}");
}
for w in &ls.writes {
let _ = writeln!(out, " write {:?} -> {:?}", w.src, w.dst);
}
}
out
}
pub(crate) fn tape_rule_records(
&self,
discrete_forcing: &HashSet<String>,
) -> (Vec<TapeRuleRecord>, TapeBuildReport) {
let (prog, report) = self.tape(discrete_forcing);
let report = (*report).clone();
let (rules, reasons): (&[RuleInfo], &[Option<String>]) = match &prog {
Ok(prog) => (&prog.rules, &[]),
Err(declined) => (&declined.rules, &declined.reasons),
};
let records = rules
.iter()
.enumerate()
.map(|(k, r)| TapeRuleRecord {
name: r.name.clone(),
kind: r.kind_name(),
cadence: r.cadence_name(),
fallback_reason: reasons.get(k).cloned().flatten(),
})
.collect();
(records, report)
}
pub fn debug_dump_fuse_stats(&self) {
let (prog, report) = self.tape(&HashSet::new());
eprintln!("{report}");
let Ok(prog) = prog else {
return;
};
let mut by_regs: Vec<&FusedSpec> = prog.fused.iter().collect();
by_regs.sort_by_key(|f| std::cmp::Reverse(f.n_regs));
let max_regs = by_regs.first().map(|f| f.n_regs).unwrap_or(0);
let tot_runs: usize = prog.fused.iter().map(|f| f.schedule.n_runs).sum();
let tot_micro: usize = prog.fused.iter().map(|f| f.micro.len()).sum();
let tot_inputs: usize = prog.fused.iter().map(|f| f.inputs.len()).sum();
let tot_outputs: usize = prog.fused.iter().map(|f| f.outputs.len()).sum();
let tot_elem_ops: usize = prog.fused.iter().map(|f| f.micro.len() * f.n_elems()).sum();
let tot_out_elems: usize = prog
.fused
.iter()
.map(|f| f.outputs.len() * f.n_elems())
.sum();
eprintln!(
"element-ops per call: {} micro x elems | output elems stored {}",
tot_elem_ops, tot_out_elems
);
eprintln!(
"fused groups: {} | micro ops {} | inputs {} | outputs {} | runs {} | max n_regs {}",
prog.fused.len(),
tot_micro,
tot_inputs,
tot_outputs,
tot_runs,
max_regs
);
eprintln!("micro-op histogram (element-weighted, post-superop):");
for (k, n) in prog.fuse_stats.micro_hist.iter().take(24) {
eprintln!(" {n:>10} {k}");
}
eprintln!("single-use adjacency (element-weighted, pre-superop):");
for (k, n) in prog.fuse_stats.adj_hist.iter().take(24) {
eprintln!(" {n:>10} {k}");
}
eprintln!("arith chain lengths (element-weighted, pre-superop):");
for (len, n) in prog.fuse_stats.chain_hist.iter() {
eprintln!(" {n:>10} len={len}");
}
for f in by_regs.iter().take(12) {
eprintln!(
" group: shape={:?} micro={} regs={} inputs={} (shifted {}) scalars={} \
outputs={} runs={} folded_gathers={}",
&f.shape[..],
f.micro.len(),
f.n_regs,
f.inputs.len(),
f.inputs.iter().filter(|i| i.shifted_ix.is_some()).count(),
f.scalars.len(),
f.outputs.len(),
f.schedule.n_runs,
f.n_folded_gathers
);
}
}
}