Optimization: differentiate function calls (#724)

* rename FuncBody -> CompiledFunc

* make CompiledFunc use u32 internally

This is to make it possible to use CompiledFunc in wasmi bytecode.

* implement custom Default impl for CodeMap

* rename insts -> instrs in CodeMap

"insts" is ambiguous in Wasm with "instances" therefore we always want to use "instr" to refer to an instruction.

* refactor InstructionsRef for special 0 value case

* add FuncHeader::new

* cleanup code slightly

* rename parameter & minor refactoring

* fix internal doc links with FuncBody -> CompiledFunc

* improve docs in Engine

* rename resolve_inst -> resolve_instr

* add docs to some EngineInner APIs

* add doc link

* allocate CompiledFunc prior to translation

Now instead of allocation CompiledFunc to the `wasmi` Engine as the result of validation and translation of a Wasm function, the CompiledFunc is already allocated and initialized instead.
This allows to access the CompiledFunc index for function calls in other translation units which finally provides the possibility to optimize calls to internal functions.
Also this will make it simpler to run function translations concurrently if we ever decide to want this feature.

* add new CallInternal instruction

* remove unnecessary import

* add docs

Maybe we should specialize Call to CallImported and mandate that it is only used for imported functions. This could further optimize the call infrastructure in the engine by relying on this.

* fix internal doc links

* apply rustfmt

* add ReturnCallInternal instruction

crates/wasmi/src/engine/bytecode/mod.rs

@@ -20,7 +20,7 @@ pub use self::utils::{
     SignatureIdx,
     TableIdx,
 };
-use super::{const_pool::ConstRef, TranslationError};
+use super::{const_pool::ConstRef, CompiledFunc, TranslationError};
 use core::fmt::Debug;
 use wasmi_core::F32;
 
@@ -77,8 +77,29 @@ pub enum Instruction {
     ConsumeFuel(BlockFuel),
     Return(DropKeep),
     ReturnIfNez(DropKeep),
+    /// Tail calls an internal (compiled) function.
+    ///
+    /// # Note
+    ///
+    /// This instruction can be used for calls to functions that are engine internal
+    /// (or compiled) and acts as an optimization for those common cases.
+    ///
+    /// # Encoding
+    ///
+    /// This [`Instruction`] must be followed by an [`Instruction::Return`] that
+    /// encodes the [`DropKeep`] parameter. Note that the [`Instruction::Return`]
+    /// only acts as a storage for the parameter of the [`Instruction::ReturnCall`]
+    /// and will never be executed by itself.
+    ReturnCallInternal(CompiledFunc),
     /// Tail calling `func`.
     ///
+    /// # Note
+    ///
+    /// Since [`Instruction::ReturnCallInternal`] should be used for all functions internal
+    /// (or compiled) to the engine this instruction should mainly be used for tail calling
+    /// imported functions. However, it is a general form that can technically be used
+    /// for both.
+    ///
     /// # Encoding
     ///
     /// This [`Instruction`] must be followed by an [`Instruction::Return`] that
@@ -96,7 +117,21 @@ pub enum Instruction {
     /// and [`Instruction::TableGet`] only act as a storage for parameters to the
     /// [`Instruction::ReturnCallIndirect`] and will never be executed by themselves.
     ReturnCallIndirect(SignatureIdx),
+    /// Calls an internal (compiled) function.
+    ///
+    /// # Note
+    ///
+    /// This instruction can be used for calls to functions that are engine internal
+    /// (or compiled) and acts as an optimization for those common cases.
+    CallInternal(CompiledFunc),
     /// Calls the function.
+    ///
+    /// # Note
+    ///
+    /// Since [`Instruction::CallInternal`] should be used for all functions internal
+    /// (or compiled) to the engine this instruction should mainly be used for calling
+    /// imported functions. However, it is a general form that can technically be used
+    /// for both.
     Call(FuncIdx),
     /// Calling a function indirectly.
     ///

crates/wasmi/src/engine/code_map.rs

@@ -4,25 +4,59 @@ use super::Instruction;
 use alloc::vec::Vec;
 use wasmi_arena::ArenaIndex;
 
-/// A reference to a Wasm function body stored in the [`CodeMap`].
-#[derive(Debug, Copy, Clone)]
-pub struct FuncBody(usize);
+/// A reference to a compiled function stored in the [`CodeMap`] of an [`Engine`](crate::Engine).
+#[derive(Debug, Copy, Clone, PartialEq, Eq)]
+pub struct CompiledFunc(u32);
 
-impl ArenaIndex for FuncBody {
+impl ArenaIndex for CompiledFunc {
     fn into_usize(self) -> usize {
-        self.0
+        self.0 as usize
     }
 
-    fn from_usize(value: usize) -> Self {
-        FuncBody(value)
+    fn from_usize(index: usize) -> Self {
+        let index = u32::try_from(index)
+            .unwrap_or_else(|_| panic!("out of bounds compiled func index: {index}"));
+        CompiledFunc(index)
     }
 }
 
 /// A reference to the instructions of a compiled Wasm function.
 #[derive(Debug, Copy, Clone)]
 pub struct InstructionsRef {
     /// The start index in the instructions array.
-    start: usize,
+    index: usize,
+}
+
+impl InstructionsRef {
+    /// Creates a new valid [`InstructionsRef`] for the given `index`.
+    ///
+    /// # Note
+    ///
+    /// The `index` denotes the index of the first instruction in the sequence
+    /// of instructions denoted by [`InstructionsRef`].
+    ///
+    /// # Panics
+    ///
+    /// If `index` is 0 since the zero index is reserved for uninitialized [`InstructionsRef`].
+    fn new(index: usize) -> Self {
+        assert_ne!(index, 0, "must initialize with a proper non-zero index");
+        Self { index }
+    }
+
+    /// Creates a new uninitialized [`InstructionsRef`].
+    fn uninit() -> Self {
+        Self { index: 0 }
+    }
+
+    /// Returns `true` if the [`InstructionsRef`] refers to an uninitialized sequence of instructions.
+    fn is_uninit(self) -> bool {
+        self.index == 0
+    }
+
+    /// Returns the `usize` value of the underlying index.
+    fn to_usize(self) -> usize {
+        self.index
+    }
 }
 
 /// Meta information about a compiled function.
@@ -37,6 +71,32 @@ pub struct FuncHeader {
 }
 
 impl FuncHeader {
+    /// Create a new initialized [`FuncHeader`].
+    pub fn new(iref: InstructionsRef, len_locals: usize, local_stack_height: usize) -> Self {
+        let max_stack_height = local_stack_height
+            .checked_add(len_locals)
+            .unwrap_or_else(|| panic!("invalid maximum stack height for function"));
+        Self {
+            iref,
+            len_locals,
+            max_stack_height,
+        }
+    }
+
+    /// Create a new uninitialized [`FuncHeader`].
+    pub fn uninit() -> Self {
+        Self {
+            iref: InstructionsRef::uninit(),
+            len_locals: 0,
+            max_stack_height: 0,
+        }
+    }
+
+    /// Returns `true` if the [`FuncHeader`] is uninitialized.
+    pub fn is_uninit(&self) -> bool {
+        self.iref.is_uninit()
+    }
+
     /// Returns a reference to the instructions of the function.
     pub fn iref(&self) -> InstructionsRef {
         self.iref
@@ -59,7 +119,7 @@ impl FuncHeader {
 }
 
 /// Datastructure to efficiently store Wasm function bodies.
-#[derive(Debug, Default)]
+#[derive(Debug)]
 pub struct CodeMap {
     /// The headers of all compiled functions.
     headers: Vec<FuncHeader>,
@@ -72,63 +132,91 @@ pub struct CodeMap {
     ///
     /// Also this improves efficiency of deallocating the [`CodeMap`]
     /// and generally improves data locality.
-    insts: Vec<Instruction>,
+    instrs: Vec<Instruction>,
+}
+
+impl Default for CodeMap {
+    fn default() -> Self {
+        Self {
+            headers: Vec::new(),
+            // The first instruction always is a simple trapping instruction
+            // so that we safely can use `InstructionsRef(0)` as an uninitialized
+            // index value for compiled functions that have yet to be
+            // initialized with their actual function bodies.
+            instrs: vec![Instruction::Unreachable],
+        }
+    }
 }
 
 impl CodeMap {
-    /// Allocates a new function body to the [`CodeMap`].
+    /// Allocates a new uninitialized [`CompiledFunc`] to the [`CodeMap`].
+    ///
+    /// # Note
+    ///
+    /// The uninitialized [`CompiledFunc`] must be initialized using
+    /// [`CodeMap::init_func`] before it is executed.
+    pub fn alloc_func(&mut self) -> CompiledFunc {
+        let header_index = self.headers.len();
+        self.headers.push(FuncHeader::uninit());
+        CompiledFunc::from_usize(header_index)
+    }
+
+    /// Initializes the [`CompiledFunc`].
     ///
-    /// Returns a reference to the allocated function body that can
-    /// be used with [`CodeMap::header`] in order to resolve its
-    /// instructions.
-    pub fn alloc<I>(&mut self, len_locals: usize, max_stack_height: usize, insts: I) -> FuncBody
-    where
+    /// # Panics
+    ///
+    /// - If `func` is an invalid [`CompiledFunc`] reference for this [`CodeMap`].
+    /// - If `func` refers to an already initialized [`CompiledFunc`].
+    pub fn init_func<I>(
+        &mut self,
+        func: CompiledFunc,
+        len_locals: usize,
+        local_stack_height: usize,
+        instrs: I,
+    ) where
         I: IntoIterator<Item = Instruction>,
     {
-        let start = self.insts.len();
-        self.insts.extend(insts);
-        let iref = InstructionsRef { start };
-        let header = FuncHeader {
-            iref,
-            len_locals,
-            max_stack_height: len_locals + max_stack_height,
-        };
-        let header_index = self.headers.len();
-        self.headers.push(header);
-        FuncBody(header_index)
+        assert!(
+            self.header(func).is_uninit(),
+            "func {func:?} is already initialized"
+        );
+        let start = self.instrs.len();
+        self.instrs.extend(instrs);
+        let iref = InstructionsRef::new(start);
+        self.headers[func.into_usize()] = FuncHeader::new(iref, len_locals, local_stack_height);
     }
 
     /// Returns an [`InstructionPtr`] to the instruction at [`InstructionsRef`].
     #[inline]
     pub fn instr_ptr(&self, iref: InstructionsRef) -> InstructionPtr {
-        InstructionPtr::new(self.insts[iref.start..].as_ptr())
+        InstructionPtr::new(self.instrs[iref.to_usize()..].as_ptr())
     }
 
-    /// Returns the [`FuncHeader`] of the [`FuncBody`].
-    pub fn header(&self, func_body: FuncBody) -> &FuncHeader {
-        &self.headers[func_body.0]
+    /// Returns the [`FuncHeader`] of the [`CompiledFunc`].
+    pub fn header(&self, func_body: CompiledFunc) -> &FuncHeader {
+        &self.headers[func_body.into_usize()]
     }
 
-    /// Resolves the instruction at `index` of the compiled [`FuncBody`].
+    /// Resolves the instruction at `index` of the compiled [`CompiledFunc`].
     #[cfg(test)]
-    pub fn get_instr(&self, func_body: FuncBody, index: usize) -> Option<&Instruction> {
+    pub fn get_instr(&self, func_body: CompiledFunc, index: usize) -> Option<&Instruction> {
         let header = self.header(func_body);
-        let start = header.iref.start;
+        let start = header.iref.to_usize();
         let end = self.instr_end(func_body);
-        let instrs = &self.insts[start..end];
+        let instrs = &self.instrs[start..end];
         instrs.get(index)
     }
 
-    /// Returns the `end` index of the instructions of [`FuncBody`].
+    /// Returns the `end` index of the instructions of [`CompiledFunc`].
     ///
     /// This is important to synthesize how many instructions there are in
-    /// the function referred to by [`FuncBody`].
+    /// the function referred to by [`CompiledFunc`].
     #[cfg(test)]
-    pub fn instr_end(&self, func_body: FuncBody) -> usize {
+    pub fn instr_end(&self, func_body: CompiledFunc) -> usize {
         self.headers
-            .get(func_body.0 + 1)
-            .map(|header| header.iref.start)
-            .unwrap_or(self.insts.len())
+            .get(func_body.into_usize() + 1)
+            .map(|header| header.iref.to_usize())
+            .unwrap_or(self.instrs.len())
     }
 }
 

crates/wasmi/src/engine/executor.rs

@@ -1,4 +1,4 @@
-use super::{bytecode::BranchOffset, const_pool::ConstRef, ConstPoolView};
+use super::{bytecode::BranchOffset, const_pool::ConstRef, CompiledFunc, ConstPoolView};
 use crate::{
     core::TrapCode,
     engine::{
@@ -238,12 +238,16 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
                         return Ok(WasmOutcome::Return);
                     }
                 }
+                Instr::ReturnCallInternal(compiled_func) => {
+                    self.visit_return_call_internal(compiled_func)?
+                }
                 Instr::ReturnCall(func) => {
                     forward_call!(self.visit_return_call(func))
                 }
                 Instr::ReturnCallIndirect(func_type) => {
                     forward_call!(self.visit_return_call_indirect(func_type))
                 }
+                Instr::CallInternal(compiled_func) => self.visit_call_internal(compiled_func)?,
                 Instr::Call(func) => forward_call!(self.visit_call(func)),
                 Instr::CallIndirect(func_type) => {
                     forward_call!(self.visit_call_indirect(func_type))
@@ -629,6 +633,29 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
         }
     }
 
+    /// Calls the given internal [`CompiledFunc`].
+    ///
+    /// This also prepares the instruction pointer and stack pointer for
+    /// the function call so that the stack and execution state is synchronized
+    /// with the outer structures.
+    #[inline(always)]
+    fn call_func_internal(&mut self, func: CompiledFunc, kind: CallKind) -> Result<(), TrapCode> {
+        self.next_instr_at(match kind {
+            CallKind::Nested => 1,
+            CallKind::Tail => 2,
+        });
+        self.sync_stack_ptr();
+        if matches!(kind, CallKind::Nested) {
+            self.call_stack
+                .push(FuncFrame::new(self.ip, self.cache.instance()))?;
+        }
+        let header = self.code_map.header(func);
+        self.value_stack.prepare_wasm_call(header)?;
+        self.sp = self.value_stack.stack_ptr();
+        self.ip = self.code_map.instr_ptr(header.iref());
+        Ok(())
+    }
+
     /// Returns to the caller.
     ///
     /// This also modifies the stack as the caller would expect it
@@ -965,6 +992,13 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
         self.next_instr()
     }
 
+    #[inline(always)]
+    fn visit_return_call_internal(&mut self, compiled_func: CompiledFunc) -> Result<(), TrapCode> {
+        let drop_keep = self.fetch_drop_keep(1);
+        self.sp.drop_keep(drop_keep);
+        self.call_func_internal(compiled_func, CallKind::Tail)
+    }
+
     #[inline(always)]
     fn visit_return_call(&mut self, func_index: FuncIdx) -> Result<CallOutcome, TrapCode> {
         let drop_keep = self.fetch_drop_keep(1);
@@ -984,6 +1018,11 @@ impl<'ctx, 'engine> Executor<'ctx, 'engine> {
         self.execute_call_indirect(3, table, func_index, func_type, CallKind::Tail)
     }
 
+    #[inline(always)]
+    fn visit_call_internal(&mut self, compiled_func: CompiledFunc) -> Result<(), TrapCode> {
+        self.call_func_internal(compiled_func, CallKind::Nested)
+    }
+
     #[inline(always)]
     fn visit_call(&mut self, func_index: FuncIdx) -> Result<CallOutcome, TrapCode> {
         let callee = self.cache.get_func(self.ctx, func_index);