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};
@@ -58,6 +58,7 @@ impl<'b, B: BlackBoxFunctionSolver<F>, F: AcirField> BrilligSolver<'b, F, B> {
 
     /// Constructs a solver for a Brillig block given the bytecode and initial
     /// witness.
+    #[allow(clippy::too_many_arguments)]
     pub(crate) fn new_call(
         initial_witness: &WitnessMap<F>,
         memory: &HashMap<BlockId, MemoryOpSolver<F>>,
@@ -66,9 +67,16 @@ impl<'b, B: BlackBoxFunctionSolver<F>, F: AcirField> BrilligSolver<'b, F, B> {
         bb_solver: &'b B,
         acir_index: usize,
         brillig_function_id: BrilligFunctionId,
+        profiling_active: bool,
     ) -> Result<Self, OpcodeResolutionError<F>> {
-        let vm =
-            Self::setup_brillig_vm(initial_witness, memory, inputs, brillig_bytecode, bb_solver)?;
+        let vm = Self::setup_brillig_vm(
+            initial_witness,
+            memory,
+            inputs,
+            brillig_bytecode,
+            bb_solver,
+            profiling_active,
+        )?;
         Ok(Self { vm, acir_index, function_id: brillig_function_id })
     }
 
@@ -78,6 +86,7 @@ impl<'b, B: BlackBoxFunctionSolver<F>, F: AcirField> BrilligSolver<'b, F, B> {
         inputs: &[BrilligInputs<F>],
         brillig_bytecode: &'b [BrilligOpcode<F>],
         bb_solver: &'b B,
+        profiling_active: bool,
     ) -> Result<VM<'b, F, B>, OpcodeResolutionError<F>> {
         // Set input values
         let mut calldata: Vec<F> = Vec::new();
@@ -125,7 +134,7 @@ impl<'b, B: BlackBoxFunctionSolver<F>, F: AcirField> BrilligSolver<'b, F, B> {
 
         // Instantiate a Brillig VM given the solved calldata
         // along with the Brillig bytecode.
-        let vm = VM::new(calldata, brillig_bytecode, vec![], bb_solver);
+        let vm = VM::new(calldata, brillig_bytecode, vec![], bb_solver, profiling_active);
         Ok(vm)
     }
 
@@ -203,6 +212,25 @@ impl<'b, B: BlackBoxFunctionSolver<F>, F: AcirField> BrilligSolver<'b, F, B> {
         self,
         witness: &mut WitnessMap<F>,
         outputs: &[BrilligOutputs],
+    ) -> Result<(), OpcodeResolutionError<F>> {
+        assert!(!self.vm.is_profiling_active(), "Expected VM profiling to not be active");
+        self.finalize_inner(witness, outputs)
+    }
+
+    pub(crate) fn finalize_with_profiling(
+        mut self,
+        witness: &mut WitnessMap<F>,
+        outputs: &[BrilligOutputs],
+    ) -> Result<BrilligProfilingSamples, OpcodeResolutionError<F>> {
+        assert!(self.vm.is_profiling_active(), "Expected VM profiling to be active");
+        self.finalize_inner(witness, outputs)?;
+        Ok(self.vm.take_profiling_samples())
+    }
+
+    fn finalize_inner(
+        &self,
+        witness: &mut WitnessMap<F>,
+        outputs: &[BrilligOutputs],
     ) -> Result<(), OpcodeResolutionError<F>> {
         // Finish the Brillig execution by writing the outputs to the witness map
         let vm_status = self.vm.get_status();

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,14 @@
     }
 }
 
+pub type ProfilingSamples = Vec<ProfilingSample>;
+
+#[derive(Default)]
+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 +206,10 @@
     unconstrained_functions: &'a [BrilligBytecode<F>],
 
     assertion_payloads: &'a [(OpcodeLocation, AssertionPayload<F>)],
+
+    profiling_active: bool,
+
+    profiling_samples: ProfilingSamples,
 }
 
 impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
