feat(profiler)!: New flamegraph command that profiles the opcodes executed

acvm-repo/acvm/src/pwg/brillig.rs

@@ -12,7 +12,7 @@ use acir::{
     AcirField,
 };
 use acvm_blackbox_solver::BlackBoxFunctionSolver;
-use brillig_vm::{FailureReason, MemoryValue, VMStatus, VM};
+use brillig_vm::{BrilligProfilingSamples, FailureReason, MemoryValue, VMStatus, VM};
 use serde::{Deserialize, Serialize};
 
 use crate::{pwg::OpcodeNotSolvable, OpcodeResolutionError};
@@ -200,16 +200,16 @@ impl<'b, B: BlackBoxFunctionSolver<F>, F: AcirField> BrilligSolver<'b, F, B> {
     }
 
     pub(crate) fn finalize(
-        self,
+        mut self,
         witness: &mut WitnessMap<F>,
         outputs: &[BrilligOutputs],
-    ) -> Result<(), OpcodeResolutionError<F>> {
+    ) -> Result<BrilligProfilingSamples, OpcodeResolutionError<F>> {
         // Finish the Brillig execution by writing the outputs to the witness map
         let vm_status = self.vm.get_status();
         match vm_status {
             VMStatus::Finished { return_data_offset, return_data_size } => {
                 self.write_brillig_outputs(witness, return_data_offset, return_data_size, outputs)?;
-                Ok(())
+                Ok(self.vm.take_profiling_samples())
             }
             _ => panic!("Brillig VM has not completed execution"),
         }

acvm-repo/acvm/src/pwg/mod.rs

@@ -58,7 +58,7 @@
     RequiresForeignCall(ForeignCallWaitInfo<F>),
 
     /// The ACVM has encountered a request for an ACIR [call][acir::circuit::Opcode]
     /// to execute a separate ACVM instance. The result of the ACIR call must be passd back to the ACVM.
     ///
     /// Once this is done, the ACVM can be restarted to solve the remaining opcodes.
     RequiresAcirCall(AcirCallWaitInfo<F>),
@@ -168,6 +168,13 @@
     }
 }
 
+pub type ProfilingSamples = Vec<ProfilingSample>;
+
+pub struct ProfilingSample {
+    pub call_stack: Vec<OpcodeLocation>,
+    pub brillig_function_id: Option<BrilligFunctionId>,
+}
+
 pub struct ACVM<'a, F, B: BlackBoxFunctionSolver<F>> {
     status: ACVMStatus<F>,
 
@@ -198,6 +205,8 @@
     unconstrained_functions: &'a [BrilligBytecode<F>],
 
     assertion_payloads: &'a [(OpcodeLocation, AssertionPayload<F>)],
+
+    profiling_samples: ProfilingSamples,
 }
 
 impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
@@ -222,6 +231,7 @@
             acir_call_results: Vec::default(),
             unconstrained_functions,
             assertion_payloads,
+            profiling_samples: Vec::new(),
         }
     }
 
@@ -247,11 +257,11 @@
     }
 
     /// Finalize the ACVM execution, returning the resulting [`WitnessMap`].
-    pub fn finalize(self) -> WitnessMap<F> {
+    pub fn finalize(self) -> (WitnessMap<F>, ProfilingSamples) {
         if self.status != ACVMStatus::Solved {
             panic!("ACVM execution is not complete: ({})", self.status);
         }
-        self.witness_map
+        (self.witness_map, self.profiling_samples)
     }
 
     /// Updates the current status of the VM.
@@ -519,7 +529,19 @@
             }
             BrilligSolverStatus::Finished => {
                 // Write execution outputs
-                solver.finalize(&mut self.witness_map, outputs)?;
+                let profiling_info = solver.finalize(&mut self.witness_map, outputs)?;
+                profiling_info.into_iter().for_each(|sample| {
+                    let mapped = sample.call_stack.into_iter().map(|loc| OpcodeLocation::Brillig {
+                        acir_index: self.instruction_pointer,
+                        brillig_index: loc,
+                    });
+                    self.profiling_samples.push(ProfilingSample {
+                        call_stack: std::iter::once(OpcodeLocation::Acir(self.instruction_pointer))
+                            .chain(mapped)
+                            .collect(),
+                        brillig_function_id: Some(*id),
+                    });
+                });
                 Ok(None)
             }
         }

