chore: Pretty print timings for codegen (#4659)

# Description

## Problem\*

Resolves #4658

## Summary\*

I was already timing things for a separate issue so I decided to push
this PR. I felt this provided a nice temporary solution. If we deem that
the `tracing` crate can still be used for pretty printing the timings of
each codegen pass we can switch to that in a follow-up, and for now this
provides an easy way to check the timings of each codegen pass.

Example output:
```
After Defunctionalization:: 0 ms
After Removing Paired rc_inc & rc_decs:: 0 ms
After Inlining:: 0 ms
After Mem2Reg:: 0 ms
After Assert Constant:: 0 ms
After Unrolling:: 13 ms
After Simplifying:: 0 ms
After Flattening:: 9 ms
After Removing Bit Shifts:: 0 ms
After Mem2Reg:: 19 ms
After Constant Folding:: 3 ms
After Constraint Folding:: 3 ms
After Dead Instruction Elimination:: 0 ms
SSA to Brillig: 0 ms
SSA to ACIR: 14 ms
``` 

## Additional Context



## Documentation\*

Check one:
- [ ] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[For Experimental Features]** Documentation to be submitted in a
separate PR.

# PR Checklist\*

- [ ] I have tested the changes locally.
- [ ] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.

---------

Co-authored-by: jfecher <jake@aztecprotocol.com>

compiler/noirc_driver/src/lib.rs

@@ -68,6 +68,10 @@ pub struct CompileOptions {
     #[arg(long)]
     pub print_acir: bool,
 
+    /// Pretty print benchmark times of each code generation pass
+    #[arg(long, hide = true)]
+    pub benchmark_codegen: bool,
+
     /// Treat all warnings as errors
     #[arg(long, conflicts_with = "silence_warnings")]
     pub deny_warnings: bool,
@@ -478,8 +482,13 @@ pub fn compile_no_check(
         return Ok(cached_program.expect("cache must exist for hashes to match"));
     }
     let visibility = program.return_visibility;
-    let (circuit, debug, input_witnesses, return_witnesses, warnings) =
-        create_circuit(program, options.show_ssa, options.show_brillig, options.force_brillig)?;
+    let (circuit, debug, input_witnesses, return_witnesses, warnings) = create_circuit(
+        program,
+        options.show_ssa,
+        options.show_brillig,
+        options.force_brillig,
+        options.benchmark_codegen,
+    )?;
 
     let abi =
         abi_gen::gen_abi(context, &main_function, input_witnesses, return_witnesses, visibility);

compiler/noirc_evaluator/Cargo.toml

@@ -17,4 +17,5 @@ thiserror.workspace = true
 num-bigint = "0.4"
 im = { version = "15.1", features = ["serde"] }
 serde.workspace = true
-tracing.workspace = true
+tracing.workspace = true
+chrono = "0.4.37"

compiler/noirc_evaluator/src/ssa.rs

@@ -38,15 +38,16 @@ pub mod ssa_gen;
 /// convert the final SSA into ACIR and return it.
 pub(crate) fn optimize_into_acir(
     program: Program,
-    print_ssa_passes: bool,
+    print_passes: bool,
     print_brillig_trace: bool,
     force_brillig_output: bool,
+    print_timings: bool,
 ) -> Result<GeneratedAcir, RuntimeError> {
     let abi_distinctness = program.return_distinctness;
 
     let ssa_gen_span = span!(Level::TRACE, "ssa_generation");
     let ssa_gen_span_guard = ssa_gen_span.enter();
-    let ssa = SsaBuilder::new(program, print_ssa_passes, force_brillig_output)?
+    let ssa = SsaBuilder::new(program, print_passes, force_brillig_output, print_timings)?
         .run_pass(Ssa::defunctionalize, "After Defunctionalization:")
         .run_pass(Ssa::remove_paired_rc, "After Removing Paired rc_inc & rc_decs:")
         .run_pass(Ssa::inline_functions, "After Inlining:")
@@ -64,13 +65,28 @@ pub(crate) fn optimize_into_acir(
         .run_pass(Ssa::dead_instruction_elimination, "After Dead Instruction Elimination:")
         .finish();
 
-    let brillig = ssa.to_brillig(print_brillig_trace);
+    let brillig = time("SSA to Brillig", print_timings, || ssa.to_brillig(print_brillig_trace));
 
     drop(ssa_gen_span_guard);
 
     let last_array_uses = ssa.find_last_array_uses();
 
-    ssa.into_acir(brillig, abi_distinctness, &last_array_uses)
+    time("SSA to ACIR", print_timings, || {
+        ssa.into_acir(brillig, abi_distinctness, &last_array_uses)
+    })
+}
+
+// Helper to time SSA passes
+fn time<T>(name: &str, print_timings: bool, f: impl FnOnce() -> T) -> T {
+    let start_time = chrono::Utc::now().time();
+    let result = f();
+
+    if print_timings {
+        let end_time = chrono::Utc::now().time();
+        println!("{name}: {} ms", (end_time - start_time).num_milliseconds());
+    }
+
+    result
 }
 
 /// Compiles the [`Program`] into [`ACIR`][acvm::acir::circuit::Circuit].
@@ -83,6 +99,7 @@ pub fn create_circuit(
     enable_ssa_logging: bool,
     enable_brillig_logging: bool,
     force_brillig_output: bool,
+    print_codegen_timings: bool,
 ) -> Result<(Circuit, DebugInfo, Vec<Witness>, Vec<Witness>, Vec<SsaReport>), RuntimeError> {
     let debug_variables = program.debug_variables.clone();
     let debug_types = program.debug_types.clone();
@@ -94,6 +111,7 @@ pub fn create_circuit(
         enable_ssa_logging,
         enable_brillig_logging,
         force_brillig_output,
+        print_codegen_timings,
     )?;
     let opcodes = generated_acir.take_opcodes();
     let current_witness_index = generated_acir.current_witness_index().0;
@@ -176,16 +194,18 @@ fn split_public_and_private_inputs(
 struct SsaBuilder {
     ssa: Ssa,
     print_ssa_passes: bool,
+    print_codegen_timings: bool,
 }
 
 impl SsaBuilder {
     fn new(
         program: Program,
         print_ssa_passes: bool,
         force_brillig_runtime: bool,
+        print_codegen_timings: bool,
     ) -> Result<SsaBuilder, RuntimeError> {
         let ssa = ssa_gen::generate_ssa(program, force_brillig_runtime)?;
-        Ok(SsaBuilder { print_ssa_passes, ssa }.print("Initial SSA:"))
+        Ok(SsaBuilder { print_ssa_passes, print_codegen_timings, ssa }.print("Initial SSA:"))
     }
 
     fn finish(self) -> Ssa {
@@ -194,7 +214,7 @@ impl SsaBuilder {
 
     /// Runs the given SSA pass and prints the SSA afterward if `print_ssa_passes` is true.
     fn run_pass(mut self, pass: fn(Ssa) -> Ssa, msg: &str) -> Self {
-        self.ssa = pass(self.ssa);
+        self.ssa = time(msg, self.print_codegen_timings, || pass(self.ssa));
         self.print(msg)
     }
 
@@ -204,7 +224,7 @@ impl SsaBuilder {
         pass: fn(Ssa) -> Result<Ssa, RuntimeError>,
         msg: &str,
     ) -> Result<Self, RuntimeError> {
-        self.ssa = pass(self.ssa)?;
+        self.ssa = time(msg, self.print_codegen_timings, || pass(self.ssa))?;
         Ok(self.print(msg))
     }
 

cspell.json

@@ -36,6 +36,7 @@
         "castable",
         "catmcgee",
         "Celo",
+        "chrono",
         "chumsky",
         "codegen",
         "codegenned",