@@ -222,9 +234,16 @@
             acir_call_results: Vec::default(),
             unconstrained_functions,
             assertion_payloads,
+            profiling_active: false,
+            profiling_samples: Vec::new(),
         }
     }
 
+    // Enable profiling
+    pub fn with_profiler(&mut self, profiling_active: bool) {
+        self.profiling_active = profiling_active;
+    }
+
     /// Returns a reference to the current state of the ACVM's [`WitnessMap`].
     ///
     /// Once execution has completed, the witness map can be extracted using [`ACVM::finalize`]
@@ -246,6 +265,10 @@
         self.instruction_pointer
     }
 
+    pub fn take_profiling_samples(&mut self) -> ProfilingSamples {
+        std::mem::take(&mut self.profiling_samples)
+    }
+
     /// Finalize the ACVM execution, returning the resulting [`WitnessMap`].
     pub fn finalize(self) -> WitnessMap<F> {
         if self.status != ACVMStatus::Solved {
@@ -503,6 +526,7 @@
                 self.backend,
                 self.instruction_pointer,
                 *id,
+                self.profiling_active,
             )?,
         };
 
@@ -519,7 +543,28 @@
             }
             BrilligSolverStatus::Finished => {
                 // Write execution outputs
-                solver.finalize(&mut self.witness_map, outputs)?;
+                if self.profiling_active {
+                    let profiling_info =
+                        solver.finalize_with_profiling(&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),
+                        });
+                    });
+                } else {
+                    solver.finalize(&mut self.witness_map, outputs)?;
+                }
+
                 Ok(None)
             }
         }
@@ -575,6 +620,7 @@
             self.backend,
             self.instruction_pointer,
             *id,
+            self.profiling_active,
         );
         match solver {
             Ok(solver) => StepResult::IntoBrillig(solver),

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,10 @@
     black_box_solver: &'a B,
     // The solver for big integers
     bigint_solver: BrilligBigintSolver,
+    // Flag that determines whether we want to profile VM.
+    profiling_active: bool,
+    // Samples for profiling the VM execution.
+    profiling_samples: BrilligProfilingSamples,
 }
 
 impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> VM<'a, F, B> {
@@ -97,6 +110,7 @@
         bytecode: &'a [Opcode<F>],
         foreign_call_results: Vec<ForeignCallResult<F>>,
         black_box_solver: &'a B,
+        profiling_active: bool,
     ) -> Self {
         Self {
             calldata,
@@ -109,9 +123,19 @@
             call_stack: Vec::new(),
             black_box_solver,
             bigint_solver: Default::default(),
+            profiling_active,
+            profiling_samples: Vec::with_capacity(bytecode.len()),
         }
     }
 
+    pub fn is_profiling_active(&self) -> bool {
+        self.profiling_active
+    }
+
+    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 +220,15 @@
 
     /// Process a single opcode and modify the program counter.
     pub fn process_opcode(&mut self) -> VMStatus<F> {
+        if self.profiling_active {
+            let call_stack: Vec<usize> = self.get_call_stack();
+            self.profiling_samples.push(BrilligProfilingSample { call_stack });
+        }
+
+        self.process_opcode_internal()
+    }
+
+    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 +380,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()
             }
@@ -813,7 +846,7 @@
         }];
 
         // Start VM
-        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::Finished { return_data_offset: 0, return_data_size: 0 });
@@ -863,7 +896,7 @@
             Opcode::JumpIf { condition: destination, location: 6 },
         ];
 
-        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -885,7 +918,7 @@
     }
 
     #[test]
     fn jmpifnot_opcode() {
         let calldata: Vec<FieldElement> = vec![1u128.into(), 2u128.into()];
 
         let opcodes = vec![
@@ -921,7 +954,7 @@
             },
         ];
 
-        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -982,7 +1015,7 @@
             },
             Opcode::Stop { return_data_offset: 1, return_data_size: 1 },
         ];
