diff --git a/src/backend/localify.rs b/src/backend/localify.rs new file mode 100644 index 0000000..2260961 --- /dev/null +++ b/src/backend/localify.rs @@ -0,0 +1,258 @@ +//! Localification: a simple form of register allocation that picks +//! locations for SSA values in Wasm locals. + +use crate::backend::treeify::Trees; +use crate::cfg::CFGInfo; +use crate::entity::{EntityVec, PerEntity}; +use crate::ir::{Block, FunctionBody, Local, Type, Value, ValueDef}; +use smallvec::{smallvec, SmallVec}; +use std::collections::{HashMap, HashSet}; + +#[derive(Clone, Debug, Default)] +pub struct Localifier { + values: PerEntity>, + locals: EntityVec, +} + +impl Localifier { + pub fn compute(body: &FunctionBody, cfg: &CFGInfo, trees: &Trees) -> Self { + Context::new(body, cfg, trees).compute() + } +} + +struct Context<'a> { + body: &'a FunctionBody, + cfg: &'a CFGInfo, + trees: &'a Trees, + results: Localifier, + + // Affinities (blockparam value to input arg value). + affinities: HashMap>, + + // Livein to each block from dominators. + livein_values: PerEntity>, + livein_locals: PerEntity>, +} + +impl<'a> Context<'a> { + fn new(body: &'a FunctionBody, cfg: &'a CFGInfo, trees: &'a Trees) -> Self { + let mut results = Localifier::default(); + // Create locals for function args. + for &(ty, value) in &body.blocks[body.entry].params { + let param_local = results.locals.push(ty); + results.values[value] = smallvec![param_local]; + } + // Compute affinities. + let mut affinities = HashMap::default(); + for block in body.blocks.values() { + block.terminator.visit_targets(|target| { + for (&arg, ¶m) in target + .args + .iter() + .zip(body.blocks[target.block].params.iter().map(|(_, val)| val)) + { + affinities + .entry(param) + .or_insert_with(|| smallvec![]) + .push(arg); + } + }); + } + + Self { + body, + cfg, + trees, + results, + affinities, + livein_values: PerEntity::default(), + livein_locals: PerEntity::default(), + } + } + + fn compute(mut self) -> Localifier { + // Create domtree preorder for traversal (we iterate in + // reverse preorder below). + let mut order = vec![]; + order.push(self.body.entry); + let mut i = 0; + while i < order.len() { + for child in self.cfg.dom_children(order[i]) { + order.push(child); + } + i += 1; + } + + for &block in order.iter().rev() { + self.process(block); + } + + debug_assert!(self.livein_values[self.body.entry].is_empty()); + debug_assert!(self.livein_locals[self.body.entry].is_empty()); + + self.results + } + + fn process(&mut self, block: Block) { + let mut live_values = HashSet::new(); + let mut live_locals = HashSet::new(); + + // Collect liveins of all dominated blocks; this is our initial live-set. + for child in self.cfg.dom_children(block) { + for &livein_value in &self.livein_values[child] { + log::trace!( + "localify: block {} gets livein value {} from block {}", + block, + livein_value, + child + ); + live_values.insert(livein_value); + } + for &livein_local in &self.livein_locals[child] { + live_locals.insert(livein_local); + } + } + + log::trace!( + "localify: process block {}: liveout values {:?} locals {:?}", + block, + live_values, + live_locals + ); + + // For each use/def in reverse order, update live-set; on last + // use (first observed), allocate a local. + fn handle_use( + body: &FunctionBody, + u: Value, + live_values: &mut HashSet, + live_locals: &mut HashSet, + results: &mut Localifier, + affinities: &HashMap>, + ) { + let u = body.resolve_alias(u); + if live_values.insert(u) { + // If there is already an allocation (e.g. for a + // function parameter), return. + if !results.values[u].is_empty() { + return; + } + + // Need to create an allocation. + let def = &body.values[u]; + match def { + &ValueDef::Alias(_value) => { + unreachable!(); + } + &ValueDef::PickOutput(value, idx, _) => { + handle_use(body, value, live_values, live_locals, results, affinities); + results.values[u] = smallvec![results.values[value][idx]]; + } + &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => { + let locals = def + .tys() + .iter() + .map(|&ty| { + // Try to find a local of the right type that is not live. + let affinities = affinities.get(&u).map(|v| &v[..]).unwrap_or(&[]); + let local = affinities + .iter() + .filter_map(|&aff_val| { + let local = *results.values[aff_val].get(0)?; + Some((local, results.locals[local])) + }) + .chain(results.locals.entries().map(|(local, &ty)| (local, ty))) + .filter_map(|(local, local_ty)| { + if local_ty == ty && live_locals.insert(local) { + Some(local) + } else { + None + } + }) + .next(); + local.unwrap_or_else(|| { + let local = results.locals.push(ty); + live_locals.insert(local); + local + }) + }) + .collect::>(); + results.values[u] = locals; + } + &ValueDef::Placeholder(_) | &ValueDef::None => unreachable!(), + } + } + } + + fn handle_def( + body: &FunctionBody, + d: Value, + live_values: &mut HashSet, + live_locals: &mut HashSet, + results: &Localifier, + ) { + if let ValueDef::Alias(..) = &body.values[d] { + return; + } + + if live_values.remove(&d) { + for &local in &results.values[d] { + live_locals.remove(&local); + } + } + } + + self.body.blocks[block].terminator.visit_uses(|u| { + handle_use( + self.body, + u, + &mut live_values, + &mut live_locals, + &mut self.results, + &self.affinities, + ) + }); + + for &inst in self.body.blocks[block].insts.iter().rev() { + handle_def( + self.body, + inst, + &mut live_values, + &mut live_locals, + &self.results, + ); + self.body.values[inst].visit_uses(|u| { + // If treeified, then don't process use. + if !self.trees.owner.contains_key(&u) { + handle_use( + self.body, + u, + &mut live_values, + &mut live_locals, + &mut self.results, + &self.affinities, + ) + } + }); + } + + for &(_, param) in &self.body.blocks[block].params { + handle_def( + self.body, + param, + &mut live_values, + &mut live_locals, + &self.results, + ); + } + + log::trace!( + "localify: process block {}: livein values {:?} locals {:?}", + block, + live_values, + live_locals + ); + self.livein_locals[block] = live_locals; + self.livein_values[block] = live_values; + } +} diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 9059176..8cf332e 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -1,26 +1,28 @@ //! Backend: IR to Wasm. use crate::cfg::CFGInfo; -use crate::entity::{EntityRef, PerEntity}; -use crate::ir::{Block, FunctionBody, Value, ValueDef}; -use crate::passes::rpo::{RPOIndex, RPO}; +use crate::ir::FunctionBody; +use crate::passes::rpo::RPO; use anyhow::Result; pub mod stackify; use stackify::{Context as StackifyContext, WasmBlock}; pub mod treeify; use treeify::Trees; +pub mod localify; +use localify::Localifier; pub struct WasmBackend<'a> { body: &'a FunctionBody, rpo: RPO, trees: Trees, ctrl: Vec>, + locals: Localifier, } impl<'a> WasmBackend<'a> { pub fn new(body: &'a FunctionBody) -> Result> { - log::debug!("Backend compiling:\n{}\n", body.display("| ")); + log::debug!("Backend compiling:\n{}\n", body.display_verbose("| ")); let cfg = CFGInfo::new(body); let rpo = RPO::compute(body); log::debug!("RPO:\n{:?}\n", rpo); @@ -28,11 +30,14 @@ impl<'a> WasmBackend<'a> { log::debug!("Trees:\n{:?}\n", trees); let ctrl = StackifyContext::new(body, &cfg, &rpo)?.compute(); log::debug!("Ctrl:\n{:?}\n", ctrl); + let locals = Localifier::compute(body, &cfg, &trees); + log::debug!("Locals:\n{:?}\n", locals); Ok(WasmBackend { body, rpo, trees, ctrl, + locals, }) } diff --git a/src/backend/treeify.rs b/src/backend/treeify.rs index 30b03f6..e25b385 100644 --- a/src/backend/treeify.rs +++ b/src/backend/treeify.rs @@ -28,7 +28,7 @@ impl Trees { for (value, def) in body.values.entries() { match def { - &ValueDef::Operator(_, ref args, ref tys) => { + &ValueDef::Operator(_, ref args, _) => { // For each of the args, if the value is produced // by a single-output op and is movable, and is // not already recorded in `multi_use`, place it