diff --git a/Cargo.lock b/Cargo.lock index 99af652..7d2be98 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2963,6 +2963,7 @@ dependencies = [ "wasmtime-wasi", "wasmtime-wasi-http", "wasmtime-wizer", + "wat", ] [[package]] diff --git a/Cargo.toml b/Cargo.toml index fa1655f..c65e61c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -29,6 +29,9 @@ sha2 = "0.10.8" sqlite = "0.36.0" serde = { version = "1.0.197", features = ["derive"] } +[dev-dependencies] +wat = "1.245.1" + [lints.clippy] # The default set of lints in Clippy is viewed as "too noisy" right now so # they're all turned off by default. Selective lints are then enabled below as diff --git a/src/constant_offsets.rs b/src/constant_offsets.rs index eefce8a..e0f6772 100644 --- a/src/constant_offsets.rs +++ b/src/constant_offsets.rs @@ -3,7 +3,7 @@ //! these offsets into loads/stores where possible. use fxhash::{FxHashMap, FxHashSet}; -use std::collections::{BTreeMap, VecDeque}; +use std::collections::{BTreeMap, BTreeSet, VecDeque}; use waffle::{ cfg::CFGInfo, entity::PerEntity, pool::ListRef, Block, FunctionBody, Operator, Terminator, Type, Value, ValueDef, @@ -46,6 +46,10 @@ pub fn run(func: &mut FunctionBody, cfg: &CFGInfo) { // Compute a fixpoint analysis: which values are some original SSA // value plus an offset? let mut values: PerEntity = PerEntity::default(); + // A value may be used across a join without being a block parameter + // there. Revisit its users when its abstract value changes, even if + // neither predecessor dominates the join. + let mut value_dep_blocks: FxHashMap> = FxHashMap::default(); let mut workqueue: VecDeque = VecDeque::new(); let mut workqueue_set: FxHashSet = FxHashSet::default(); @@ -61,9 +65,21 @@ pub fn run(func: &mut FunctionBody, cfg: &CFGInfo) { while let Some(block) = workqueue.pop_front() { log::trace!("processing {block}"); workqueue_set.remove(&block); + if visited.insert(block) { + for &inst in &func.blocks[block].insts { + func.values[inst].visit_uses(&func.arg_pool, |value| { + value_dep_blocks.entry(value).or_default().insert(block); + }); + } + func.blocks[block].terminator.visit_uses(|value| { + value_dep_blocks.entry(value).or_default().insert(block); + }); + } + let mut changed_values = vec![]; for &inst in &func.blocks[block].insts { log::trace!("block {} value {}: {:?}", block, inst, func.values[inst]); + let old = values[inst]; match &func.values[inst] { ValueDef::BlockParam(..) => { unreachable!(); @@ -135,29 +151,43 @@ pub fn run(func: &mut FunctionBody, cfg: &CFGInfo) { } } log::trace!(" -> values[{}] = {:?}", inst, values[inst]); + if old != values[inst] { + changed_values.push(inst); + } } func.blocks[block].terminator.visit_targets(|target| { - let mut changed = false; let succ_params = &func.blocks[target.block].params; for (&arg, &(_, blockparam)) in target.args.iter().zip(succ_params.iter()) { let new = AbsValue::meet(values[arg], values[blockparam]); - log::trace!(" -> block {} target {}: arg {} to blockparam {}: value {:?} -> {:?}", - block, target.block, arg, blockparam, values[blockparam], new); - changed |= new != values[blockparam]; + log::trace!( + " -> block {} target {}: arg {} to blockparam {}: value {:?} -> {:?}", + block, + target.block, + arg, + blockparam, + values[blockparam], + new + ); + if new != values[blockparam] { + changed_values.push(blockparam); + } values[blockparam] = new; } - if changed || - visited.insert(target.block) || - (block != target.block && cfg.dominates(block, target.block)) { - log::trace!(" -> at least one blockparam changed, or we dominate target block; enqueuing {}", - target.block); - if workqueue_set.insert(target.block) { - workqueue.push_back(target.block); - } + if !visited.contains(&target.block) && workqueue_set.insert(target.block) { + workqueue.push_back(target.block); } }); + for value in changed_values { + if let Some(deps) = value_dep_blocks.get(&value) { + for &block in deps { + if workqueue_set.insert(block) { + workqueue.push_back(block); + } + } + } + } } // Constant-fold conditional branches. @@ -343,3 +373,50 @@ pub fn run(func: &mut FunctionBody, cfg: &CFGInfo) { block_def.insts = new_insts; } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_loop_carried_offset_after_diamond() { + let bytes = wat::parse_str(include_str!("../tests/constant-offsets.wat")).unwrap(); + let mut module = waffle::Module::from_wasm_bytes(&bytes, &Default::default()).unwrap(); + module.expand_all_funcs().unwrap(); + for (_, decl) in module.funcs.entries_mut() { + if let waffle::FuncDecl::Body(_, _, body) = decl { + body.optimize(&waffle::OptOptions { + gvn: false, + cprop: false, + redundant_blockparams: true, + }); + let cfg = CFGInfo::new(body); + run(body, &cfg); + } + } + let optimized = module.to_wasm_bytes().unwrap(); + let engine = wasmtime::Engine::default(); + for (name, bytes) in [("original", &bytes), ("optimized", &optimized)] { + let module = wasmtime::Module::new(&engine, bytes).unwrap(); + let mut store = wasmtime::Store::new(&engine, ()); + let instance = wasmtime::Instance::new(&mut store, &module, &[]).unwrap(); + let compute = instance + .get_typed_func::<(i32, i32, i32), i32>(&mut store, "compute") + .unwrap(); + for base in [4096i32, 991, 0, i32::MIN, i32::MAX] { + for count in [0, 1, 2, 7, 32] { + for branch in [0, 1] { + let expected = base + .wrapping_sub(16 + 8 * (count + 1)) + .wrapping_add(if branch == 0 { 29 } else { 17 }); + assert_eq!( + compute.call(&mut store, (base, count, branch)).unwrap(), + expected, + "{name}: base={base}, count={count}, branch={branch}" + ); + } + } + } + } + } +} diff --git a/tests/constant-offsets.wat b/tests/constant-offsets.wat new file mode 100644 index 0000000..25017c0 --- /dev/null +++ b/tests/constant-offsets.wat @@ -0,0 +1,45 @@ +(module + (memory 1) + (data (i32.const 0) "\11\00\00\00\1d\00\00\00") + (func (export "compute") (param $base i32) (param $count i32) (param $branch i32) (result i32) + (local $pointer i32) + local.get $base + i32.const 16 + i32.sub + local.set $pointer + block $done + local.get $count + i32.eqz + br_if $done + ;; This extra diamond exposes the missed update across the later join. + local.get $branch + if (result i32) + i32.const 3 + else + i32.const 5 + end + drop + loop $decrement + local.get $pointer + i32.const 8 + i32.sub + local.set $pointer + local.get $count + i32.const 1 + i32.sub + local.tee $count + br_if $decrement + end + end + local.get $branch + if (result i32) + i32.const 0 + i32.load + else + i32.const 4 + i32.load + end + local.get $pointer + i32.const 8 + i32.sub + i32.add))