chore: collapsable metrics and join tables (#274)

* chore: collapsable metrics and join tables

* chore: switch to special primitives workflow

* chore: markdown newline

.github/workflows/db.yml

@@ -7,6 +7,7 @@ on:
     branches: ["main"]
     paths:
       - "stark-backend/**"
+      - "primitives/**"
       - "db/**"
 
 concurrency:

.github/workflows/unit.yml → .github/workflows/primitives.yml

@@ -1,15 +1,14 @@
-name: Default Unit Tests
+name: Primitives Tests
 
 on:
   push:
     branches: ["main"]
   pull_request:
     branches: ["main"]
-    paths-ignore:
+    paths:
       - "stark-backend/**"
-      - "db/**"
-      - "vm/**"
-      - "compiler/**"
+      - "primitives/**"
+      - "poseidon2-air/**"
 
 concurrency:
   group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.sha }}
@@ -30,6 +29,12 @@ jobs:
         with:
           cache-on-failure: true
 
-      - name: Run unit tests
+      - name: Run tests for primitives
+        working-directory: primitives
         run: |
-          cargo test --workspace --exclude afs-stark-backend --exclude afs-page --exclude stark-vm --exclude afs-compiler --exclude afs-recursion
+          cargo test --features parallel
+
+      - name: Run tests for poseidon2-air
+        working-directory: poseidon2-air
+        run: |
+          cargo test --features parallel

.github/workflows/recursion.yml

@@ -7,6 +7,7 @@ on:
     branches: ["main"]
     paths:
       - "stark-backend/**"
+      - "primitives/**"
       - "vm/**"
       - "compiler/**"
       - "recursion/**"

.github/workflows/vm.yml

@@ -7,6 +7,7 @@ on:
     branches: ["main"]
     paths:
       - "stark-backend/**"
+      - "primitives/**"
       - "vm/**"
       - "compiler/**"
 

benchmark/src/commands/benchmark.rs

@@ -24,7 +24,7 @@ use crate::{
         },
         tracing::{clear_tracing_log, extract_event_data_from_log, extract_timing_data_from_log},
     },
-    workflow::metrics::{BenchmarkMetrics, CustomMetrics},
+    workflow::metrics::BenchmarkMetrics,
     AFI_FILE_PATH, DB_FILE_PATH, MULTITIER_TABLE_ID, TABLE_ID, TMP_FOLDER, TMP_RESULT_MD,
     TMP_TRACING_LOG,
 };
@@ -148,7 +148,7 @@ pub fn benchmark_execute(
             perm_trace_gen_ms,
             calc_quotient_values_ms,
             trace: trace_metrics,
-            custom: CustomMetrics::default(),
+            custom: String::default(), // nothing for now
         };
         write!(File::create(TMP_RESULT_MD.as_str())?, "{}", metrics)?;
 

benchmark/src/commands/vm/benchmark_helpers.rs

@@ -31,7 +31,7 @@ use tracing::info_span;
 use crate::{
     config::benchmark_data::{BenchmarkSetup, BACKEND_TIMING_FILTERS, BACKEND_TIMING_HEADERS},
     utils::tracing::{clear_tracing_log, extract_timing_data_from_log, setup_benchmark_tracing},
-    workflow::metrics::{BenchmarkMetrics, CustomMetrics},
+    workflow::metrics::{BenchmarkMetrics, VmCustomMetrics},
     TMP_RESULT_MD, TMP_TRACING_LOG,
 };
 
