feat: Add flamegraph support to transpiler (#216)

This PR adds flamegraph generation support to the transpiler. I've
manually verified it by integrating it to the airbender-platform
(locally for now).

It does so without depending on `risc_v_simulator` for that, since one
of the goals is to remove `risc_v_simulator` completely as obsolete.

Once this is merged, I'll start working on removing `risc_v_simulator`
from the repository.

Author: popzxc

Cargo.toml

@@ -175,13 +175,16 @@ gpu_prover = { path = "./gpu_prover", default-features = false }
 
 serde = { version = "1", default-features = false, features = ["derive", "alloc"] }
 clap = { version = "4.5.21", features = ["derive"] }
+addr2line_new = { package = "addr2line", version = "0.25" }
+object = "0.37"
 # rand = {version = "0.8", default-features = false, features = ["std_rng"] }
 rand = { version = "0.9", default-features = false }
 unroll = "0.1"
 seq-macro = "0.3"
 super-seq-macro = "0.3"
 arrayvec = { version = "0.7", default-features = false }
 itertools = { version = "0.14" }
+inferno = "0.12"
 log = "0.4"
 sha3 = { version = "*", default-features = false }
 

riscv_transpiler/Cargo.toml

@@ -18,6 +18,9 @@ risc_v_simulator = { workspace = true, features = ["delegation"] }
 worker = { workspace = true }
 serde = { workspace = true }
 seq-macro = { workspace = true }
+inferno = { workspace = true, optional = true }
+addr2line_new = { workspace = true, optional = true }
+object = { workspace = true, optional = true }
 serde-big-array = "*"
 keccak = "*"
 
@@ -27,6 +30,7 @@ capstone = {version = "0.13", optional = true }
 
 [features]
 jit = ["dynasmrt", "riscv-decode", "capstone"]
+flamegraph = ["dep:inferno", "dep:addr2line_new", "dep:object"]
 default = ["jit"]
 
 [dev-dependencies]

riscv_transpiler/src/vm/flamegraph/collapse.rs

@@ -0,0 +1,109 @@
+use std::collections::HashMap;
+
+use super::symbolizer::Addr2LineContext;
+
+/// Converts raw sampled stacks into `inferno`'s collapsed stack format.
+pub(super) fn build_collapsed_stack_lines(
+    frames: &[(u32, Vec<u32>)],
+    symbolizer: &Addr2LineContext<'_>,
+) -> Vec<String> {
+    // Hot PCs appear in many samples, so we resolve symbols once per address.
+    let mut symbol_cache: HashMap<u32, Vec<String>> = HashMap::new();
+    let mut collapsed_line_counts: HashMap<String, usize> = HashMap::new();
+    // Temporary merged stack for one sample before collapsing into a line.
+    let mut buffer = Vec::with_capacity(64);
+
+    for (pc, callsites) in frames.iter() {
+        buffer.clear();
+
+        if let Some(names) = try_frames_for_pc(symbolizer, &mut symbol_cache, *pc) {
+            append_frames_with_overlap(&mut buffer, names);
+        }
+
+        // Stack unwinding and DWARF frame expansion can produce overlapping
+        // frame sequences. We merge overlaps to avoid duplicate path segments.
+        for callsite_pc in callsites.iter().copied().skip(1) {
+            let Some(names) = try_frames_for_pc(symbolizer, &mut symbol_cache, callsite_pc) else {
+                continue;
+            };
+            append_frames_with_overlap(&mut buffer, names);
+        }
+
+        if buffer.is_empty() {
+            continue;
+        }
+
+        // `inferno` collapsed format expects root-first paths separated by `;`.
+        // Our merged buffer is leaf-first, so we reverse it at serialization.
+        let mut line = String::with_capacity(buffer.len() * 16 + 12);
+        for (idx, el) in buffer.iter().rev().enumerate() {
+            if idx > 0 {
+                line.push(';');
+            }
+            line.push_str(el);
+        }
+
+        *collapsed_line_counts.entry(line).or_default() += 1;
+    }
+
+    let mut remapped = Vec::with_capacity(collapsed_line_counts.len());
+    for (line, count) in collapsed_line_counts.into_iter() {
+        let mut line_with_count = line;
+        line_with_count.push(' ');
+        line_with_count.push_str(&count.to_string());
+        remapped.push(line_with_count);
+    }
+
+    remapped
+}
+
+#[inline(always)]
+fn try_frames_for_pc<'a>(
+    symbolizer: &Addr2LineContext<'_>,
+    symbol_cache: &'a mut HashMap<u32, Vec<String>>,
+    pc: u32,
+) -> Option<&'a [String]> {
+    // Unaligned PCs are not valid instruction addresses in this VM and are
+    // usually artifacts of incomplete stack data.
+    if pc % 4 != 0 {
+        return None;
+    }
+
+    if symbol_cache.contains_key(&pc) == false {
+        let frames = symbolizer.collect_frames(pc);
+        symbol_cache.insert(pc, frames);
+    }
+
+    let names = symbol_cache
+        .get(&pc)
+        .expect("symbol cache must contain a value");
+    if names.is_empty() {
+        None
+    } else {
+        Some(names)
+    }
+}
+
+#[inline(always)]
+fn append_frames_with_overlap(buffer: &mut Vec<String>, names: &[String]) {
+    if names.is_empty() {
+        return;
+    }
+
+    if buffer.is_empty() {
+        buffer.extend(names.iter().cloned());
+        return;
+    }
+
+    let max_overlap = buffer.len().min(names.len());
+    // Keep only the non-overlapping suffix from `names` so each logical stack
+    // segment appears once in the merged path.
+    let overlap = (1..=max_overlap)
+        .rev()
+        .find(|overlap_len| buffer[(buffer.len() - *overlap_len)..] == names[..*overlap_len])
+        .unwrap_or(0);
+
+    if overlap < names.len() {
+        buffer.extend(names[overlap..].iter().cloned());
+    }
+}

