Skip to content

Commit 74aab36

Browse files
avi-starkwareclaude
andcommitted
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. The profile-collection pattern is callback-driven so the per-call key (tx hash, etc.) stays out of cairo-native. - metadata::profiler::Profile is now pub (was a private type alias). - AotContractExecutor::run_with_libfunc_profile<H, F> (gated on with-libfunc-profiling, in new file src/executor/libfunc_profile.rs) wraps run with the bookkeeping the runtime needs: 1. Acquires a process-wide PROFILE_LOCK so concurrent profile calls serialize on the global trace-id symbol. The lock is recovered if poisoned. 2. Looks up the profile-id symbol before touching any global state. Absent symbol -> typed Error::UnexpectedValue rather than panic. 3. Allocates a trace ID, inserts a slot in LIBFUNC_PROFILE, points the global at the new ID. 4. Calls run; on success drains the slot and hands the Profile to on_profile. On failure the callback is not invoked. 5. A ProfilerGuard restores the previous trace ID on success and unwind. - ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) is a new variant that bundles an AOT executor with the Sierra program it was built from. From<AotWithProgram> is provided. - ContractExecutor::run dispatches the new variant via run_with_libfunc_profile with a no-op profile callback. - ContractExecutor::run_with_profile<H, F> is the profile-capturing counterpart of run; for non-AotWithProgram variants it falls through to run (callback never fires). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
1 parent 4f53f27 commit 74aab36

4 files changed

Lines changed: 230 additions & 3 deletions

File tree

src/executor.rs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,12 @@
33
//! This module provides methods to execute the programs, either via JIT or compiled ahead
44
//! of time. It also provides a cache to avoid recompiling previously compiled programs.
55
6+
#[cfg(feature = "with-libfunc-profiling")]
7+
pub use self::contract_executor::AotWithProgram;
8+
// Re-exported so libfunc-profiling consumers (e.g. blockifier) can refer to the
9+
// program type without adding a direct `cairo-lang-sierra` dependency.
10+
#[cfg(feature = "with-libfunc-profiling")]
11+
pub use cairo_lang_sierra::program::Program;
612
#[cfg(feature = "sierra-emu")]
713
pub use self::contract_executor::EmuContractInfo;
814
pub use self::{
@@ -46,6 +52,8 @@ mod aot;
4652
mod contract;
4753
mod contract_executor;
4854
mod jit;
55+
#[cfg(feature = "with-libfunc-profiling")]
56+
mod libfunc_profile;
4957

5058
#[cfg(target_arch = "aarch64")]
5159
global_asm!(include_str!("arch/aarch64.s"));

src/executor/contract_executor.rs

Lines changed: 72 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -5,21 +5,23 @@
55
//! `H: StarknetSyscallHandler` -- sierra-emu and cairo-native re-export the trait from
66
//! `cairo-native-syscalls`, so no adapter is needed.
77
8-
#[cfg(feature = "sierra-emu")]
8+
#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))]
99
use cairo_lang_sierra::program::Program;
1010
#[cfg(feature = "sierra-emu")]
1111
use cairo_lang_starknet_classes::compiler_version::VersionId;
1212
#[cfg(feature = "sierra-emu")]
1313
use cairo_lang_starknet_classes::contract_class::ContractEntryPoints;
1414
use starknet_types_core::felt::Felt;
15-
#[cfg(feature = "sierra-emu")]
15+
#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))]
1616
use std::sync::Arc;
1717

1818
#[cfg(feature = "sierra-emu")]
1919
use crate::error::Error;
2020
use crate::error::Result;
2121
use crate::execution_result::ContractExecutionResult;
2222
use crate::executor::AotContractExecutor;
23+
#[cfg(feature = "with-libfunc-profiling")]
24+
use crate::metadata::profiler::Profile;
2325
use crate::starknet::StarknetSyscallHandler;
2426
use crate::utils::BuiltinCosts;
2527

@@ -33,6 +35,8 @@ pub enum ContractExecutor {
3335
Aot(AotContractExecutor),
3436
#[cfg(feature = "sierra-emu")]
3537
Emu(EmuContractInfo),
38+
#[cfg(feature = "with-libfunc-profiling")]
39+
AotWithProgram(AotWithProgram),
3640
}
3741

3842
/// Inputs required to construct a `sierra_emu::VirtualMachine` for the `Emu` variant.
@@ -44,6 +48,16 @@ pub struct EmuContractInfo {
4448
pub sierra_version: VersionId,
4549
}
4650

51+
/// AOT executor paired with the Sierra program it was built from. Required by
52+
/// [`ContractExecutor::run_with_profile`] so libfunc samples can be resolved against
53+
/// the program's declarations.
54+
#[cfg(feature = "with-libfunc-profiling")]
55+
#[derive(Debug)]
56+
pub struct AotWithProgram {
57+
pub executor: AotContractExecutor,
58+
pub program: Arc<Program>,
59+
}
60+
4761
impl From<AotContractExecutor> for ContractExecutor {
4862
fn from(value: AotContractExecutor) -> Self {
4963
Self::Aot(value)
@@ -57,6 +71,13 @@ impl From<EmuContractInfo> for ContractExecutor {
5771
}
5872
}
5973

74+
#[cfg(feature = "with-libfunc-profiling")]
75+
impl From<AotWithProgram> for ContractExecutor {
76+
fn from(value: AotWithProgram) -> Self {
77+
Self::AotWithProgram(value)
78+
}
79+
}
80+
6081
impl ContractExecutor {
6182
/// Run the contract entry point identified by `selector`.
6283
///
@@ -105,6 +126,55 @@ impl ContractExecutor {
105126
builtin_stats: Default::default(),
106127
})
107128
}
129+
#[cfg(feature = "with-libfunc-profiling")]
130+
ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) => executor
131+
.run_with_libfunc_profile(
132+
program,
133+
selector,
134+
args,
135+
gas,
136+
builtin_costs,
137+
syscall_handler,
138+
// Profile is collected and dropped on this path. Use
139+
// `run_with_profile` to capture it.
140+
|_profile| {},
141+
),
142+
}
143+
}
144+
145+
/// Like [`Self::run`] but, for the `AotWithProgram` variant, hands the captured
146+
/// libfunc profile -- together with the `Arc<Program>` the executor was paired with --
147+
/// to `on_profile` after the call returns successfully. The program is included so
148+
/// callers don't have to keep their own copy around just to resolve libfunc samples.
149+
/// For other variants this is identical to `run` and `on_profile` is never invoked.
150+
#[cfg(feature = "with-libfunc-profiling")]
151+
pub fn run_with_profile<H, F>(
152+
&self,
153+
selector: Felt,
154+
args: &[Felt],
155+
gas: u64,
156+
builtin_costs: Option<BuiltinCosts>,
157+
syscall_handler: H,
158+
on_profile: F,
159+
) -> Result<ContractExecutionResult>
160+
where
161+
H: StarknetSyscallHandler,
162+
F: FnOnce(Profile, Arc<Program>),
163+
{
164+
match self {
165+
ContractExecutor::AotWithProgram(AotWithProgram { executor, program }) => {
166+
let program_for_cb = Arc::clone(program);
167+
executor.run_with_libfunc_profile(
168+
program,
169+
selector,
170+
args,
171+
gas,
172+
builtin_costs,
173+
syscall_handler,
174+
move |profile| on_profile(profile, program_for_cb),
175+
)
176+
}
177+
_ => self.run(selector, args, gas, builtin_costs, syscall_handler),
108178
}
109179
}
110180
}