acvm-repo/acvm/tests/solver.rs

@@ -53,7 +53,7 @@
     assert_eq!(solver_status, ACVMStatus::Solved, "should be fully solved");
 
     // ACVM should be able to be finalized in `Solved` state.
-    let witness_stack = acvm.finalize();
+    let witness_stack = acvm.finalize().0;
 
     assert_eq!(witness_stack.get(&Witness(3)).unwrap(), &Bls12FieldElement::from(5u128));
 }
@@ -785,7 +785,7 @@
         ACVM::new(&StubbedBlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
-    let witness_map = acvm.finalize();
+    let witness_map = acvm.finalize().0;
 
     assert_eq!(witness_map[&Witness(8)], FieldElement::from(6u128));
 }
@@ -891,7 +891,7 @@
         ACVM::new(&Bn254BlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
-    let witness_map = acvm.finalize();
+    let witness_map = acvm.finalize().0;
 
     Ok(outputs
         .iter()
@@ -1027,7 +1027,7 @@
         ACVM::new(&StubbedBlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
-    let witness_map = acvm.finalize();
+    let witness_map = acvm.finalize().0;
     Ok(output_witnesses
         .iter()
         .map(|witness| *witness_map.get(witness).expect("all witnesses to be set"))
@@ -1070,7 +1070,7 @@
         ACVM::new(&StubbedBlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
-    let witness_map = acvm.finalize();
+    let witness_map = acvm.finalize().0;
 
     Ok(witness_map[&Witness(3)])
 }
@@ -1326,9 +1326,9 @@
     let message = format!("not injective:\n{:?}\n{:?}", &inputs, &distinct_inputs);
     let outputs_not_equal =
         solve_array_input_blackbox_call(inputs, num_outputs, num_bits, op.clone())
             .expect("injectivity test operations to have valid input")
             != solve_array_input_blackbox_call(distinct_inputs, num_outputs, num_bits, op)
                 .expect("injectivity test operations to have valid input");
     (equal_inputs || outputs_not_equal, message)
 }
 

acvm-repo/acvm_js/src/execute.rs

@@ -276,7 +276,7 @@ impl<'a, B: BlackBoxFunctionSolver<FieldElement>> ProgramExecutor<'a, B> {
                 }
             }
 
-            Ok(acvm.finalize())
+            Ok(acvm.finalize().0)
         })
     }
 }

acvm-repo/brillig_vm/src/lib.rs

@@ -63,6 +63,15 @@
     },
 }
 
+// A sample for each opcode that was executed.
+pub type BrilligProfilingSamples = Vec<BrilligProfilingSample>;
+
+#[derive(Debug, PartialEq, Eq, Clone)]
+pub struct BrilligProfilingSample {
+    // The call stack when processing a given opcode.
+    pub call_stack: Vec<usize>,
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 /// VM encapsulates the state of the Brillig VM during execution.
 pub struct VM<'a, F, B: BlackBoxFunctionSolver<F>> {
@@ -88,6 +97,8 @@
     black_box_solver: &'a B,
     // The solver for big integers
     bigint_solver: BrilligBigintSolver,
+    // Samples for profiling the VM execution.
+    profiling_samples: BrilligProfilingSamples,
 }
 
 impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> VM<'a, F, B> {
@@ -109,9 +120,14 @@
             call_stack: Vec::new(),
             black_box_solver,
             bigint_solver: Default::default(),
+            profiling_samples: Vec::with_capacity(bytecode.len()),
         }
     }
 
