Add flamegraph option to benchmarks

jakzale · Dec 7, 2023 · 4857a46 · 4857a46
1 parent 56af135
commit 4857a46
Show file tree

Hide file tree

Showing 4 changed files with 73 additions and 28 deletions.
diff --git a/.gitignore b/.gitignore
@@ -12,5 +12,8 @@ Cargo.lock
 *.txt
 .DS_Store
 
+# flamegraphs, probably
 **/*.svg
-flamegraph.svg
+
+# tracing-flame output
+**.folded
diff --git a/README.md b/README.md
@@ -39,9 +39,19 @@ Note: requires nightly Rust
 -   `cargo run --release -p jolt-core -- --name <bench_name>`
 -   `cargo run --release -- -p jolt-core --name <bench_name> --chart`: Display performance gant chart
 -   `cargo run --release -p jolt-core --features ark-msm -- --name <bench_name>`: Run without MSM small field optimizations
--   `sudo cargo flamegraph`
 
-_Note on flamegraphing: Turn off the parallel feature in Cargo.toml (`multicore`) and / or `export RAYON_NUM_THREADS=1` to make flamegraph more interpretable._
+## Flamegraph
+
+Requires `inferno`:
+
+```
+cargo install inferno
+```
+
+Then run 
+```
+cargo run -p jolt-core --release -- --name rv32 --format flamegraph && cat tracing.folded | inferno-flamegraph > tracing-flamegraph.svg
+```
 
 ## Disclaimer
 

diff --git a/jolt-core/Cargo.toml b/jolt-core/Cargo.toml
@@ -40,13 +40,14 @@ ark-serialize = { version = "0.4.2", default-features = false, features = [
     "derive",
 ] }
 
-criterion = { version = "0.3.1", features = ["html_reports"] }
+criterion = { version = "0.5.1", features = ["html_reports"] }
 num-integer = "0.1.45"
 seq-macro = "0.3.3"
 ark-curve25519 = "0.4.0"
 tracing = "0.1.37"
 tracing-subscriber = "0.3.17"
 tracing-texray = "0.2.0"
+tracing-flame = "0.2.0"
 clap = { version = "4.3.10", features = ["derive"] }
 hashbrown = "0.14.0"
 ff = "0.13.0"
@@ -61,9 +62,6 @@ enum_dispatch = "0.3.12"
 common = { path = "../common" }
 eyre = "0.6.9"
 
-[dev-dependencies]
-criterion = "0.3.1"
-
 [lib]
 name = "liblasso"
 path = "src/lib.rs"

diff --git a/jolt-core/src/main.rs b/jolt-core/src/main.rs
@@ -1,38 +1,72 @@
+use clap::{Parser, ValueEnum};
 use liblasso::benches::bench::{benchmarks, BenchType};
-use tracing_subscriber::{self, fmt::format::FmtSpan};
-
-use clap::Parser;
+use std::{fs::File, io::BufWriter};
+use tracing_flame::FlameLayer;
+use tracing_subscriber::{self, fmt, fmt::format::FmtSpan, prelude::*, registry::Registry};
 
 /// Search for a pattern in a file and display the lines that contain it.
 #[derive(Parser, Debug)]
 struct Cli {
-    /// Whether to present in chart format
-    #[clap(long, short, action)]
-    chart: bool,
+    /// Output format
+    #[clap(short, long, value_enum, default_value_t = Format::Default)]
+    format: Format,
 
     /// Type of benchmark to run
     #[clap(long, value_enum)]
     name: BenchType,
 }
 
+#[derive(Debug, Clone, ValueEnum)]
+enum Format {
+    Default,
+    Texray,
+    Flamegraph,
+}
+
+fn setup_global_subscriber() -> impl Drop {
+    let fmt_layer = fmt::Layer::default();
+
+    let (flame_layer, _guard) = FlameLayer::with_file("./tracing.folded").unwrap();
+
+    tracing_subscriber::registry()
+        .with(fmt_layer)
+        .with(flame_layer)
+        .init();
+
+    _guard
+}
+
 fn main() {
     let args = Cli::parse();
-    if args.chart {
-        tracing_texray::init();
-        for (span, bench) in benchmarks(args.name).into_iter() {
-            tracing_texray::examine(span.to_owned()).in_scope(bench);
+    match args.format {
+        Format::Default => {
+            let collector = tracing_subscriber::fmt()
+                .with_max_level(tracing::Level::TRACE)
+                .with_span_events(FmtSpan::CLOSE)
+                .finish();
+            tracing::subscriber::set_global_default(collector)
+                .expect("setting tracing default failed");
+            for (span, bench) in benchmarks(args.name).into_iter() {
+                span.to_owned().in_scope(|| {
+                    bench();
+                    tracing::info!("Bench Complete");
+                });
+            }
+        }
+        Format::Texray => {
+            tracing_texray::init();
+            for (span, bench) in benchmarks(args.name).into_iter() {
+                tracing_texray::examine(span.to_owned()).in_scope(bench);
+            }
         }
-    } else {
-        let collector = tracing_subscriber::fmt()
-            .with_max_level(tracing::Level::TRACE)
-            .with_span_events(FmtSpan::CLOSE)
-            .finish();
-        tracing::subscriber::set_global_default(collector).expect("setting tracing default failed");
-        for (span, bench) in benchmarks(args.name).into_iter() {
-            span.to_owned().in_scope(|| {
-                bench();
-                tracing::info!("Bench Complete");
-            });
+        Format::Flamegraph => {
+            let _guard = setup_global_subscriber();
+            for (span, bench) in benchmarks(args.name).into_iter() {
+                span.to_owned().in_scope(|| {
+                    bench();
+                    tracing::info!("Bench Complete");
+                });
+            }
         }
     }
 }