Implement libfunc profiling (squash) (#1253)

This commit is a squash of the whole #946 PR
(#946)

Author: JulianGCalderon

.github/workflows/starknet-blocks.yml

@@ -32,7 +32,7 @@ jobs:
         with:
           repository: lambdaclass/starknet-replay
           path: starknet-replay
-          ref: 1b8e2e0be21a8df9f5f6b8f8514d1a40b456ef58
+          ref: d36491aa5fca3f48b4d7fb25eba599603ff48225
       # We need native to use the linux deps ci action
       - name: Checkout Native
         uses: actions/checkout@v4
@@ -43,8 +43,7 @@ jobs:
         with:
           repository: lambdaclass/sequencer
           path: sequencer
-          ref: 40331042c1149f5cb84b27f9dd8d47994a010bbe
-
+          ref: 14be65ca995ac702bad26ac20f2a522d9515f70a
       - name: Cache RPC Calls
         uses: actions/cache@v4.2.0
         with:

Cargo.toml

@@ -60,6 +60,7 @@ scarb = ["build-cli", "dep:scarb-ui", "dep:scarb-metadata"]
 with-cheatcode = []
 with-debug-utils = []
 with-mem-tracing = []
+with-libfunc-profiling = []
 with-segfault-catcher = []
 with-trace-dump = ["dep:sierra-emu"]
 

src/bin/cairo-native-run.rs

@@ -3,14 +3,20 @@ use cairo_lang_compiler::{
     compile_prepared_db, db::RootDatabase, project::setup_project, CompilerConfig,
 };
 use cairo_lang_runner::short_string::as_cairo_short_string;
+#[cfg(feature = "with-libfunc-profiling")]
+use cairo_lang_sierra::ids::ConcreteLibfuncId;
 use cairo_lang_sierra_to_casm::metadata::MetadataComputationConfig;
+#[cfg(feature = "with-libfunc-profiling")]
+use cairo_native::metadata::profiler::LibfuncProfileData;
 use cairo_native::{
     context::NativeContext,
     executor::{AotNativeExecutor, JitNativeExecutor},
     metadata::gas::GasMetadata,
     starknet_stub::StubSyscallHandler,
 };
 use clap::{Parser, ValueEnum};
+#[cfg(feature = "with-libfunc-profiling")]
+use std::collections::HashMap;
 use std::path::PathBuf;
 use tracing_subscriber::{EnvFilter, FmtSubscriber};
 use utils::{find_function, result_to_runresult};
@@ -46,6 +52,11 @@ struct Args {
     #[arg(short = 'O', long, default_value_t = 0)]
     opt_level: u8,
 
+    #[cfg(feature = "with-libfunc-profiling")]
+    #[arg(long)]
+    /// The output path for the libfunc profilling results
+    profiler_output: Option<PathBuf>,
+
     #[cfg(feature = "with-trace-dump")]
     #[arg(long)]
     /// The output path for the execution trace
@@ -131,6 +142,18 @@ fn main() -> anyhow::Result<()> {
                 }
             }
 
+            #[cfg(feature = "with-libfunc-profiling")]
+            {
+                use cairo_native::metadata::profiler::ProfilerBinding;
+
+                if let Some(trace_id) =
+                    executor.find_symbol_ptr(ProfilerBinding::ProfileId.symbol())
+                {
+                    let trace_id = trace_id.cast::<u64>();
+                    unsafe { *trace_id = 0 };
+                }
+            }
+
             Box::new(move |function_id, args, gas, syscall_handler| {
                 executor.invoke_dynamic_with_syscall_handler(
                     function_id,
@@ -153,6 +176,16 @@ fn main() -> anyhow::Result<()> {
         );
     }
 
+    #[cfg(feature = "with-libfunc-profiling")]
+    {
+        use cairo_native::metadata::profiler::{ProfilerImpl, LIBFUNC_PROFILE};
+
+        LIBFUNC_PROFILE
+            .lock()
+            .unwrap()
+            .insert(0, ProfilerImpl::new());
+    }
+
     let gas_metadata =
         GasMetadata::new(&sierra_program, Some(MetadataComputationConfig::default())).unwrap();
 
@@ -188,6 +221,66 @@ fn main() -> anyhow::Result<()> {
         println!("Remaining gas: {gas}");
     }
 
+    #[cfg(feature = "with-libfunc-profiling")]
+    {
+        use std::{fs::File, io::Write};
+
+        let profile = cairo_native::metadata::profiler::LIBFUNC_PROFILE
+            .lock()
+            .unwrap();
+
+        assert_eq!(profile.values().len(), 1);
+
+        let profile = profile.values().next().unwrap();
+
+        if let Some(profiler_output_path) = args.profiler_output {
+            let mut output = File::create(profiler_output_path)?;
+
+            let raw_profile = profile.get_profile(&sierra_program);
+            let mut processed_profile = process_profile(raw_profile);
+
+            processed_profile.sort_by_key(|LibfuncProfileSummary { libfunc_idx, .. }| {
+                sierra_program
+                    .libfunc_declarations
+                    .iter()
+                    .enumerate()
+                    .find_map(|(i, x)| (x.id == *libfunc_idx).then_some(i))
+                    .unwrap()
+            });
+
+            for LibfuncProfileSummary {
+                libfunc_idx,
+                samples,
+                total_time,
+                average_time,
+                std_deviation,
+                quartiles,
+            } in processed_profile
+            {
+                writeln!(output, "{libfunc_idx}")?;
+                writeln!(output, "    Total Samples:          {samples}")?;
+
+                let (Some(total_time), Some(average_time), Some(std_deviation), Some(quartiles)) =
+                    (total_time, average_time, std_deviation, quartiles)
+                else {
+                    writeln!(output, "    Total Execution Time:   none")?;
+                    writeln!(output, "    Average Execution Time: none")?;
+                    writeln!(output, "    Standard Deviation:     none")?;
+                    writeln!(output, "    Quartiles:              none")?;
+                    writeln!(output)?;
+
+                    continue;
+                };
+
+                writeln!(output, "    Total Execution Time:   {total_time}")?;
+                writeln!(output, "    Average Execution Time: {average_time}")?;
+                writeln!(output, "    Standard Deviation:     {std_deviation}")?;
+                writeln!(output, "    Quartiles:              {quartiles:?}")?;
+                writeln!(output)?;
+            }
+        }
+    }
+
     #[cfg(feature = "with-trace-dump")]
     if let Some(trace_output) = args.trace_output {
         let traces = cairo_native::metadata::trace_dump::trace_dump_runtime::TRACE_DUMP
@@ -205,3 +298,92 @@ fn main() -> anyhow::Result<()> {
 
     Ok(())
 }
+
+#[cfg(feature = "with-libfunc-profiling")]
+struct LibfuncProfileSummary {
+    pub libfunc_idx: ConcreteLibfuncId,
+    pub samples: u64,
+    pub total_time: Option<u64>,
+    pub average_time: Option<f64>,
+    pub std_deviation: Option<f64>,
+    pub quartiles: Option<[u64; 5]>,
+}
+
+#[cfg(feature = "with-libfunc-profiling")]
+fn process_profile(
+    profiles: HashMap<ConcreteLibfuncId, LibfuncProfileData>,
+) -> Vec<LibfuncProfileSummary> {
+    profiles
+        .into_iter()
+        .map(
+            |(
+                libfunc_idx,
+                LibfuncProfileData {
+                    mut deltas,
+                    extra_counts,
+                },
+            )| {
+                // if no deltas were registered, we only return the libfunc's calls amount
+                if deltas.is_empty() {
+                    return LibfuncProfileSummary {
+                        libfunc_idx,
+                        samples: extra_counts,
+                        total_time: None,
+                        average_time: None,
+                        std_deviation: None,
+                        quartiles: None,
+                    };
+                }
+
+                deltas.sort();
+
+                // Drop outliers.
+                {
+                    let q1 = deltas[deltas.len() / 4];
+                    let q3 = deltas[3 * deltas.len() / 4];
+                    let iqr = q3 - q1;
+
+                    let q1_thr = q1.saturating_sub(iqr + iqr / 2);
+                    let q3_thr = q3 + (iqr + iqr / 2);
+
+                    deltas.retain(|x| *x >= q1_thr && *x <= q3_thr);
+                }
+
+                // Compute the quartiles.
+                let quartiles = [
+                    *deltas.first().unwrap(),
+                    deltas[deltas.len() / 4],
+                    deltas[deltas.len() / 2],
+                    deltas[3 * deltas.len() / 4],
+                    *deltas.last().unwrap(),
+                ];
+
+                // Compute the average.
+                let average = deltas.iter().copied().sum::<u64>() as f64 / deltas.len() as f64;
+
+                // Compute the standard deviation.
+                let std_dev = {
+                    let sum = deltas
+                        .iter()
+                        .copied()
+                        .map(|x| x as f64)
+                        .map(|x| (x - average))
+                        .map(|x| x * x)
+                        .sum::<f64>();
+                    sum / (deltas.len() as u64 + extra_counts) as f64
+                };
+
+                LibfuncProfileSummary {
+                    libfunc_idx,
+                    samples: deltas.len() as u64 + extra_counts,
+                    total_time: Some(
+                        deltas.iter().sum::<u64>() + (extra_counts as f64 * average).round() as u64,
+                    ),
+                    average_time: Some(average),
+                    std_deviation: Some(std_dev),
+                    quartiles: Some(quartiles),
+                }
+            },
+        )
+        .collect::<Vec<_>>()
+}

src/compiler.rs

@@ -550,26 +550,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() {
@@ -620,6 +600,46 @@ 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-libfunc-profiling")]
+                        profiler: match libfunc {
+                            CoreConcreteLibfunc::FunctionCall(_) => {
+                                // Tail-recursive function calls are broken beacuse a stack of timestamps is required,
+                                // which would invalidate tail recursion. Also, since each libfunc is measured individually,
+                                // it doesn't make sense to take function calls into account, 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,
@@ -651,6 +671,11 @@ fn compile_func(
                         libfunc_name
                     );
 
+                    #[cfg(feature = "with-libfunc-profiling")]
+                    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

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

src/executor/aot.rs

@@ -60,6 +60,9 @@ impl AotNativeExecutor {
         #[cfg(feature = "with-trace-dump")]
         crate::metadata::trace_dump::setup_runtime(|name| executor.find_symbol_ptr(name));
 
+        #[cfg(feature = "with-libfunc-profiling")]
+        crate::metadata::profiler::setup_runtime(|name| executor.find_symbol_ptr(name));
+
         executor
     }
 

src/executor/contract.rs

@@ -330,6 +330,9 @@ impl AotContractExecutor {
         #[cfg(feature = "with-trace-dump")]
         crate::metadata::trace_dump::setup_runtime(|name| executor.find_symbol_ptr(name));
 
+        #[cfg(feature = "with-libfunc-profiling")]
+        crate::metadata::profiler::setup_runtime(|name| executor.find_symbol_ptr(name));
+
         Ok(Some(executor))
     }
 

src/executor/jit.rs

@@ -71,6 +71,9 @@ impl<'m> JitNativeExecutor<'m> {
         #[cfg(feature = "with-trace-dump")]
         crate::metadata::trace_dump::setup_runtime(|name| executor.find_symbol_ptr(name));
 
+        #[cfg(feature = "with-libfunc-profiling")]
+        crate::metadata::profiler::setup_runtime(|name| executor.find_symbol_ptr(name));
+
         Ok(executor)
     }