Skip to content

Commit

Permalink
Refactor Instance management within the CallStack (take 2) (#1065)
Browse files Browse the repository at this point in the history
* initial WIP work

* refactor InstanceStack::push

Also remove faulty debug_assert.

* fix after merge

* add TopVec type

* add missing import

* use TopVec in InstanceStack

* add InstanceStack::reset

* rename IndexedInstance -> InstanceAndHeight

* rename field: calls -> frames

* apply clippy suggestion

* improve docs

* rename index -> height

* return new Instance in CallStack::pop

* remove unused TopVec methods

* add inline annotations

* remove no longer needed import

* improve docs

* rename top -> last as in Rust's Vec

* improve docs

* remove commented out code

* fix typo

* rename TopVec -> HeadVec

* refactor CallStack to no longer use InstanceStack

This also fixes a bug with instance reuse of tail calls when the call stack was temporarily empty while merging the tail call frames.

* add some inline annotations

* put cached memory bytes directly into Executor

This has fewer indirections when using the default linear memory.
Local perf measurements indicated up to 20% improvements for memory-intense workloads.

* fix doc link and cleanup code

* fix code from merge

* add StackOffsets abstraction

* add changed_instance field to CallFrame

This allows to minimize overhead when returning from a CallFrame that did not adjust the Instance.

* add inline annotations

* add missing docs to default_memory method

* add CachedMemory type

* add safety comments

* replace InstanceCache with finer grained caches

We now have a single MemoryCache and a single GlobalCache each caching the linear memory and global variable at index 0 respectively.
The linear memory at index 0 makes sense to cache since Wasmi does not yet support multi-memory. The global at index 0 makes sense since it usually serves as the shadow stack pointer.

* add hints to global variable executors

* define From for bytecode index type to convert to u32

* use macro to generate get_entity funcs

* inline(always) for instance getter

* fix doc links

* apply clippy suggestions

* apply rustfmt

* use old docs for memory field

* remove invalid code after merge

* remove inline annotations not present before this PR

* replace some #[inline(always)] with #[inline]

* move HeadVec to wasmi_collections

* use Option::replace where applicable

* add basic derives to HeadVec

* swap field ordering in InstanceAndHeight

* remove InstanceAndHeight due to redundant information

The CallFrame::changed_instance (bool) field already contains all the necessary information for when to pop an Instance.
  • Loading branch information
Robbepop authored Jun 15, 2024
1 parent f11e850 commit 3297811
Show file tree
Hide file tree
Showing 9 changed files with 184 additions and 61 deletions.
77 changes: 77 additions & 0 deletions crates/collections/src/head_vec.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,77 @@
use core::mem;
use std::vec::Vec;

/// A [`Vec`]-like data structure with fast access to the last item.
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct HeadVec<T> {
/// The top (or last) item in the [`HeadVec`].
head: Option<T>,
/// The rest of the items in the [`HeadVec`] excluding the last item.
rest: Vec<T>,
}

impl<T> Default for HeadVec<T> {
#[inline]
fn default() -> Self {
Self {
head: None,
rest: Vec::new(),
}
}
}

impl<T> HeadVec<T> {
/// Removes all items from the [`HeadVec`].
#[inline]
pub fn clear(&mut self) {
self.head = None;
self.rest.clear();
}

/// Returns the number of items stored in the [`HeadVec`].
#[inline]
pub fn len(&self) -> usize {
match self.head {
Some(_) => 1 + self.rest.len(),
None => 0,
}
}

/// Returns `true` if the [`HeadVec`] contains no items.
#[inline]
pub fn is_empty(&self) -> bool {
self.len() == 0
}

/// Returns a shared reference to the last item in the [`HeadVec`] if any.
///
/// Returns `None` if the [`HeadVec`] is empty.
#[inline]
pub fn last(&self) -> Option<&T> {
self.head.as_ref()
}

/// Returns an exclusive reference to the last item in the [`HeadVec`] if any.
///
/// Returns `None` if the [`HeadVec`] is empty.
#[inline]
pub fn last_mut(&mut self) -> Option<&mut T> {
self.head.as_mut()
}

/// Pushes a new `value` onto the [`HeadVec`].
#[inline]
pub fn push(&mut self, value: T) {
let prev_head = self.head.replace(value);
if let Some(prev_head) = prev_head {
self.rest.push(prev_head);
}
}

/// Pops the last `value` from the [`HeadVec`] if any.
#[inline]
pub fn pop(&mut self) -> Option<T> {
let new_top = self.rest.pop();
mem::replace(&mut self.head, new_top)
}
}
2 changes: 2 additions & 0 deletions crates/collections/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -38,6 +38,7 @@ extern crate std;

pub mod arena;
pub mod hash;
mod head_vec;
pub mod map;
pub mod set;
pub mod string_interner;
Expand All @@ -48,6 +49,7 @@ mod tests;
#[doc(inline)]
pub use self::{
arena::{Arena, ComponentVec, DedupArena},
head_vec::HeadVec,
map::Map,
set::Set,
string_interner::StringInterner,
Expand Down
3 changes: 1 addition & 2 deletions crates/wasmi/src/engine/executor/instrs.rs
Original file line number Diff line number Diff line change
Expand Up @@ -152,8 +152,7 @@ impl<'engine> Executor<'engine> {
#[inline(always)]
fn instance(call_stack: &CallStack) -> &Instance {
call_stack
.peek()
.map(CallFrame::instance)
.instance()
.expect("missing instance for non-empty call stack")
}

Expand Down
31 changes: 19 additions & 12 deletions crates/wasmi/src/engine/executor/instrs/call.rs
Original file line number Diff line number Diff line change
Expand Up @@ -214,7 +214,6 @@ impl<'engine> Executor<'engine> {
&mut self,
results: RegisterSpan,
func: &CompiledFuncEntity,
instance: &Instance,
) -> Result<CallFrame, Error> {
// We have to reinstantiate the `self.sp` [`FrameRegisters`] since we just called
// [`ValueStack::alloc_call_frame`] which might invalidate all live [`FrameRegisters`].
Expand All @@ -227,7 +226,7 @@ impl<'engine> Executor<'engine> {
self.sp = unsafe { this.stack_ptr_at(caller.base_offset()) };
})?;
let instr_ptr = InstructionPtr::new(func.instrs().as_ptr());
let frame = CallFrame::new(instr_ptr, offsets, results, *instance);
let frame = CallFrame::new(instr_ptr, offsets, results);
if <C as CallContext>::HAS_PARAMS {
self.copy_call_params(&mut uninit_params);
}
Expand Down Expand Up @@ -300,8 +299,7 @@ impl<'engine> Executor<'engine> {
mut instance: Option<Instance>,
) -> Result<(), Error> {
let func = self.code_map.get(Some(store.fuel_mut()), func)?;
let instance = instance.get_or_insert_with(|| *Self::instance(self.call_stack));
let mut called = self.dispatch_compiled_func::<C>(results, func, instance)?;
let mut called = self.dispatch_compiled_func::<C>(results, func)?;
match <C as CallContext>::KIND {
CallKind::Nested => {
// We need to update the instruction pointer of the caller call frame.
Expand All @@ -318,11 +316,16 @@ impl<'engine> Executor<'engine> {
// on the value stack which is what the function expects. After this operation we ensure
// that `self.sp` is adjusted via a call to `init_call_frame` since it may have been
// invalidated by this method.
unsafe { Stack::merge_call_frames(self.call_stack, self.value_stack, &mut called) };
let caller_instance = unsafe {
Stack::merge_call_frames(self.call_stack, self.value_stack, &mut called)
};
if let Some(caller_instance) = caller_instance {
instance.get_or_insert(caller_instance);
}
}
}
self.init_call_frame(&called);
self.call_stack.push(called)?;
self.call_stack.push(called, instance)?;
Ok(())
}

Expand Down Expand Up @@ -495,11 +498,15 @@ impl<'engine> Executor<'engine> {
) -> Result<(), Error> {
let (len_params, len_results) = self.func_types.len_params_results(host_func.ty_dedup());
let max_inout = len_params.max(len_results);
let instance = *self
.call_stack
.instance()
.expect("need to have an instance on the call stack");
// We have to reinstantiate the `self.sp` [`FrameRegisters`] since we just called
// [`ValueStack::reserve`] which might invalidate all live [`FrameRegisters`].
let caller = match <C as CallContext>::KIND {
CallKind::Nested => self.call_stack.peek().copied(),
CallKind::Tail => self.call_stack.pop(),
CallKind::Tail => self.call_stack.pop().map(|(frame, _instance)| frame),
}
.expect("need to have a caller on the call stack");
let buffer = self.value_stack.extend_by(max_inout, |this| {
Expand All @@ -513,13 +520,13 @@ impl<'engine> Executor<'engine> {
if matches!(<C as CallContext>::KIND, CallKind::Nested) {
self.update_instr_ptr_at(1);
}
self.dispatch_host_func::<T>(store, host_func, caller)
self.dispatch_host_func::<T>(store, host_func, &instance)
.map_err(|error| match self.call_stack.is_empty() {
true => error,
false => ResumableHostError::new(error, *func, results).into(),
})?;
self.memory.update(&mut store.inner, caller.instance());
self.global.update(&mut store.inner, caller.instance());
self.memory.update(&mut store.inner, &instance);
self.global.update(&mut store.inner, &instance);
let results = results.iter(len_results);
let returned = self.value_stack.drop_return(max_inout);
for (result, value) in results.zip(returned) {
Expand All @@ -542,14 +549,14 @@ impl<'engine> Executor<'engine> {
&mut self,
store: &mut Store<T>,
host_func: HostFuncEntity,
caller: CallFrame,
instance: &Instance,
) -> Result<(usize, usize), Error> {
dispatch_host_func(
store,
self.func_types,
self.value_stack,
host_func,
Some(caller.instance()),
Some(instance),
)
}

Expand Down
10 changes: 6 additions & 4 deletions crates/wasmi/src/engine/executor/instrs/return_.rs
Original file line number Diff line number Diff line change
Expand Up @@ -25,17 +25,19 @@ impl<'engine> Executor<'engine> {
/// from the returning callee to the caller.
#[inline(always)]
fn return_impl(&mut self, store: &mut StoreInner) -> ReturnOutcome {
let returned = self
let (returned, popped_instance) = self
.call_stack
.pop()
.expect("the executing call frame is always on the stack");
self.value_stack.truncate(returned.frame_offset());
let new_instance = popped_instance.and_then(|_| self.call_stack.instance());
if let Some(new_instance) = new_instance {
self.global.update(store, new_instance);
self.memory.update(store, new_instance);
}
match self.call_stack.peek() {
Some(caller) => {
Self::init_call_frame_impl(self.value_stack, &mut self.sp, &mut self.ip, caller);
let instance = caller.instance();
self.memory.update(store, instance);
self.global.update(store, instance);
ReturnOutcome::Wasm
}
None => ReturnOutcome::Host,
Expand Down
14 changes: 8 additions & 6 deletions crates/wasmi/src/engine/executor/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -215,12 +215,14 @@ impl<'engine> EngineExecutor<'engine> {
unsafe { uninit_params.init_next(value) };
}
uninit_params.init_zeroes();
self.stack.calls.push(CallFrame::new(
InstructionPtr::new(compiled_func.instrs().as_ptr()),
offsets,
RegisterSpan::new(Register::from_i16(0)),
instance,
))?;
self.stack.calls.push(
CallFrame::new(
InstructionPtr::new(compiled_func.instrs().as_ptr()),
offsets,
RegisterSpan::new(Register::from_i16(0)),
),
Some(instance),
)?;
self.execute_func(store)?;
}
FuncEntity::Host(host_func) => {
Expand Down
Loading

0 comments on commit 3297811

Please sign in to comment.