perf: reduce block overhead (#3)

explodingcamera · web-flow · commit bf30b77a5ce9 · 2024-02-24T19:17:46.000+01:00
diff --git a/BENCHMARKS.md b/BENCHMARKS.md
@@ -31,11 +31,11 @@ All runtimes are compiled with the following settings:
 | Benchmark    | Native   | TinyWasm\* | Wasmi     | Wasmer (Single Pass) |
 | ------------ | -------- | ---------- | --------- | -------------------- |
 | `fib`        | \*\*     | ` 44.11µs` | `49.46µs` | ` 50.65µs`           |
-| `fib-rec`    | `0.26ms` | ` 24.91ms` | ` 4.62ms` | `  0.49ms`           |
+| `fib-rec`    | `0.26ms` | ` 20.99ms` | ` 4.64ms` | `  0.50ms`           |
 | `argon2id`   | `0.53ms` | `109.38ms` | `45.85ms` | `  4.82ms`           |
-| `selfhosted` | `0.05ms` | `  2.07ms` | ` 4.26ms` | `260.32ms`           |
+| `selfhosted` | `0.05ms` | `  1.97ms` | ` 4.26ms` | `260.32ms`           |
 
-_\* converting WASM to TinyWasm bytecode is not included. 7.2.ms is the time it takes to convert `tinywasm.wasm` to TinyWasm bytecode._
+_\* converting WASM to TinyWasm bytecode is not included. I takes ~7ms to convert `tinywasm.wasm` to TinyWasm bytecode._
 
 _\*\* essentially instant as it gets computed at compile time._
 
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -30,6 +30,7 @@ test=false
 color-eyre="0.6"
 tinywasm={path="crates/tinywasm", features=["unsafe"]}
 wat={version="1.0"}
+pretty_env_logger="0.5"
 
 [profile.bench]
 opt-level=3
diff --git a/crates/parser/src/conversion.rs b/crates/parser/src/conversion.rs
@@ -291,8 +291,6 @@ pub(crate) fn process_operators(
             }
             End => {
                 if let Some(label_pointer) = labels_ptrs.pop() {
-                    log::debug!("ending block: {:?}", instructions[label_pointer]);
-
                     let current_instr_ptr = instructions.len();
 
                     // last_label_pointer is Some if we're ending a block
diff --git a/crates/tinywasm/src/error.rs b/crates/tinywasm/src/error.rs
@@ -24,17 +24,17 @@ pub enum Error {
     FuncDidNotReturn,
 
     /// The stack is empty
-    StackUnderflow,
+    ValueStackUnderflow,
 
     /// The label stack is empty
-    LabelStackUnderflow,
+    BlockStackUnderflow,
+
+    /// The call stack is empty
+    CallStackUnderflow,
 
     /// An invalid label type was encountered
     InvalidLabelType,
 
-    /// The call stack is empty
-    CallStackEmpty,
-
     /// The store is not the one that the module instance was instantiated in
     InvalidStore,
 
@@ -189,13 +189,13 @@ impl Display for Error {
 
             Self::Trap(trap) => write!(f, "trap: {}", trap),
             Self::Linker(err) => write!(f, "linking error: {}", err),
-            Self::CallStackEmpty => write!(f, "call stack empty"),
+            Self::CallStackUnderflow => write!(f, "call stack empty"),
             Self::InvalidLabelType => write!(f, "invalid label type"),
             Self::Other(message) => write!(f, "unknown error: {}", message),
             Self::UnsupportedFeature(feature) => write!(f, "unsupported feature: {}", feature),
             Self::FuncDidNotReturn => write!(f, "function did not return"),
-            Self::LabelStackUnderflow => write!(f, "label stack underflow"),
-            Self::StackUnderflow => write!(f, "stack underflow"),
+            Self::BlockStackUnderflow => write!(f, "label stack underflow"),
+            Self::ValueStackUnderflow => write!(f, "value stack underflow"),
             Self::InvalidStore => write!(f, "invalid store"),
         }
     }
diff --git a/crates/tinywasm/src/func.rs b/crates/tinywasm/src/func.rs
@@ -58,7 +58,7 @@ impl FuncHandle {
 
         // 6. Let f be the dummy frame
         let call_frame =
-            CallFrame::new(wasm_func.clone(), func_inst.owner, params.iter().map(|v| RawWasmValue::from(*v)));
+            CallFrame::new(wasm_func.clone(), func_inst.owner, params.iter().map(|v| RawWasmValue::from(*v)), 0);
 
         // 7. Push the frame f to the call stack
         // & 8. Push the values to the stack (Not needed since the call frame owns the values)
diff --git a/crates/tinywasm/src/instance.rs b/crates/tinywasm/src/instance.rs
@@ -61,7 +61,7 @@ impl ModuleInstance {
         // don't need to create a auxiliary frame etc.
 
         let idx = store.next_module_instance_idx();
-        log::error!("Instantiating module at index {}", idx);
+        log::info!("Instantiating module at index {}", idx);
         let imports = imports.unwrap_or_default();
 
         let mut addrs = imports.link(store, &module, idx)?;
diff --git a/crates/tinywasm/src/runtime/interpreter/macros.rs b/crates/tinywasm/src/runtime/interpreter/macros.rs
@@ -11,7 +11,7 @@
 // from a function, so we need to check if the label stack is empty
 macro_rules! break_to {
     ($cf:ident, $stack:ident, $break_to_relative:ident) => {{
-        if $cf.break_to(*$break_to_relative, &mut $stack.values).is_none() {
+        if $cf.break_to(*$break_to_relative, &mut $stack.values, &mut $stack.blocks).is_none() {
             if $stack.call_stack.is_empty() {
                 return Ok(ExecResult::Return);
             } else {
@@ -53,7 +53,6 @@ macro_rules! mem_load {
 
         const LEN: usize = core::mem::size_of::<$load_type>();
         let val = mem_ref.load_as::<LEN, $load_type>(addr, $arg.align as usize)?;
-        // let loaded_value = mem_ref.load_as::<$load_type>(addr, $arg.align as usize)?;
         $stack.values.push((val as $target_type).into());
     }};
 }
diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -4,7 +4,7 @@ use core::ops::{BitAnd, BitOr, BitXor, Neg};
 use tinywasm_types::{ElementKind, ValType};
 
 use super::{InterpreterRuntime, Stack};
-use crate::runtime::{BlockType, CallFrame, LabelFrame};
+use crate::runtime::{BlockFrame, BlockType, CallFrame};
 use crate::{cold, log, unlikely};
 use crate::{Error, FuncContext, ModuleInstance, Result, Store, Trap};
 
@@ -32,8 +32,13 @@ impl InterpreterRuntime {
             match exec_one(&mut cf, stack, store, &current_module) {
                 // Continue execution at the new top of the call stack
                 Ok(ExecResult::Call) => {
+                    let old = cf.block_ptr;
                     cf = stack.call_stack.pop()?;
 
+                    if old > cf.block_ptr {
+                        stack.blocks.truncate(old);
+                    }
+
                     // keeping the pointer seperate from the call frame is about 2% faster
                     // than storing it in the call frame
                     if cf.func_instance.1 != current_module.id() {
@@ -123,7 +128,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             };
 
             let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params);
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len());
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -180,7 +185,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             }
 
             let params = stack.values.pop_n_rev(wasm_func.ty.params.len())?;
-            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params);
+            let call_frame = CallFrame::new(wasm_func, func_inst.owner, params, stack.blocks.len());
 
             // push the call frame
             cf.instr_ptr += 1; // skip the call instruction
@@ -194,8 +199,8 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
         If(args, else_offset, end_offset) => {
             // truthy value is on the top of the stack, so enter the then block
             if stack.values.pop_t::<i32>()? != 0 {
-                cf.enter_label(
-                    LabelFrame::new(
+                cf.enter_block(
+                    BlockFrame::new(
                         cf.instr_ptr,
                         cf.instr_ptr + *end_offset,
                         stack.values.len(), // - params,
@@ -204,13 +209,14 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                         module,
                     ),
                     &mut stack.values,
+                    &mut stack.blocks,
                 );
                 return Ok(ExecResult::Ok);
             }
 
             // falsy value is on the top of the stack
             if let Some(else_offset) = else_offset {
-                let label = LabelFrame::new(
+                let label = BlockFrame::new(
                     cf.instr_ptr + *else_offset,
                     cf.instr_ptr + *end_offset,
                     stack.values.len(), // - params,
@@ -219,15 +225,15 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                     module,
                 );
                 cf.instr_ptr += *else_offset;
-                cf.enter_label(label, &mut stack.values);
+                cf.enter_block(label, &mut stack.values, &mut stack.blocks);
             } else {
                 cf.instr_ptr += *end_offset;
             }
         }
 
         Loop(args, end_offset) => {
-            cf.enter_label(
-                LabelFrame::new(
+            cf.enter_block(
+                BlockFrame::new(
                     cf.instr_ptr,
                     cf.instr_ptr + *end_offset,
                     stack.values.len(), // - params,
@@ -236,12 +242,13 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                     module,
                 ),
                 &mut stack.values,
+                &mut stack.blocks,
             );
         }
 
         Block(args, end_offset) => {
-            cf.enter_label(
-                LabelFrame::new(
+            cf.enter_block(
+                BlockFrame::new(
                     cf.instr_ptr,
                     cf.instr_ptr + *end_offset,
                     stack.values.len(), // - params,
@@ -250,6 +257,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
                     module,
                 ),
                 &mut stack.values,
+                &mut stack.blocks,
             );
         }
 
@@ -291,10 +299,10 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
         },
 
         EndFunc => {
-            assert!(
-                cf.labels.len() == 0,
-                "endfunc: block frames not empty, this should have been validated by the parser"
-            );
+            if stack.blocks.len() != cf.block_ptr {
+                cold();
+                panic!("endfunc: block frames not empty, this should have been validated by the parser");
+            }
 
             match stack.call_stack.is_empty() {
                 true => return Ok(ExecResult::Return),
@@ -304,7 +312,7 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
 
         // We're essentially using else as a EndBlockFrame instruction for if blocks
         Else(end_offset) => {
-            let Some(block) = cf.labels.pop() else {
+            let Some(block) = stack.blocks.pop() else {
                 cold();
                 panic!("else: no label to end, this should have been validated by the parser");
             };
@@ -316,16 +324,28 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
 
         EndBlockFrame => {
             // remove the label from the label stack
-            let Some(block) = cf.labels.pop() else {
+            let Some(block) = stack.blocks.pop() else {
                 cold();
-                panic!("end: no label to end, this should have been validated by the parser");
+                panic!("end blockframe: no label to end, this should have been validated by the parser");
             };
-            stack.values.truncate_keep(block.stack_ptr, block.results)
+
+            stack.values.truncate_keep(block.stack_ptr, block.results);
         }
 
         LocalGet(local_index) => stack.values.push(cf.get_local(*local_index as usize)),
         LocalSet(local_index) => cf.set_local(*local_index as usize, stack.values.pop()?),
-        LocalTee(local_index) => cf.set_local(*local_index as usize, *stack.values.last()?),
+        LocalTee(local_index) => {
+            let last_val = match stack.values.last() {
+                Ok(val) => val,
+                Err(_) => {
+                    log::error!("index: {}", local_index);
+                    log::error!("stack: {:?}", stack.values);
+
+                    panic!();
+                }
+            };
+            cf.set_local(*local_index as usize, *last_val)
+        }
 
         GlobalGet(global_index) => {
             let idx = module.resolve_global_addr(*global_index);
diff --git a/crates/tinywasm/src/runtime/stack.rs b/crates/tinywasm/src/runtime/stack.rs
@@ -1,20 +1,21 @@
-mod blocks;
+mod block_stack;
 mod call_stack;
 mod value_stack;
 
 use self::{call_stack::CallStack, value_stack::ValueStack};
-pub(crate) use blocks::{BlockType, LabelFrame};
+pub(crate) use block_stack::{BlockFrame, BlockStack, BlockType};
 pub(crate) use call_stack::CallFrame;
 
 /// A WebAssembly Stack
 #[derive(Debug)]
 pub struct Stack {
     pub(crate) values: ValueStack,
+    pub(crate) blocks: BlockStack,
     pub(crate) call_stack: CallStack,
 }
 
 impl Stack {
     pub(crate) fn new(call_frame: CallFrame) -> Self {
-        Self { values: ValueStack::default(), call_stack: CallStack::new(call_frame) }
+        Self { values: ValueStack::default(), blocks: BlockStack::default(), call_stack: CallStack::new(call_frame) }
     }
 }
diff --git a/crates/tinywasm/src/runtime/stack/block_stack.rs b/crates/tinywasm/src/runtime/stack/block_stack.rs
@@ -1,41 +1,36 @@
+use crate::{unlikely, ModuleInstance};
 use alloc::vec::Vec;
 use tinywasm_types::BlockArgs;
 
-use crate::{unlikely, ModuleInstance};
-
-#[derive(Debug, Clone)]
-pub(crate) struct Labels(Vec<LabelFrame>); // TODO: maybe Box<[LabelFrame]> by analyzing the lable count when parsing the module?
-
-impl Labels {
-    #[inline]
-    pub(crate) fn new() -> Self {
-        // this is somehow a lot faster than Vec::with_capacity(128) or even using Default::default() in the benchmarks
-        Self(Vec::new())
-    }
+#[derive(Debug, Clone, Default)]
+pub(crate) struct BlockStack(Vec<BlockFrame>); // TODO: maybe Box<[LabelFrame]> by analyzing the lable count when parsing the module?
 
+impl BlockStack {
     #[inline]
     pub(crate) fn len(&self) -> usize {
         self.0.len()
     }
 
     #[inline]
-    pub(crate) fn push(&mut self, label: LabelFrame) {
-        self.0.push(label);
+    pub(crate) fn push(&mut self, block: BlockFrame) {
+        self.0.push(block);
     }
 
     #[inline]
     /// get the label at the given index, where 0 is the top of the stack
-    pub(crate) fn get_relative_to_top(&self, index: usize) -> Option<&LabelFrame> {
+    pub(crate) fn get_relative_to(&self, index: usize, offset: usize) -> Option<&BlockFrame> {
+        let len = self.0.len() - offset;
+
         // the vast majority of wasm functions don't use break to return
-        if unlikely(index >= self.0.len()) {
+        if unlikely(index >= len) {
             return None;
         }
 
         Some(&self.0[self.0.len() - index - 1])
     }
 
     #[inline]
-    pub(crate) fn pop(&mut self) -> Option<LabelFrame> {
+    pub(crate) fn pop(&mut self) -> Option<BlockFrame> {
         self.0.pop()
     }
 
@@ -47,7 +42,7 @@ impl Labels {
 }
 
 #[derive(Debug, Clone)]
-pub(crate) struct LabelFrame {
+pub(crate) struct BlockFrame {
     // position of the instruction pointer when the block was entered
     pub(crate) instr_ptr: usize,
     // position of the end instruction of the block
@@ -60,7 +55,7 @@ pub(crate) struct LabelFrame {
     pub(crate) ty: BlockType,
 }
 
-impl LabelFrame {
+impl BlockFrame {
     #[inline]
     pub(crate) fn new(
         instr_ptr: usize,
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
diff --git a/examples/wasm-rust.rs b/examples/wasm-rust.rs

Original file line number	Diff line number	Diff line change
`@@ -291,8 +291,6 @@ pub(crate) fn process_operators(`
`291`	`291`	`}`
`292`	`292`	`End => {`
`293`	`293`	`if let Some(label_pointer) = labels_ptrs.pop() {`
`294`		`- log::debug!("ending block: {:?}", instructions[label_pointer]);`
`295`		`-`
`296`	`294`	`let current_instr_ptr = instructions.len();`
`297`	`295`
`298`	`296`	`// last_label_pointer is Some if we're ending a block`