+    pub fn take_profiling_samples(&mut self) -> BrilligProfilingSamples {
+        std::mem::take(&mut self.profiling_samples)
+    }
+
     /// Updates the current status of the VM.
     /// Returns the given status.
     fn status(&mut self, status: VMStatus<F>) -> VMStatus<F> {
@@ -196,6 +212,13 @@
 
     /// Process a single opcode and modify the program counter.
     pub fn process_opcode(&mut self) -> VMStatus<F> {
+        let call_stack: Vec<usize> = self.get_call_stack();
+        self.profiling_samples.push(BrilligProfilingSample { call_stack });
+
+        self.process_opcode_internal()
+    }
+
+    pub fn process_opcode_internal(&mut self) -> VMStatus<F> {
         let opcode = &self.bytecode[self.program_counter];
         match opcode {
             Opcode::BinaryFieldOp { op, lhs, rhs, destination: result } => {
@@ -347,7 +370,7 @@
                 self.set_program_counter(*location)
             }
             Opcode::Const { destination, value, bit_size } => {
                 // Consts are not checked in runtime to fit in the bit size, since they can safely be checked statically.
                 self.memory.write(*destination, MemoryValue::new_from_field(*value, *bit_size));
                 self.increment_program_counter()
             }
@@ -885,7 +908,7 @@
     }
 
     #[test]
     fn jmpifnot_opcode() {
         let calldata: Vec<FieldElement> = vec![1u128.into(), 2u128.into()];
 
         let opcodes = vec![
@@ -1101,7 +1124,7 @@
     }
 
     #[test]
     fn cmov_opcode() {
         let calldata: Vec<FieldElement> =
             vec![(0u128).into(), (1u128).into(), (2u128).into(), (3u128).into()];
 
@@ -1355,7 +1378,7 @@
     }
 
     #[test]
     fn iconst_opcode() {
         let opcodes = &[
             Opcode::Const {
                 destination: MemoryAddress::direct(0),

tooling/acvm_cli/src/cli/execute_cmd.rs

@@ -75,5 +75,6 @@ pub(crate) fn execute_program_from_witness(
         &Bn254BlackBoxSolver,
         &mut DefaultForeignCallExecutor::new(true, None, None, None),
     )
+    .map(|(witness_stack, _)| witness_stack)
     .map_err(CliError::CircuitExecutionError)
 }

tooling/debugger/src/context.rs

@@ -590,7 +590,7 @@ impl<'a, B: BlackBoxFunctionSolver<FieldElement>> DebugContext<'a, B> {
         let caller_acvm = caller_frame.acvm;
         let callee_acvm = std::mem::replace(&mut self.acvm, caller_acvm);
         self.current_circuit_id = caller_frame.circuit_id;
-        let call_solved_witness = callee_acvm.finalize();
+        let call_solved_witness = callee_acvm.finalize().0;
 
         let ACVMStatus::RequiresAcirCall(call_info) = self.acvm.get_status() else {
             unreachable!("Resolving an ACIR call, the caller is in an invalid state");
@@ -851,7 +851,7 @@ impl<'a, B: BlackBoxFunctionSolver<FieldElement>> DebugContext<'a, B> {
     }
 
     pub fn finalize(mut self) -> WitnessStack<FieldElement> {
-        let last_witness_map = self.acvm.finalize();
+        let last_witness_map = self.acvm.finalize().0;
         self.witness_stack.push(0, last_witness_map);
         self.witness_stack
     }

tooling/nargo/src/ops/execute.rs

@@ -3,7 +3,9 @@ use acvm::acir::circuit::{
     OpcodeLocation, Program, ResolvedAssertionPayload, ResolvedOpcodeLocation,
 };
 use acvm::acir::native_types::WitnessStack;
-use acvm::pwg::{ACVMStatus, ErrorLocation, OpcodeNotSolvable, OpcodeResolutionError, ACVM};
+use acvm::pwg::{
+    ACVMStatus, ErrorLocation, OpcodeNotSolvable, OpcodeResolutionError, ProfilingSamples, ACVM,
+};
 use acvm::{acir::circuit::Circuit, acir::native_types::WitnessMap};
 use acvm::{AcirField, BlackBoxFunctionSolver};
 
@@ -62,7 +64,7 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>, E: ForeignCallExecutor<F>>
     fn execute_circuit(
         &mut self,
         initial_witness: WitnessMap<F>,
-    ) -> Result<WitnessMap<F>, NargoError<F>> {
+    ) -> Result<(WitnessMap<F>, ProfilingSamples), NargoError<F>> {
         let circuit = &self.functions[self.current_function_index];
         let mut acvm = ACVM::new(
             self.blackbox_solver,
@@ -155,7 +157,8 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>, E: ForeignCallExecutor<F>>
                     // Execute the ACIR call
                     let acir_to_call = &self.functions[call_info.id.as_usize()];
                     let initial_witness = call_info.initial_witness;
-                    let call_solved_witness = self.execute_circuit(initial_witness)?;
+                    // TODO: Profiling among multiple circuits is not supported
+                    let (call_solved_witness, _) = self.execute_circuit(initial_witness)?;
 
                     // Set tracking index back to the parent function after ACIR call execution
                     self.current_function_index = acir_function_caller;
@@ -194,15 +197,15 @@ pub fn execute_program<F: AcirField, B: BlackBoxFunctionSolver<F>, E: ForeignCal
     initial_witness: WitnessMap<F>,
     blackbox_solver: &B,
     foreign_call_executor: &mut E,
-) -> Result<WitnessStack<F>, NargoError<F>> {
+) -> Result<(WitnessStack<F>, ProfilingSamples), NargoError<F>> {
     let mut executor = ProgramExecutor::new(
         &program.functions,
         &program.unconstrained_functions,
         blackbox_solver,
         foreign_call_executor,
     );
-    let main_witness = executor.execute_circuit(initial_witness)?;
+    let (main_witness, profiling_samples) = executor.execute_circuit(initial_witness)?;
     executor.witness_stack.push(0, main_witness);
 
-    Ok(executor.finalize())
+    Ok((executor.finalize(), profiling_samples))
 }

tooling/nargo/src/ops/test.rs

@@ -63,7 +63,7 @@ pub fn run_test<B: BlackBoxFunctionSolver<FieldElement>>(
                     test_function,
                     compiled_program.abi,
                     compiled_program.debug,
-                    circuit_execution,
+                    circuit_execution.map(|(witness_stack, _)| witness_stack),
                 )
             } else {
                 #[cfg(target_arch = "wasm32")]
@@ -97,6 +97,7 @@ pub fn run_test<B: BlackBoxFunctionSolver<FieldElement>>(
                                     package_name.clone(),
                                 ),
                             )
+                            .map(|(witness_stack, _)| witness_stack)
                             .map_err(|err| err.to_string())
                         };
                     let fuzzer = FuzzedExecutor::new(compiled_program.into(), executor, runner);

tooling/nargo_cli/src/cli/execute_cmd.rs

@@ -133,7 +133,8 @@ pub(crate) fn execute_program(
             root_path,
             package_name,
         ),
-    );
+    )
+    .map(|(witness_stack, _)| witness_stack);
     match solved_witness_stack_err {
         Ok(solved_witness_stack) => Ok(solved_witness_stack),
         Err(err) => {

tooling/nargo_cli/tests/stdlib-props.rs

@@ -95,6 +95,7 @@ fn run_snippet_proptest(
             &blackbox_solver,
             &mut *foreign_call_executor,
         )
+        .map(|(witness_stack, _)| witness_stack)
         .expect("failed to execute");
 
         let main_witness = witness_stack.peek().expect("should have return value on witness stack");

tooling/profiler/Cargo.toml

@@ -18,8 +18,10 @@ name = "noir-profiler"
 path = "src/main.rs"
 
 [dependencies]
+bn254_blackbox_solver.workspace = true
 color-eyre.workspace = true
 clap.workspace = true
+fxhash.workspace = true
 noirc_artifacts.workspace = true
 const_format.workspace = true
 serde.workspace = true
@@ -28,7 +30,9 @@ fm.workspace = true
 inferno = "0.11.19"
 im.workspace = true
 acir.workspace = true
+nargo.workspace = true
 noirc_errors.workspace = true
+noirc_abi.workspace = true
 
 # Logs
 tracing-subscriber.workspace = true