feat(profiler): Reduce memory in Brillig execution flamegraph (#6538)

vezenovm · TomAFrench · web-flow · commit 1cad7c887893 · 2024-11-18T15:46:05.000Z
Co-authored-by: Tom French &lt;15848336+TomAFrench@users.noreply.github.com&gt;
diff --git a/tooling/profiler/Cargo.toml b/tooling/profiler/Cargo.toml
@@ -44,6 +44,3 @@ noirc_abi.workspace = true
 noirc_driver.workspace = true
 tempfile.workspace = true
 
-[features]
-default = ["bn254"]
-bn254 = ["acir/bn254"]
diff --git a/tooling/profiler/src/cli/execution_flamegraph_cmd.rs b/tooling/profiler/src/cli/execution_flamegraph_cmd.rs
@@ -1,13 +1,12 @@
 use std::path::{Path, PathBuf};
 
 use acir::circuit::OpcodeLocation;
-use acir::FieldElement;
 use clap::Args;
 use color_eyre::eyre::{self, Context};
 
-use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator, Sample};
+use crate::flamegraph::{BrilligExecutionSample, FlamegraphGenerator, InfernoFlamegraphGenerator};
 use crate::fs::{read_inputs_from_file, read_program_from_file};
-use crate::opcode_formatter::AcirOrBrilligOpcode;
+use crate::opcode_formatter::format_brillig_opcode;
 use bn254_blackbox_solver::Bn254BlackBoxSolver;
 use nargo::ops::DefaultForeignCallExecutor;
 use noirc_abi::input_parser::Format;
@@ -51,19 +50,21 @@ fn run_with_generator(
     let initial_witness = program.abi.encode(&inputs_map, None)?;
 
     println!("Executing");
-    let (_, profiling_samples) = nargo::ops::execute_program_with_profiling(
+    let (_, mut profiling_samples) = nargo::ops::execute_program_with_profiling(
         &program.bytecode,
         initial_witness,
         &Bn254BlackBoxSolver,
         &mut DefaultForeignCallExecutor::new(true, None, None, None),
     )?;
     println!("Executed");
 
-    let profiling_samples: Vec<Sample<FieldElement>> = profiling_samples
-        .into_iter()
+    println!("Collecting {} samples", profiling_samples.len());
+
+    let profiling_samples: Vec<BrilligExecutionSample> = profiling_samples
+        .iter_mut()
         .map(|sample| {
-            let call_stack = sample.call_stack;
-            let brillig_function_id = sample.brillig_function_id;
+            let call_stack = std::mem::take(&mut sample.call_stack);
+            let brillig_function_id = std::mem::take(&mut sample.brillig_function_id);
             let last_entry = call_stack.last();
             let opcode = brillig_function_id
                 .and_then(|id| program.bytecode.unconstrained_functions.get(id.0 as usize))
@@ -74,8 +75,8 @@ fn run_with_generator(
                         None
                     }
                 })
-                .map(|opcode| AcirOrBrilligOpcode::Brillig(opcode.clone()));
-            Sample { opcode, call_stack, count: 1, brillig_function_id }
+                .map(format_brillig_opcode);
+            BrilligExecutionSample { opcode, call_stack, brillig_function_id }
         })
         .collect();
 
diff --git a/tooling/profiler/src/cli/gates_flamegraph_cmd.rs b/tooling/profiler/src/cli/gates_flamegraph_cmd.rs
@@ -6,10 +6,10 @@ use color_eyre::eyre::{self, Context};
 
 use noirc_artifacts::debug::DebugArtifact;
 
-use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator, Sample};
+use crate::flamegraph::{CompilationSample, FlamegraphGenerator, InfernoFlamegraphGenerator};
 use crate::fs::read_program_from_file;
 use crate::gates_provider::{BackendGatesProvider, GatesProvider};
