mod defunctionalize;
mod fold;
mod splice_fuse;
mod inline_tiny;
mod inline_leaf;
mod inline_recursive;
mod popcount_leaf;
mod symmetry;
mod dce;
mod propagate;
mod gvn;
mod loop_carried_cse;
mod float_induction_sr;
mod licm;
mod bound_version;
mod closed_form;
mod ctfe;
mod deforest;
#[cfg(feature = "codegen")]
mod unroll;
#[cfg(feature = "codegen")]
mod scalarize;
#[cfg(feature = "codegen")]
mod affine_scalarize;
#[cfg(feature = "codegen")]
mod loop_split;
mod affine;
mod abstract_interp;
pub use abstract_interp::{
lin_to_rust, oracle_analyze, oracle_analyze_with, oracle_analyze_with_entry_guards, OracleFacts,
ScalarKind, VarProvenFacts,
};
pub mod egraph;
pub mod supercompile;
pub mod effects;
pub mod bta;
pub mod partial_eval;
use std::collections::HashMap;
use crate::arena::Arena;
use crate::ast::stmt::{Expr, Stmt};
use crate::intern::{Interner, Symbol};
use crate::optimization::{
admits_pinned, FiredOptimizations, HotswapConfig, Opt, OptimizationConfig, Pin, Tier,
};
thread_local! {
static ACTIVE_CONFIG: std::cell::Cell<OptimizationConfig> =
std::cell::Cell::new(OptimizationConfig::all_on());
}
pub fn set_active_config(cfg: OptimizationConfig) {
ACTIVE_CONFIG.with(|c| c.set(cfg));
}
thread_local! {
static PROVEN_RAW_INT_OPS: std::cell::RefCell<std::collections::HashSet<usize>> =
std::cell::RefCell::new(std::collections::HashSet::new());
}
pub(crate) fn clear_proven_raw_int_ops() {
PROVEN_RAW_INT_OPS.with(|s| s.borrow_mut().clear());
}
pub(crate) fn mark_proven_raw_int_op(e: &crate::ast::stmt::Expr) {
PROVEN_RAW_INT_OPS.with(|s| {
s.borrow_mut().insert(e as *const crate::ast::stmt::Expr as usize);
});
}
pub fn expr_proven_raw_int_op(e: &crate::ast::stmt::Expr) -> bool {
PROVEN_RAW_INT_OPS.with(|s| s.borrow().contains(&(e as *const crate::ast::stmt::Expr as usize)))
}
pub fn active_config() -> OptimizationConfig {
ACTIVE_CONFIG.with(|c| c.get())
}
thread_local! {
static FIRED: std::cell::Cell<Option<u64>> = const { std::cell::Cell::new(None) };
static PREEMPTED: std::cell::RefCell<Option<Vec<(Opt, Opt)>>> =
const { std::cell::RefCell::new(None) };
}
pub fn begin_fired_trace() {
FIRED.with(|f| f.set(Some(0)));
PREEMPTED.with(|p| *p.borrow_mut() = Some(Vec::new()));
}
pub fn end_fired_trace() -> Option<FiredOptimizations> {
FIRED.with(|f| f.take().map(FiredOptimizations::from_bits))
}
#[inline]
pub fn fired_trace_active() -> bool {
FIRED.with(|f| f.get().is_some())
}
#[inline]
pub fn mark_fired(opt: Opt) {
FIRED.with(|f| {
if let Some(bits) = f.get() {
f.set(Some(bits | opt.bit()));
}
});
}
#[inline]
pub fn mark_preempted(winner: Opt, loser: Opt) {
if !fired_trace_active() {
return;
}
let cfg = active_config();
if !cfg.is_on(winner) || !cfg.is_on(loser) {
return;
}
PREEMPTED.with(|p| {
if let Some(v) = p.borrow_mut().as_mut() {
v.push((winner, loser));
}
});
}
pub fn end_preempted_trace() -> Vec<(Opt, Opt)> {
let mut edges = PREEMPTED.with(|p| p.borrow_mut().take().unwrap_or_default());
edges.sort_by_key(|&(w, l)| (w as u8, l as u8));
edges.dedup();
edges
}
#[inline]
pub(crate) fn run_traced<'a>(
opt: Opt,
input: Vec<Stmt<'a>>,
pass: impl FnOnce(Vec<Stmt<'a>>) -> Vec<Stmt<'a>>,
) -> Vec<Stmt<'a>> {
if !fired_trace_active() {
return pass(input);
}
let before = format!("{input:?}");
let out = pass(input);
if format!("{out:?}") != before {
mark_fired(opt);
}
out
}
pub fn optimize_for_projection<'a>(
stmts: Vec<Stmt<'a>>,
expr_arena: &'a Arena<Expr<'a>>,
stmt_arena: &'a Arena<Stmt<'a>>,
interner: &mut Interner,
cfg: &OptimizationConfig,
) -> Vec<Stmt<'a>> {
let bta_cache = bta::analyze_with_sccs(&stmts, interner);
let mut current = stmts;
if cfg.is_on(Opt::Specialize) {
let mut variant_count: HashMap<Symbol, usize> = HashMap::new();
for _ in 0..16 {
let folded = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
let propagated = propagate::propagate_stmts(folded, expr_arena, stmt_arena, interner);
let (specialized, changes) = partial_eval::specialize_stmts_with_state(
propagated, expr_arena, stmt_arena, interner, &mut variant_count,
Some(&bta_cache),
);
current = specialized;
if changes == 0 {
break;
}
}
}
if cfg.is_on(Opt::Comptime) {
current = ctfe::ctfe_stmts(current, expr_arena, stmt_arena, interner);
}
fold::fold_stmts(current, expr_arena, stmt_arena, interner)
}
pub fn optimize_for_run<'a>(
stmts: Vec<Stmt<'a>>,
expr_arena: &'a Arena<Expr<'a>>,
stmt_arena: &'a Arena<Stmt<'a>>,
interner: &mut Interner,
cfg: &OptimizationConfig,
) -> Vec<Stmt<'a>> {
optimize_for_run_tiered(
stmts,
expr_arena,
stmt_arena,
interner,
cfg,
&HotswapConfig::default(),
Tier::T3,
)
}
fn pe_iters(hs: &HotswapConfig, tier: Tier) -> usize {
if matches!(hs.pin(Opt::Specialize), Some(Pin::Eager) | Some(Pin::At(_))) {
return 16;
}
match tier {
Tier::T0 | Tier::T1 => 0,
Tier::T2 => 4,
Tier::T3 => 16,
}
}
pub fn optimize_for_run_tiered<'a>(
stmts: Vec<Stmt<'a>>,
expr_arena: &'a Arena<Expr<'a>>,
stmt_arena: &'a Arena<Stmt<'a>>,
interner: &mut Interner,
cfg: &OptimizationConfig,
hs: &HotswapConfig,
tier: Tier,
) -> Vec<Stmt<'a>> {
set_active_config(*cfg);
clear_proven_raw_int_ops();
if cfg.is_all_off() || tier == Tier::T0 {
return stmts;
}
const MAX_RUN_OPT_STMTS: usize = 5_000;
if stmts.len() > MAX_RUN_OPT_STMTS {
return stmts;
}
let bta_cache = bta::analyze_with_sccs(&stmts, interner);
let mut current = stmts;
if admits_pinned(cfg, hs, tier, Opt::Inline) {
current = run_traced(Opt::Inline, current, |c| {
inline_tiny::inline_tiny_fns(c, expr_arena, stmt_arena, interner)
});
current = run_traced(Opt::Inline, current, |c| {
inline_leaf::inline_leaf_fns(c, expr_arena, stmt_arena, interner)
});
}
if admits_pinned(cfg, hs, tier, Opt::Specialize) {
let mut variant_count: HashMap<Symbol, usize> = HashMap::new();
let mut specialized_any = false;
let cap = pe_iters(hs, tier);
let mut capped_early = false;
for i in 0..cap {
let folded = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
let propagated = propagate::propagate_stmts(folded, expr_arena, stmt_arena, interner);
let (specialized, changes) = partial_eval::specialize_stmts_with_state(
propagated, expr_arena, stmt_arena, interner, &mut variant_count,
Some(&bta_cache),
);
current = specialized;
if changes == 0 {
break;
}
specialized_any = true;
capped_early = i + 1 == cap;
}
if specialized_any {
mark_fired(Opt::Specialize);
}
if capped_early && tier < Tier::T3 {
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
}
if admits_pinned(cfg, hs, tier, Opt::Comptime) {
current = run_traced(Opt::Comptime, current, |c| {
ctfe::ctfe_stmts(c, expr_arena, stmt_arena, interner)
});
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
}
if admits_pinned(cfg, hs, tier, Opt::LoopCse) {
let (c, fired) =
loop_carried_cse::loop_carried_cse_stmts(current, expr_arena, stmt_arena, interner);
current = c;
if fired {
mark_fired(Opt::LoopCse);
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
}
if admits_pinned(cfg, hs, tier, Opt::FloatStrength) {
let (c, fired) = float_induction_sr::float_induction_sr_stmts(
current, expr_arena, stmt_arena, interner,
);
current = c;
if fired {
mark_fired(Opt::FloatStrength);
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
}
if admits_pinned(cfg, hs, tier, Opt::Cse) {
current = run_traced(Opt::Cse, current, |c| {
gvn::cse_stmts(c, expr_arena, stmt_arena, interner)
});
}
#[cfg(feature = "codegen")]
if admits_pinned(cfg, hs, tier, Opt::Affine) {
let (a, ca) =
affine_scalarize::affine_scalarize_seqs(current, expr_arena, stmt_arena, interner);
current = a;
if ca {
mark_fired(Opt::Affine);
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
}
#[cfg(feature = "codegen")]
{
let mut changed = false;
if admits_pinned(cfg, hs, tier, Opt::Unroll) {
let (u, c1) = unroll::unroll_stmts_run(current, expr_arena, stmt_arena, interner);
current = u;
if c1 {
mark_fired(Opt::Unroll);
}
changed |= c1;
}
if admits_pinned(cfg, hs, tier, Opt::Scalarize) {
let (s, c2) = scalarize::scalarize_seqs(current, expr_arena, stmt_arena, interner);
current = s;
if c2 {
mark_fired(Opt::Scalarize);
}
changed |= c2;
}
if changed {
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
}
if admits_pinned(cfg, hs, tier, Opt::LoopHoist) {
current = run_traced(Opt::LoopHoist, current, |c| {
licm::licm_stmts_run(c, expr_arena, stmt_arena, interner)
});
}
if admits_pinned(cfg, hs, tier, Opt::ClosedForm) {
current = run_traced(Opt::ClosedForm, current, |c| {
closed_form::closed_form_stmts(c, expr_arena, stmt_arena, interner)
});
}
if admits_pinned(cfg, hs, tier, Opt::Fuse) {
current = run_traced(Opt::Fuse, current, |c| {
deforest::deforest_stmts(c, expr_arena, stmt_arena, interner)
});
}
if admits_pinned(cfg, hs, tier, Opt::Oracle) {
current = run_traced(Opt::Oracle, current, |c| {
abstract_interp::abstract_interp_stmts(c, expr_arena, stmt_arena)
});
}
if admits_pinned(cfg, hs, tier, Opt::DeadCode) {
current = run_traced(Opt::DeadCode, current, |c| {
dce::eliminate_dead_code(c, stmt_arena, expr_arena)
});
}
const MAX_RUN_INLINE_STMTS: usize = 1_500;
if admits_pinned(cfg, hs, tier, Opt::Unfold) && current.len() <= MAX_RUN_INLINE_STMTS {
let unrolled = run_traced(Opt::Unfold, current, |c| {
inline_recursive::inline_recursive_fns_run(c, expr_arena, stmt_arena, interner)
});
let folded = fold::fold_stmts(unrolled, expr_arena, stmt_arena, interner);
let propagated = propagate::propagate_stmts(folded, expr_arena, stmt_arena, interner);
current = dce::eliminate_dead_code(propagated, stmt_arena, expr_arena);
}
current
}
pub fn optimize_program<'a>(
stmts: Vec<Stmt<'a>>,
expr_arena: &'a Arena<Expr<'a>>,
stmt_arena: &'a Arena<Stmt<'a>>,
interner: &mut Interner,
cfg: &OptimizationConfig,
) -> Vec<Stmt<'a>> {
set_active_config(*cfg);
clear_proven_raw_int_ops();
let stmts = splice_fuse::fuse_place_splices(stmts, expr_arena, stmt_arena, interner);
let bta_cache = bta::analyze_with_sccs(&stmts, interner);
let mut current = stmts;
if cfg.is_on(Opt::Specialize) {
let mut variant_count: HashMap<Symbol, usize> = HashMap::new();
let mut specialized_any = false;
for _ in 0..16 {
let folded = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
let propagated = propagate::propagate_stmts(folded, expr_arena, stmt_arena, interner);
let (specialized, changes) = partial_eval::specialize_stmts_with_state(
propagated, expr_arena, stmt_arena, interner, &mut variant_count,
Some(&bta_cache),
);
current = specialized;
if changes == 0 {
break;
}
specialized_any = true;
}
if specialized_any {
mark_fired(Opt::Specialize);
}
}
if cfg.is_on(Opt::Comptime) {
current = run_traced(Opt::Comptime, current, |c| {
ctfe::ctfe_stmts(c, expr_arena, stmt_arena, interner)
});
}
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
if cfg.is_on(Opt::Symmetry) {
current = run_traced(Opt::Symmetry, current, |c| {
symmetry::break_symmetry_stmts(c, expr_arena, stmt_arena, interner)
});
}
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
#[cfg(feature = "codegen")]
if cfg.is_on(Opt::Unroll) {
let (unrolled, changed) = unroll::unroll_stmts(current, expr_arena, stmt_arena, interner);
current = unrolled;
if changed {
mark_fired(Opt::Unroll);
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
}
if cfg.is_on(Opt::Defunctionalize) {
current = run_traced(Opt::Defunctionalize, current, |c| {
defunctionalize::defunctionalize_stmts(c, expr_arena, stmt_arena, interner)
});
}
if cfg.is_on(Opt::Saturate) {
current = run_traced(Opt::Saturate, current, |c| {
egraph::convert::egraph_stmts(c, expr_arena, stmt_arena, interner)
});
}
if cfg.is_on(Opt::Cse) {
current = run_traced(Opt::Cse, current, |c| {
gvn::cse_stmts(c, expr_arena, stmt_arena, interner)
});
}
if cfg.is_on(Opt::LoopHoist) {
current = run_traced(Opt::LoopHoist, current, |c| {
licm::licm_stmts(c, expr_arena, stmt_arena, interner)
});
}
if cfg.is_on(Opt::ClosedForm) {
current = run_traced(Opt::ClosedForm, current, |c| {
closed_form::closed_form_stmts(c, expr_arena, stmt_arena, interner)
});
}
if cfg.is_on(Opt::Fuse) {
current = run_traced(Opt::Fuse, current, |c| {
deforest::deforest_stmts(c, expr_arena, stmt_arena, interner)
});
}
#[cfg(feature = "codegen")]
{
current = run_traced(Opt::LoopSplit, current, |c| {
loop_split::loop_split_stmts(c, expr_arena, stmt_arena, interner, cfg)
});
current = fold::fold_stmts(current, expr_arena, stmt_arena, interner);
current = propagate::propagate_stmts(current, expr_arena, stmt_arena, interner);
}
if cfg.is_on(Opt::Oracle) {
current = run_traced(Opt::Oracle, current, |c| {
abstract_interp::abstract_interp_stmts(c, expr_arena, stmt_arena)
});
}
let dced = if cfg.is_on(Opt::DeadCode) {
run_traced(Opt::DeadCode, current, |c| {
dce::eliminate_dead_code(c, stmt_arena, expr_arena)
})
} else {
current
};
let current = if cfg.is_on(Opt::Supercompile) {
run_traced(Opt::Supercompile, dced, |c| {
supercompile::supercompile_stmts(c, expr_arena, stmt_arena, interner)
})
} else {
dced
};
let current = run_traced(Opt::Popcount, current, |c| {
popcount_leaf::popcount_leaf_stmts(c, expr_arena, stmt_arena, interner, cfg)
});
let unrolled = if cfg.is_on(Opt::Unfold) {
run_traced(Opt::Unfold, current, |c| {
inline_recursive::inline_recursive_fns(c, expr_arena, stmt_arena, interner)
})
} else {
current
};
let folded = fold::fold_stmts(unrolled, expr_arena, stmt_arena, interner);
let propagated = propagate::propagate_stmts(folded, expr_arena, stmt_arena, interner);
#[cfg(feature = "codegen")]
let propagated = if cfg.is_on(Opt::LoopSplit) {
run_traced(Opt::LoopSplit, propagated, |c| {
bound_version::bound_version_stmts(c, expr_arena, stmt_arena)
})
} else {
propagated
};
if cfg.is_on(Opt::DeadCode) {
run_traced(Opt::DeadCode, propagated, |c| {
dce::eliminate_dead_code(c, stmt_arena, expr_arena)
})
} else {
propagated
}
}