src/executor/libfunc_profile.rs

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
//! Profiling-instrumented run wrapper around [`AotContractExecutor::run`].
2+
//!
3+
//! Available under the `with-libfunc-profiling` feature (gated at the `mod`
4+
//! declaration in `src/executor.rs`).
5+
6+
use std::sync::atomic::{AtomicU64, Ordering};
7+
use std::sync::{Arc, Mutex};
8+
9+
use cairo_lang_sierra::program::Program;
10+
use starknet_types_core::felt::Felt;
11+
12+
use crate::error::{Error, Result};
13+
use crate::execution_result::ContractExecutionResult;
14+
use crate::executor::AotContractExecutor;
15+
use crate::metadata::profiler::{Profile, ProfilerBinding, ProfilerImpl, LIBFUNC_PROFILE};
16+
use crate::starknet::StarknetSyscallHandler;
17+
use crate::utils::BuiltinCosts;
18+
19+
/// Process-wide lock that serializes calls into [`AotContractExecutor::run_with_libfunc_profile`].
20+
/// The profiler hot-swaps a process-global symbol (`cairo_native__profiler__profile_id`);
21+
/// concurrent callers would race on that write and on the [`LIBFUNC_PROFILE`] slot bookkeeping.
22+
static PROFILE_LOCK: Mutex<()> = Mutex::new(());
23+
24+
impl AotContractExecutor {
25+
/// Run the entrypoint with libfunc-level profiling instrumentation.
26+
///
27+
/// Wraps [`AotContractExecutor::run`] with the bookkeeping the
28+
/// `with-libfunc-profiling` runtime needs:
29+
///
30+
/// 1. Acquires [`PROFILE_LOCK`] so concurrent profile calls serialize on the
31+
/// global trace-id symbol. The lock is recovered if poisoned.
32+
/// 2. Looks up the executor's `cairo_native__profiler__profile_id` symbol. If
33+
/// absent (the .so was compiled without profiling instrumentation) the call
34+
/// returns an error before touching any global state.
35+
/// 3. Allocates a unique trace ID and inserts an empty `ProfilerImpl` slot in
36+
/// [`LIBFUNC_PROFILE`]; points the profile-id symbol at the new ID, saving
37+
/// the previous value.
38+
/// 4. Calls `run`. Per-statement samples accumulate in the slot via the runtime
39+
/// `push_stmt` callback.
40+
/// 5. Drains the slot. On success (and only on success) hands the resulting
41+
/// [`Profile`] to `on_profile`; on failure the callback is not invoked
42+
/// (partial profiles aren't meaningful).
43+
/// 6. A [`ProfilerGuard`] restores the previous trace ID and clears the slot on
44+
/// both the success and unwind paths.
45+
///
46+
/// `program` must be the Sierra program this executor was compiled from; it's used
47+
/// by `get_profile` to map runtime libfunc IDs back to declarations.
48+
#[allow(clippy::too_many_arguments)]
49+
pub fn run_with_libfunc_profile<H, F>(
50+
&self,
51+
program: &Arc<Program>,
52+
selector: Felt,
53+
args: &[Felt],
54+
gas: u64,
55+
builtin_costs: Option<BuiltinCosts>,
56+
syscall_handler: H,
57+
on_profile: F,
58+
) -> Result<ContractExecutionResult>
59+
where
60+
H: StarknetSyscallHandler,
61+
F: FnOnce(Profile),
62+
{
63+
// Serialize against concurrent profile calls. Recover from a poisoned lock --
64+
// we don't have invariants on the protected state itself; the lock only gates
65+
// access to the global trace-id symbol.
66+
let _profile_lock = PROFILE_LOCK.lock().unwrap_or_else(|e| e.into_inner());
67+
68+
// Look up the profile-id symbol before touching any global state. If the
69+
// executor wasn't compiled with libfunc-profiling instrumentation, the
70+
// symbol is absent -- return a typed error rather than panicking.
71+
let trace_id_ptr = self
72+
.find_symbol_ptr(ProfilerBinding::ProfileId.symbol())
73+
.ok_or_else(|| {
74+
Error::UnexpectedValue(format!(
75+
"AOT executor missing libfunc-profiling symbol `{}`; \
76+
was the program compiled with libfunc-profiling enabled?",
77+
ProfilerBinding::ProfileId.symbol()
78+
))
79+
})?
80+
.cast::<u64>();
81+
82+
static COUNTER: AtomicU64 = AtomicU64::new(0);
83+
let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
84+
85+
LIBFUNC_PROFILE
86+
.lock()
87+
.unwrap_or_else(|e| e.into_inner())
88+
.insert(counter, ProfilerImpl::new());
89+
90+
// SAFETY: the pointer targets a memref-global emitted into the executor's
91+
// shared library; the executor outlives the call. `PROFILE_LOCK` serializes
92+
// us against any other writer, and the JIT/AOT code reads through the same
93+
// address. Reads/writes are aligned `u64`s.
94+
let old_trace_id = unsafe { *trace_id_ptr };
95+
unsafe {
96+
*trace_id_ptr = counter;
97+
}
98+
99+
let _guard = ProfilerGuard {
100+
trace_id_ptr,
101+
old_trace_id,
102+
counter,
103+
};
104+
105+
let result = self.run(selector, args, gas, builtin_costs, syscall_handler);
106+
107+
// Drain the slot. `ProfilerGuard::drop` would also remove it; doing it here
108+
// means we hold the lock for the shortest time and can hand the profile to
109+
// the callback. Tolerate a poisoned mutex (we'd lose the profile, not state).
110+
let drained = LIBFUNC_PROFILE
111+
.lock()
112+
.unwrap_or_else(|e| e.into_inner())
113+
.remove(&counter);
114+
115+
// Only call the user's callback when `run` succeeded -- a partial profile
116+
// captured against an aborted execution wouldn't be meaningful.
117+
if let (Some(profiler), Ok(_)) = (drained, &result) {
118+
on_profile(profiler.get_profile(program));
119+
}
120+
121+
result
122+
}
123+
}
124+
125+
/// RAII cleanup for the profiler globals. Restores `*trace_id_ptr` on success or
126+
/// unwind. The [`LIBFUNC_PROFILE`] slot at `counter` is normally drained on the
127+
/// success path; this guard removes it if it's still occupied (panic case).
128+
struct ProfilerGuard {
129+
trace_id_ptr: *mut u64,
130+
old_trace_id: u64,
131+
counter: u64,
132+
}
133+
134+
impl Drop for ProfilerGuard {
135+
fn drop(&mut self) {
136+
// SAFETY: same provenance as the construction site. `PROFILE_LOCK` is held
137+
// by the enclosing scope (still in flight while we drop) so no other thread
138+
// races us.
139+
unsafe {
140+
*self.trace_id_ptr = self.old_trace_id;
141+
}
142+
// Tolerate a poisoned mutex silently -- Drop must not panic. Slot leak on
143+
// poison is intentional and matches the behavior of other Drop impls in
144+
// this crate; the alternative (panic in Drop) is worse.
145+
if let Ok(mut profile) = LIBFUNC_PROFILE.lock() {
146+
profile.remove(&self.counter);
147+
}
148+
}
149+
}

src/metadata/profiler.rs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -330,7 +330,7 @@ impl ProfilerMeta {
330330
/// Represents the entire profile of the execution.
331331
///
332332
/// It maps the libfunc ID to a libfunc profile.
333-
type Profile = HashMap<ConcreteLibfuncId, LibfuncProfileData>;
333+
pub type Profile = HashMap<ConcreteLibfuncId, LibfuncProfileData>;
334334

335335
/// Represents the profile data for a particular libfunc.
336336
#[derive(Default)]

0 commit comments

Comments
 (0)