Support vtune profiling of trampolines too (#3687)

* Provide helpers for demangling function names

* Profile trampolines in vtune too

* get rid of mapping

* avoid code duplication with jitdump_linux

* maintain previous default display name for wasm functions

* no dash, grrr

* Remove unused profiling error type

Cargo.lock

@@ -3339,7 +3339,6 @@ dependencies = [
  "backtrace",
  "bincode",
  "cfg-if 1.0.0",
- "cpp_demangle",
  "indexmap",
  "lazy_static",
  "libc",
@@ -3349,7 +3348,6 @@ dependencies = [
  "psm",
  "rayon",
  "region",
- "rustc-demangle",
  "serde",
  "target-lexicon",
  "tempfile",
@@ -3563,6 +3561,7 @@ dependencies = [
  "cpp_demangle",
  "gimli",
  "ittapi-rs",
+ "log",
  "object",
  "region",
  "rustc-demangle",

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 crate::{demangling::demangle_function_name_or_index, CompiledModule};
-use std::error::Error;
-use std::fmt;
 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.
+    /// Notify the profiler about a single dynamically-generated trampoline (for host function)
-    fn trampoline_load(&self, file: &object::File<'_>);
+    /// 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) {
+    let mut debug_name = String::new();
-        Some(s) => rustc_demangle::try_demangle(s)
+    demangle_function_name_or_index(
-            .map(|demangle| demangle.to_string())
+        &mut debug_name,
-            .or_else(|_| cpp_demangle::Symbol::new(s).map(|sym| sym.to_string()))
+        module.func_names.get(&index).map(|s| s.as_str()),
-            .unwrap_or_else(|_| s.clone()),
+        index.index(),
-        None => format!("wasm::wasm-function[{}]", 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<'_>) {
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32) {
-        self.state.lock().unwrap().trampoline_load(file)
+        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<'_>) {
+    fn load_single_trampoline(
-        use object::{ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
+        &mut self,
-        let pid = process::id();
+        name: &str,
-        let tid = pid;
+        addr: *const u8,
-
+        size: usize,
-        let text_base = match image.sections().find(|s| s.kind() == SectionKind::Text) {
+        pid: u32,
-            Some(section) => match section.data() {
+        tid: u32,
-                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(
+        self.dump_code_load_record(name, addr, size, timestamp, pid, tid);
-                    &name,
-                    addr as *const u8,
-                    size as usize,
-                    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,
+    ) {
+    }
 }

crates/jit/src/profiling/vtune_linux.rs

@@ -9,10 +9,10 @@
 //! ### Profile
 //!
 //! ```ignore
-//! amplxe-cl -run-pass-thru=--no-altstack -v -collect hotspots target/debug/wasmtime --vtune test.wasm
+//! vtune -run-pass-thru=--no-altstack -v -collect hotspots target/debug/wasmtime --vtune test.wasm
 //! ```
 //!
-//! Note: `amplxe-cl` is a command-line tool for VTune which must [be
+//! Note: `vtune` is a command-line tool for VTune which must [be
 //! installed](https://www.intel.com/content/www/us/en/developer/tools/oneapi/vtune-profiler.html#standalone)
 //! for this to work.
 
@@ -20,10 +20,9 @@ use crate::{CompiledModule, ProfilingAgent};
 use anyhow::Result;
 use core::ptr;
 use ittapi_rs::*;
-use std::collections::HashMap;
 use std::ffi::CString;
 use std::sync::{atomic, Mutex};
-use wasmtime_environ::DefinedFuncIndex;
+use wasmtime_environ::EntityRef;
 
 /// Interface for driving the ittapi for VTune support
 pub struct VTuneAgent {
@@ -34,19 +33,13 @@ pub struct VTuneAgent {
 
 /// Interface for driving vtune
 #[derive(Clone, Debug, Default)]
-struct State {
+struct State;
-    /// Unique identifier for the jitted function
-    method_id: HashMap<(usize, DefinedFuncIndex), u32>,
-}
 
 impl VTuneAgent {
-    /// Intialize a VTuneAgent and write out the header
+    /// Initialize a VTuneAgent.
     pub fn new() -> Result<Self> {
-        let state = State {
-            method_id: HashMap::new(),
-        };
         Ok(VTuneAgent {
-            state: Mutex::new(state),
+            state: Mutex::new(State),
         })
     }
 }
@@ -58,24 +51,15 @@ impl Drop for VTuneAgent {
 }
 
 impl State {
-    /// Return the unique method ID for use with the ittapi
+    /// Return a method ID for use with the ittapi.
-    pub fn get_method_id(&mut self, module_id: usize, func_idx: DefinedFuncIndex) -> u32 {
+    fn get_method_id(&self) -> u32 {
-        let method_id: u32;
+        unsafe { iJIT_GetNewMethodID() }
-        unsafe {
-            method_id = iJIT_GetNewMethodID();
-        }
-        assert_eq!(
-            self.method_id.insert((module_id, func_idx), method_id),
-            None
-        );
-        method_id
     }
 
-    /// Load module
+    /// Notify vtune about a newly tracked code region.
-    pub fn event_load(
+    fn event_load(
         &mut self,
         method_id: u32,
-        filename: &str,
         module_name: &str,
         method_name: &str,
         addr: *const u8,
@@ -94,13 +78,16 @@ impl State {
             class_file_name: CString::new(module_name)
                 .expect("CString::new failed")
                 .into_raw(),
-            source_file_name: CString::new(filename)
+            source_file_name: CString::new("<unknown wasm filename>")
                 .expect("CString::new failed")
                 .into_raw(),
         };
         let jmethod_ptr = &mut jmethod as *mut _ as *mut _;
         unsafe {
-            println!("EventLoad: NotifyEvent Called {}", method_id);
+            log::trace!(
+                "NotifyEvent: method load (single method with id {})",
+                method_id
+            );
             let _ret = iJIT_NotifyEvent(
                 iJIT_jvm_event_iJVM_EVENT_TYPE_METHOD_LOAD_FINISHED,
                 jmethod_ptr as *mut ::std::os::raw::c_void,
@@ -111,7 +98,7 @@ impl State {
     /// Shutdown module
     fn event_shutdown(&mut self) -> () {
         unsafe {
-            println!("Drop was called!!!!!!\n");
+            log::trace!("NotifyEvent shutdown (whole module)");
             let _ret = iJIT_NotifyEvent(iJIT_jvm_event_iJVM_EVENT_TYPE_SHUTDOWN, ptr::null_mut());
         }
     }
@@ -121,8 +108,11 @@ 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<'_>) {
+    fn load_single_trampoline(&self, name: &str, addr: *const u8, size: usize, pid: u32, tid: u32) {
-        // TODO: needs an implementation
+        self.state
+            .lock()
+            .unwrap()
+            .load_single_trampoline(name, addr, size, pid, tid);
     }
 }
 
@@ -132,29 +122,52 @@ impl State {
         static MODULE_ID: atomic::AtomicUsize = atomic::AtomicUsize::new(0);
         let global_module_id = MODULE_ID.fetch_add(1, atomic::Ordering::SeqCst);
 
-        for (idx, func) in module.finished_functions() {
-            let (addr, len) = unsafe { ((*func).as_ptr() as *const u8, (*func).len()) };
-            let default_filename = "wasm_file";
-            let default_module_name = String::from("wasm_module");
         let module_name = module
             .module()
             .name
             .as_ref()
-                .unwrap_or(&default_module_name);
+            .cloned()
+            .unwrap_or_else(|| format!("wasm_module_{}", global_module_id));
+
+        for (idx, func) in module.finished_functions() {
+            let (addr, len) = unsafe { ((*func).as_ptr() as *const u8, (*func).len()) };
             let method_name = super::debug_name(module.module(), idx);
-            let method_id = self.get_method_id(global_module_id, idx);
+            let method_id = self.get_method_id();
-            println!(
+            log::trace!(
-                "Event Load: ({}) {:?}::{:?} Addr:{:?}\n",
+                "new function ({}) {:?}::{:?} @ {:?}\n",
-                method_id, module_name, method_name, addr
-            );
-            self.event_load(
                 method_id,
-                default_filename,
                 module_name,
-                &method_name,
+                method_name,
-                addr,
+                addr
-                len,
             );
+            self.event_load(method_id, &module_name, &method_name, addr, len);
         }
+
+        // Note: these are the trampolines into exported functions.
+        for (idx, func, len) in module.trampolines() {
+            let idx = idx.index();
+            let (addr, len) = (func as usize as *const u8, len);
+            let method_name = format!("wasm::trampoline[{}]", idx,);
+            let method_id = self.get_method_id();
+            log::trace!(
+                "new trampoline ({}) for exported signature {} @ {:?}\n",
+                method_id,
+                idx,
+                addr
+            );
+            self.event_load(method_id, &module_name, &method_name, addr, len);
+        }
+    }
+
+    fn load_single_trampoline(
+        &mut self,
+        name: &str,
+        addr: *const u8,
+        size: usize,
+        _pid: u32,
+        _tid: u32,
+    ) {
+        let method_id = self.get_method_id();
+        self.event_load(method_id, "wasm trampoline for Func::new", name, addr, size);
     }
 }

crates/wasmtime/Cargo.toml

@@ -26,8 +26,6 @@ region = "2.2.0"
 libc = "0.2"
 cfg-if = "1.0"
 backtrace = "0.3.61"
-rustc-demangle = "0.1.16"
-cpp_demangle = "0.3.2"
 log = "0.4.8"
 wat = { version = "1.0.36", optional = true }
 serde = { version = "1.0.94", features = ["derive"] }

crates/wasmtime/src/trampoline/func.rs

@@ -6,7 +6,7 @@ use std::any::Any;
 use std::panic::{self, AssertUnwindSafe};
 use std::sync::Arc;
 use wasmtime_environ::{EntityIndex, Module, ModuleType, PrimaryMap, SignatureIndex};
-use wasmtime_jit::{CodeMemory, MmapVec};
+use wasmtime_jit::{CodeMemory, MmapVec, ProfilingAgent};
 use wasmtime_runtime::{
     Imports, InstanceAllocationRequest, InstanceAllocator, InstanceHandle,
     OnDemandInstanceAllocator, StorePtr, VMContext, VMFunctionBody, VMSharedSignatureIndex,
@@ -66,6 +66,39 @@ unsafe extern "C" fn stub_fn<F>(
     }
 }
 
+#[cfg(compiler)]
+fn register_trampolines(profiler: &dyn ProfilingAgent, image: &object::File<'_>) {
+    use object::{Object as _, ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
+    let pid = std::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;
+            profiler.load_single_trampoline(name, addr as *const u8, size as usize, pid, tid);
+        }
+    }
+}
+
 #[cfg(compiler)]
 pub fn create_function<F>(
     ft: &FuncType,
@@ -87,7 +120,8 @@ 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);
+
+    register_trampolines(engine.config().profiler.as_ref(), &code.obj);
 
     // Extract the host/wasm trampolines from the results of compilation since
     // we know their start/length.

crates/wasmtime/src/trap.rs

@@ -4,6 +4,7 @@ use backtrace::Backtrace;
 use std::fmt;
 use std::sync::Arc;
 use wasmtime_environ::TrapCode as EnvTrapCode;
+use wasmtime_jit::{demangle_function_name, demangle_function_name_or_index};
 
 /// A struct representing an aborted instruction execution, with a message
 /// indicating the cause.
@@ -327,17 +328,8 @@ impl fmt::Display for Trap {
                 write!(f, "{:#6x} - ", offset)?;
             }
 
-            let demangle =
+            let write_raw_func_name = |f: &mut fmt::Formatter<'_>| {
-                |f: &mut fmt::Formatter<'_>, name: &str| match rustc_demangle::try_demangle(name) {
+                demangle_function_name_or_index(f, frame.func_name(), frame.func_index() as usize)
-                    Ok(name) => write!(f, "{}", name),
-                    Err(_) => match cpp_demangle::Symbol::new(name) {
-                        Ok(name) => write!(f, "{}", name),
-                        Err(_) => write!(f, "{}", name),
-                    },
-                };
-            let write_raw_func_name = |f: &mut fmt::Formatter<'_>| match frame.func_name() {
-                Some(name) => demangle(f, name),
-                None => write!(f, "<wasm function {}>", frame.func_index()),
             };
             if frame.symbols().is_empty() {
                 write!(f, "{}!", name)?;
@@ -351,7 +343,7 @@ impl fmt::Display for Trap {
                         // ...
                     }
                     match symbol.name() {
-                        Some(name) => demangle(f, name)?,
+                        Some(name) => demangle_function_name(f, name)?,
                         None if i == 0 => write_raw_func_name(f)?,
                         None => write!(f, "<inlined function>")?,
                     }