portal/waffle
1
0
Fork 0
Code Issues Pull requests Actions Packages Projects Releases Wiki Activity

add traits, passes, and utils from more_waffle

Browse source
This commit is contained in:
Graham Kelly 2024-06-03 17:26:56 -04:00
parent 284d141c36
commit 7cf77d466c
8 changed files with 546 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
Cargo.toml
View file

@ -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 <chris@cfallin.org>"]
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"

34
src/ir/module.rs
View file

@ -205,6 +205,40 @@ impl<'a> Module<'a> {
}
}
pub fn take_per_func_body<F: FnMut(&mut Self,&mut FunctionBody)>(&mut self, mut f: F) {
for func_decl in self.funcs.iter().collect::<Vec<_>>() {
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<F: FnMut(&mut FunctionBody) -> 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<F: FnMut(&mut Self,&mut FunctionBody) -> Result<(),E>,E>(&mut self, mut f: F) -> Result<(),E>{
for func_decl in self.funcs.iter().collect::<Vec<_>>() {
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).

1
src/lib.rs
View file

@ -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;

119
src/op_traits.rs
View file

@ -2151,3 +2151,122 @@ pub fn op_rematerialize(op: &Operator) -> bool {
_ => false,
}
}
pub fn rewrite_mem<V,E>(
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,
}
}

2
src/passes.rs
View file

@ -7,3 +7,5 @@ pub mod maxssa;
pub mod resolve_aliases;
pub mod ssa;
pub mod trace;
pub mod mem_fusing;
pub mod unmem;

194
src/passes/mem_fusing.rs Normal file
View file

@ -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<String, ExportKind> {
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<Self> {
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::<Vec<_>>();
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::<Vec<_>>() {
// 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(());
}

152
src/passes/unmem.rs Normal file
View file

@ -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<Item = (usize, u8)>, 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::<Vec<_>>()
.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::<Vec<_>>() {
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::<Vec<_>>() {
// if self.cfg.do_lower(&i.module, &i.name) {
if i.kind == ImportKind::Memory(mem) {
return true;
}
// }
}
false
}
}

39
src/util.rs Normal file
View file

@ -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);
}