Load generated trampolines into jitdump when profiling

crates/jit/src/instantiate.rs

@@ -397,13 +397,17 @@ impl CompiledModule {
  }
 
  /// Returns the per-signature trampolines for this module.
- pub fn trampolines(&self) -> impl Iterator<Item = (SignatureIndex, VMTrampoline)> + '_ {
+ pub fn trampolines(&self) -> impl Iterator<Item = (SignatureIndex, VMTrampoline, usize)> + '_ {
  let code = self.code();
  self.trampolines.iter().map(move |info| {
- (info.signature, unsafe {
- let ptr = &code[info.start as usize];
- std::mem::transmute::<*const u8, VMTrampoline>(ptr)
- })
+ (
+ info.signature,
+ unsafe {
+ let ptr = &code[info.start as usize];
+ std::mem::transmute::<*const u8, VMTrampoline>(ptr)
+ },
+ info.length as usize,
+ )
  })
  }
 

crates/jit/src/profiling.rs

@@ -29,7 +29,10 @@ pub use vtune::VTuneAgent;
 /// Common interface for profiling tools.
 pub trait ProfilingAgent: Send + Sync + 'static {
  /// Notify the profiler of a new module loaded into memory
- fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>) -> ();
+ fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>);
+ /// Notify the profiler that the object file provided contains
+ /// dynamically-generated trampolines which are now being loaded.
+ fn trampoline_load(&self, file: &object::File<'_>);
 }
 
 /// Default agent for unsupported profiling build.
@@ -54,7 +57,8 @@ impl Error for NullProfilerAgentError {
 }
 
 impl ProfilingAgent for NullProfilerAgent {
- fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) -> () {}
+ fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
+ fn trampoline_load(&self, _file: &object::File<'_>) {}
 }
 
 #[allow(dead_code)]

crates/jit/src/profiling/jitdump_disabled.rs

@@ -20,4 +20,5 @@ impl JitDumpAgent {
 
 impl ProfilingAgent for JitDumpAgent {
  fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
+ fn trampoline_load(&self, _file: &object::File<'_>) {}
 }

crates/jit/src/profiling/jitdump_linux.rs

@@ -23,6 +23,7 @@ use std::ptr;
 use std::sync::Mutex;
 use std::{borrow, mem, process};
 use target_lexicon::Architecture;
+use wasmtime_environ::EntityRef;
 
 use object::elf;
 
@@ -207,6 +208,9 @@ impl ProfilingAgent for JitDumpAgent {
  fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>) {
  self.state.lock().unwrap().module_load(module, dbg_image);
  }
+ fn trampoline_load(&self, file: &object::File<'_>) {
+ self.state.lock().unwrap().trampoline_load(file)
+ }
 }
 
 impl State {
@@ -292,7 +296,7 @@ impl State {
  }
 
  /// Sent when a method is compiled and loaded into memory by the VM.
- pub fn module_load(&mut self, module: &CompiledModule, dbg_image: Option<&[u8]>) -> () {
+ pub fn module_load(&mut self, module: &CompiledModule, dbg_image: Option<&[u8]>) {
  let pid = process::id();
  let tid = pid; // ThreadId does appear to track underlying thread. Using PID.
 
@@ -311,6 +315,52 @@ impl State {
  self.dump_code_load_record(&name, addr, len, timestamp, pid, tid);
  }
  }
+ for (idx, func, len) in module.trampolines() {
+ let (addr, len) = (func as usize as *const u8, len);
+ let timestamp = self.get_time_stamp();
+ let name = format!("wasm::trampoline[{}]", idx.index());
+ self.dump_code_load_record(&name, addr, len, timestamp, pid, tid);
+ }
+ }
+
+ fn trampoline_load(&mut self, image: &object::File<'_>) {
+ use object::{ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
+ let pid = process::id();
+ let tid = pid;
+
+ let text_base = match image.sections().find(|s| s.kind() == SectionKind::Text) {
+ Some(section) => match section.data() {
+ Ok(data) => data.as_ptr() as usize,
+ Err(_) => return,
+ },
+ None => return,
+ };
+
+ for sym in image.symbols() {
+ if !sym.is_definition() {
+ continue;
+ }
+ if sym.kind() != SymbolKind::Text {
+ continue;
+ }
+ let address = sym.address();
+ let size = sym.size();
+ if address == 0 || size == 0 {
+ continue;
+ }
+ if let Ok(name) = sym.name() {
+ let addr = text_base + address as usize;
+ let timestamp = self.get_time_stamp();
+ self.dump_code_load_record(
+ &name,
+ addr as *const u8,
+ size as usize,
+ timestamp,
+ pid,
+ tid,
+ );
+ }
+ }
  }
 
  fn dump_code_load_record(
@@ -321,7 +371,7 @@ impl State {
  timestamp: u64,
  pid: u32,
  tid: u32,
- ) -> () {
+ ) {
  let name_len = method_name.len() + 1;
  let size_limit = mem::size_of::<CodeLoadRecord>();
 

crates/jit/src/profiling/vtune_disabled.rs

@@ -20,4 +20,5 @@ impl VTuneAgent {
 
 impl ProfilingAgent for VTuneAgent {
  fn module_load(&self, _module: &crate::CompiledModule, _dbg_image: Option<&[u8]>) {}
+ fn trampoline_load(&self, _file: &object::File<'_>) {}
 }

crates/jit/src/profiling/vtune_linux.rs

@@ -112,6 +112,9 @@ impl ProfilingAgent for VTuneAgent {
  fn module_load(&self, module: &CompiledModule, dbg_image: Option<&[u8]>) {
  self.state.lock().unwrap().module_load(module, dbg_image);
  }
+ fn trampoline_load(&self, _file: &object::File<'_>) {
+ // TODO: needs an implementation
+ }
 }
 
 impl State {

crates/wasmtime/src/module.rs

@@ -500,7 +500,9 @@ impl Module {
  let signatures = Arc::new(SignatureCollection::new_for_module(
  engine.signatures(),
  &types.wasm_signatures,
- modules.iter().flat_map(|m| m.trampolines()),
+ modules
+ .iter()
+ .flat_map(|m| m.trampolines().map(|(idx, f, _)| (idx, f))),
  ));
 
  let module = modules.remove(main_module);

crates/wasmtime/src/trampoline/func.rs

@@ -86,6 +86,7 @@ where
  // also take care of unwind table registration.
  let mut code_memory = CodeMemory::new(obj);
  let code = code_memory.publish()?;
+ engine.config().profiler.trampoline_load(&code.obj);
 
  // Extract the host/wasm trampolines from the results of compilation since
  // we know their start/length.