-        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -1033,7 +1066,7 @@
             },
             Opcode::Stop { return_data_offset: 1, return_data_size: 1 },
         ];
-        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -1079,7 +1112,7 @@
             },
             Opcode::Mov { destination: MemoryAddress::direct(2), source: MemoryAddress::direct(0) },
         ];
-        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -1101,7 +1134,7 @@
     }
 
     #[test]
     fn cmov_opcode() {
         let calldata: Vec<FieldElement> =
             vec![(0u128).into(), (1u128).into(), (2u128).into(), (3u128).into()];
 
@@ -1144,7 +1177,7 @@
                 condition: MemoryAddress::direct(1),
             },
         ];
-        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -1240,7 +1273,7 @@
             .chain(cast_opcodes)
             .chain([equal_opcode, not_equal_opcode, less_than_opcode, less_than_equal_opcode])
             .collect();
-        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         // Calldata copy
         let status = vm.process_opcode();
@@ -1355,7 +1388,7 @@
     }
 
     #[test]
     fn iconst_opcode() {
         let opcodes = &[
             Opcode::Const {
                 destination: MemoryAddress::direct(0),
@@ -1368,7 +1401,7 @@
                 value: FieldElement::from(27_usize),
             },
         ];
-        let mut vm = VM::new(vec![], opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(vec![], opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         let status = vm.process_opcode();
         assert_eq!(status, VMStatus::InProgress);
@@ -1592,7 +1625,7 @@
         calldata: Vec<F>,
         opcodes: &[Opcode<F>],
     ) -> VM<'_, F, StubbedBlackBoxSolver> {
-        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, opcodes, vec![], &StubbedBlackBoxSolver, false);
         brillig_execute(&mut vm);
         assert_eq!(vm.call_stack, vec![]);
         vm
@@ -2271,7 +2304,7 @@
             },
         ];
 
-        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver);
+        let mut vm = VM::new(calldata, &opcodes, vec![], &StubbedBlackBoxSolver, false);
 
         vm.process_opcode();
         vm.process_opcode();

compiler/noirc_evaluator/src/brillig/brillig_ir.rs

@@ -233,7 +233,8 @@ pub(crate) mod tests {
         calldata: Vec<FieldElement>,
         bytecode: &[BrilligOpcode<FieldElement>],
     ) -> (VM<'_, FieldElement, DummyBlackBoxSolver>, usize, usize) {
-        let mut vm = VM::new(calldata, bytecode, vec![], &DummyBlackBoxSolver);
+        let profiling_active = false;
+        let mut vm = VM::new(calldata, bytecode, vec![], &DummyBlackBoxSolver, profiling_active);
 
         let status = vm.process_opcodes();
         if let VMStatus::Finished { return_data_offset, return_data_size } = status {
@@ -301,6 +302,7 @@ pub(crate) mod tests {
             &bytecode,
             vec![ForeignCallResult { values: vec![ForeignCallParam::Array(number_sequence)] }],
             &DummyBlackBoxSolver,
+            false,
         );
         let status = vm.process_opcodes();
         assert_eq!(status, VMStatus::Finished { return_data_offset: 0, return_data_size: 0 });

compiler/noirc_evaluator/src/ssa/acir_gen/acir_ir/acir_variable.rs

@@ -2195,7 +2195,8 @@ fn execute_brillig<F: AcirField>(
     // We pass a stubbed solver here as a concrete solver implies a field choice which conflicts with this function
     // being generic.
     let solver = acvm::blackbox_solver::StubbedBlackBoxSolver;
-    let mut vm = VM::new(calldata, code, Vec::new(), &solver);
+    let profiling_active = false;
+    let mut vm = VM::new(calldata, code, Vec::new(), &solver, profiling_active);
 
     // Run the Brillig VM on these inputs, bytecode, etc!
     let vm_status = vm.process_opcodes();