profiler-html.rs

//! Example program of parsing and displaying pulley profiles.
//!
//! To use this program first build Wasmtime with support for profiling Pulley:
//!
//! ```text
//! $ cargo build --release --features profile-pulley
//! ```
//!
//! Next record a profile
//!
//! ```text
//! $ ./target/release/wasmtime run --profile pulley --target pulley64 \
//!   your_wasm_file.wasm
//! ```
//!
//! This will emit `pulley-$pid.data` to the current working directory. That
//! file is then fed to this program:
//!
//! ```text
//! $ cargo run -p pulley-interpreter --example profiler-html --all-features \
//!     ./pulley-$pid.data
//! ```
//!
//! This will print all functions and their disassemblies to stdout. Functions
//! are annotated with the % of samples that fell in that function. Instructions
//! in functions are annotated with the % of samples in that function that fell
//! on that instruction. Functions are dropped if their sample rate is below the
//! CLI threshold and instructions are un-annotated if they're below the
//! threshold.

use anyhow::{Context, Result, bail};
use clap::Parser;
use pulley_interpreter::decode::{Decoder, OpVisitor};
use pulley_interpreter::disas::Disassembler;
use pulley_interpreter::profile::{Event, decode};
use std::collections::BTreeMap;
use std::io::Write;
use std::path::PathBuf;
use termcolor::{Color, ColorChoice, ColorSpec, StandardStream, WriteColor};

#[derive(Parser)]
struct ProfilerHtml {
    /// The profile data to load which was generated by a `--profile pulley` run
    /// of Wasmtime previously.
    profile: PathBuf,

    /// The minimum threshold to display a function or annotate an instruction.
    #[clap(long, default_value = "0.5")]
    threshold: f32,

    /// Whether or not to show instruction disassemblies.
    #[clap(long)]
    instructions: Option<bool>,
}

struct Function<'a> {
    addr: u64,
    hits: u64,
    name: &'a str,
    body: &'a [u8],
    instructions: BTreeMap<u32, u32>,
}

fn main() -> Result<()> {
    let args = ProfilerHtml::parse();
    let profile = std::fs::read(&args.profile)
        .with_context(|| format!("failed to read {:?}", args.profile))?;

    // All known functions and the total of all samples taken.
    let mut functions = BTreeMap::new();
    let mut total = 0;

    let mut found_samples = false;
    for event in decode(&profile) {
        match event? {
            Event::Function(addr, name, body) => {
                let prev = functions.insert(
                    addr,
                    Function {
                        addr,
                        name,
                        body,
                        hits: 0,
                        instructions: BTreeMap::new(),
                    },
                );
                assert!(prev.is_none());
            }
            Event::Samples(samples) => {
                found_samples = true;
                for sample in samples {
                    let addr = sample.0;
                    let (_, function) = functions.range_mut(..=addr).next_back().unwrap();
                    assert!(addr >= function.addr);
                    assert!(addr < function.addr + (function.body.len() as u64));

                    total += 1;
                    function.hits += 1;
                    *function
                        .instructions
                        .entry(u32::try_from(addr - function.addr).unwrap())
                        .or_insert(0) += 1;
                }
            }
        }
    }

    if functions.is_empty() {
        bail!("no functions found in profile");
    }
    if !found_samples {
        bail!("no samples found in profile");
    }

    let mut funcs = functions
        .into_iter()
        .map(|(_, func)| func)
        .collect::<Vec<_>>();
    funcs.sort_by_key(|f| f.hits);

    let mut term = StandardStream::stdout(ColorChoice::Auto);
    let mut reset = ColorSpec::new();
    reset.set_reset(true);

    for mut func in funcs {
        let func_pct = (func.hits as f32) / (total as f32) * 100.0;
        if func_pct < args.threshold {
            continue;
        }
        writeln!(
            term,
            "{:6.02}% {}",
            (func.hits as f32) / (total as f32) * 100.0,
            func.name,
        )?;

        if !args.instructions.unwrap_or(true) {
            continue;
        }

        let mut disas = Disassembler::new(func.body);
        disas.hexdump(false);
        disas.offsets(false);
        disas.br_tables(false);
        let mut decoder = Decoder::new();
        let mut prev = 0;
        let mut offset = 0;
        let mut remaining = func.body.len();

        let min_instruction = func
            .instructions
            .iter()
            .map(|(_, hits)| *hits)
            .min()
            .unwrap_or(0);
        let max_instruction = func
            .instructions
            .iter()
            .map(|(_, hits)| *hits)
            .max()
            .unwrap_or(0);

        while !disas.bytecode().as_slice().is_empty() {
            decoder.decode_one(&mut disas)?;
            let instr = &disas.disas()[prev..].trim();
            let hits = func.instructions.remove(&offset).unwrap_or(0);
            let pct = (hits as f32) / (func.hits as f32) * 100.;
            if pct < args.threshold {
                term.set_color(&reset)?;
                writeln!(term, "\t        {:6x}: {instr}", u64::from(offset))?;
            } else {
                // Attempt to do a bit of a gradient from red-to-green from
                // least-hit to most-hit instruction. Note that un-annotated
                // instructions will have no color still (e.g. they aren't
                // green).
                let mut color = ColorSpec::new();
                color.set_bold(hits == max_instruction);
                let pct_r =
                    (hits - min_instruction) as f32 / (max_instruction - min_instruction) as f32;

                let r = ((0xff as f32) * pct_r) as u8;
                let g = ((0xff as f32) * (1. - pct_r)) as u8;
                let b = 0;
                color.set_fg(Some(Color::Rgb(r, g, b)));
                term.set_color(&color)?;
                writeln!(term, "\t{pct:6.02}% {:6x}: {instr}", u64::from(offset))?;
            }
            offset += u32::try_from(remaining - disas.bytecode().as_slice().len()).unwrap();
            remaining = disas.bytecode().as_slice().len();
            prev = disas.disas().len();
        }

        term.set_color(&reset)?;

        assert!(func.instructions.is_empty(), "{:?}", func.instructions);
    }

    Ok(())
}