chore: improve simd

Signed-off-by: Henry Gressmann <mail@henrygressmann.de>

6 files changed, 98 insertions, 33 deletions

diff --git a/crates/tinywasm/src/runtime/interpreter/mod.rs b/crates/tinywasm/src/runtime/interpreter/mod.rs
index 10d90d9..d1bad62 100644
--- a/crates/tinywasm/src/runtime/interpreter/mod.rs
+++ b/crates/tinywasm/src/runtime/interpreter/mod.rs
@@ -295,7 +295,7 @@ impl<'store, 'stack> Executor<'store, 'stack> {
             // custom instructions
             LocalGet2(a, b) => self.exec_local_get2(*a, *b),
             LocalGet3(a, b, c) => self.exec_local_get3(*a, *b, *c),
-            LocalTeeGet(a, b) => self.exec_local_tee_get(*a, *b),
+            LocalTeeGet(a, b) => self.exec_local_tee_get(*a, *b)?,
             LocalGetSet(a, b) => self.exec_local_get_set(*a, *b),
             I64XorConstRotl(rotate_by) => self.exec_i64_xor_const_rotl(*rotate_by)?,
             I32LocalGetConstAdd(local, val) => self.exec_i32_local_get_const_add(*local, *val),
@@ -317,13 +317,21 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     fn exec_end_block(&mut self) -> Result<()> {
         let block = self.stack.blocks.pop()?;
         self.stack.values.truncate_keep(block.stack_ptr, block.results as u32);
+
+        #[cfg(feature = "simd")]
+        self.stack.values.truncate_keep_simd(block.simd_stack_ptr, block.simd_results as u32);
         Ok(())
     }
 
     #[inline(always)]
     fn exec_else(&mut self, end_offset: u32) -> Result<()> {
         let block = self.stack.blocks.pop()?;
+
         self.stack.values.truncate_keep(block.stack_ptr, block.results as u32);
+
+        #[cfg(feature = "simd")]
+        self.stack.values.truncate_keep_simd(block.simd_stack_ptr, block.simd_results as u32);
+
         self.cf.instr_ptr += end_offset as usize;
         Ok(())
     }
@@ -448,14 +456,14 @@ impl<'store, 'stack> Executor<'store, 'stack> {
     }
 
     #[inline(always)]
-    fn exec_local_tee_get(&mut self, a: u32, b: u32) {
-        let last =
-            self.stack.values.last().expect("localtee: stack is empty. this should have been validated by the parser");
+    fn exec_local_tee_get(&mut self, a: u32, b: u32) -> Result<()> {
+        let last = self.stack.values.last()?;
         self.cf.set_local(a, *last);
         self.stack.values.push(match a == b {
             true => *last,
             false => self.cf.get_local(b),
         });
+        Ok(())
     }
 
     #[inline(always)]
diff --git a/crates/tinywasm/src/runtime/mod.rs b/crates/tinywasm/src/runtime/mod.rs
index 705b085..d4572dd 100644
--- a/crates/tinywasm/src/runtime/mod.rs
+++ b/crates/tinywasm/src/runtime/mod.rs
@@ -3,7 +3,10 @@ mod stack;
 
 mod raw;
 
-#[cfg(all(nightly, feature = "simd"))]
+#[cfg(all(not(nightly), feature = "simd"))]
+compile_error!("`simd` feature requires nightly");
+
+#[cfg(feature = "simd")]
 mod raw_simd;
 
 use crate::Result;
diff --git a/crates/tinywasm/src/runtime/stack/block_stack.rs b/crates/tinywasm/src/runtime/stack/block_stack.rs
index 6c3a844..7382e15 100644
--- a/crates/tinywasm/src/runtime/stack/block_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/block_stack.rs
@@ -56,12 +56,16 @@ pub(crate) struct BlockFrame {
     pub(crate) end_instr_offset: u32, // position of the end instruction of the block
 
     pub(crate) stack_ptr: u32, // position of the stack pointer when the block was entered
+    pub(crate) results: u8,
+    pub(crate) params: u8,
 
-    #[cfg(all(nightly, feature = "simd"))]
+    #[cfg(feature = "simd")]
     pub(crate) simd_stack_ptr: u32, // position of the large stack pointer when the block was entered
+    #[cfg(feature = "simd")]
+    pub(crate) simd_results: u8,
+    #[cfg(feature = "simd")]
+    pub(crate) simd_params: u8,
 
-    pub(crate) results: u8,
-    pub(crate) params: u8,
     pub(crate) ty: BlockType,
 }
 
diff --git a/crates/tinywasm/src/runtime/stack/call_stack.rs b/crates/tinywasm/src/runtime/stack/call_stack.rs
index 316e3ed..b3a78ed 100644
--- a/crates/tinywasm/src/runtime/stack/call_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/call_stack.rs
@@ -62,10 +62,7 @@ impl CallFrame {
     pub(crate) fn fetch_instr(&self) -> &Instruction {
         match self.func_instance.instructions.get(self.instr_ptr) {
             Some(instr) => instr,
-            None => {
-                cold();
-                panic!("Instruction pointer out of bounds");
-            }
+            None => unreachable!("Instruction pointer out of bounds"),
         }
     }
 
@@ -87,7 +84,7 @@ impl CallFrame {
                 self.instr_ptr = break_to.instr_ptr;
 
                 // We also want to push the params to the stack
-                values.break_to(break_to.stack_ptr, break_to.params);
+                values.break_to_params(&break_to);
 
                 // check if we're breaking to the loop
                 if break_to_relative != 0 {
@@ -100,7 +97,7 @@ impl CallFrame {
             BlockType::Block | BlockType::If | BlockType::Else => {
                 // this is a block, so we want to jump to the next instruction after the block ends
                 // We also want to push the block's results to the stack
-                values.break_to(break_to.stack_ptr, break_to.results);
+                values.break_to_results(&break_to);
 
                 // (the inst_ptr will be incremented by 1 before the next instruction is executed)
                 self.instr_ptr = break_to.instr_ptr + break_to.end_instr_offset as usize;
diff --git a/crates/tinywasm/src/runtime/stack/value_stack.rs b/crates/tinywasm/src/runtime/stack/value_stack.rs
index 811cdd0..5794268 100644
--- a/crates/tinywasm/src/runtime/stack/value_stack.rs
+++ b/crates/tinywasm/src/runtime/stack/value_stack.rs
@@ -2,16 +2,18 @@ use crate::{cold, runtime::RawWasmValue, unlikely, Error, Result};
 use alloc::vec::Vec;
 use tinywasm_types::{ValType, WasmValue};
 
+use super::BlockFrame;
+
 pub(crate) const MIN_VALUE_STACK_SIZE: usize = 1024 * 128;
 
-#[cfg(all(nightly, feature = "simd"))]
+#[cfg(feature = "simd")]
 pub(crate) const MIN_SIMD_VALUE_STACK_SIZE: usize = 1024 * 32;
 
 #[derive(Debug)]
 pub(crate) struct ValueStack {
     stack: Vec<RawWasmValue>,
 
-    #[cfg(all(nightly, feature = "simd"))]
+    #[cfg(feature = "simd")]
     simd_stack: Vec<RawSimdWasmValue>,
 }
 
@@ -20,7 +22,7 @@ impl Default for ValueStack {
         Self {
             stack: Vec::with_capacity(MIN_VALUE_STACK_SIZE),
 
-            #[cfg(all(nightly, feature = "simd"))]
+            #[cfg(feature = "simd")]
             simd_stack: Vec::with_capacity(MIN_SIMD_VALUE_STACK_SIZE),
         }
     }
@@ -29,7 +31,16 @@ impl Default for ValueStack {
 impl ValueStack {
     #[inline]
     pub(crate) fn extend_from_typed(&mut self, values: &[WasmValue]) {
+        #[cfg(not(feature = "simd"))]
         self.stack.extend(values.iter().map(|v| RawWasmValue::from(*v)));
+
+        #[cfg(feature = "simd")]
+        {
+            values.iter().for_each(|v| match v {
+                WasmValue::V128(v) => self.simd_stack.push(*v),
+                v => self.stack.push(RawWasmValue::from(*v)),
+            });
+        }
     }
 
     #[inline(always)]
@@ -65,18 +76,13 @@ impl ValueStack {
 
     #[inline]
     pub(crate) fn truncate_keep(&mut self, n: u32, end_keep: u32) {
-        let total_to_keep = n + end_keep;
-        let len = self.stack.len() as u32;
-        assert!(len >= total_to_keep, "RawWasmValueotal to keep should be less than or equal to self.top");
-
-        if len <= total_to_keep {
-            return; // No need to truncate if the current size is already less than or equal to total_to_keep
-        }
+        truncate_keep(&mut self.stack, n, end_keep);
+    }
 
-        let items_to_remove = len - total_to_keep;
-        let remove_start_index = (len - items_to_remove - end_keep) as usize;
-        let remove_end_index = (len - end_keep) as usize;
-        self.stack.drain(remove_start_index..remove_end_index);
+    #[cfg(feature = "simd")]
+    #[inline]
+    pub(crate) fn truncate_keep_simd(&mut self, n: u32, end_keep: u32) {
+        truncate_keep(&mut self.simd_stack, n, end_keep);
     }
 
     #[inline(always)]
@@ -124,14 +130,44 @@ impl ValueStack {
 
     #[inline]
     pub(crate) fn pop_params(&mut self, types: &[ValType]) -> Result<Vec<WasmValue>> {
-        Ok(self.pop_n_rev(types.len())?.zip(types.iter()).map(|(v, ty)| v.attach_type(*ty)).collect())
+        #[cfg(not(feature = "simd"))]
+        return Ok(self.pop_n_rev(types.len())?.zip(types.iter()).map(|(v, ty)| v.attach_type(*ty)).collect());
+
+        #[cfg(feature = "simd")]
+        {
+            let mut values = Vec::with_capacity(types.len());
+            for ty in types {
+                match ty {
+                    ValType::V128 => values.push(WasmValue::V128(self.simd_stack.pop().unwrap())),
+                    ty => values.push(self.pop()?.attach_type(*ty)),
+                }
+            }
+            Ok(values)
+        }
+    }
+
+    #[inline]
+    pub(crate) fn break_to_results(&mut self, bf: &BlockFrame) {
+        let end = self.stack.len() - bf.results as usize;
+        self.stack.drain(bf.stack_ptr as usize..end);
+
+        #[cfg(feature = "simd")]
+        {
+            let end = self.simd_stack.len() - bf.simd_results as usize;
+            self.simd_stack.drain(bf.simd_stack_ptr as usize..end);
+        }
     }
 
     #[inline]
-    pub(crate) fn break_to(&mut self, new_stack_size: u32, result_count: u8) {
-        let start = new_stack_size as usize;
-        let end = self.stack.len() - result_count as usize;
-        self.stack.drain(start..end);
+    pub(crate) fn break_to_params(&mut self, bf: &BlockFrame) {
+        let end = self.stack.len() - bf.params as usize;
+        self.stack.drain(bf.stack_ptr as usize..end);
+
+        #[cfg(feature = "simd")]
+        {
+            let end = self.simd_stack.len() - bf.simd_params as usize;
+            self.simd_stack.drain(bf.simd_stack_ptr as usize..end);
+        }
     }
 
     #[inline]
@@ -152,6 +188,22 @@ impl ValueStack {
     }
 }
 
+#[inline(always)]
+fn truncate_keep<T>(data: &mut Vec<T>, n: u32, end_keep: u32) {
+    let total_to_keep = n + end_keep;
+    let len = data.len() as u32;
+    assert!(len >= total_to_keep, "RawWasmValueotal to keep should be less than or equal to self.top");
+
+    if len <= total_to_keep {
+        return; // No need to truncate if the current size is already less than or equal to total_to_keep
+    }
+
+    let items_to_remove = len - total_to_keep;
+    let remove_start_index = (len - items_to_remove - end_keep) as usize;
+    let remove_end_index = (len - end_keep) as usize;
+    data.drain(remove_start_index..remove_end_index);
+}
+
 #[cfg(test)]
 mod tests {
     use super::*;
diff --git a/crates/types/src/value.rs b/crates/types/src/value.rs
index 0f127a8..2248c17 100644
--- a/crates/types/src/value.rs
+++ b/crates/types/src/value.rs
@@ -18,6 +18,7 @@ pub enum WasmValue {
     F64(f64),
     // /// A 128-bit vector
     V128(u128),
+
     RefExtern(ExternAddr),
     RefFunc(FuncAddr),
     RefNull(ValType),