riscv_transpiler/src/vm/flamegraph/config.rs

@@ -0,0 +1,42 @@
+use std::path::PathBuf;
+
+/// Runtime knobs for VM flamegraph generation.
+///
+/// We intentionally separate collection (cheap, during execution) from
+/// symbolization/rendering (heavier, after execution). These values tune both
+/// phases.
+#[derive(Clone, Debug)]
+pub struct FlamegraphConfig {
+    /// ELF/object file that provides symbols and DWARF info for PC resolution.
+    pub symbols_path: PathBuf,
+    /// Destination SVG written by `inferno`.
+    pub output_path: PathBuf,
+    /// Controls whether stacks are rendered in reverse order.
+    pub reverse_graph: bool,
+    /// Collect one sample every `frequency_recip` VM cycles.
+    ///
+    /// Larger values reduce runtime overhead but may hide short-lived frames.
+    pub frequency_recip: usize,
+}
+
+impl FlamegraphConfig {
+    pub fn new(symbols_path: PathBuf, output_path: PathBuf) -> Self {
+        // Defaults bias toward low overhead while keeping a readable graph.
+        Self {
+            symbols_path,
+            output_path,
+            reverse_graph: false,
+            frequency_recip: 100,
+        }
+    }
+}
+
+/// Sampling counters that help estimate profiler effectiveness.
+///
+/// `samples_total` tracks how many sampling points were attempted, while
+/// `samples_collected` tracks how many produced a non-empty stack trace.
+#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)]
+pub struct FlamegraphSampleStats {
+    pub samples_total: usize,
+    pub samples_collected: usize,
+}

riscv_transpiler/src/vm/flamegraph/mod.rs

@@ -0,0 +1,10 @@
+mod collapse;
+mod config;
+mod profiler;
+mod ram;
+mod stacktrace;
+mod symbolizer;
+
+pub use self::config::{FlamegraphConfig, FlamegraphSampleStats};
+pub use self::profiler::VmFlamegraphProfiler;
+pub use self::ram::FlamegraphReadableRam;

riscv_transpiler/src/vm/flamegraph/profiler.rs

