Implement libfunc profiler.

Cargo.toml

@@ -50,6 +50,7 @@ scarb = ["build-cli", "dep:scarb-ui", "dep:scarb-metadata"]
 with-cheatcode = []
 with-debug-utils = []
 with-mem-tracing = []
+with-profiler = []
 with-runtime = ["dep:cairo-native-runtime"]
 
 # the aquamarine dep is only used in docs and cannot be detected as used by cargo udeps

programs/examples/brainfuck.cairo

@@ -1,7 +1,8 @@
+use core::num::traits::WrappingSub;
+use core::num::traits::WrappingAdd;
 use core::{
-    traits::Default,
-    array::{ArrayTrait, SpanTrait}, debug::PrintTrait, dict::Felt252DictTrait,
-    integer::{u8_wrapping_add, u8_wrapping_sub}, option::OptionTrait, traits::Into,
+    traits::Default, array::{ArrayTrait, SpanTrait}, debug::PrintTrait, dict::Felt252DictTrait,
+    option::OptionTrait, traits::Into,
 };
 
 fn generate_jump_table(program: @Array<u8>) -> Felt252Dict<u32> {
@@ -50,9 +51,9 @@ fn run_program(program: @Array<u8>, input: Option<Span<u8>>) {
         } else if op_code == '<' {
             mp -= 1;
         } else if op_code == '+' {
-            memory.insert(mp, u8_wrapping_add(memory.get(mp), 1));
+            memory.insert(mp, memory.get(mp).wrapping_add(1));
         } else if op_code == '-' {
-            memory.insert(mp, u8_wrapping_sub(memory.get(mp), 1));
+            memory.insert(mp, memory.get(mp).wrapping_sub(1));
         } else if op_code == '.' {
             memory.get(mp).print();
         } else if op_code == ',' {

src/bin/cairo-native-run.rs

@@ -44,6 +44,10 @@ struct Args {
     /// Optimization level, Valid: 0, 1, 2, 3. Values higher than 3 are considered as 3.
     #[arg(short = 'O', long, default_value_t = 0)]
     opt_level: u8,
+
+    #[cfg(feature = "with-profiler")]
+    #[arg(long)]
+    profiler_output: Option<PathBuf>,
 }
 
 fn main() -> anyhow::Result<()> {
@@ -138,5 +142,104 @@ fn main() -> anyhow::Result<()> {
         println!("Remaining gas: {gas}");
     }
 
+    #[cfg(feature = "with-profiler")]
+    if let Some(profiler_output) = args.profiler_output {
+        use cairo_lang_sierra::{ids::ConcreteLibfuncId, program::Statement};
+        use std::{collections::HashMap, fs::File, io::Write};
+
+        let mut trace = HashMap::<ConcreteLibfuncId, (Vec<u64>, u64)>::new();
+
+        for (statement_idx, tick_delta) in cairo_native::metadata::profiler::ProfilerImpl::take() {
+            if let Statement::Invocation(invocation) = &sierra_program.statements[statement_idx.0] {
+                let (tick_deltas, extra_count) =
+                    trace.entry(invocation.libfunc_id.clone()).or_default();
+
+                if tick_delta != u64::MAX {
+                    tick_deltas.push(tick_delta);
+                } else {
+                    *extra_count += 1;
+                }
+            }
+        }
+
+        let mut trace = trace
+            .into_iter()
+            .map(|(libfunc_id, (mut tick_deltas, extra_count))| {
+                tick_deltas.sort();
+
+                // Drop outliers.
+                {
+                    let q1 = tick_deltas[tick_deltas.len() / 4];
+                    let q3 = tick_deltas[3 * tick_deltas.len() / 4];
+                    let iqr = q3 - q1;
+
+                    let q1_thr = q1.saturating_sub(iqr + iqr / 2);
+                    let q3_thr = q3 + (iqr + iqr / 2);
+
+                    tick_deltas.retain(|x| *x >= q1_thr && *x <= q3_thr);
+                }
+
+                // Compute the quartiles.
+                let quartiles = [
+                    *tick_deltas.first().unwrap(),
+                    tick_deltas[tick_deltas.len() / 4],
+                    tick_deltas[tick_deltas.len() / 2],
+                    tick_deltas[3 * tick_deltas.len() / 4],
+                    *tick_deltas.last().unwrap(),
+                ];
+
+                // Compuite the average.
+                let average =
+                    tick_deltas.iter().copied().sum::<u64>() as f64 / tick_deltas.len() as f64;
+
+                // Compute the standard deviation.
+                let std_dev = {
+                    let sum = tick_deltas
+                        .iter()
+                        .copied()
+                        .map(|x| x as f64)
+                        .map(|x| (x - average))
+                        .map(|x| x * x)
+                        .sum::<f64>();
+                    sum / (tick_deltas.len() as u64 + extra_count) as f64
+                };
+
+                (
+                    libfunc_id,
+                    (
+                        tick_deltas.len() as u64 + extra_count,
+                        tick_deltas.iter().sum::<u64>()
+                            + (extra_count as f64 * average).round() as u64,
+                        quartiles,
+                        average,
+                        std_dev,
+                    ),
+                )
+            })
+            .collect::<Vec<_>>();
+
+        // Sort libfuncs by the order in which they are declared.
+        trace.sort_by_key(|(libfunc_id, _)| {
+            sierra_program
+                .libfunc_declarations
+                .iter()
+                .enumerate()
+                .find_map(|(i, x)| (&x.id == libfunc_id).then_some(i))
+                .unwrap()
+        });
+
+        let mut output = File::create(profiler_output)?;
+
+        for (libfunc_id, (n_samples, sum, quartiles, average, std_dev)) in trace {
+            writeln!(output, "{libfunc_id}")?;
+            writeln!(output, "    Samples  : {n_samples}")?;
+            writeln!(output, "    Sum      : {sum}")?;
+            writeln!(output, "    Average  : {average}")?;
+            writeln!(output, "    Deviation: {std_dev}")?;
+            writeln!(output, "    Quartiles: {quartiles:?}")?;
+            writeln!(output)?;
+        }
+    }
+
     Ok(())
 }

src/compiler.rs

@@ -524,26 +524,6 @@ fn compile_func(
 
                     let (state, _) = edit_state::take_args(state, invocation.args.iter())?;
 
-                    let helper = LibfuncHelper {
-                        module,
-                        init_block: &pre_entry_block,
-                        region: &region,
-                        blocks_arena: &blocks_arena,
-                        last_block: Cell::new(block),
-                        branches: generate_branching_targets(
-                            &blocks,
-                            statements,
-                            statement_idx,
-                            invocation,
-                            &state,
-                        ),
-                        results: invocation
-                            .branches
-                            .iter()
-                            .map(|x| vec![Cell::new(None); x.results.len()])
-                            .collect::<Vec<_>>(),
-                    };
-
                     let libfunc = registry.get_libfunc(&invocation.libfunc_id)?;
                     if is_recursive {
                         if let Some(target) = libfunc.is_function_call() {
@@ -594,6 +574,45 @@ fn compile_func(
                         }
                     }
 
+                    #[allow(unused_mut)]
+                    let mut helper = LibfuncHelper {
+                        module,
+                        init_block: &pre_entry_block,
+                        region: &region,
+                        blocks_arena: &blocks_arena,
+                        last_block: Cell::new(block),
+                        branches: generate_branching_targets(
+                            &blocks,
+                            statements,
+                            statement_idx,
+                            invocation,
+                            &state,
+                        ),
+                        results: invocation
+                            .branches
+                            .iter()
+                            .map(|x| vec![Cell::new(None); x.results.len()])
+                            .collect::<Vec<_>>(),
+
+                        #[cfg(feature = "with-profiler")]
+                        profiler: match libfunc {
+                            CoreConcreteLibfunc::FunctionCall(_) => {
+                                // Tail-recursive function calls are broken. Also it'd include the entire function which
+                                // doesn't make sense, therefore it's ignored on purpose.
+                                None
+                            }
+                            _ => match metadata.remove::<crate::metadata::profiler::ProfilerMeta>()
+                            {
+                                Some(profiler_meta) => {
+                                    let t0 = profiler_meta
+                                        .measure_timestamp(context, block, location)?;
+                                    Some((profiler_meta, statement_idx, t0))
+                                }
+                                None => None,
+                            },
+                        },
+                    };
+
                     libfunc.build(
                         context,
                         registry,
@@ -604,6 +623,11 @@ fn compile_func(
                     )?;
                     assert!(block.terminator().is_some());
 
+                    #[cfg(feature = "with-profiler")]
+                    if let Some((profiler_meta, _, _)) = helper.profiler.take() {
+                        metadata.insert(profiler_meta);
+                    }
+
                     if let Some(tailrec_meta) = metadata.remove::<TailRecursionMeta>() {
                         if let Some(return_block) = tailrec_meta.return_target() {
                             tailrec_state = Some((tailrec_meta.depth_counter(), return_block));

src/context.rs

@@ -167,6 +167,12 @@ impl NativeContext {
         // already some metadata of the same type.
         metadata.insert(gas_metadata);
 
+        #[cfg(feature = "with-profiler")]
+        metadata.insert(crate::metadata::profiler::ProfilerMeta::new(
+            &self.context,
+            &module,
+        )?);
+
         // Create the Sierra program registry
         let registry = ProgramRegistry::<CoreType, CoreLibfunc>::new(program)?;
 

src/executor/jit.rs

@@ -85,6 +85,9 @@ impl<'m> JitNativeExecutor<'m> {
         let set_builtin_costs_fnptr: extern "C" fn(*const u64) -> *const u64 =
             unsafe { std::mem::transmute(self.engine.lookup("cairo_native__set_costs_builtin")) };
 
+        #[cfg(feature = "with-profiler")]
+        self.setup_profiling();
+
         super::invoke_dynamic(
             &self.registry,
             self.find_function_ptr(function_id),
@@ -112,6 +115,9 @@ impl<'m> JitNativeExecutor<'m> {
         let set_builtin_costs_fnptr: extern "C" fn(*const u64) -> *const u64 =
             unsafe { std::mem::transmute(self.engine.lookup("cairo_native__set_costs_builtin")) };
 
+        #[cfg(feature = "with-profiler")]
+        self.setup_profiling();
+
         super::invoke_dynamic(
             &self.registry,
             self.find_function_ptr(function_id),
@@ -138,6 +144,9 @@ impl<'m> JitNativeExecutor<'m> {
         let set_builtin_costs_fnptr: extern "C" fn(*const u64) -> *const u64 =
             unsafe { std::mem::transmute(self.engine.lookup("cairo_native__set_costs_builtin")) };
 
+        #[cfg(feature = "with-profiler")]
+        self.setup_profiling();
+
         ContractExecutionResult::from_execution_result(super::invoke_dynamic(
             &self.registry,
             self.find_function_ptr(function_id),
@@ -178,4 +187,14 @@ impl<'m> JitNativeExecutor<'m> {
             .get_function(function_id)
             .map(|func| &func.signature)?)
     }
+
+    #[cfg(feature = "with-profiler")]
+    fn setup_profiling(&self) {
+        unsafe {
+            let callback_ptr: *mut extern "C" fn(u64, u64) =
+                self.engine.lookup("__profiler_callback").cast();
+
+            *callback_ptr = crate::metadata::profiler::ProfilerImpl::callback;
+        }
+    }
 }