Generate trampolines based on signatures (#947)

* Generate trampolines based on signatures

Instead of generating a trampoline-per-function generate a
trampoline-per-signature. This should hopefully greatly increase the
cache hit rate on trampolines within a module and avoid generating a
function-per-function.

* Update crates/runtime/src/traphandlers.rs

Co-Authored-By: Sergei Pepyakin <s.pepyakin@gmail.com>

Co-authored-by: Sergei Pepyakin <s.pepyakin@gmail.com>

crates/api/src/callable.rs

@@ -160,7 +160,7 @@ impl WrappedCallable for WasmtimeFn {
         let exec_code_buf = self
             .store
             .compiler_mut()
-            .get_published_trampoline(body, &signature, value_size)
+            .get_published_trampoline(&signature, value_size)
             .map_err(|e| Trap::new(format!("trampoline error: {:?}", e)))?;
 
         // Call the trampoline.
@@ -169,6 +169,7 @@ impl WrappedCallable for WasmtimeFn {
                 vmctx,
                 ptr::null_mut(),
                 exec_code_buf,
+                body,
                 values_vec.as_mut_ptr() as *mut u8,
             )
         } {

crates/jit/src/compiler.rs

@@ -22,6 +22,7 @@ use wasmtime_environ::{
 };
 use wasmtime_runtime::{
     InstantiationError, SignatureRegistry, TrapRegistration, TrapRegistry, VMFunctionBody,
+    VMSharedSignatureIndex,
 };
 
 /// Select which kind of compilation to use.
@@ -51,7 +52,7 @@ pub struct Compiler {
 
     code_memory: CodeMemory,
     trap_registry: TrapRegistry,
-    trampoline_park: HashMap<*const VMFunctionBody, *const VMFunctionBody>,
+    trampoline_park: HashMap<VMSharedSignatureIndex, *const VMFunctionBody>,
     signatures: SignatureRegistry,
     strategy: CompilationStrategy,
     cache_config: CacheConfig,
@@ -200,37 +201,31 @@ impl Compiler {
     /// Create a trampoline for invoking a function.
     pub(crate) fn get_trampoline(
         &mut self,
-        callee_address: *const VMFunctionBody,
         signature: &ir::Signature,
         value_size: usize,
     ) -> Result<*const VMFunctionBody, SetupError> {
-        use std::collections::hash_map::Entry::{Occupied, Vacant};
-        Ok(match self.trampoline_park.entry(callee_address) {
-            Occupied(entry) => *entry.get(),
-            Vacant(entry) => {
-                let body = make_trampoline(
-                    &*self.isa,
-                    &mut self.code_memory,
-                    &mut self.fn_builder_ctx,
-                    callee_address,
-                    signature,
-                    value_size,
-                )?;
-
-                entry.insert(body);
-                body
-            }
-        })
+        let index = self.signatures.register(signature);
+        if let Some(trampoline) = self.trampoline_park.get(&index) {
+            return Ok(*trampoline);
+        }
+        let body = make_trampoline(
+            &*self.isa,
+            &mut self.code_memory,
+            &mut self.fn_builder_ctx,
+            signature,
+            value_size,
+        )?;
+        self.trampoline_park.insert(index, body);
+        return Ok(body);
     }
 
     /// Create and publish a trampoline for invoking a function.
     pub fn get_published_trampoline(
         &mut self,
-        callee_address: *const VMFunctionBody,
         signature: &ir::Signature,
         value_size: usize,
     ) -> Result<*const VMFunctionBody, SetupError> {
-        let result = self.get_trampoline(callee_address, signature, value_size)?;
+        let result = self.get_trampoline(signature, value_size)?;
         self.publish_compiled_code();
         Ok(result)
     }
@@ -256,7 +251,6 @@ fn make_trampoline(
     isa: &dyn TargetIsa,
     code_memory: &mut CodeMemory,
     fn_builder_ctx: &mut FunctionBuilderContext,
-    callee_address: *const VMFunctionBody,
     signature: &ir::Signature,
     value_size: usize,
 ) -> Result<*const VMFunctionBody, SetupError> {
@@ -272,6 +266,9 @@ fn make_trampoline(
     // Add the caller `vmctx` parameter.
     wrapper_sig.params.push(ir::AbiParam::new(pointer_type));
 
+    // Add the `callee_address` parameter.
+    wrapper_sig.params.push(ir::AbiParam::new(pointer_type));
+
     // Add the `values_vec` parameter.
     wrapper_sig.params.push(ir::AbiParam::new(pointer_type));
 
@@ -287,9 +284,9 @@ fn make_trampoline(
         builder.switch_to_block(block0);
         builder.seal_block(block0);
 
-        let (vmctx_ptr_val, caller_vmctx_ptr_val, values_vec_ptr_val) = {
+        let (vmctx_ptr_val, caller_vmctx_ptr_val, callee_value, values_vec_ptr_val) = {
             let params = builder.func.dfg.block_params(block0);
-            (params[0], params[1], params[2])
+            (params[0], params[1], params[2], params[3])
         };
 
         // Load the argument values out of `values_vec`.
@@ -318,10 +315,6 @@ fn make_trampoline(
 
         let new_sig = builder.import_signature(signature.clone());
 
-        // TODO: It's possible to make this a direct call. We just need Cranelift
-        // to support functions declared with an immediate integer address.
-        // ExternalName::Absolute(u64). Let's do it.
-        let callee_value = builder.ins().iconst(pointer_type, callee_address as i64);
         let call = builder
             .ins()
             .call_indirect(new_sig, callee_value, &callee_args);

crates/runtime/signalhandlers/Trampolines.cpp

@@ -7,20 +7,25 @@ int WasmtimeCallTrampoline(
     void **buf_storage,
     void *vmctx,
     void *caller_vmctx,
-    void (*body)(void*, void*, void*),
+    void (*trampoline)(void*, void*, void*, void*),
+    void *body,
     void *args)
 {
   jmp_buf buf;
   if (setjmp(buf) != 0) {
     return 0;
   }
   *buf_storage = &buf;
-  body(vmctx, caller_vmctx, args);
+  trampoline(vmctx, caller_vmctx, body, args);
   return 1;
 }
 
 extern "C"
-int WasmtimeCall(void **buf_storage, void *vmctx, void *caller_vmctx, void (*body)(void*, void*)) {
+int WasmtimeCall(
+    void **buf_storage,
+    void *vmctx,
+    void *caller_vmctx,
+    void (*body)(void*, void*)) {
   jmp_buf buf;
   if (setjmp(buf) != 0) {
     return 0;

crates/runtime/src/instance.rs

@@ -347,10 +347,14 @@ impl Instance {
         &*self.host_state
     }
 
-    fn invoke_function(&self, index: FuncIndex) -> Result<(), InstantiationError> {
-        // TODO: Check that the callee's calling convention matches what we expect.
+    /// Invoke the WebAssembly start function of the instance, if one is present.
+    fn invoke_start_function(&self) -> Result<(), InstantiationError> {
+        let start_index = match self.module.start_func {
+            Some(idx) => idx,
+            None => return Ok(()),
+        };
 
-        let (callee_address, callee_vmctx) = match self.module.defined_func_index(index) {
+        let (callee_address, callee_vmctx) = match self.module.defined_func_index(start_index) {
             Some(defined_index) => {
                 let body = *self
                     .finished_functions
@@ -359,8 +363,8 @@ impl Instance {
                 (body as *const _, self.vmctx_ptr())
             }
             None => {
-                assert_lt!(index.index(), self.module.imported_funcs.len());
-                let import = self.imported_function(index);
+                assert_lt!(start_index.index(), self.module.imported_funcs.len());
+                let import = self.imported_function(start_index);
                 (import.body, import.vmctx)
             }
         };
@@ -370,15 +374,6 @@ impl Instance {
             .map_err(InstantiationError::StartTrap)
     }
 
-    /// Invoke the WebAssembly start function of the instance, if one is present.
-    fn invoke_start_function(&self) -> Result<(), InstantiationError> {
-        if let Some(start_index) = self.module.start_func {
-            self.invoke_function(start_index)
-        } else {
-            Ok(())
-        }
-    }
-
     /// Return the offset from the vmctx pointer to its containing Instance.
     pub(crate) fn vmctx_offset() -> isize {
         offset_of!(Self, vmctx) as isize

crates/runtime/src/traphandlers.rs

@@ -17,6 +17,7 @@ extern "C" {
         jmp_buf: *mut *const u8,
         vmctx: *mut u8,
         caller_vmctx: *mut u8,
+        trampoline: *const VMFunctionBody,
         callee: *const VMFunctionBody,
         values_vec: *mut u8,
     ) -> i32;
@@ -133,13 +134,23 @@ impl fmt::Display for Trap {
 
 impl std::error::Error for Trap {}
 
-/// Call the wasm function pointed to by `callee`. `values_vec` points to
-/// a buffer which holds the incoming arguments, and to which the outgoing
-/// return values will be written.
-#[no_mangle]
-pub unsafe extern "C" fn wasmtime_call_trampoline(
+/// Call the wasm function pointed to by `callee`.
+///
+/// * `vmctx` - the callee vmctx argument
+/// * `caller_vmctx` - the caller vmctx argument
+/// * `trampoline` - the jit-generated trampoline whose ABI takes 4 values, the
+///   callee vmctx, the caller vmctx, the `callee` argument below, and then the
+///   `values_vec` argument.
+/// * `callee` - the third argument to the `trampoline` function
+/// * `values_vec` - points to a buffer which holds the incoming arguments, and to
+///   which the outgoing return values will be written.
+///
+/// Wildly unsafe because it calls raw function pointers and reads/writes raw
+/// function pointers.
+pub unsafe fn wasmtime_call_trampoline(
     vmctx: *mut VMContext,
     caller_vmctx: *mut VMContext,
+    trampoline: *const VMFunctionBody,
     callee: *const VMFunctionBody,
     values_vec: *mut u8,
 ) -> Result<(), Trap> {
@@ -148,6 +159,7 @@ pub unsafe extern "C" fn wasmtime_call_trampoline(
             cx.jmp_buf.as_ptr(),
             vmctx as *mut u8,
             caller_vmctx as *mut u8,
+            trampoline,
             callee,
             values_vec,
         )
@@ -156,8 +168,7 @@ pub unsafe extern "C" fn wasmtime_call_trampoline(
 
 /// Call the wasm function pointed to by `callee`, which has no arguments or
 /// return values.
-#[no_mangle]
-pub unsafe extern "C" fn wasmtime_call(
+pub unsafe fn wasmtime_call(
     vmctx: *mut VMContext,
     caller_vmctx: *mut VMContext,
     callee: *const VMFunctionBody,