Compile fewer trampolines with module linking

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);
  }
  }