feat: improve memory perf + basic bulk memory access

Signed-off-by: Henry Gressmann <[email protected]>

Author: explodingcamera

crates/parser/src/conversion.rs

@@ -332,6 +332,12 @@ pub fn process_operators<'a>(
             GlobalSet { global_index } => Instruction::GlobalSet(global_index),
             MemorySize { mem, mem_byte } => Instruction::MemorySize(mem, mem_byte),
             MemoryGrow { mem, mem_byte } => Instruction::MemoryGrow(mem, mem_byte),
+
+            MemoryCopy { dst_mem, src_mem } => Instruction::MemoryCopy(src_mem, dst_mem),
+            MemoryFill { mem } => Instruction::MemoryFill(mem),
+            MemoryInit { data_index, mem } => Instruction::MemoryInit(data_index, mem),
+            DataDrop { data_index } => Instruction::DataDrop(data_index),
+
             I32Const { value } => Instruction::I32Const(value),
             I64Const { value } => Instruction::I64Const(value),
             F32Const { value } => Instruction::F32Const(f32::from_bits(value.bits())),

crates/tinywasm/Cargo.toml

@@ -33,6 +33,7 @@ default=["std", "parser", "logging"]
 logging=["log", "tinywasm-types/logging", "tinywasm-parser?/logging"]
 std=["tinywasm-parser?/std", "tinywasm-types/std"]
 parser=["tinywasm-parser"]
+unsafe=[]
 
 [[test]]
 name="generate-charts"

crates/tinywasm/src/lib.rs

@@ -1,11 +1,11 @@
 #![no_std]
-#![forbid(unsafe_code)]
 #![doc(test(
     no_crate_inject,
     attr(deny(warnings, rust_2018_idioms), allow(dead_code, unused_assignments, unused_variables))
 ))]
 #![warn(missing_docs, missing_debug_implementations, rust_2018_idioms, unreachable_pub)]
 #![cfg_attr(nightly, feature(error_in_core))]
+#![cfg_attr(not(feature = "unsafe"), deny(unsafe_code))]
 
 //! A tiny WebAssembly Runtime written in Rust
 //!

crates/tinywasm/src/runtime/interpreter/macros.rs

@@ -14,33 +14,29 @@ macro_rules! mem_load {
         // TODO: there could be a lot of performance improvements here
         let mem_idx = $module.resolve_mem_addr($arg.mem_addr);
         let mem = $store.get_mem(mem_idx as usize)?;
+        let mem_ref = mem.borrow_mut();
 
         let addr = $stack.values.pop()?.raw_value();
-
         let addr = $arg.offset.checked_add(addr).ok_or_else(|| {
             Error::Trap(crate::Trap::MemoryOutOfBounds {
                 offset: $arg.offset as usize,
                 len: core::mem::size_of::<$load_type>(),
-                max: mem.borrow().max_pages(),
+                max: mem_ref.max_pages(),
             })
         })?;
 
         let addr: usize = addr.try_into().ok().ok_or_else(|| {
             Error::Trap(crate::Trap::MemoryOutOfBounds {
                 offset: $arg.offset as usize,
                 len: core::mem::size_of::<$load_type>(),
-                max: mem.borrow().max_pages(),
+                max: mem_ref.max_pages(),
             })
         })?;
 
-        let val: [u8; core::mem::size_of::<$load_type>()] = {
-            let mem = mem.borrow_mut();
-            let val = mem.load(addr, $arg.align as usize, core::mem::size_of::<$load_type>())?;
-            val.try_into().expect("slice with incorrect length")
-        };
-
-        let loaded_value = <$load_type>::from_le_bytes(val);
-        $stack.values.push((loaded_value as $target_type).into());
+        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());
     }};
 }
 

crates/tinywasm/src/runtime/interpreter/mod.rs