-use crate::opcode_formatter::AcirOrBrilligOpcode;
+use crate::opcode_formatter::format_acir_opcode;
 
 #[derive(Debug, Clone, Args)]
 pub(crate) struct GatesFlamegraphCommand {
@@ -83,8 +83,8 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
             .into_iter()
             .zip(bytecode.opcodes)
             .enumerate()
-            .map(|(index, (gates, opcode))| Sample {
-                opcode: Some(AcirOrBrilligOpcode::Acir(opcode)),
+            .map(|(index, (gates, opcode))| CompilationSample {
+                opcode: Some(format_acir_opcode(&opcode)),
                 call_stack: vec![OpcodeLocation::Acir(index)],
                 count: gates,
                 brillig_function_id: None,
@@ -106,10 +106,7 @@ fn run_with_provider<Provider: GatesProvider, Generator: FlamegraphGenerator>(
 
 #[cfg(test)]
 mod tests {
-    use acir::{
-        circuit::{Circuit, Program},
-        AcirField,
-    };
+    use acir::circuit::{Circuit, Program};
     use color_eyre::eyre::{self};
     use fm::codespan_files::Files;
     use noirc_artifacts::program::ProgramArtifact;
@@ -143,9 +140,9 @@ mod tests {
     struct TestFlamegraphGenerator {}
 
     impl super::FlamegraphGenerator for TestFlamegraphGenerator {
-        fn generate_flamegraph<'files, F: AcirField>(
+        fn generate_flamegraph<'files, S: Sample>(
             &self,
-            _samples: Vec<Sample<F>>,
+            _samples: Vec<S>,
             _debug_symbols: &DebugInfo,
             _files: &'files impl Files<'files, FileId = fm::FileId>,
             _artifact_name: &str,
diff --git a/tooling/profiler/src/cli/opcodes_flamegraph_cmd.rs b/tooling/profiler/src/cli/opcodes_flamegraph_cmd.rs
@@ -7,9 +7,9 @@ use color_eyre::eyre::{self, Context};
 
 use noirc_artifacts::debug::DebugArtifact;
 
-use crate::flamegraph::{FlamegraphGenerator, InfernoFlamegraphGenerator, Sample};
+use crate::flamegraph::{CompilationSample, FlamegraphGenerator, InfernoFlamegraphGenerator};
 use crate::fs::read_program_from_file;
-use crate::opcode_formatter::AcirOrBrilligOpcode;
+use crate::opcode_formatter::{format_acir_opcode, format_brillig_opcode};
 
 #[derive(Debug, Clone, Args)]
 pub(crate) struct OpcodesFlamegraphCommand {
@@ -59,8 +59,8 @@ fn run_with_generator<Generator: FlamegraphGenerator>(
             .opcodes
             .iter()
             .enumerate()
-            .map(|(index, opcode)| Sample {
-                opcode: Some(AcirOrBrilligOpcode::Acir(opcode.clone())),
+            .map(|(index, opcode)| CompilationSample {
+                opcode: Some(format_acir_opcode(opcode)),
                 call_stack: vec![OpcodeLocation::Acir(index)],
                 count: 1,
                 brillig_function_id: None,
@@ -96,8 +96,8 @@ fn run_with_generator<Generator: FlamegraphGenerator>(
                 .bytecode
                 .into_iter()
                 .enumerate()
-                .map(|(brillig_index, opcode)| Sample {
-                    opcode: Some(AcirOrBrilligOpcode::Brillig(opcode)),
+                .map(|(brillig_index, opcode)| CompilationSample {
+                    opcode: Some(format_brillig_opcode(&opcode)),
                     call_stack: vec![OpcodeLocation::Brillig {
                         acir_index: acir_opcode_index,
                         brillig_index,
@@ -146,7 +146,7 @@ mod tests {
             brillig::{BrilligBytecode, BrilligFunctionId},
             Circuit, Opcode, Program,
         },
-        AcirField, FieldElement,
+        FieldElement,
     };
     use color_eyre::eyre::{self};
     use fm::codespan_files::Files;
@@ -160,9 +160,9 @@ mod tests {
     struct TestFlamegraphGenerator {}
 
     impl super::FlamegraphGenerator for TestFlamegraphGenerator {
-        fn generate_flamegraph<'files, F: AcirField>(
+        fn generate_flamegraph<'files, S: Sample>(
             &self,
-            _samples: Vec<Sample<F>>,
+            _samples: Vec<S>,
             _debug_symbols: &DebugInfo,
             _files: &'files impl Files<'files, FileId = fm::FileId>,
             _artifact_name: &str,
diff --git a/tooling/profiler/src/flamegraph.rs b/tooling/profiler/src/flamegraph.rs
@@ -3,7 +3,6 @@ use std::{collections::BTreeMap, io::BufWriter};
 
 use acir::circuit::brillig::BrilligFunctionId;
 use acir::circuit::OpcodeLocation;
-use acir::AcirField;
 use color_eyre::eyre::{self};
 use fm::codespan_files::Files;
 use fxhash::FxHashMap as HashMap;
@@ -13,18 +12,66 @@ use noirc_errors::reporter::line_and_column_from_span;
 use noirc_errors::Location;
 use noirc_evaluator::brillig::ProcedureId;
 
-use crate::opcode_formatter::AcirOrBrilligOpcode;
+pub(crate) trait Sample {
+    fn count(&self) -> usize;
 
-use super::opcode_formatter::format_opcode;
+    fn brillig_function_id(&self) -> Option<BrilligFunctionId>;
+
+    fn call_stack(&self) -> &[OpcodeLocation];
+
+    fn opcode(self) -> Option<String>;
+}
 
 #[derive(Debug)]
-pub(crate) struct Sample<F: AcirField> {
-    pub(crate) opcode: Option<AcirOrBrilligOpcode<F>>,
+pub(crate) struct CompilationSample {
+    pub(crate) opcode: Option<String>,
     pub(crate) call_stack: Vec<OpcodeLocation>,
     pub(crate) count: usize,
     pub(crate) brillig_function_id: Option<BrilligFunctionId>,
 }
 
+impl Sample for CompilationSample {
+    fn count(&self) -> usize {
+        self.count
+    }
+
+    fn brillig_function_id(&self) -> Option<BrilligFunctionId> {
+        self.brillig_function_id
+    }
+
+    fn call_stack(&self) -> &[OpcodeLocation] {
+        &self.call_stack
+    }
+
+    fn opcode(self) -> Option<String> {
+        self.opcode
+    }
+}
+
+pub(crate) struct BrilligExecutionSample {
+    pub(crate) opcode: Option<String>,
+    pub(crate) call_stack: Vec<OpcodeLocation>,
+    pub(crate) brillig_function_id: Option<BrilligFunctionId>,
+}
+
+impl Sample for BrilligExecutionSample {
+    fn count(&self) -> usize {
+        1
+    }
+
+    fn brillig_function_id(&self) -> Option<BrilligFunctionId> {
+        self.brillig_function_id
+    }
+
+    fn call_stack(&self) -> &[OpcodeLocation] {
+        &self.call_stack
+    }
+
+    fn opcode(self) -> Option<String> {
+        self.opcode
+    }
+}
+
 #[derive(Debug, Default)]
 pub(crate) struct FoldedStackItem {
     pub(crate) total_samples: usize,
@@ -33,9 +80,9 @@ pub(crate) struct FoldedStackItem {
 
 pub(crate) trait FlamegraphGenerator {
     #[allow(clippy::too_many_arguments)]
-    fn generate_flamegraph<'files, F: AcirField>(
+    fn generate_flamegraph<'files, S: Sample>(
         &self,
-        samples: Vec<Sample<F>>,
+        samples: Vec<S>,
         debug_symbols: &DebugInfo,
         files: &'files impl Files<'files, FileId = fm::FileId>,
         artifact_name: &str,
@@ -49,9 +96,9 @@ pub(crate) struct InfernoFlamegraphGenerator {
 }
 
 impl FlamegraphGenerator for InfernoFlamegraphGenerator {
-    fn generate_flamegraph<'files, F: AcirField>(
+    fn generate_flamegraph<'files, S: Sample>(
         &self,
-        samples: Vec<Sample<F>>,
+        samples: Vec<S>,
         debug_symbols: &DebugInfo,
         files: &'files impl Files<'files, FileId = fm::FileId>,
         artifact_name: &str,
@@ -82,8 +129,8 @@ impl FlamegraphGenerator for InfernoFlamegraphGenerator {
     }
 }
 
-fn generate_folded_sorted_lines<'files, F: AcirField>(
-    samples: Vec<Sample<F>>,
+fn generate_folded_sorted_lines<'files, S: Sample>(
+    samples: Vec<S>,
     debug_symbols: &DebugInfo,
     files: &'files impl Files<'files, FileId = fm::FileId>,
 ) -> Vec<String> {
@@ -92,15 +139,15 @@ fn generate_folded_sorted_lines<'files, F: AcirField>(
 
     let mut resolution_cache: HashMap<OpcodeLocation, Vec<String>> = HashMap::default();
     for sample in samples {
-        let mut location_names = Vec::with_capacity(sample.call_stack.len());
-        for opcode_location in sample.call_stack {
+        let mut location_names = Vec::with_capacity(sample.call_stack().len());
+        for opcode_location in sample.call_stack() {
             let callsite_labels = resolution_cache
-                .entry(opcode_location)
+                .entry(*opcode_location)
                 .or_insert_with(|| {
                     find_callsite_labels(
                         debug_symbols,
-                        &opcode_location,
-                        sample.brillig_function_id,
+                        opcode_location,
+                        sample.brillig_function_id(),
                         files,
                     )
                 })
@@ -109,11 +156,14 @@ fn generate_folded_sorted_lines<'files, F: AcirField>(
             location_names.extend(callsite_labels);
         }
 
-        if let Some(opcode) = &sample.opcode {
-            location_names.push(format_opcode(opcode));
+        // We move `sample` by calling `sample.opcode()` so we want to fetch the sample count here.
+        let count = sample.count();
+
+        if let Some(opcode) = sample.opcode() {
+            location_names.push(opcode);
         }
 
-        add_locations_to_folded_stack_items(&mut folded_stack_items, location_names, sample.count);
+        add_locations_to_folded_stack_items(&mut folded_stack_items, location_names, count);
     }
 
     to_folded_sorted_lines(&folded_stack_items, Default::default())
@@ -251,7 +301,7 @@ mod tests {
     use noirc_errors::{debug_info::DebugInfo, Location, Span};
     use std::{collections::BTreeMap, path::Path};
 
-    use crate::{flamegraph::Sample, opcode_formatter::AcirOrBrilligOpcode};
+    use crate::{flamegraph::CompilationSample, opcode_formatter::format_acir_opcode};
 
     use super::generate_folded_sorted_lines;
 
@@ -338,25 +388,25 @@ mod tests {
             BTreeMap::default(),
         );
 
-        let samples: Vec<Sample<FieldElement>> = vec![
-            Sample {
-                opcode: Some(AcirOrBrilligOpcode::Acir(AcirOpcode::AssertZero(
+        let samples: Vec<CompilationSample> = vec![
+            CompilationSample {
+                opcode: Some(format_acir_opcode(&AcirOpcode::AssertZero::<FieldElement>(
                     Expression::default(),
                 ))),
                 call_stack: vec![OpcodeLocation::Acir(0)],
                 count: 10,
                 brillig_function_id: None,
             },
-            Sample {
-                opcode: Some(AcirOrBrilligOpcode::Acir(AcirOpcode::AssertZero(
+            CompilationSample {
+                opcode: Some(format_acir_opcode(&AcirOpcode::AssertZero::<FieldElement>(
                     Expression::default(),
                 ))),
                 call_stack: vec![OpcodeLocation::Acir(1)],
                 count: 20,
                 brillig_function_id: None,
             },
-            Sample {
-                opcode: Some(AcirOrBrilligOpcode::Acir(AcirOpcode::MemoryInit {
+            CompilationSample {
+                opcode: Some(format_acir_opcode(&AcirOpcode::MemoryInit::<FieldElement> {
                     block_id: BlockId(0),
                     init: vec![],
                     block_type: acir::circuit::opcodes::BlockType::Memory,
diff --git a/tooling/profiler/src/opcode_formatter.rs b/tooling/profiler/src/opcode_formatter.rs
@@ -2,12 +2,6 @@ use acir::brillig::{BinaryFieldOp, BinaryIntOp, BlackBoxOp, Opcode as BrilligOpc
 use acir::circuit::{directives::Directive, opcodes::BlackBoxFuncCall, Opcode as AcirOpcode};
 use acir::AcirField;
 
-#[derive(Debug)]
-pub(crate) enum AcirOrBrilligOpcode<F: AcirField> {
-    Acir(AcirOpcode<F>),
-    Brillig(BrilligOpcode<F>),
-}
-
 fn format_blackbox_function<F>(call: &BlackBoxFuncCall<F>) -> String {
     match call {
         BlackBoxFuncCall::AES128Encrypt { .. } => "aes128_encrypt".to_string(),
@@ -136,11 +130,10 @@ fn format_brillig_opcode_kind<F>(opcode: &BrilligOpcode<F>) -> String {
     }
 }
 
-pub(crate) fn format_opcode<F: AcirField>(opcode: &AcirOrBrilligOpcode<F>) -> String {
-    match opcode {
-        AcirOrBrilligOpcode::Acir(opcode) => format!("acir::{}", format_acir_opcode_kind(opcode)),
-        AcirOrBrilligOpcode::Brillig(opcode) => {
-            format!("brillig::{}", format_brillig_opcode_kind(opcode))
-        }
-    }
+pub(crate) fn format_acir_opcode<F: AcirField>(opcode: &AcirOpcode<F>) -> String {
+    format!("acir::{}", format_acir_opcode_kind(opcode))
+}
+
+pub(crate) fn format_brillig_opcode<F: AcirField>(opcode: &BrilligOpcode<F>) -> String {
+    format!("brillig::{}", format_brillig_opcode_kind(opcode))
 }
