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

…cuted (#6327) # Description ## Problem\* Resolves  Initial profiling to get better numbers for the opcodes actually being executed by the Brillig VM. ## Summary\* cc @sirasistant for a lot of the initial code in this PR. The noir-profiler currently only looks at compiled artifacts to generate a flamegraph showing where constraints counts and/or opcodes counts are most prominent in a program. This PR adds execution profiling sampling to be part of the ACVM/Brillig VM and the noir-profiler. This will enable us to see if there are opcodes which are being unnecessarily executed (such as array copies). This initial PR simply adds the `execution-flamegraph` command while a follow-up will add extra metadata to look at Brillig procedures. After compiling `execution_success/sha256` with `--force-brillig` we can run the following command: ``` noir-profiler execution-opcodes -a ./target/sha256.json -p ./Prover.toml -o ./target/ ``` Result for execution flamegraph: <img width="1591" alt="Screenshot 2024-10-23 at 7 17 48 PM" src="https://github.com/user-attachments/assets/c8a7f601-4431-4e86-a3e0-160facb0ec08"> To compare here is the result of `noir-profiler opcodes-flamegraph`: <img width="1589" alt="Screenshot 2024-10-23 at 7 18 12 PM" src="https://github.com/user-attachments/assets/c352ee09-26de-43da-9fdb-7c02558dfacf"> ## Additional Context This is a minor breaking change as I changed the names of the commands in the `noir-profiler` to not all be suffixed with `flamegraph`. ## Documentation\* Check one: - [ ] No documentation needed. - [ ] Documentation included in this PR. - [X] **[For Experimental Features]** Documentation to be submitted in a separate PR. # PR Checklist\* - [X] I have tested the changes locally. - [X] I have formatted the changes with [Prettier](https://prettier.io/) and/or `cargo fmt` on default settings. --------- Co-authored-by: Akosh Farkash <[email protected]>
noir-lang · Oct 24, 2024 · 4d87c9a · 4d87c9a
1 parent 0de3241
commit 4d87c9a
Show file tree

Hide file tree

Showing 15 changed files with 385 additions and 65 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/acvm-repo/acvm/src/pwg/brillig.rs b/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();

diff --git a/acvm-repo/acvm/src/pwg/mod.rs b/acvm-repo/acvm/src/pwg/mod.rs
@@ -168,6 +168,14 @@ impl<F> From<InvalidInputBitSize> for OpcodeResolutionError<F> {
     }
 }
 
+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 @@ pub struct ACVM<'a, F, B: BlackBoxFunctionSolver<F>> {
     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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
             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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
         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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
                 self.backend,
                 self.instruction_pointer,
                 *id,
+                self.profiling_active,
             )?,
         };
 
@@ -519,7 +543,28 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
             }
             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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> ACVM<'a, F, B> {
             self.backend,
             self.instruction_pointer,
             *id,
+            self.profiling_active,
         );
         match solver {
             Ok(solver) => StepResult::IntoBrillig(solver),

diff --git a/acvm-repo/brillig_vm/src/lib.rs b/acvm-repo/brillig_vm/src/lib.rs
@@ -63,6 +63,15 @@ pub enum VMStatus<F> {
     },
 }
 
+// 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 @@ pub struct VM<'a, F, B: BlackBoxFunctionSolver<F>> {
     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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> VM<'a, F, B> {
         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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> VM<'a, F, B> {
             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 @@ impl<'a, F: AcirField, B: BlackBoxFunctionSolver<F>> VM<'a, F, B> {
 
     /// 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 } => {
@@ -813,7 +846,7 @@ mod tests {
         }];
 
         // 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 @@ mod tests {
             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);
@@ -921,7 +954,7 @@ mod tests {
             },
         ];
 
-        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 @@ mod tests {
             },
             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 @@ mod tests {
             },
             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 @@ mod tests {
             },
             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);
@@ -1144,7 +1177,7 @@ mod tests {
                 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 @@ mod tests {
             .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();
@@ -1368,7 +1401,7 @@ mod tests {
                 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 @@ mod tests {
         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 @@ mod tests {
             },
         ];
 
-        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();

diff --git a/compiler/noirc_evaluator/src/brillig/brillig_ir.rs b/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 });

diff --git a/compiler/noirc_evaluator/src/ssa/acir_gen/acir_ir/acir_variable.rs b/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();