@@ -23,7 +23,7 @@ use macros::*;
 use traits::*;
 
 impl InterpreterRuntime {
-    #[inline(always)] // a small 2-3% performance improvement in some cases
+    // #[inline(always)] // a small 2-3% performance improvement in some cases
     pub(crate) fn exec(&self, store: &mut Store, stack: &mut Stack) -> Result<()> {
         // The current call frame, gets updated inside of exec_one
         let mut cf = stack.call_stack.pop()?;
@@ -388,6 +388,42 @@ fn exec_one(cf: &mut CallFrame, stack: &mut Stack, store: &mut Store, module: &M
             }
         }
 
+        // Bulk memory operations
+        MemoryCopy(from, to) => {
+            let size = stack.values.pop_t::<i32>()?;
+            let src = stack.values.pop_t::<i32>()?;
+            let dst = stack.values.pop_t::<i32>()?;
+
+            let mem = store.get_mem(module.resolve_mem_addr(*from) as usize)?;
+            let mut mem = mem.borrow_mut();
+
+            if from == to {
+                // copy within the same memory
+                mem.copy_within(dst as usize, src as usize, size as usize)?;
+            } else {
+                // copy between two memories
+                let mem2 = store.get_mem(module.resolve_mem_addr(*to) as usize)?;
+                let mut mem2 = mem2.borrow_mut();
+                mem2.copy_from_slice(dst as usize, mem.load(src as usize, 0, size as usize)?)?;
+            }
+        }
+
+        MemoryFill(addr) => {
+            let size = stack.values.pop_t::<i32>()?;
+            let val = stack.values.pop_t::<i32>()?;
+            let dst = stack.values.pop_t::<i32>()?;
+
+            let mem = store.get_mem(module.resolve_mem_addr(*addr) as usize)?;
+            let mut mem = mem.borrow_mut();
+            mem.fill(dst as usize, size as usize, val as u8)?;
+        }
+
+        // MemoryInit(data_index, mem_index) => {}
+        // DataDrop(data_index) => {
+        //     // let data_idx = module.resolve_data_addr(*data_index);
+        //     // let data = store.get_data(data_idx as usize)?;
+        //     // data.borrow_mut().drop()?;
+        // }
         I32Store(arg) => mem_store!(i32, arg, stack, store, module),
         I64Store(arg) => mem_store!(i64, arg, stack, store, module),
         F32Store(arg) => mem_store!(f32, arg, stack, store, module),

crates/tinywasm/src/runtime/stack/call_stack.rs

@@ -58,7 +58,6 @@ pub(crate) struct CallFrame {
 }
 
 impl CallFrame {
-    // TOOD: perf: this is called a lot, and it's a bit slow
     /// Push a new label to the label stack and ensure the stack has the correct values
     pub(crate) fn enter_label(&mut self, label_frame: LabelFrame, stack: &mut super::ValueStack) {
         if label_frame.params > 0 {
@@ -77,14 +76,14 @@ impl CallFrame {
         // will increment it by 1 since we're changing the "current" instr_ptr
         match break_to.ty {
             BlockType::Loop => {
+                // this is a loop, so we want to jump back to the start of the loop
                 self.instr_ptr = break_to.instr_ptr;
 
+                // We also want to push the params to the stack
+                value_stack.break_to(break_to.stack_ptr, break_to.params);
+
                 // check if we're breaking to the loop
                 if break_to_relative != 0 {
-                    // this is a loop, so we want to jump back to the start of the loop
-                    // We also want to push the params to the stack
-                    value_stack.break_to(break_to.stack_ptr, break_to.params);
-
                     // we also want to trim the label stack to the loop (but not including the loop)
                     self.labels.truncate(self.labels.len() - break_to_relative as usize);
                     return Some(());

crates/tinywasm/src/runtime/stack/value_stack.rs

@@ -1,7 +1,6 @@
 use core::ops::Range;
 
 use crate::{cold, runtime::RawWasmValue, unlikely, Error, Result};
-use alloc::vec;
 use alloc::vec::Vec;
 use tinywasm_types::{ValType, WasmValue};
 
@@ -15,7 +14,7 @@ pub(crate) struct ValueStack {
 
 impl Default for ValueStack {
     fn default() -> Self {
-        Self { stack: vec![RawWasmValue::default(); MIN_VALUE_STACK_SIZE] }
+        Self { stack: Vec::with_capacity(MIN_VALUE_STACK_SIZE) }
     }
 }
 

crates/tinywasm/src/runtime/value.rs

@@ -18,6 +18,7 @@ impl Debug for RawWasmValue {
 }
 
 impl RawWasmValue {
+    #[inline(always)]
     pub fn raw_value(&self) -> u64 {
         self.0
     }