Skip to content

Commit

Permalink
Cache module instantiations per-thread
Browse files Browse the repository at this point in the history 
This commit continues to rejigger the API of `watt` by having a
top-level `Instance` type now instead of a suite of top-level functions.
By using a top-level `struct` we can store an internal identifier
(`usize`) which is used to key a thread-local cache for wasm blobs. This
should allow us to share resources like module instantiations across
macro invocations, ensuring that we only instantiate modules
once-per-process.

Closes #16
  • Loading branch information
@alexcrichton
alexcrichton committed Oct 31, 2019
1 parent 640edfd commit 8fd98d4
Show file tree
Hide file tree
Showing 4 changed files with 283 additions and 172 deletions.
4 changes: 2 additions & 2 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -126,11 +126,11 @@ extern crate proc_macro;

use proc_macro::TokenStream;

static WASM: &[u8] = include_bytes!("my_macro.wasm");
static WASM: watt::Instance = watt::Instance::new(include_bytes!("my_macro.wasm"));

#[proc_macro]
pub fn my_macro(input: TokenStream) -> TokenStream {
watt::proc_macro("my_macro", input, WASM)
WASM.proc_macro("my_macro", input)
}
```

Expand Down
6 changes: 3 additions & 3 deletions demo/wa/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,11 +2,11 @@ extern crate proc_macro;

use proc_macro::TokenStream;

static WASM: &[u8] = include_bytes! {
static WASM: watt::Instance = watt::Instance::new(include_bytes! {
"../../impl/target/wasm32-unknown-unknown/release/watt_demo.wasm"
};
});

#[proc_macro_derive(Demo)]
pub fn demo(input: TokenStream) -> TokenStream {
watt::proc_macro_derive("demo", input, WASM)
WASM.proc_macro_derive("demo", input)
}
196 changes: 133 additions & 63 deletions src/exec.rs
View file
Original file line number Diff line number Diff line change
@@ -1,82 +1,152 @@
use crate::data::Data;
use crate::import;
use crate::Instance;
use proc_macro::TokenStream;
use std::cell::RefCell;
use std::collections::hash_map::{HashMap, Entry};

pub fn proc_macro(fun: &str, inputs: Vec<TokenStream>, wasm: &[u8]) -> TokenStream {
_proc_macro(fun, inputs, wasm)
pub fn proc_macro(fun: &str, inputs: Vec<TokenStream>, instance: &Instance) -> TokenStream {
_proc_macro(fun, inputs, instance)
}

#[cfg(jit)]
fn _proc_macro(fun: &str, inputs: Vec<TokenStream>, wasm: &[u8]) -> TokenStream {
fn _proc_macro(fun: &str, inputs: Vec<TokenStream>, instance: &crate::Instance) -> TokenStream {
use crate::runtime::*;

let engine = Engine::new();
let mut store = Store::new(&engine);
let module = Module::new(&store, wasm);
let imports = import::extern_vals(&module, &mut store);
let module_instance = Instance::new(&store, &module, &imports).unwrap();
let main = module
.exports()
.iter()
.position(|p| p.name() == fun)
.unwrap();
let exports = module_instance.exports();
let main = exports[main].func().unwrap();
let memory = exports.iter().filter_map(|e| e.memory()).next().unwrap();

let _guard = Data::guard();
let args = Data::with(|d| {
inputs
.into_iter()
.map(|input| Val::i32(d.tokenstream.push(input) as i32))
.collect::<Vec<_>>()
});

current_memory::set(&memory, || {
let values = main.call(&args).unwrap();
let handle = values.into_iter().next().unwrap();
let handle = handle.as_i32().unwrap() as u32;
Data::with(|d| d.tokenstream[handle].clone())
struct ThreadState {
_engine: Engine,
store: Store,
modules: HashMap<usize, Module>,
instances: HashMap<usize, Instance>,
}

std::thread_local! {
static STATE: RefCell<ThreadState> = {
let engine = Engine::new();
let store = Store::new(&engine);
RefCell::new(ThreadState {
_engine: engine,
store,
modules: HashMap::new(),
instances: HashMap::new(),
})
};
}

impl ThreadState {
pub fn instance(&mut self, instance: &crate::Instance) -> (&Module, &Instance) {
let id = instance.id();
let entry = match self.instances.entry(id) {
Entry::Occupied(e) => return (&self.modules[&id], &*e.into_mut()),
Entry::Vacant(v) => v,
};

let module = Module::new(&self.store, instance.wasm);
let imports = import::extern_vals(&module, &mut self.store);
let module_instance = Instance::new(&self.store, &module, &imports).unwrap();
self.modules.insert(id, module);
let instance = entry.insert(module_instance);
(&self.modules[&id], &*instance)
}
}

STATE.with(|state| {
let mut state = state.borrow_mut();
let (module, instance) = state.instance(instance);
let main = module
.exports()
.iter()
.position(|p| p.name() == fun)
.unwrap();
let exports = instance.exports();
let main = exports[main].func().unwrap();
let memory = exports.iter().filter_map(|e| e.memory()).next().unwrap();

let _guard = Data::guard();
let args = Data::with(|d| {
inputs
.into_iter()
.map(|input| Val::i32(d.tokenstream.push(input) as i32))
.collect::<Vec<_>>()
});

current_memory::set(&memory, || {
let values = main.call(&args).unwrap();
let handle = values.into_iter().next().unwrap();
let handle = handle.as_i32().unwrap() as u32;
Data::with(|d| d.tokenstream[handle].clone())
})
})
}

#[cfg(not(jit))]
fn _proc_macro(fun: &str, inputs: Vec<TokenStream>, wasm: &[u8]) -> TokenStream {
fn _proc_macro(fun: &str, inputs: Vec<TokenStream>, instance: &Instance) -> TokenStream {
use crate::runtime::{
decode_module, get_export, init_store, instantiate_module, invoke_func, ExternVal, Value,
decode_module, get_export, init_store, instantiate_module, invoke_func,
runtime::ModuleInst, ExternVal, Store, Value,
};
use std::io::Cursor;
use std::rc::Rc;

let cursor = Cursor::new(wasm);
let module = decode_module(cursor).unwrap();
#[cfg(watt_debug)]
crate::debug::print_module(&module);

let mut store = init_store();
let extern_vals = import::extern_vals(&module, &mut store);
let module_instance = instantiate_module(&mut store, module, &extern_vals).unwrap();
let main = match get_export(&module_instance, fun) {
Ok(ExternVal::Func(main)) => main,
_ => unimplemented!("unresolved macro: {:?}", fun),
};
struct ThreadState {
store: Store,
instances: HashMap<usize, Rc<ModuleInst>>,
}

let _guard = Data::guard();
let args = Data::with(|d| {
inputs
.into_iter()
.map(|input| Value::I32(d.tokenstream.push(input)))
.collect()
});

let res = invoke_func(&mut store, main, args);
let values = match res {
Ok(values) => values,
Err(err) => panic!("{:?}", err),
};
let handle = values.into_iter().next().unwrap();
let handle = match handle {
Value::I32(handle) => handle,
_ => unimplemented!("unexpected macro return type"),
};
Data::with(|d| d.tokenstream[handle].clone())
std::thread_local! {
static STATE: RefCell<ThreadState> = {
RefCell::new(ThreadState {
store: init_store(),
instances: HashMap::new(),
})
};
}

impl ThreadState {
pub fn instance(&mut self, instance: &crate::Instance) -> &ModuleInst {
let id = instance.id();
let entry = match self.instances.entry(id) {
Entry::Occupied(e) => return e.into_mut(),
Entry::Vacant(v) => v,
};

let cursor = Cursor::new(instance.wasm);
let module = decode_module(cursor).unwrap();
#[cfg(watt_debug)]
crate::debug::print_module(&module);
let extern_vals = import::extern_vals(&module, &mut self.store);
let module_instance =
instantiate_module(&mut self.store, module, &extern_vals).unwrap();
entry.insert(module_instance)
}
}

STATE.with(|state| {
let mut state = state.borrow_mut();
let instance = state.instance(instance);
let main = match get_export(instance, fun) {
Ok(ExternVal::Func(main)) => main,
_ => unimplemented!("unresolved macro: {:?}", fun),
};

let _guard = Data::guard();
let args = Data::with(|d| {
inputs
.into_iter()
.map(|input| Value::I32(d.tokenstream.push(input)))
.collect()
});

let res = invoke_func(&mut state.store, main, args);
let values = match res {
Ok(values) => values,
Err(err) => panic!("{:?}", err),
};
let handle = values.into_iter().next().unwrap();
let handle = match handle {
Value::I32(handle) => handle,
_ => unimplemented!("unexpected macro return type"),
};
Data::with(|d| d.tokenstream[handle].clone())
})
}
Loading

0 comments on commit 8fd98d4

Please sign in to comment.