diff --git a/src/backend/localify.rs b/src/backend/localify.rs index 2770575..8b1008a 100644 --- a/src/backend/localify.rs +++ b/src/backend/localify.rs @@ -74,7 +74,7 @@ impl<'a, V: Visitor> BlockVisitor<'a, V> { self.visitor.pre_term(); for &inst in self.body.blocks[block].insts.iter().rev() { - if self.trees.owner.contains_key(&inst) { + if self.trees.owner.contains_key(&inst) || self.trees.remat.contains(&inst) { continue; } self.visitor.post_inst(inst); diff --git a/src/backend/mod.rs b/src/backend/mod.rs index 8a649a5..9b42cce 100644 --- a/src/backend/mod.rs +++ b/src/backend/mod.rs @@ -142,7 +142,7 @@ impl<'a> WasmFuncBackend<'a> { for &inst in &self.body.blocks[*block].insts { // If this value is "owned", do nothing: it will be lowered in // the one place it's used. - if self.trees.owner.contains_key(&inst) { + if self.trees.owner.contains_key(&inst) || self.trees.remat.contains(&inst) { continue; } if let &ValueDef::Operator(..) = &self.body.values[inst] { @@ -180,12 +180,16 @@ impl<'a> WasmFuncBackend<'a> { fn lower_value(&self, value: Value, func: &mut wasm_encoder::Function) { log::trace!("lower_value: value {}", value); let value = self.body.resolve_alias(value); - let local = match &self.body.values[value] { - &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => self.locals.values[value][0], - &ValueDef::PickOutput(orig_value, idx, _) => self.locals.values[orig_value][idx], - _ => unreachable!(), - }; - func.instruction(&wasm_encoder::Instruction::LocalGet(local.index() as u32)); + if self.trees.remat.contains(&value) { + self.lower_inst(value, /* root = */ false, func); + } else { + let local = match &self.body.values[value] { + &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => self.locals.values[value][0], + &ValueDef::PickOutput(orig_value, idx, _) => self.locals.values[orig_value][idx], + _ => unreachable!(), + }; + func.instruction(&wasm_encoder::Instruction::LocalGet(local.index() as u32)); + } } fn lower_set_value(&self, value: Value, func: &mut wasm_encoder::Function) { @@ -205,7 +209,7 @@ impl<'a> WasmFuncBackend<'a> { &ValueDef::Operator(ref op, ref args, ref tys) => { for &arg in &args[..] { let arg = self.body.resolve_alias(arg); - if self.trees.owner.contains_key(&arg) { + if self.trees.owner.contains_key(&arg) || self.trees.remat.contains(&arg) { log::trace!(" -> arg {} is owned", arg); self.lower_inst(arg, /* root = */ false, func); } else { diff --git a/src/backend/treeify.rs b/src/backend/treeify.rs index 4693ce0..550bc8b 100644 --- a/src/backend/treeify.rs +++ b/src/backend/treeify.rs @@ -4,7 +4,7 @@ use crate::entity::EntityRef; use crate::ir::{FunctionBody, Value, ValueDef}; use crate::Operator; -use std::collections::{HashMap, HashSet}; +use fxhash::{FxHashMap as HashMap, FxHashSet as HashSet}; use std::convert::TryFrom; /// One "argument slot" of an operator defining a value. @@ -19,22 +19,43 @@ pub struct Trees { /// For a given value that is defined by an operator, which /// Values, if any, live at each slot? pub owned: HashMap, + /// Values that are regenerated every time they are used. + pub remat: HashSet, +} + +fn is_remat(op: &Operator) -> bool { + // Only ops with no args can be always-rematerialized. + match op { + Operator::I32Const { .. } + | Operator::I64Const { .. } + | Operator::F32Const { .. } + | Operator::F64Const { .. } => true, + _ => false, + } } impl Trees { pub fn compute(body: &FunctionBody) -> Trees { - let mut owner = HashMap::new(); - let mut owned = HashMap::new(); - let mut multi_use = HashSet::new(); + let mut owner = HashMap::default(); + let mut owned = HashMap::default(); + let mut remat = HashSet::default(); + let mut multi_use = HashSet::default(); for (value, def) in body.values.entries() { match def { - &ValueDef::Operator(_, ref args, _) => { + &ValueDef::Operator(op, ref args, _) => { // Ignore operators with invalid args: these must // always be unreachable. if args.iter().any(|arg| arg.is_invalid()) { continue; } + // If this is an always-rematerialized operator, + // mark it as such and continue. + if is_remat(&op) { + remat.insert(value); + continue; + } + // 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 @@ -75,7 +96,11 @@ impl Trees { }); } - Trees { owner, owned } + Trees { + owner, + owned, + remat, + } } fn is_single_output_op(body: &FunctionBody, value: Value) -> Option {