@@ -228,25 +228,19 @@ pub fn vm_benchmark_execute_and_prove<const WORD_SIZE: usize>(
         .map(|(k, v)| (k, v.to_string()))
         .collect();
 
-    let sorted_opcode_counts: Vec<(String, String)> = opcode_counts
-        .clone()
+    let opcode_counts: Vec<(String, usize)> = opcode_counts
         .into_iter()
-        .sorted_by(|a, b| a.1.cmp(&b.1))
-        .map(|(k, v)| (k, v.to_string()))
+        .sorted_by(|a, b| b.1.cmp(&a.1))
         .collect();
 
-    let sorted_dsl_counts: Vec<(String, String)> = dsl_counts
-        .clone()
+    let dsl_counts: Vec<(String, usize)> = dsl_counts
         .into_iter()
-        .sorted_by(|a, b| a.1.cmp(&b.1))
-        .map(|(k, v)| (k, v.to_string()))
+        .sorted_by(|a, b| b.1.cmp(&a.1))
         .collect();
 
-    let sorted_opcode_trace_cells: Vec<(String, String)> = opcode_trace_cells
-        .clone()
+    let opcode_trace_cells: Vec<(String, usize)> = opcode_trace_cells
         .into_iter()
-        .sorted_by(|a, b| a.1.cmp(&b.1))
-        .map(|(k, v)| (k, v.to_string()))
+        .sorted_by(|a, b| b.1.cmp(&a.1))
         .collect();
 
     let metrics = BenchmarkMetrics {
@@ -256,11 +250,11 @@ pub fn vm_benchmark_execute_and_prove<const WORD_SIZE: usize>(
         perm_trace_gen_ms,
         calc_quotient_values_ms,
         trace: trace_metrics,
-        custom: CustomMetrics {
+        custom: VmCustomMetrics {
             vm_metrics,
-            opcode_counts: sorted_opcode_counts,
-            dsl_counts: sorted_dsl_counts,
-            opcode_trace_cells: sorted_opcode_trace_cells,
+            opcode_counts,
+            dsl_counts,
+            opcode_trace_cells,
         },
     };
 

benchmark/src/workflow/metrics.rs

@@ -1,11 +1,14 @@
-use std::{collections::BTreeMap, fmt};
+use std::{
+    collections::BTreeMap,
+    fmt::{self, Display},
+};
 
 use afs_stark_backend::prover::metrics::{format_number_with_underscores, TraceMetrics};
 use serde::{Deserialize, Serialize};
 
 /// Reusable struct for storing benchmark metrics
-#[derive(Clone, Debug, Serialize, Deserialize)]
-pub struct BenchmarkMetrics {
+#[derive(Clone, Serialize, Deserialize)]
+pub struct BenchmarkMetrics<CustomMetrics> {
     /// Benchmark name
     pub name: String,
     // Timings:
@@ -21,16 +24,8 @@ pub struct BenchmarkMetrics {
     pub custom: CustomMetrics,
 }
 
-#[derive(Clone, Debug, Default, Serialize, Deserialize)]
-pub struct CustomMetrics {
-    pub vm_metrics: BTreeMap<String, String>,
-    pub opcode_counts: Vec<(String, String)>,
-    pub dsl_counts: Vec<(String, String)>,
-    pub opcode_trace_cells: Vec<(String, String)>,
-}
-
 // Implement the Display trait for BenchmarkMetrics to create a markdown table
-impl fmt::Display for BenchmarkMetrics {
+impl<CustomMetrics: Display> Display for BenchmarkMetrics<CustomMetrics> {
     fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
         writeln!(f, "## Benchmark for {}", self.name)?;
         // Write the markdown table header
@@ -81,48 +76,65 @@ impl fmt::Display for BenchmarkMetrics {
             )?;
         }
 
-        if !self.custom.vm_metrics.is_empty() {
-            writeln!(f)?;
-            writeln!(f, "### Custom metrics")?;
-            writeln!(f, "| Name | Value |")?;
-            writeln!(f, "|------|-------|")?;
-            for (name, value) in self.custom.vm_metrics.iter() {
-                writeln!(f, "| {:<20} | {:<10} |", name, value)?;
-            }
-        }
+        self.custom.fmt(f)?;
 
-        if !self.custom.opcode_counts.is_empty() {
-            writeln!(f)?;
-            writeln!(f, "### Opcode counts")?;
-            writeln!(f, "| Name | Count |")?;
-            writeln!(f, "|------|-------|")?;
-            for (name, value) in self.custom.opcode_counts.iter() {
-                writeln!(f, "| {:<20} | {:<10} |", name, value)?;
-            }
+        Ok(())
+    }
+}
+
+#[derive(Clone, Debug, Default, Serialize, Deserialize)]
+pub struct VmCustomMetrics {
+    pub vm_metrics: BTreeMap<String, String>,
+    pub opcode_counts: Vec<(String, usize)>,
+    pub dsl_counts: Vec<(String, usize)>,
+    pub opcode_trace_cells: Vec<(String, usize)>,
+}
+
+impl Display for VmCustomMetrics {
+    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
+        writeln!(f, "<details>")?;
+        writeln!(f, "<summary>")?;
+        writeln!(f)?;
+        writeln!(f, "### Custom VM metrics")?;
+        writeln!(f)?;
+        writeln!(f, "</summary>")?;
+        writeln!(f)?;
+
+        writeln!(f, "| Name | Value |")?;
+        writeln!(f, "|------|-------|")?;
+        for (name, value) in self.vm_metrics.iter() {
+            writeln!(f, "| {:<20} | {:<10} |", name, value)?;
         }
 
-        if !self.custom.dsl_counts.is_empty() {
-            writeln!(f)?;
-            writeln!(
-                f,
-                "### DSL counts - how many isa instructions each DSL instruction generates"
-            )?;
-            writeln!(f, "| Name | Count |")?;
-            writeln!(f, "|------|-------|")?;
-            for (name, value) in self.custom.dsl_counts.iter() {
-                writeln!(f, "| {:<20} | {:<10} |", name, value)?;
+        writeln!(f)?;
+        writeln!(f, "#### Opcode metrics")?;
+        writeln!(f, "| Name | Frequency | Trace Cells Contributed |")?;
+        writeln!(f, "|------|-------|-----|")?;
+        for (name, value) in self.opcode_counts.iter() {
+            let cell_count = *self
+                .opcode_trace_cells
+                .iter()
+                .find_map(|(k, v)| if k == name { Some(v) } else { None })
+                .unwrap_or(&0);
+            writeln!(f, "| {:<20} | {:<10} | {:<10} |", name, value, cell_count)?;
+        }
+        for (name, value) in self.opcode_trace_cells.iter() {
+            if !self.opcode_counts.iter().any(|(k, _)| k == name) {
+                // this should never happen
+                writeln!(f, "| {:<20} | 0 | {:<10} |", name, value)?;
             }
         }
 
-        if !self.custom.opcode_trace_cells.is_empty() {
-            writeln!(f)?;
-            writeln!(f, "### Opcode trace cells")?;
-            writeln!(f, "| Name | Count |")?;
-            writeln!(f, "|------|-------|")?;
-            for (name, value) in self.custom.opcode_trace_cells.iter() {
-                writeln!(f, "| {:<20} | {:<10} |", name, value)?;
-            }
+        writeln!(f)?;
+        writeln!(f, "### DSL counts")?;
+        writeln!(f, "How many opcodes each DSL instruction generates:")?;
+        writeln!(f, "| Name | Count |")?;
+        writeln!(f, "|------|-------|")?;
+        for (name, value) in self.dsl_counts.iter() {
+            writeln!(f, "| {:<20} | {:<10} |", name, value)?;
         }
+
+        writeln!(f, "</details>")?;
         Ok(())
     }
 }

compiler/tests/array.rs

@@ -1,8 +1,8 @@
 use afs_compiler::{
-    asm::{AsmBuilder, AsmCompiler},
+    asm::AsmBuilder,
     ir::{Array, Config, Ext, ExtConst, Felt, Usize, Var},
     prelude::{Builder, MemIndex, MemVariable, Ptr, Variable},
-    util::{display_program, execute_program},
+    util::execute_program,
 };
 use afs_derive::DslVariable;
 use p3_baby_bear::BabyBear;