From 7cf77d466ccd2e37f52609be8a3a27d058984c1c Mon Sep 17 00:00:00 2001 From: Graham Kelly Date: Mon, 3 Jun 2024 17:26:56 -0400 Subject: [PATCH] add traits, passes, and utils from `more_waffle` --- Cargo.toml | 7 +- src/ir/module.rs | 34 +++++++ src/lib.rs | 1 + src/op_traits.rs | 119 ++++++++++++++++++++++++ src/passes.rs | 2 + src/passes/mem_fusing.rs | 194 +++++++++++++++++++++++++++++++++++++++ src/passes/unmem.rs | 152 ++++++++++++++++++++++++++++++ src/util.rs | 39 ++++++++ 8 files changed, 546 insertions(+), 2 deletions(-) create mode 100644 src/passes/mem_fusing.rs create mode 100644 src/passes/unmem.rs create mode 100644 src/util.rs diff --git a/Cargo.toml b/Cargo.toml index 86813ba..e990887 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] -name = "waffle" -version = "0.0.27" +name = "portal-pc-waffle" +version = "0.0.27+portal" description = "Wasm Analysis Framework For Lightweight Experiments" authors = ["Chris Fallin "] license = "Apache-2.0 WITH LLVM-exception" @@ -29,3 +29,6 @@ wasm-smith = { version = "0.202", optional = true } [features] default = [] fuzzing = ["libfuzzer-sys", "wasm-smith"] + +[lib] +name = "waffle" \ No newline at end of file diff --git a/src/ir/module.rs b/src/ir/module.rs index acd45d1..493a55b 100644 --- a/src/ir/module.rs +++ b/src/ir/module.rs @@ -205,6 +205,40 @@ impl<'a> Module<'a> { } } + pub fn take_per_func_body(&mut self, mut f: F) { + for func_decl in self.funcs.iter().collect::>() { + let mut x = std::mem::take(&mut self.funcs[func_decl]); + if let Some(body) = x.body_mut(){ + f(self,body); + } + self.funcs[func_decl] = x; + } + } + + pub fn try_per_func_body Result<(),E>,E>(&mut self, mut f: F) -> Result<(),E>{ + for func_decl in self.funcs.values_mut() { + if let Some(body) = func_decl.body_mut() { + f(body)?; + } + } + Ok(()) + } + + pub fn try_take_per_func_body Result<(),E>,E>(&mut self, mut f: F) -> Result<(),E>{ + for func_decl in self.funcs.iter().collect::>() { + let mut x = std::mem::take(&mut self.funcs[func_decl]); + let mut y = None; + if let Some(body) = x.body_mut(){ + y = Some(f(self,body)); + } + self.funcs[func_decl] = x; + if let Some(z) = y{ + z?; + } + } + Ok(()) + } + pub fn expand_func<'b>(&'b mut self, id: Func) -> Result<&'b mut FuncDecl<'a>> { if let FuncDecl::Lazy(..) = self.funcs[id] { // End the borrow. This is cheap (a slice copy). diff --git a/src/lib.rs b/src/lib.rs index 2332409..74d91b6 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -23,6 +23,7 @@ pub use ops::{Ieee32, Ieee64, MemoryArg, Operator}; mod interp; pub use interp::*; +pub mod util; pub use passes::basic_opt::OptOptions; diff --git a/src/op_traits.rs b/src/op_traits.rs index 182a4e2..600b708 100644 --- a/src/op_traits.rs +++ b/src/op_traits.rs @@ -2151,3 +2151,122 @@ pub fn op_rematerialize(op: &Operator) -> bool { _ => false, } } + +pub fn rewrite_mem( + o: &mut Operator, + v: &mut [V], + mut go: impl FnMut(&mut crate::Memory, Option<&mut V>) -> Result<(),E>, +) -> Result<(),E>{ + match o { + Operator::I32Load { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::F32Load { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::F64Load { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Load8S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Load8U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Load16S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Load16U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load8S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load8U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load16S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load16U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load32S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Load32U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Store { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Store { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::F32Store { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::F64Store { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Store8 { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I32Store16 { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Store8 { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Store16 { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::I64Store32 { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::MemorySize { mem } => go(mem, None), + Operator::MemoryGrow { mem } => go(mem, None), + Operator::MemoryCopy { dst_mem, src_mem } => { + go(dst_mem, Some(&mut v[0]))?; + go(src_mem, Some(&mut v[1]))?; + Ok(()) + } + Operator::MemoryFill { mem } => go(mem, Some(&mut v[0])), + + Operator::V128Load { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load8x8S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load8x8U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load16x4S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load16x4U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load32x2S { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load32x2U { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load8Splat { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load16Splat { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load32Splat { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load64Splat { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load32Zero { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load64Zero { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Store { memory } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load8Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load16Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load32Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Load64Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Store8Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Store16Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Store32Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + Operator::V128Store64Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])), + _ => Ok(()) + } +} +pub fn mem_count(o: &Operator) -> usize { + match o { + Operator::I32Load { memory } => 1, + Operator::I64Load { memory } => 1, + Operator::F32Load { memory } => 1, + Operator::F64Load { memory } => 1, + Operator::I32Load8S { memory } => 1, + Operator::I32Load8U { memory } => 1, + Operator::I32Load16S { memory } => 1, + Operator::I32Load16U { memory } => 1, + Operator::I64Load8S { memory } => 1, + Operator::I64Load8U { memory } => 1, + Operator::I64Load16S { memory } => 1, + Operator::I64Load16U { memory } => 1, + Operator::I64Load32S { memory } => 1, + Operator::I64Load32U { memory } => 1, + Operator::I32Store { memory } => 1, + Operator::I64Store { memory } => 1, + Operator::F32Store { memory } => 1, + Operator::F64Store { memory } => 1, + Operator::I32Store8 { memory } => 1, + Operator::I32Store16 { memory } => 1, + Operator::I64Store8 { memory } => 1, + Operator::I64Store16 { memory } => 1, + Operator::I64Store32 { memory } => 1, + Operator::MemorySize { mem } => 1, + Operator::MemoryGrow { mem } => 1, + Operator::MemoryCopy { dst_mem, src_mem } => 2, + Operator::MemoryFill { mem } => 1, + + Operator::V128Load { memory } => 1, + Operator::V128Load8x8S { memory } => 1, + Operator::V128Load8x8U { memory } => 1, + Operator::V128Load16x4S { memory } => 1, + Operator::V128Load16x4U { memory } => 1, + Operator::V128Load32x2S { memory } => 1, + Operator::V128Load32x2U { memory } => 1, + Operator::V128Load8Splat { memory } => 1, + Operator::V128Load16Splat { memory } => 1, + Operator::V128Load32Splat { memory } => 1, + Operator::V128Load64Splat { memory } => 1, + Operator::V128Load32Zero { memory } => 1, + Operator::V128Load64Zero { memory } => 1, + Operator::V128Store { memory } => 1, + Operator::V128Load8Lane { memory, .. } => 1, + Operator::V128Load16Lane { memory, .. } => 1, + Operator::V128Load32Lane { memory, .. } => 1, + Operator::V128Load64Lane { memory, .. } => 1, + Operator::V128Store8Lane { memory, .. } => 1, + Operator::V128Store16Lane { memory, .. } => 1, + Operator::V128Store32Lane { memory, .. } => 1, + Operator::V128Store64Lane { memory, .. } => 1, + _ => 0, + } +} \ No newline at end of file diff --git a/src/passes.rs b/src/passes.rs index 46aa889..126ab14 100644 --- a/src/passes.rs +++ b/src/passes.rs @@ -7,3 +7,5 @@ pub mod maxssa; pub mod resolve_aliases; pub mod ssa; pub mod trace; +pub mod mem_fusing; +pub mod unmem; \ No newline at end of file diff --git a/src/passes/mem_fusing.rs b/src/passes/mem_fusing.rs new file mode 100644 index 0000000..0838710 --- /dev/null +++ b/src/passes/mem_fusing.rs @@ -0,0 +1,194 @@ +use std::{collections::BTreeMap, convert::Infallible, iter::empty}; + +use anyhow::Context; +// use libc::name_t; + +// use crate::append_before; +use crate::{ + entity::{EntityRef, EntityVec}, + ExportKind, Func, FuncDecl, FunctionBody, Memory, Module, Operator, Type, ValueDef, +}; +// use itertools::Itertools; + +pub struct Fuse { + pub resolve: Func, + pub grow: Func, + pub size: Func, +} +pub fn get_exports(m: &Module) -> BTreeMap { + let mut b = BTreeMap::new(); + for e in m.exports.iter() { + let e = e.clone(); + b.insert(e.name, e.kind); + } + return b; +} +pub fn finalize(m: &mut Module) { + let mem = m.memories[Memory::new(0)].clone(); + m.memories = EntityVec::default(); + m.memories.push(mem); +} +impl Fuse { + pub fn new(m: &Module) -> Option { + let e = get_exports(m); + let Some(ExportKind::Func(a)) = e.get("sk%resolve") else { + return None; + }; + let a = *a; + let Some(ExportKind::Func(b)) = e.get("sk%grow") else { + return None; + }; + let b = *b; + let Some(ExportKind::Func(c)) = e.get("sk%size") else { + return None; + }; + let c = *c; + return Some(Fuse { + resolve: a, + grow: b, + size: c, + }); + } + pub fn finalize(self, m: &mut Module) { + let mem = m.memories[Memory::new(0)].clone(); + let l = m.memories.len() - 1; + m.memories = EntityVec::default(); + m.memories.push(mem); + let v = vec![self.resolve, self.grow, self.size]; + let mut new = vec![]; + for f in v.clone() { + let n = m.funcs[f].clone(); + let s = n.sig(); + let name = n.name().to_owned(); + // let n = m.funcs.push(n); + let mut b = FunctionBody::new(&m, s); + let mut p = b.blocks[b.entry] + .params + .iter() + .map(|a| a.1) + .collect::>(); + let vz = b.arg_pool.from_iter(empty()); + let tz = b.type_pool.from_iter(empty()); + let ti = b.type_pool.from_iter(vec![Type::I32].into_iter()); + let i = b.add_value(ValueDef::Operator( + Operator::I32Const { value: l as u32 }, + vz, + ti, + )); + b.append_to_block(b.entry, i); + let i = b.arg_pool.from_iter(vec![p[p.len() - 1], i].into_iter()); + let i = b.add_value(ValueDef::Operator(Operator::I32Add, i, ti)); + let l = p.len(); + p[l - 1] = i; + b.append_to_block(b.entry, i); + b.set_terminator(b.entry, crate::Terminator::ReturnCall { func: f, args: p }); + let n = FuncDecl::Body(s, name, b); + let n = m.funcs.push(n); + new.push(n); + } + for x in m.exports.iter_mut() { + let ExportKind::Func(xf) = &mut x.kind else { + continue; + }; + for (o, n) in v.iter().zip(new.iter()) { + if xf == o { + *xf = *n + } + } + } + } + pub fn process(&self, f: &mut FunctionBody) { + let vz = f.arg_pool.from_iter(empty()); + let tz = f.type_pool.from_iter(empty()); + let ti = f.type_pool.from_iter(vec![Type::I32].into_iter()); + let mut ka = BTreeMap::new(); + for k in f.blocks.iter().collect::>() { + // eprintln!("dbg: {v}"); + // let k = f.value_blocks[v]; + for v in std::mem::take(&mut f.blocks[k].insts) { + let mut w = f.values[v].clone(); + // let vi = v; + if let ValueDef::Operator(a, b, c) = &mut w { + let mut bp = f.arg_pool[*b].to_vec(); + match a.clone() { + Operator::MemorySize { mem } => { + if mem.index() != 0 { + let ia = f.add_value(ValueDef::Operator( + Operator::I32Const { + value: mem.index() as u32, + }, + vz, + ti, + )); + // append_before(f, ia, vi, k); + f.append_to_block(k, ia); + *a = Operator::Call { + function_index: self.size, + }; + bp.push(ia); + } + } + Operator::MemoryGrow { mem } => { + if mem.index() != 0 { + let ia = f.add_value(ValueDef::Operator( + Operator::I32Const { + value: mem.index() as u32, + }, + vz, + ti, + )); + // append_before(f, ia, vi, k); + f.append_to_block(k, ia); + *a = Operator::Call { + function_index: self.grow, + }; + bp.push(ia); + } + } + _ => crate::op_traits::rewrite_mem(a, &mut bp, |m, v| { + if m.index() != 0 { + let ia = f.add_value(ValueDef::Operator( + Operator::I32Const { + value: m.index() as u32, + }, + vz, + ti, + )); + f.append_to_block(k, ia); + // append_before(f, ia, vi, k); + if let Some(v) = v { + let w = f.arg_pool.from_iter(vec![*v, ia].into_iter()); + let x = f.add_value(ValueDef::Operator( + Operator::Call { + function_index: self.resolve, + }, + w, + ti, + )); + f.append_to_block(k, x); + // crate::append_before(f, x, vi, k); + *v = x; + } + *m = Memory::new(0); + } + Ok::<(),Infallible>(()) + }).unwrap(), + } + *b = *ka + .entry(bp.clone()) + .or_insert_with(|| f.arg_pool.from_iter(bp.into_iter())); + } + f.values[v] = w; + f.append_to_block(k, v); + } + } + } +} +pub fn fuse(m: &mut Module) -> anyhow::Result<()> { + let f = Fuse::new(m).context("in getting the fuse funcs")?; + crate::passes::unmem::metafuse_all(m, &mut crate::passes::unmem::All {}); + // crate::passes::splice::splice_module(m)?; + m.per_func_body(|b| f.process(b)); + f.finalize(m); + return Ok(()); +} diff --git a/src/passes/unmem.rs b/src/passes/unmem.rs new file mode 100644 index 0000000..03f5ca9 --- /dev/null +++ b/src/passes/unmem.rs @@ -0,0 +1,152 @@ +use crate::util::{add_start, new_sig}; +use std::{collections::BTreeMap, iter::empty}; +use crate::{ + Block, ExportKind, Func, FuncDecl, FunctionBody, Import, ImportKind, Memory, MemoryArg, + MemoryData, Module, Operator, Signature, SignatureData, Table, Type, Value, ValueDef, +}; + +// use super::reload::{ImportedCfg, ImportedMemoriesToFunc}; + +pub fn fuse_iter(m: &mut Module, x: impl Iterator, mem: Memory) { + let null = new_sig( + m, + SignatureData { + params: vec![], + returns: vec![], + }, + ); + let mut b = FunctionBody::new(m, null); + let vz = b.arg_pool.from_iter(empty()); + let tz = b.type_pool.from_iter(empty()); + let ti = b.type_pool.from_iter(vec![Type::I32].into_iter()); + let ia = b.add_value(ValueDef::Operator(Operator::I32Const { value: 0 }, vz, ti)); + b.append_to_block(b.entry, ia); + for (a, c) in x { + // let ia = b.add_value(ValueDef::Operator(Operator::I32Const { value: a as u32 }, vz, ti)); + // b.append_to_block(b.entry, ia); + let ic = b.add_value(ValueDef::Operator( + Operator::I32Const { value: c as u32 }, + vz, + ti, + )); + b.append_to_block(b.entry, ic); + let ia = b.add_value(ValueDef::Operator( + Operator::I32Const { value: a as u32 }, + vz, + ti, + )); + b.append_to_block(b.entry, ia); + let vs = b.arg_pool.from_iter(vec![ia, ic].into_iter()); + let j = b.add_value(ValueDef::Operator( + Operator::I32Store8 { + memory: MemoryArg { + align: 0, + offset: 0, + memory: mem, + }, + }, + vs, + tz, + )); + b.append_to_block(b.entry, j); + } + b.set_terminator(b.entry, crate::Terminator::Return { values: vec![] }); + let f = m.funcs.push(FuncDecl::Body(null, format!("z"), b)); + add_start(m, f); +} +pub fn metafuse_iter(m: &mut Module, x: &[(usize, u8)], mem: Memory) { + for w in x.chunks(4096) { + fuse_iter(m, w.iter().map(|a| *a), mem); + } +} +pub fn metafuse(m: &mut Module, mem: Memory, dat: MemoryData) { + let null = new_sig( + m, + SignatureData { + params: vec![], + returns: vec![], + }, + ); + let mut v = vec![]; + for s in dat.segments.iter() { + v.extend(s.data.iter().enumerate().map(|(a, b)| (a + s.offset, *b))); + } + metafuse_iter(m, &v, mem); + let mut b = FunctionBody::new(m, null); + let vz = b.arg_pool.from_iter(empty()); + let tz = b.type_pool.from_iter(empty()); + let ti = b.type_pool.from_iter(vec![Type::I32].into_iter()); + let ia = b.add_value(ValueDef::Operator( + Operator::I32Const { + value: dat.initial_pages as u32, + }, + vz, + ti, + )); + b.append_to_block(b.entry, ia); + let vs = b.arg_pool.from_iter(vec![ia].into_iter()); + let ib = b.add_value(ValueDef::Operator( + Operator::MemoryGrow { mem: mem }, + vs, + tz, + )); + b.append_to_block(b.entry, ib); + b.set_terminator(b.entry, crate::Terminator::Return { values: vec![] }); + let f = m.funcs.push(FuncDecl::Body(null, format!("z"), b)); + add_start(m, f); +} +pub fn metafuse_all(m: &mut Module, cfg: &mut impl Cfg) { + let mut b = BTreeMap::new(); + for mem in m + .memories + .entries_mut() + .collect::>() + .into_iter() + .rev() + { + let memory64 = mem.1.memory64; + b.insert( + mem.0, + std::mem::replace( + mem.1, + MemoryData { + initial_pages: 0, + maximum_pages: None, + segments: vec![], + memory64, + }, + ), + ); + } + for mem in m.memories.iter().collect::>() { + if !cfg.unmemmable(m, mem) { + m.memories[mem] = b.remove(&mem).unwrap(); + } + } + for (c, d) in b.into_iter() { + metafuse(m, c, d); + } +} +pub trait Cfg { + fn unmemmable(&mut self, module: &mut Module, mem: Memory) -> bool; +} +pub struct All {} +impl Cfg for All { + fn unmemmable(&mut self, module: &mut Module, mem: Memory) -> bool { + true + } +} + +pub struct ImportsOnly{} +impl Cfg for ImportsOnly{ + fn unmemmable(&mut self, module: &mut Module, mem: Memory) -> bool { + for i in module.imports.iter().map(|a| a.clone()).collect::>() { + // if self.cfg.do_lower(&i.module, &i.name) { + if i.kind == ImportKind::Memory(mem) { + return true; + } + // } + } + false + } +} diff --git a/src/util.rs b/src/util.rs new file mode 100644 index 0000000..d6c6d1c --- /dev/null +++ b/src/util.rs @@ -0,0 +1,39 @@ +use crate::*; +pub fn add_start(m: &mut Module, tf: Func) { + let s = SignatureData { + params: vec![], + returns: vec![], + }; + let s = new_sig(m, s); + let mut f = FunctionBody::new(&m, s); + let vz = f.arg_pool.from_iter(std::iter::empty()); + let t = m.funcs[tf].sig(); + let t = m.signatures[t].clone().returns; + let tz = f.type_pool.from_iter(t.into_iter()); + let v = f.add_value(ValueDef::Operator( + Operator::Call { function_index: tf }, + vz, + tz, + )); + f.append_to_block(f.entry, v); + f.set_terminator( + f.entry, + match m.start_func { + Some(a) => Terminator::ReturnCall { + func: a, + args: vec![], + }, + None => Terminator::Return { values: vec![] }, + }, + ); + let f = m.funcs.push(FuncDecl::Body(s, format!("start"), f)); + m.start_func = Some(f); +} +pub fn new_sig(m: &mut Module, s: SignatureData) -> Signature { + for (a, b) in m.signatures.entries() { + if *b == s { + return a; + } + } + return m.signatures.push(s); +} \ No newline at end of file