@@ -0,0 +1,101 @@
+use super::collapse::build_collapsed_stack_lines;
+use super::config::{FlamegraphConfig, FlamegraphSampleStats};
+use super::ram::FlamegraphReadableRam;
+use super::stacktrace::collect_stacktrace_raw;
+use super::symbolizer::Addr2LineContext;
+use crate::vm::{Counters, State};
+
+/// Coordinates the flamegraph pipeline:
+/// 1) collect lightweight raw samples during execution,
+/// 2) symbolize and render once execution finishes.
+pub struct VmFlamegraphProfiler {
+    config: FlamegraphConfig,
+    symbol_binary: Vec<u8>,
+    raw_frames: Vec<(u32, Vec<u32>)>,
+    stats: FlamegraphSampleStats,
+}
+
+impl VmFlamegraphProfiler {
+    pub fn new(config: FlamegraphConfig) -> std::io::Result<Self> {
+        // Zero would both disable progress and cause division-by-zero.
+        if config.frequency_recip == 0 {
+            return Err(std::io::Error::new(
+                std::io::ErrorKind::InvalidInput,
+                "frequency_recip must be greater than zero",
+            ));
+        }
+
+        let symbol_binary = std::fs::read(&config.symbols_path)?;
+
+        Ok(Self {
+            config,
+            symbol_binary,
+            raw_frames: Vec::new(),
+            stats: FlamegraphSampleStats::default(),
+        })
+    }
+
+    pub fn stats(&self) -> FlamegraphSampleStats {
+        self.stats
+    }
+
+    #[inline(always)]
+    pub fn sample_cycle<C: Counters, R: FlamegraphReadableRam>(
+        &mut self,
+        state: &State<C>,
+        ram: &R,
+        cycle: usize,
+    ) {
+        // Sampling is on the VM hot path, so we keep this branch and data
+        // collection minimal and defer expensive work to finalization.
+        if cycle % self.config.frequency_recip != 0 {
+            return;
+        }
+
+        self.stats.samples_total += 1;
+
+        let (pc, frames) = collect_stacktrace_raw(state, ram);
+        if frames.is_empty() == false {
+            // Empty stacks are expected when we cannot reconstruct a valid frame
+            // chain; they are tracked via stats but not emitted.
+            self.stats.samples_collected += 1;
+            self.raw_frames.push((pc, frames));
+        }
+    }
+
+    pub fn write_flamegraph(&mut self) -> std::io::Result<()> {
+        // Symbolization is deferred to here to keep execution-time sampling
+        // overhead predictable and low.
+        let symbolizer = Addr2LineContext::new(&self.symbol_binary)?;
+
+        let collapsed_lines = build_collapsed_stack_lines(&self.raw_frames, &symbolizer);
+
+        let collapsed_lines = if collapsed_lines.is_empty() {
+            // Produce a minimal graph instead of failing when no usable samples
+            // were collected.
+            vec![String::from("no_samples 1")]
+        } else {
+            collapsed_lines
+        };
+
+        let output_file = std::fs::File::create(&self.config.output_path)?;
+        let mut options = inferno::flamegraph::Options::default();
+        options.reverse_stack_order = self.config.reverse_graph;
+        inferno::flamegraph::from_lines(
+            &mut options,
+            collapsed_lines.iter().map(String::as_str),
+            output_file,
+        )
+        .map_err(|error| {
+            std::io::Error::new(
+                std::io::ErrorKind::Other,
+                format!("while attempting to generate flamegraph: {error}"),
+            )
+        })?;
+
+        // The profiler can be reused across VM runs with the same config.
+        self.raw_frames.clear();
+
+        Ok(())
+    }
+}

riscv_transpiler/src/vm/flamegraph/ram.rs

@@ -0,0 +1,47 @@
+use crate::vm::{RamPeek, RamWithRomRegion};
+
+/// Minimal RAM contract needed by the flamegraph unwinder.
+///
+/// `RamPeek` is intentionally low-level and uses debug assertions for bounds,
+/// so flamegraph collection adds an explicit checked read API to keep profiling
+/// robust even when frame-pointer metadata is malformed.
+pub trait FlamegraphReadableRam: RamPeek {
+    fn total_words_for_flamegraph(&self) -> usize;
+
+    #[inline(always)]
+    fn try_peek_word_for_flamegraph(&self, address: u32) -> Option<u32> {
+        if address % 4 != 0 {
+            return None;
+        }
+
+        let word_idx = (address / 4) as usize;
+        if word_idx >= self.total_words_for_flamegraph() {
+            return None;
+        }
+
+        Some(self.peek_word(address))
+    }
+}
+
+impl<const N: usize> FlamegraphReadableRam for [u32; N] {
+    #[inline(always)]
+    fn total_words_for_flamegraph(&self) -> usize {
+        N
+    }
+}
+
+impl FlamegraphReadableRam for [u32] {
+    #[inline(always)]
+    fn total_words_for_flamegraph(&self) -> usize {
+        self.len()
+    }
+}
+
+impl<const ROM_BOUND_SECOND_WORD_BITS: usize> FlamegraphReadableRam
+    for RamWithRomRegion<ROM_BOUND_SECOND_WORD_BITS>
+{
+    #[inline(always)]
+    fn total_words_for_flamegraph(&self) -> usize {
+        self.backing.len()
+    }
+}