Pre-generate trampoline functions (#957)

* Refactor wasmtime_runtime::Export

Instead of an enumeration with variants that have data fields have an
enumeration where each variant has a struct, and each struct has the
data fields. This allows us to store the structs in the `wasmtime` API
and avoid lots of `panic!` calls and various extraneous matches.

* Pre-generate trampoline functions

The `wasmtime` crate supports calling arbitrary function signatures in
wasm code, and to do this it generates "trampoline functions" which have
a known ABI that then internally convert to a particular signature's ABI
and call it. These trampoline functions are currently generated
on-the-fly and are cached in the global `Store` structure. This,
however, is suboptimal for a few reasons:

* Due to how code memory is managed each trampoline resides in its own
  64kb allocation of memory. This means if you have N trampolines you're
  using N * 64kb of memory, which is quite a lot of overhead!

* Trampolines are never free'd, even if the referencing module goes
  away. This is similar to #925.

* Trampolines are a source of shared state which prevents `Store` from
  being easily thread safe.

This commit refactors how trampolines are managed inside of the
`wasmtime` crate and jit/runtime internals. All trampolines are now
allocated in the same pass of `CodeMemory` that the main module is
allocated into. A trampoline is generated per-signature in a module as
well, instead of per-function. This cache of trampolines is stored
directly inside of an `Instance`. Trampolines are stored based on
`VMSharedSignatureIndex` so they can be looked up from the internals of
the `ExportFunction` value.

The `Func` API has been updated with various bits and pieces to ensure
the right trampolines are registered in the right places. Overall this
should ensure that all trampolines necessary are generated up-front
rather than lazily. This allows us to remove the trampoline cache from
the `Compiler` type, and move one step closer to making `Compiler`
threadsafe for usage across multiple threads.

Note that as one small caveat the `Func::wrap*` family of functions
don't need to generate a trampoline at runtime, they actually generate
the trampoline at compile time which gets passed in.

Also in addition to shuffling a lot of code around this fixes one minor
bug found in `code_memory.rs`, where `self.position` was loaded before
allocation, but the allocation may push a new chunk which would cause
`self.position` to be zero instead.

* Pass the `SignatureRegistry` as an argument to where it's needed.

This avoids the need for storing it in an `Arc`.

* Ignore tramoplines for functions with lots of arguments

Co-authored-by: Dan Gohman <sunfish@mozilla.com>

crates/api/src/trampoline/create_handle.rs

@@ -3,17 +3,20 @@
 use crate::runtime::Store;
 use anyhow::Result;
 use std::any::Any;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::wasm::DefinedFuncIndex;
 use wasmtime_environ::Module;
-use wasmtime_runtime::{Imports, InstanceHandle, VMFunctionBody};
+use wasmtime_runtime::{
+    Imports, InstanceHandle, VMFunctionBody, VMSharedSignatureIndex, VMTrampoline,
+};
 
 pub(crate) fn create_handle(
     module: Module,
     store: &Store,
     finished_functions: PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
+    trampolines: HashMap<VMSharedSignatureIndex, VMTrampoline>,
     state: Box<dyn Any>,
 ) -> Result<InstanceHandle> {
     let imports = Imports::new(
@@ -38,6 +41,7 @@ pub(crate) fn create_handle(
             Arc::new(module),
             store.compiler().trap_registry().register_traps(Vec::new()),
             finished_functions.into_boxed_slice(),
+            trampolines,
             imports,
             &data_initializers,
             signatures.into_boxed_slice(),

crates/api/src/trampoline/func.rs

@@ -5,12 +5,14 @@ use crate::{Callable, FuncType, Store, Trap, Val};
 use anyhow::{bail, Result};
 use std::any::Any;
 use std::cmp;
+use std::collections::HashMap;
+use std::mem;
 use std::panic::{self, AssertUnwindSafe};
 use std::rc::Rc;
 use wasmtime_environ::entity::{EntityRef, PrimaryMap};
 use wasmtime_environ::ir::types;
 use wasmtime_environ::isa::TargetIsa;
-use wasmtime_environ::wasm::{DefinedFuncIndex, FuncIndex};
+use wasmtime_environ::wasm::FuncIndex;
 use wasmtime_environ::{
     ir, settings, CompiledFunction, CompiledFunctionUnwindInfo, Export, Module,
 };
@@ -21,7 +23,7 @@ use wasmtime_jit::trampoline::{
     binemit, pretty_error, Context, FunctionBuilder, FunctionBuilderContext,
 };
 use wasmtime_jit::{native, CodeMemory};
-use wasmtime_runtime::{InstanceHandle, VMContext, VMFunctionBody};
+use wasmtime_runtime::{InstanceHandle, VMContext, VMFunctionBody, VMTrampoline};
 
 struct TrampolineState {
     func: Rc<dyn Callable + 'static>,
@@ -145,7 +147,7 @@ fn make_trampoline(
     stub_sig.params.push(ir::AbiParam::new(pointer_type));
 
     // Compute the size of the values vector. The vmctx and caller vmctx are passed separately.
-    let value_size = 16;
+    let value_size = mem::size_of::<u128>();
     let values_vec_len = ((value_size as usize)
         * cmp::max(signature.params.len() - 2, signature.returns.len()))
         as u32;
@@ -264,10 +266,12 @@ pub fn create_handle_with_function(
 
     let mut fn_builder_ctx = FunctionBuilderContext::new();
     let mut module = Module::new();
-    let mut finished_functions: PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]> =
-        PrimaryMap::new();
+    let mut finished_functions = PrimaryMap::new();
+    let mut trampolines = HashMap::new();
     let mut code_memory = CodeMemory::new();
 
+    // First up we manufacture a trampoline which has the ABI specified by `ft`
+    // and calls into `stub_fn`...
     let sig_id = module.local.signatures.push(sig.clone());
     let func_id = module.local.functions.push(sig_id);
     module
@@ -280,16 +284,31 @@ pub fn create_handle_with_function(
         func_id.index() as u32,
         &sig,
     );
-    code_memory.publish();
-
     finished_functions.push(trampoline);
 
-    let trampoline_state = TrampolineState::new(func.clone(), code_memory);
+    // ... and then we also need a trampoline with the standard "trampoline ABI"
+    // which enters into the ABI specified by `ft`. Note that this is only used
+    // if `Func::call` is called on an object created by `Func::new`.
+    let trampoline = wasmtime_jit::make_trampoline(
+        &*isa,
+        &mut code_memory,
+        &mut fn_builder_ctx,
+        &sig,
+        mem::size_of::<u128>(),
+    )?;
+    let sig_id = store.compiler().signatures().register(&sig);
+    trampolines.insert(sig_id, trampoline);
 
+    // Next up we wrap everything up into an `InstanceHandle` by publishing our
+    // code memory (makes it executable) and ensuring all our various bits of
+    // state make it into the instance constructors.
+    code_memory.publish();
+    let trampoline_state = TrampolineState::new(func.clone(), code_memory);
     create_handle(
         module,
         store,
         finished_functions,
+        trampolines,
         Box::new(trampoline_state),
     )
 }
@@ -297,6 +316,7 @@ pub fn create_handle_with_function(
 pub unsafe fn create_handle_with_raw_function(
     ft: &FuncType,
     func: *mut [VMFunctionBody],
+    trampoline: VMTrampoline,
     store: &Store,
     state: Box<dyn Any>,
 ) -> Result<InstanceHandle> {
@@ -314,6 +334,7 @@ pub unsafe fn create_handle_with_raw_function(
 
     let mut module = Module::new();
     let mut finished_functions = PrimaryMap::new();
+    let mut trampolines = HashMap::new();
 
     let sig_id = module.local.signatures.push(sig.clone());
     let func_id = module.local.functions.push(sig_id);
@@ -321,6 +342,8 @@ pub unsafe fn create_handle_with_raw_function(
         .exports
         .insert("trampoline".to_string(), Export::Function(func_id));
     finished_functions.push(func);
+    let sig_id = store.compiler().signatures().register(&sig);
+    trampolines.insert(sig_id, trampoline);
 
-    create_handle(module, store, finished_functions, state)
+    create_handle(module, store, finished_functions, trampolines, state)
 }

crates/api/src/trampoline/global.rs

@@ -30,6 +30,12 @@ pub fn create_global(store: &Store, gt: &GlobalType, val: Val) -> Result<Instanc
         "global".to_string(),
         wasmtime_environ::Export::Global(global_id),
     );
-    let handle = create_handle(module, store, PrimaryMap::new(), Box::new(()))?;
+    let handle = create_handle(
+        module,
+        store,
+        PrimaryMap::new(),
+        Default::default(),
+        Box::new(()),
+    )?;
     Ok(handle)
 }

crates/api/src/trampoline/memory.rs

@@ -23,5 +23,11 @@ pub fn create_handle_with_memory(store: &Store, memory: &MemoryType) -> Result<I
         wasmtime_environ::Export::Memory(memory_id),
     );
 
-    create_handle(module, store, PrimaryMap::new(), Box::new(()))
+    create_handle(
+        module,
+        store,
+        PrimaryMap::new(),
+        Default::default(),
+        Box::new(()),
+    )
 }

crates/api/src/trampoline/mod.rs

@@ -14,16 +14,21 @@ use super::{Callable, FuncType, GlobalType, MemoryType, Store, TableType, Val};
 use anyhow::Result;
 use std::any::Any;
 use std::rc::Rc;
-use wasmtime_runtime::VMFunctionBody;
+use wasmtime_runtime::{VMFunctionBody, VMTrampoline};
 
 pub fn generate_func_export(
     ft: &FuncType,
     func: &Rc<dyn Callable + 'static>,
     store: &Store,
-) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
+) -> Result<(
+    wasmtime_runtime::InstanceHandle,
+    wasmtime_runtime::ExportFunction,
+)> {
     let instance = create_handle_with_function(ft, func, store)?;
-    let export = instance.lookup("trampoline").expect("trampoline export");
-    Ok((instance, export))
+    match instance.lookup("trampoline").expect("trampoline export") {
+        wasmtime_runtime::Export::Function(f) => Ok((instance, f)),
+        _ => unreachable!(),
+    }
 }
 
 /// Note that this is `unsafe` since `func` must be a valid function pointer and
@@ -32,38 +37,59 @@ pub fn generate_func_export(
 pub unsafe fn generate_raw_func_export(
     ft: &FuncType,
     func: *mut [VMFunctionBody],
+    trampoline: VMTrampoline,
     store: &Store,
     state: Box<dyn Any>,
-) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
-    let instance = func::create_handle_with_raw_function(ft, func, store, state)?;
-    let export = instance.lookup("trampoline").expect("trampoline export");
-    Ok((instance, export))
+) -> Result<(
+    wasmtime_runtime::InstanceHandle,
+    wasmtime_runtime::ExportFunction,
+)> {
+    let instance = func::create_handle_with_raw_function(ft, func, trampoline, store, state)?;
+    match instance.lookup("trampoline").expect("trampoline export") {
+        wasmtime_runtime::Export::Function(f) => Ok((instance, f)),
+        _ => unreachable!(),
+    }
 }
 
 pub fn generate_global_export(
     store: &Store,
     gt: &GlobalType,
     val: Val,
-) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
+) -> Result<(
+    wasmtime_runtime::InstanceHandle,
+    wasmtime_runtime::ExportGlobal,
+)> {
     let instance = create_global(store, gt, val)?;
-    let export = instance.lookup("global").expect("global export");
-    Ok((instance, export))
+    match instance.lookup("global").expect("global export") {
+        wasmtime_runtime::Export::Global(g) => Ok((instance, g)),
+        _ => unreachable!(),
+    }
 }
 
 pub fn generate_memory_export(
     store: &Store,
     m: &MemoryType,
-) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
+) -> Result<(
+    wasmtime_runtime::InstanceHandle,
+    wasmtime_runtime::ExportMemory,
+)> {
     let instance = create_handle_with_memory(store, m)?;
-    let export = instance.lookup("memory").expect("memory export");
-    Ok((instance, export))
+    match instance.lookup("memory").expect("memory export") {
+        wasmtime_runtime::Export::Memory(m) => Ok((instance, m)),
+        _ => unreachable!(),
+    }
 }
 
 pub fn generate_table_export(
     store: &Store,
     t: &TableType,
-) -> Result<(wasmtime_runtime::InstanceHandle, wasmtime_runtime::Export)> {
+) -> Result<(
+    wasmtime_runtime::InstanceHandle,
+    wasmtime_runtime::ExportTable,
+)> {
     let instance = create_handle_with_table(store, t)?;
-    let export = instance.lookup("table").expect("table export");
-    Ok((instance, export))
+    match instance.lookup("table").expect("table export") {
+        wasmtime_runtime::Export::Table(t) => Ok((instance, t)),
+        _ => unreachable!(),
+    }
 }

crates/api/src/trampoline/table.rs

@@ -26,5 +26,11 @@ pub fn create_handle_with_table(store: &Store, table: &TableType) -> Result<Inst
         wasmtime_environ::Export::Table(table_id),
     );
 
-    create_handle(module, store, PrimaryMap::new(), Box::new(()))
+    create_handle(
+        module,
+        store,
+        PrimaryMap::new(),
+        Default::default(),
+        Box::new(()),
+    )
 }