Compile fewer trampolines with module linking (#2774)

Previously each module in a module-linking-using-module would compile
all the trampolines for all signatures for all modules. In forest-like
situations with lots of modules this would cause quite a few trampolines
to get compiled. The original intention was to have one global list of
trampolines for all modules in the module-linking graph that they could
all share. With the current design of module linking, however, the
intention is for modules to be relatively isolated from one another
which would make achieving this difficult.

In lieu of total sharing (which would be good for the global scope
anyway but we also don't do that right now) this commit implements an
alternative strategy where each module simply compiles its own
trampolines that it itself can reach. This should mean that
module-linking modules behave more similarly to standalone modules in
terms of trampoline duplication. If we ever do global trampoline
deduplication we can likely batch this all together into one, but for
now this should fix the performance issues seen in fuzzing.

Closes #2525

crates/jit/src/instantiate.rs

@@ -205,7 +205,7 @@ pub struct CompiledModule {
     artifacts: CompilationArtifacts,
     code: Arc<ModuleCode>,
     finished_functions: FinishedFunctions,
-    trampolines: PrimaryMap<SignatureIndex, VMTrampoline>,
+    trampolines: Vec<(SignatureIndex, VMTrampoline)>,
 }
 
 impl CompiledModule {
@@ -292,7 +292,7 @@ impl CompiledModule {
     }
 
     /// Returns the per-signature trampolines for this module.
-    pub fn trampolines(&self) -> &PrimaryMap<SignatureIndex, VMTrampoline> {
+    pub fn trampolines(&self) -> &[(SignatureIndex, VMTrampoline)] {
         &self.trampolines
     }
 
@@ -444,7 +444,7 @@ fn build_code_memory(
         CodeMemory,
         (*const u8, usize),
         PrimaryMap<DefinedFuncIndex, *mut [VMFunctionBody]>,
-        PrimaryMap<SignatureIndex, VMTrampoline>,
+        Vec<(SignatureIndex, VMTrampoline)>,
     ),
     String,
 > {
@@ -464,12 +464,15 @@ fn build_code_memory(
         );
     }
 
-    let mut trampolines = PrimaryMap::new();
-    for (i, fat_ptr) in allocation.trampolines() {
-        let fat_ptr =
-            unsafe { std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr()) };
-        assert_eq!(trampolines.push(fat_ptr), i);
-    }
+    let trampolines = allocation
+        .trampolines()
+        .map(|(i, fat_ptr)| {
+            let fnptr = unsafe {
+                std::mem::transmute::<*const VMFunctionBody, VMTrampoline>(fat_ptr.as_ptr())
+            };
+            (i, fnptr)
+        })
+        .collect();
 
     let code_range = allocation.code_range();
 

crates/jit/src/object.rs

@@ -4,11 +4,11 @@ use super::trampoline::build_trampoline;
 use cranelift_frontend::FunctionBuilderContext;
 use object::write::Object;
 use serde::{Deserialize, Serialize};
+use std::collections::BTreeSet;
 use wasmtime_debug::DwarfSection;
-use wasmtime_environ::entity::PrimaryMap;
 use wasmtime_environ::isa::{unwind::UnwindInfo, TargetIsa};
 use wasmtime_environ::wasm::{FuncIndex, SignatureIndex};
-use wasmtime_environ::{CompiledFunctions, ModuleTranslation, TypeTables};
+use wasmtime_environ::{CompiledFunctions, ModuleTranslation, ModuleType, TypeTables};
 use wasmtime_obj::{ObjectBuilder, ObjectBuilderTarget};
 
 pub use wasmtime_obj::utils;
@@ -39,22 +39,26 @@ pub(crate) fn build_object(
             .map(|info| ObjectUnwindInfo::Func(translation.module.func_index(index), info.clone()))
     }));
 
-    let mut trampolines = PrimaryMap::with_capacity(types.native_signatures.len());
+    // Build trampolines for every signature that can be used by this module.
+    let signatures = translation
+        .module
+        .types
+        .values()
+        .filter_map(|t| match t {
+            ModuleType::Function(f) => Some(*f),
+            _ => None,
+        })
+        .collect::<BTreeSet<_>>();
+    let mut trampolines = Vec::with_capacity(signatures.len());
     let mut cx = FunctionBuilderContext::new();
-    // Build trampolines for every signature.
-    //
-    // TODO: for the module linking proposal this builds too many native
-    // signatures. This builds trampolines for all signatures for all modules
-    // for each module. That's a lot of trampolines! We should instead figure
-    // out a way to share trampolines amongst all modules when compiling
-    // module-linking modules.
-    for (i, native_sig) in types.native_signatures.iter() {
+    for i in signatures {
+        let native_sig = &types.native_signatures[i];
         let func = build_trampoline(isa, &mut cx, native_sig, std::mem::size_of::<u128>())?;
         // Preserve trampoline function unwind info.
         if let Some(info) = &func.unwind_info {
             unwind_info.push(ObjectUnwindInfo::Trampoline(i, info.clone()))
         }
-        trampolines.push(func);
+        trampolines.push((i, func));
     }
 
     let target = ObjectBuilderTarget::new(isa.triple().architecture)?;

crates/obj/src/builder.rs

@@ -257,7 +257,7 @@ pub struct ObjectBuilder<'a> {
     module: &'a Module,
     code_alignment: u64,
     compilation: &'a CompiledFunctions,
-    trampolines: PrimaryMap<SignatureIndex, CompiledFunction>,
+    trampolines: Vec<(SignatureIndex, CompiledFunction)>,
     dwarf_sections: Vec<DwarfSection>,
 }
 
@@ -271,7 +271,7 @@ impl<'a> ObjectBuilder<'a> {
             target,
             module,
             code_alignment: 1,
-            trampolines: PrimaryMap::new(),
+            trampolines: Vec::new(),
             dwarf_sections: vec![],
             compilation,
         }
@@ -284,7 +284,7 @@ impl<'a> ObjectBuilder<'a> {
 
     pub fn set_trampolines(
         &mut self,
-        trampolines: PrimaryMap<SignatureIndex, CompiledFunction>,
+        trampolines: Vec<(SignatureIndex, CompiledFunction)>,
     ) -> &mut Self {
         self.trampolines = trampolines;
         self
@@ -359,7 +359,7 @@ impl<'a> ObjectBuilder<'a> {
         }
         let mut trampolines = Vec::new();
         for (i, func) in self.trampolines.iter() {
-            let name = utils::trampoline_symbol_name(i).as_bytes().to_vec();
+            let name = utils::trampoline_symbol_name(*i).as_bytes().to_vec();
             trampolines.push(append_func(name, func));
         }
 
@@ -399,7 +399,7 @@ impl<'a> ObjectBuilder<'a> {
             }
         }
 
-        for (func, symbol) in self.trampolines.values().zip(trampolines) {
+        for ((_, func), symbol) in self.trampolines.iter().zip(trampolines) {
             let (_, off) = obj.symbol_section_and_offset(symbol).unwrap();
             for r in to_object_relocations(
                 func.relocations.iter(),

crates/wasmtime/src/store.rs

@@ -320,10 +320,11 @@ impl Store {
     }
 
     fn register_signatures(&self, module: &Module) {
-        let trampolines = module.compiled_module().trampolines();
         let mut signatures = self.signatures().borrow_mut();
-        for (index, wasm) in module.types().wasm_signatures.iter() {
-            signatures.register(wasm, trampolines[index]);
+        let types = module.types();
+        for (index, trampoline) in module.compiled_module().trampolines() {
+            let wasm = &types.wasm_signatures[*index];
+            signatures.register(wasm, *trampoline);
         }
     }