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> {
@@ -207,6 +219,7 @@
         initial_witness: WitnessMap<F>,
         unconstrained_functions: &'a [BrilligBytecode<F>],
         assertion_payloads: &'a [(OpcodeLocation, AssertionPayload<F>)],
+        profiling_active: bool,
     ) -> Self {
         let status = if opcodes.is_empty() { ACVMStatus::Solved } else { ACVMStatus::InProgress };
         ACVM {
@@ -222,6 +235,8 @@
             acir_call_results: Vec::default(),
             unconstrained_functions,
             assertion_payloads,
+            profiling_active,
+            profiling_samples: Vec::new(),
         }
     }
 
@@ -247,6 +262,13 @@
     }
 
     /// Finalize the ACVM execution, returning the resulting [`WitnessMap`].
+    pub fn finalize_with_profiling(self) -> (WitnessMap<F>, ProfilingSamples) {
+        if self.status != ACVMStatus::Solved {
+            panic!("ACVM execution is not complete: ({})", self.status);
+        }
+        (self.witness_map, self.profiling_samples)
+    }
+
     pub fn finalize(self) -> WitnessMap<F> {
         if self.status != ACVMStatus::Solved {
             panic!("ACVM execution is not complete: ({})", self.status);
@@ -503,6 +525,7 @@
                 self.backend,
                 self.instruction_pointer,
                 *id,
+                self.profiling_active,
             )?,
         };
 
@@ -519,7 +542,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 +619,7 @@
             self.backend,
             self.instruction_pointer,
             *id,
+            self.profiling_active,
         );
         match solver {
             Ok(solver) => StepResult::IntoBrillig(solver),

acvm-repo/acvm/tests/solver.rs

@@ -47,7 +47,8 @@
     ])
     .into();
 
-    let mut acvm = ACVM::new(&StubbedBlackBoxSolver, &opcodes, witness_assignments, &[], &[]);
+    let mut acvm =
+        ACVM::new(&StubbedBlackBoxSolver, &opcodes, witness_assignments, &[], &[], false);
     // use the partial witness generation solver with our acir program
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved, "should be fully solved");
@@ -164,6 +165,7 @@
         witness_assignments,
         &unconstrained_functions,
         &[],
+        false,
     );
     // use the partial witness generation solver with our acir program
     let solver_status = acvm.solve();
@@ -319,6 +321,7 @@
         witness_assignments,
         &unconstrained_functions,
         &[],
+        false,
     );
 
     // use the partial witness generation solver with our acir program
@@ -465,6 +468,7 @@
         witness_assignments,
         &unconstrained_functions,
         &[],
+        false,
     );
 
     // use the partial witness generation solver with our acir program
@@ -587,6 +591,7 @@
         witness_assignments,
         &unconstrained_functions,
         &[],
+        false,
     );
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved, "should be fully solved");
@@ -623,7 +628,7 @@
     let opcodes = vec![Opcode::AssertZero(opcode_a)];
     let unconstrained_functions = vec![];
     let mut acvm =
-        ACVM::new(&StubbedBlackBoxSolver, &opcodes, values, &unconstrained_functions, &[]);
+        ACVM::new(&StubbedBlackBoxSolver, &opcodes, values, &unconstrained_functions, &[], false);
     let solver_status = acvm.solve();
     assert_eq!(
         solver_status,
@@ -732,7 +737,7 @@
     ];
     let unconstrained_functions = vec![brillig_bytecode];
     let mut acvm =
-        ACVM::new(&StubbedBlackBoxSolver, &opcodes, values, &unconstrained_functions, &[]);
+        ACVM::new(&StubbedBlackBoxSolver, &opcodes, values, &unconstrained_functions, &[], false);
     let solver_status = acvm.solve();
     assert_eq!(
         solver_status,
@@ -781,8 +786,14 @@
 
     let opcodes = vec![init, read_op, expression];
     let unconstrained_functions = vec![];
-    let mut acvm =
-        ACVM::new(&StubbedBlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
+    let mut acvm = ACVM::new(
+        &StubbedBlackBoxSolver,
+        &opcodes,
+        initial_witness,
+        &unconstrained_functions,
+        &[],
+        false,
+    );
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
     let witness_map = acvm.finalize();
@@ -887,8 +898,14 @@
     let op = Opcode::BlackBoxFuncCall(f((inputs.clone(), outputs.clone()))?);
     let opcodes = vec![op];
     let unconstrained_functions = vec![];
-    let mut acvm =
-        ACVM::new(&Bn254BlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
+    let mut acvm = ACVM::new(
+        &Bn254BlackBoxSolver,
+        &opcodes,
+        initial_witness,
+        &unconstrained_functions,
+        &[],
+        false,
+    );
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
     let witness_map = acvm.finalize();
@@ -1023,8 +1040,14 @@
     opcodes.push(bigint_to_op);
 
     let unconstrained_functions = vec![];
-    let mut acvm =
-        ACVM::new(&StubbedBlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
+    let mut acvm = ACVM::new(
+        &StubbedBlackBoxSolver,
+        &opcodes,
+        initial_witness,
+        &unconstrained_functions,
+        &[],
+        false,
+    );
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
     let witness_map = acvm.finalize();
@@ -1066,8 +1089,14 @@
     let op = Opcode::BlackBoxFuncCall(blackbox_func_call(lhs_opt, rhs_opt)?);
     let opcodes = vec![op];
     let unconstrained_functions = vec![];
-    let mut acvm =
-        ACVM::new(&StubbedBlackBoxSolver, &opcodes, initial_witness, &unconstrained_functions, &[]);
+    let mut acvm = ACVM::new(
+        &StubbedBlackBoxSolver,
+        &opcodes,
+        initial_witness,
+        &unconstrained_functions,
+        &[],
+        false,
+    );
     let solver_status = acvm.solve();
     assert_eq!(solver_status, ACVMStatus::Solved);
     let witness_map = acvm.finalize();
@@ -1326,9 +1355,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

@@ -174,12 +174,14 @@ impl<'a, B: BlackBoxFunctionSolver<FieldElement>> ProgramExecutor<'a, B> {
         witness_stack: &'a mut WitnessStack<FieldElement>,
     ) -> Pin<Box<dyn Future<Output = Result<WitnessMap<FieldElement>, Error>> + 'a>> {
         Box::pin(async {
+            let profiling_active = false;
             let mut acvm = ACVM::new(
                 self.blackbox_solver,
                 &circuit.opcodes,
                 initial_witness,
                 self.unconstrained_functions,
                 &circuit.assert_messages,
+                profiling_active,
             );
 
             loop {