From 8dab3a4b33da7fdcf8ae3013f8112710e516f727 Mon Sep 17 00:00:00 2001 From: Avi Cohen Date: Mon, 27 Apr 2026 14:31:59 +0300 Subject: [PATCH] add run_with_libfunc_profile + AotWithProgram variant for ContractExecutor Exposes the libfunc-profiling primitives that downstream consumers (e.g. the blockifier in starkware-libs/sequencer) currently maintain locally. Makes the profile-collection pattern callback-driven so the per-call key (tx hash, etc.) stays out of cairo-native. - `Profile` (HashMap) is now public. - `AotContractExecutor::run_with_libfunc_profile` (gated on `with-libfunc-profiling`) wraps `run`: allocates a unique trace ID, points the executor's `cairo_native__profiler__profile_id` symbol at it, drains the resulting `Profile` after `run` returns, and hands it to a caller- supplied `FnOnce(Profile)`. A `ProfilerGuard` restores the previous trace ID and drops the LIBFUNC_PROFILE slot on both the success and unwind paths. - `ContractExecutor::AotWithProgram(AotWithProgram { executor, program })` is a new variant that bundles an AOT executor with the Sierra program it was built from. `From` is provided. - `ContractExecutor::run` dispatches the new variant via `run_with_libfunc_profile` with a no-op profile callback. - `ContractExecutor::run_with_profile` is the profile-capturing counterpart of `run`; for non-`AotWithProgram` variants it falls through to `run` (callback never fires). --- src/executor.rs | 4 ++ src/executor/contract_executor.rs | 75 +++++++++++++++++++++- src/executor/libfunc_profile.rs | 103 ++++++++++++++++++++++++++++++ src/metadata/profiler.rs | 2 +- 4 files changed, 180 insertions(+), 4 deletions(-) create mode 100644 src/executor/libfunc_profile.rs diff --git a/src/executor.rs b/src/executor.rs index 195971fe59..42dd2e2e25 100644 --- a/src/executor.rs +++ b/src/executor.rs @@ -11,6 +11,8 @@ pub use self::{ }; #[cfg(feature = "sierra-emu")] pub use self::contract_executor::EmuContractInfo; +#[cfg(feature = "with-libfunc-profiling")] +pub use self::contract_executor::AotWithProgram; use crate::{ arch::{AbiArgument, ValueWithInfoWrapper}, error::{panic::ToNativeAssertError, Error}, @@ -48,6 +50,8 @@ mod aot; mod contract; mod contract_executor; mod jit; +#[cfg(feature = "with-libfunc-profiling")] +mod libfunc_profile; #[cfg(target_arch = "aarch64")] global_asm!(include_str!("arch/aarch64.s")); diff --git a/src/executor/contract_executor.rs b/src/executor/contract_executor.rs index f1d03089af..2b1b3dba1d 100644 --- a/src/executor/contract_executor.rs +++ b/src/executor/contract_executor.rs @@ -3,21 +3,26 @@ //! //! The `Emu` variant is gated on the `sierra-emu` feature and uses //! [`crate::sierra_emu_bridge::SierraEmuSyscallBridge`] to thread a cairo-native syscall -//! handler through the sierra-emu VM. +//! handler through the sierra-emu VM. The `AotWithProgram` variant is gated on +//! `with-libfunc-profiling` and bundles an [`AotContractExecutor`] with the Sierra +//! `Program` it was built from so [`ContractExecutor::run_with_profile`] can resolve +//! libfunc IDs after collecting samples. -#[cfg(feature = "sierra-emu")] +#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))] use cairo_lang_sierra::program::Program; #[cfg(feature = "sierra-emu")] use cairo_lang_starknet_classes::compiler_version::VersionId; #[cfg(feature = "sierra-emu")] use cairo_lang_starknet_classes::contract_class::ContractEntryPoints; use starknet_types_core::felt::Felt; -#[cfg(feature = "sierra-emu")] +#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))] use std::sync::Arc; use crate::error::Result; use crate::execution_result::ContractExecutionResult; use crate::executor::AotContractExecutor; +#[cfg(feature = "with-libfunc-profiling")] +use crate::metadata::profiler::Profile; #[cfg(feature = "sierra-emu")] use crate::sierra_emu_bridge::SierraEmuSyscallBridge; use crate::starknet::StarknetSyscallHandler; @@ -33,6 +38,8 @@ pub enum ContractExecutor { Aot(AotContractExecutor), #[cfg(feature = "sierra-emu")] Emu(EmuContractInfo), + #[cfg(feature = "with-libfunc-profiling")] + AotWithProgram(AotWithProgram), } /// Inputs required to construct a `sierra_emu::VirtualMachine` for the `Emu` variant. @@ -44,6 +51,16 @@ pub struct EmuContractInfo { pub sierra_version: VersionId, } +/// AOT executor paired with the Sierra program it was built from. Required by +/// [`ContractExecutor::run_with_profile`] so libfunc samples can be resolved against +/// the program's declarations. +#[cfg(feature = "with-libfunc-profiling")] +#[derive(Debug)] +pub struct AotWithProgram { + pub executor: AotContractExecutor, + pub program: Arc, +} + impl From for ContractExecutor { fn from(value: AotContractExecutor) -> Self { Self::Aot(value) @@ -57,6 +74,13 @@ impl From for ContractExecutor { } } +#[cfg(feature = "with-libfunc-profiling")] +impl From for ContractExecutor { + fn from(value: AotWithProgram) -> Self { + Self::AotWithProgram(value) + } +} + impl ContractExecutor { /// Run the contract entry point identified by `selector`. /// @@ -99,6 +123,51 @@ impl ContractExecutor { builtin_stats: Default::default(), }) } + #[cfg(feature = "with-libfunc-profiling")] + ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) => executor + .run_with_libfunc_profile( + program, + selector, + args, + gas, + builtin_costs, + syscall_handler, + // Profile is collected and dropped on this path. Use + // `run_with_profile` to capture it. + |_profile| {}, + ), + } + } + + /// Like [`Self::run`] but, for the `AotWithProgram` variant, hands the captured + /// libfunc profile to `on_profile` after the call returns successfully. For other + /// variants this is identical to `run` and `on_profile` is never invoked. + #[cfg(feature = "with-libfunc-profiling")] + pub fn run_with_profile( + &self, + selector: Felt, + args: &[Felt], + gas: u64, + builtin_costs: Option, + syscall_handler: H, + on_profile: F, + ) -> Result + where + H: StarknetSyscallHandler, + F: FnOnce(Profile), + { + match self { + ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) => executor + .run_with_libfunc_profile( + program, + selector, + args, + gas, + builtin_costs, + syscall_handler, + on_profile, + ), + _ => self.run(selector, args, gas, builtin_costs, syscall_handler), } } } diff --git a/src/executor/libfunc_profile.rs b/src/executor/libfunc_profile.rs new file mode 100644 index 0000000000..c5aecadebc --- /dev/null +++ b/src/executor/libfunc_profile.rs @@ -0,0 +1,103 @@ +//! Profiling-instrumented run wrapper around [`AotContractExecutor::run`]. +//! +//! Available under the `with-libfunc-profiling` feature. + +#![cfg(feature = "with-libfunc-profiling")] + +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::Arc; + +use cairo_lang_sierra::program::Program; +use starknet_types_core::felt::Felt; + +use crate::error::Result; +use crate::execution_result::ContractExecutionResult; +use crate::executor::AotContractExecutor; +use crate::metadata::profiler::{Profile, ProfilerBinding, ProfilerImpl, LIBFUNC_PROFILE}; +use crate::starknet::StarknetSyscallHandler; +use crate::utils::BuiltinCosts; + +impl AotContractExecutor { + /// Run the entrypoint with libfunc-level profiling instrumentation. + /// + /// Wraps [`AotContractExecutor::run`] with the bookkeeping the + /// `with-libfunc-profiling` runtime needs: + /// + /// 1. Allocates a unique trace ID and inserts an empty `ProfilerImpl` slot in + /// [`LIBFUNC_PROFILE`]. + /// 2. Points the executor's `cairo_native__profiler__profile_id` symbol at the new + /// trace ID, saving the previous value. + /// 3. Calls `run`. Per-statement samples accumulate in the slot via the runtime + /// `push_stmt` callback. + /// 4. Drains the slot, calls [`ProfilerImpl::get_profile`] with `program`, and hands + /// the resulting [`Profile`] to `on_profile`. + /// 5. A [`ProfilerGuard`] restores the previous trace ID — and removes the slot if + /// the success path didn't — on both success and unwind paths. + /// + /// `program` must be the Sierra program this executor was compiled from; it's used + /// by `get_profile` to map runtime libfunc IDs back to declarations. + /// + /// Profiling is intended to run single-threaded; concurrent calls would race on the + /// global `trace_id` symbol. + pub fn run_with_libfunc_profile( + &self, + program: &Arc, + selector: Felt, + args: &[Felt], + gas: u64, + builtin_costs: Option, + syscall_handler: H, + on_profile: F, + ) -> Result + where + H: StarknetSyscallHandler, + F: FnOnce(Profile), + { + static COUNTER: AtomicU64 = AtomicU64::new(0); + let counter = COUNTER.fetch_add(1, Ordering::Relaxed); + + LIBFUNC_PROFILE.lock().unwrap().insert(counter, ProfilerImpl::new()); + + // The pointer targets a global symbol in the executor's shared library; it lives + // for the executor's lifetime. Single-threaded profiling means no concurrent writer. + let trace_id_ptr = + self.find_symbol_ptr(ProfilerBinding::ProfileId.symbol()).unwrap().cast::(); + // SAFETY: see above. Read/write to a non-null, properly-aligned `*mut u64`. + let old_trace_id = unsafe { *trace_id_ptr }; + unsafe { + *trace_id_ptr = counter; + } + + // Restore on the success path AND on unwind. On success the caller drains the + // slot below; the guard's `remove` is then a no-op. + let _guard = ProfilerGuard { trace_id_ptr, old_trace_id, counter }; + + let result = self.run(selector, args, gas, builtin_costs, syscall_handler); + + let profiler = LIBFUNC_PROFILE.lock().unwrap().remove(&counter).unwrap(); + on_profile(profiler.get_profile(program)); + + result + } +} + +/// RAII cleanup for the profiler globals. Restores `*trace_id_ptr` and drops the +/// `LIBFUNC_PROFILE` slot at `counter` if it's still occupied. +struct ProfilerGuard { + trace_id_ptr: *mut u64, + old_trace_id: u64, + counter: u64, +} + +impl Drop for ProfilerGuard { + fn drop(&mut self) { + // SAFETY: same provenance as the construction site; single-threaded use. + unsafe { + *self.trace_id_ptr = self.old_trace_id; + } + // Tolerate a poisoned mutex silently — Drop must not panic. + if let Ok(mut profile) = LIBFUNC_PROFILE.lock() { + profile.remove(&self.counter); + } + } +} diff --git a/src/metadata/profiler.rs b/src/metadata/profiler.rs index feb96daa31..8e58b4b0e6 100644 --- a/src/metadata/profiler.rs +++ b/src/metadata/profiler.rs @@ -330,7 +330,7 @@ impl ProfilerMeta { /// Represents the entire profile of the execution. /// /// It maps the libfunc ID to a libfunc profile. -type Profile = HashMap; +pub type Profile = HashMap; /// Represents the profile data for a particular libfunc. #[derive(Default)]