Support vtune profiling of trampolines too

Cargo.toml

@@ -86,7 +86,7 @@ members = [
 exclude = [
   'crates/wasi-common/WASI/tools/witx-cli',
   'docs/rust_wasi_markdown_parser'
-  ]
+]
 
 [features]
 default = ["jitdump", "wasmtime/wat", "wasmtime/parallel-compilation", "wasi-nn"]

crates/jit/Cargo.toml

@@ -26,6 +26,7 @@ ittapi-rs = { version = "0.1.6", optional = true  }
 bincode = "1.2.1"
 rustc-demangle = "0.1.16"
 cpp_demangle = "0.3.2"
+log = "0.4.8"
 
 [target.'cfg(target_os = "windows")'.dependencies]
 winapi = { version = "0.3.8", features = ["winnt", "impl-default"] }

crates/jit/src/demangling.rs

@@ -0,0 +1,27 @@
+//! Helpers for demangling function names.
+
+/// Demangles a single function name into a user-readable form.
+///
+/// Currently supported: Rust/C/C++ function names.
+pub fn demangle_function_name(writer: &mut impl std::fmt::Write, name: &str) -> std::fmt::Result {
+    if let Ok(demangled) = rustc_demangle::try_demangle(name) {
+        write!(writer, "{}", demangled)
+    } else if let Ok(demangled) = cpp_demangle::Symbol::new(name) {
+        write!(writer, "{}", demangled)
+    } else {
+        write!(writer, "{}", name)
+    }
+}
+
+/// Demangles a function name if it's provided, or returns a unified representation based on the
+/// function index otherwise.
+pub fn demangle_function_name_or_index(
+    writer: &mut impl std::fmt::Write,
+    name: Option<&str>,
+    func_id: usize,
+) -> std::fmt::Result {
+    match name {
+        Some(name) => demangle_function_name(writer, name),
+        None => write!(writer, "<wasm function {}>", func_id),
+    }
+}

crates/jit/src/lib.rs

@@ -22,6 +22,7 @@
 
 mod code_memory;
 mod debug;
+mod demangling;
 mod instantiate;
 mod link;
 mod mmap_vec;
@@ -34,6 +35,7 @@ pub use crate::instantiate::{
     SymbolizeContext, TypeTables,
 };
 pub use crate::mmap_vec::MmapVec;
+pub use demangling::*;
 pub use profiling::*;
 
 /// Version number of this crate.

crates/jit/src/profiling.rs

@@ -1,6 +1,4 @@
-use crate::CompiledModule;
-use std::error::Error;
-use std::fmt;
+use crate::{demangling::demangle_function_name_or_index, CompiledModule};
 use wasmtime_environ::{DefinedFuncIndex, EntityRef, Module};
 
 cfg_if::cfg_if! {
@@ -30,45 +28,38 @@ pub use vtune::VTuneAgent;
 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]>);
-    /// Notify the profiler that the object file provided contains
-    /// dynamically-generated trampolines which are now being loaded.
-    fn trampoline_load(&self, file: &object::File<'_>);
+
+    /// Notify the profiler about a single dynamically-generated trampoline (for host function)
+    /// that is being loaded now.`
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32);
 }
 
 /// Default agent for unsupported profiling build.
 #[derive(Debug, Default, Clone, Copy)]
 pub struct NullProfilerAgent;
 
-#[derive(Debug)]
-struct NullProfilerAgentError;
-
-impl fmt::Display for NullProfilerAgentError {
-    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
-        write!(f, "A profiler agent is not supported by this build")
-    }
-}
-
-// This is important for other errors to wrap this one.
-impl Error for NullProfilerAgentError {
-    fn source(&self) -> Option<&(dyn Error + 'static)> {
-        // Generic error, underlying cause isn't tracked.
-        None
-    }
-}
-
 impl ProfilingAgent for NullProfilerAgent {
     fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn trampoline_load(&self, _file: &object::File<'_>) {}
+    fn load_single_trampoline(
+        &self,
+        _name: &str,
+        _addr: *const u8,
+        _size: usize,
+        _pid: u32,
+        _tid: u32,
+    ) {
+    }
 }
 
 #[allow(dead_code)]
 fn debug_name(module: &Module, index: DefinedFuncIndex) -> String {
     let index = module.func_index(index);
-    match module.func_names.get(&index) {
-        Some(s) => rustc_demangle::try_demangle(s)
-            .map(|demangle| demangle.to_string())
-            .or_else(|_| cpp_demangle::Symbol::new(s).map(|sym| sym.to_string()))
-            .unwrap_or_else(|_| s.clone()),
-        None => format!("wasm::wasm-function[{}]", index.index()),
-    }
+    let mut debug_name = String::new();
+    demangle_function_name_or_index(
+        &mut debug_name,
+        module.func_names.get(&index).map(|s| s.as_str()),
+        index.index(),
+    )
+    .unwrap();
+    debug_name
 }

crates/jit/src/profiling/jitdump_disabled.rs

@@ -20,5 +20,13 @@ impl JitDumpAgent {
 
 impl ProfilingAgent for JitDumpAgent {
     fn module_load(&self, _module: &CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn trampoline_load(&self, _file: &object::File<'_>) {}
+    fn load_single_trampoline(
+        &self,
+        _name: &str,
+        _addr: *const u8,
+        _size: usize,
+        __pid: u32,
+        _tid: u32,
+    ) {
+    }
 }

crates/jit/src/profiling/jitdump_linux.rs

@@ -203,8 +203,11 @@ 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)
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32) {
+        self.state
+            .lock()
+            .unwrap()
+            .load_single_trampoline(name, addr, size, pid, tid);
     }
 }
 
@@ -303,6 +306,8 @@ impl State {
                 self.dump_code_load_record(&name, addr, len, timestamp, pid, tid);
             }
         }
+
+        // Note: these are the trampolines into exported functions.
         for (idx, func, len) in module.trampolines() {
             let (addr, len) = (func as usize as *const u8, len);
             let timestamp = self.get_time_stamp();
@@ -311,44 +316,16 @@ impl State {
         }
     }
 
-    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 load_single_trampoline(
+        &mut self,
+        name: &str,
+        addr: *const u8,
+        size: usize,
+        pid: u32,
+        tid: u32,
+    ) {
+        let timestamp = self.get_time_stamp();
+        self.dump_code_load_record(name, addr, size, timestamp, pid, tid);
     }
 
     fn dump_code_load_record(

crates/jit/src/profiling/vtune_disabled.rs

@@ -20,5 +20,13 @@ impl VTuneAgent {
 
 impl ProfilingAgent for VTuneAgent {
     fn module_load(&self, _module: &crate::CompiledModule, _dbg_image: Option<&[u8]>) {}
-    fn trampoline_load(&self, _file: &object::File<'_>) {}
+    fn load_single_trampoline(
+        &self,
+        _name: &str,
+        _addr: *const u8,
+        _size: usize,
+        __pid: u32,
+        _tid: u32,
+    ) {
+    }
 }