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,4 +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};
@@ -31,7 +31,8 @@ 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.
+    /// dynamically-generated trampolines (for host functions) which are now
+    /// being loaded.
     fn trampoline_load(&self, file: &object::File<'_>);
 }
 
@@ -64,11 +65,12 @@ impl ProfilingAgent for NullProfilerAgent {
 #[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_linux.rs

@@ -303,6 +303,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();

crates/jit/src/profiling/vtune_linux.rs

@@ -9,21 +9,22 @@
 //! ### 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.
 
 use crate::{CompiledModule, ProfilingAgent};
 use anyhow::Result;
 use core::ptr;
 use ittapi_rs::*;
+use object::Object as _;
 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 {
@@ -32,15 +33,21 @@ pub struct VTuneAgent {
     state: Mutex<State>,
 }
 
+#[derive(Clone, Debug, Hash, PartialEq, Eq)]
+enum MethodKind {
+    Func { module_id: usize, func_id: usize },
+    ExportedFuncTrampoline { module_id: usize, func_id: usize },
+}
+
 /// Interface for driving vtune
 #[derive(Clone, Debug, Default)]
 struct State {
-    /// Unique identifier for the jitted function
-    method_id: HashMap<(usize, DefinedFuncIndex), u32>,
+    /// Unique identifier for the jitted code region.
+    method_id: HashMap<MethodKind, 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(),
@@ -58,21 +65,20 @@ impl Drop for VTuneAgent {
 }
 
 impl State {
-    /// Return the unique method ID for use with the ittapi
-    pub fn get_method_id(&mut self, module_id: usize, func_idx: DefinedFuncIndex) -> u32 {
-        let method_id: u32;
-        unsafe {
-            method_id = iJIT_GetNewMethodID();
-        }
-        assert_eq!(
-            self.method_id.insert((module_id, func_idx), method_id),
-            None
-        );
+    /// Return the unique method ID for use with the ittapi.
+    fn get_method_id(&mut self, kind: MethodKind) -> u32 {
+        let method_id = Self::get_raw_method_id();
+        assert_eq!(self.method_id.insert(kind, method_id), None);
         method_id
     }
 
-    /// Load module
-    pub fn event_load(
+    /// Return a method ID for use with the ittapi, for code that has no unique identifier.
+    fn get_raw_method_id() -> u32 {
+        unsafe { iJIT_GetNewMethodID() }
+    }
+
+    /// Notify vtune about a newly tracked code region.
+    fn event_load(
         &mut self,
         method_id: u32,
         filename: &str,
@@ -100,7 +106,10 @@ impl State {
         };
         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 +120,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 +130,8 @@ 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
+    fn trampoline_load(&self, image: &object::File<'_>) {
+        self.state.lock().unwrap().trampoline_load(image);
     }
 }
 
@@ -132,20 +141,52 @@ impl State {
         static MODULE_ID: atomic::AtomicUsize = atomic::AtomicUsize::new(0);
         let global_module_id = MODULE_ID.fetch_add(1, atomic::Ordering::SeqCst);
 
+        let default_filename = "wasm_file";
+        let module_name = module
+            .module()
+            .name
+            .as_ref()
+            .map(|s| s.as_str())
+            .unwrap_or("wasm_module");
+
         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);
             let method_name = super::debug_name(module.module(), idx);
-            let method_id = self.get_method_id(global_module_id, idx);
-            println!(
-                "Event Load: ({}) {:?}::{:?} Addr:{:?}\n",
-                method_id, module_name, method_name, addr
+            let method_id = self.get_method_id(MethodKind::Func {
+                module_id: global_module_id,
+                func_id: idx.index(),
+            });
+            log::trace!(
+                "new function ({}) {:?}::{:?} @ {:?}\n",
+                method_id,
+                module_name,
+                method_name,
+                addr
+            );
+            self.event_load(
+                method_id,
+                default_filename,
+                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(MethodKind::ExportedFuncTrampoline {
+                module_id: global_module_id,
+                func_id: idx,
+            });
+            log::trace!(
+                "new export trampoline ({}) for function {} @ {:?}\n",
+                method_id,
+                idx,
+                addr
             );
             self.event_load(
                 method_id,
@@ -157,4 +198,42 @@ impl State {
             );
         }
     }
+
+    fn trampoline_load(&mut self, image: &object::File<'_>) {
+        use object::{ObjectSection, ObjectSymbol, SectionKind, SymbolKind};
+
+        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 method_id = Self::get_raw_method_id();
+                self.event_load(
+                    method_id,
+                    "wasm_file",
+                    "wasm_module",
+                    &name,
+                    addr as *const u8,
+                    size as usize,
+                );
+            }
+        }
+    }
 }

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/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 =
-                |f: &mut fmt::Formatter<'_>, name: &str| match rustc_demangle::try_demangle(name) {
-                    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()),
+            let write_raw_func_name = |f: &mut fmt::Formatter<'_>| {
+                demangle_function_name_or_index(f, frame.func_name(), frame.func_index() as usize)
             };
             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>")?,
                     }