fix: Reduce size of FnBounds (#1592)

manh9203 · web-flow · commit dc48fafc7c78 · 2025-04-22T17:37:04.000-04:00
- Extract guest symbol data into a string table and store it in `guest.syms` to reduce the size of metrics.json
- Update `flamegraph.py` to load guest symbols from `guest.syms`
- Fix stack trace unwinding in `update_current_fn()`
- Update `openvm-prof` to use `mmap` for faster reading of large JSON files
diff --git a/.gitignore b/.gitignore
@@ -23,6 +23,7 @@ rustc-*
 .bench_metrics/
 __pycache__/
 metrics.json
+guest.syms
 
 # KZG trusted setup
 **/params
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/ci/scripts/metric_unify/flamegraph.py b/ci/scripts/metric_unify/flamegraph.py
@@ -6,7 +6,19 @@
 
 from utils import FLAMEGRAPHS_DIR, get_git_root
 
-def get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics=None):
+def get_function_symbol(string_table, offset_str):
+    try:
+        offset_int = int(offset_str)
+        end = string_table.find(b'\0', offset_int)
+        if end == -1:
+            print(f"Invalid symbol offset: {offset_int}")
+            return None
+        return string_table[offset_int:end].decode()
+    except ValueError:
+        return offset_str
+
+
+def get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics=None, string_table=None):
     """
     Filters a metrics_dict obtained from json for entries that look like:
         [ { labels: [["key1", "span1;span2"], ["key2", "span3"]], "metric": metric_name, "value": 2 } ]
@@ -40,7 +52,15 @@ def get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_met
             if key not in labels:
                 filter = True
                 break
-            stack_values.append(labels[key])
+            if key == 'cycle_tracker_span':
+                if labels[key] == '' or string_table is None:
+                    stack_values.append(labels[key])
+                else:
+                    symbol_offsets = labels[key].split(';')
+                    function_symbols = [get_function_symbol(string_table, offset) for offset in symbol_offsets]
+                    stack_values.extend(function_symbols)
+            else:
+                stack_values.append(labels[key])
         if filter:
             continue
 
@@ -57,8 +77,8 @@ def get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_met
     return lines if non_zero else []
 
 
-def create_flamegraph(fname, metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics=None, reverse=False):
-    lines = get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics)
+def create_flamegraph(fname, metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics=None, reverse=False, string_table=None):
+    lines = get_stack_lines(metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics, string_table)
     if not lines:
         return
 
@@ -86,7 +106,7 @@ def create_flamegraph(fname, metrics_dict, group_by_kvs, stack_keys, metric_name
         print(f"Created flamegraph at {flamegraph_path}")
 
 
-def create_flamegraphs(metrics_file, group_by, stack_keys, metric_name, sum_metrics=None, reverse=False):
+def create_flamegraphs(metrics_file, group_by, stack_keys, metric_name, sum_metrics=None, reverse=False, string_table=None):
     fname_prefix = os.path.splitext(os.path.basename(metrics_file))[0]
 
     with open(metrics_file, 'r') as f:
@@ -104,18 +124,18 @@ def create_flamegraphs(metrics_file, group_by, stack_keys, metric_name, sum_metr
     for group_by_values in group_by_values_list:
         group_by_kvs = list(zip(group_by, group_by_values))
         fname = fname_prefix + '-' + '-'.join(group_by_values)
-        create_flamegraph(fname, metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics, reverse=reverse)
+        create_flamegraph(fname, metrics_dict, group_by_kvs, stack_keys, metric_name, sum_metrics, reverse=reverse, string_table=string_table)
 
 
-def create_custom_flamegraphs(metrics_file, group_by=["group"]):
+def create_custom_flamegraphs(metrics_file, group_by=["group"], string_table=None):
     for reverse in [False, True]:
         create_flamegraphs(metrics_file, group_by, ["cycle_tracker_span", "dsl_ir", "opcode"], "frequency",
-                           reverse=reverse)
+                           reverse=reverse, string_table=string_table)
         create_flamegraphs(metrics_file, group_by, ["cycle_tracker_span", "dsl_ir", "opcode", "air_name"], "cells_used",
-                           reverse=reverse)
+                           reverse=reverse, string_table=string_table)
         create_flamegraphs(metrics_file, group_by, ["cell_tracker_span"], "cells_used",
                            sum_metrics=["simple_advice_cells", "fixed_cells", "lookup_advice_cells"],
-                           reverse=reverse)
+                           reverse=reverse, string_table=string_table)
 
 
 def main():
@@ -127,9 +147,16 @@ def main():
 
     argparser = argparse.ArgumentParser()
     argparser.add_argument('metrics_json', type=str, help="Path to the metrics JSON")
+    argparser.add_argument('--guest-symbols', type=str, help="Path to the guest symbols file", default=None, required=False)
     args = argparser.parse_args()
 
-    create_custom_flamegraphs(args.metrics_json)
+    if args.guest_symbols:
+        with open(args.guest_symbols, 'rb') as f:
+            string_table = f.read()
+    else:
+        string_table = None
+
+    create_custom_flamegraphs(args.metrics_json, string_table=string_table)
 
 
 if __name__ == '__main__':
diff --git a/crates/prof/Cargo.toml b/crates/prof/Cargo.toml
@@ -16,3 +16,4 @@ clap = { workspace = true, features = ["derive"] }
 eyre.workspace = true
 itertools = { workspace = true, features = ["use_std"] }
 num-format = "0.4"
+memmap2 = "0.9"
diff --git a/crates/prof/src/lib.rs b/crates/prof/src/lib.rs
@@ -4,6 +4,7 @@ use aggregate::{
     EXECUTE_TIME_LABEL, PROOF_TIME_LABEL, PROVE_EXCL_TRACE_TIME_LABEL, TRACE_GEN_TIME_LABEL,
 };
 use eyre::Result;
+use memmap2::Mmap;
 
 use crate::types::{Labels, Metric, MetricDb, MetricsFile};
 
@@ -14,7 +15,8 @@ pub mod types;
 impl MetricDb {
     pub fn new(metrics_file: impl AsRef<Path>) -> Result<Self> {
         let file = File::open(metrics_file)?;
-        let metrics: MetricsFile = serde_json::from_reader(file)?;
+        let mmap = unsafe { Mmap::map(&file)? };
+        let metrics: MetricsFile = serde_json::from_slice(&mmap)?;
 
         let mut db = MetricDb::default();
 
diff --git a/crates/toolchain/transpiler/Cargo.toml b/crates/toolchain/transpiler/Cargo.toml
@@ -15,6 +15,7 @@ eyre.workspace = true
 thiserror.workspace = true
 elf = "0.7.4"
 rrs-lib.workspace = true
+rustc-demangle = { version = "0.1.24", optional = true }
 
 [features]
-function-span = []
+function-span = ["dep:rustc-demangle"]
diff --git a/crates/toolchain/transpiler/src/elf.rs b/crates/toolchain/transpiler/src/elf.rs
@@ -1,6 +1,11 @@
 // Initial version taken from https://github.com/succinctlabs/sp1/blob/v2.0.0/crates/core/executor/src/disassembler/elf.rs under MIT License
 // and https://github.com/risc0/risc0/blob/f61379bf69b24d56e49d6af96a3b284961dcc498/risc0/binfmt/src/elf.rs#L34 under Apache License
 use std::{cmp::min, collections::BTreeMap, fmt::Debug};
+#[cfg(feature = "function-span")]
+use std::{
+    collections::{hash_map::Entry, HashMap},
+    io::Write,
+};
 
 use elf::{
     abi::{EM_RISCV, ET_EXEC, PF_X, PT_LOAD},
@@ -93,18 +98,48 @@ impl Elf {
         #[cfg(feature = "function-span")]
         {
             if let Some((symtab, stringtab)) = elf.symbol_table()? {
+                let mut fn_names = Vec::new();
+                for symbol in symtab.iter() {
+                    if symbol.st_symtype() == elf::abi::STT_FUNC {
+                        let raw_name = stringtab.get(symbol.st_name as usize).unwrap().to_string();
+                        let demangled_name = rustc_demangle::demangle(&raw_name).to_string();
+                        fn_names.push((demangled_name, symbol.st_name));
+                    }
+                }
+
+                let mut buf = Vec::new();
+                let mut offsets = HashMap::new();
+                buf.push(0);
+                for (name, st_name) in fn_names {
+                    if let Entry::Vacant(e) = offsets.entry(st_name) {
+                        let offset = buf.len();
+                        e.insert(offset);
+                        buf.extend_from_slice(name.as_bytes());
+                        buf.push(0);
+                    }
+                }
+
                 for symbol in symtab.iter() {
                     if symbol.st_symtype() == elf::abi::STT_FUNC {
                         fn_bounds.insert(
                             symbol.st_value as u32,
                             FnBound {
                                 start: symbol.st_value as u32,
                                 end: (symbol.st_value + symbol.st_size - (WORD_SIZE as u64)) as u32,
-                                name: stringtab.get(symbol.st_name as usize).unwrap().to_string(),
+                                name: offsets[&symbol.st_name].to_string(),
                             },
                         );
                     }
                 }
+
+                let guest_symbols_path = std::env::var("GUEST_SYMBOLS_PATH")?;
+                let mut guest_symbols_file =
+                    std::fs::File::create(&guest_symbols_path).map_err(|e| {
+                        eyre::eyre!(
+                            "Failed to create guest symbols file at {guest_symbols_path}: {e}"
+                        )
+                    })?;
+                guest_symbols_file.write_all(buf.as_slice())?;
             } else {
                 println!("No symbol table found");
             }
diff --git a/crates/vm/src/metrics/cycle_tracker/mod.rs b/crates/vm/src/metrics/cycle_tracker/mod.rs
@@ -9,6 +9,10 @@ impl CycleTracker {
         Self::default()
     }
 
+    pub fn top(&self) -> Option<&String> {
+        self.stack.last()
+    }
+
     /// Starts a new cycle tracker span for the given name.
     /// If a span already exists for the given name, it ends the existing span and pushes a new one
     /// to the vec.
diff --git a/crates/vm/src/metrics/mod.rs b/crates/vm/src/metrics/mod.rs
@@ -106,7 +106,10 @@ impl VmMetrics {
 
     #[cfg(feature = "function-span")]
     fn update_current_fn(&mut self, pc: u32) {
-        if !self.fn_bounds.is_empty() && (pc < self.current_fn.start || pc > self.current_fn.end) {
+        if self.fn_bounds.is_empty() {
+            return;
+        }
+        if pc < self.current_fn.start || pc > self.current_fn.end {
             self.current_fn = self
                 .fn_bounds
                 .range(..=pc)
@@ -116,10 +119,16 @@ impl VmMetrics {
             if pc == self.current_fn.start {
                 self.cycle_tracker.start(self.current_fn.name.clone());
             } else {
-                self.cycle_tracker.force_end();
+                while let Some(name) = self.cycle_tracker.top() {
+                    if name == &self.current_fn.name {
+                        break;
+                    }
+                    self.cycle_tracker.force_end();
+                }
             }
         };
     }
+
     pub fn emit(&self) {
         for (name, value) in self.chip_heights.iter() {
             let labels = [("chip_name", name.clone())];
