Skip to content

Commit 18cfe2d

Browse files
avi-starkwareclaude
andcommitted
add run_with_libfunc_profile + AotWithProgram pairing
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
1 parent f62b57a commit 18cfe2d

4 files changed

Lines changed: 210 additions & 2 deletions

File tree

binaries/starknet-native-compile/Cargo.toml

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,13 @@ edition.workspace = true
77
license.workspace = true
88
repository.workspace = true
99

10+
[features]
11+
# Builds an instrumented compiler whose emitted `.so` carries the libfunc-profiling
12+
# symbols that `cairo_native::AotContractExecutor::run_with_libfunc_profile` looks up.
13+
# Install with `cargo install ... --features with-libfunc-profiling` to profile native
14+
# execution; a binary built without it produces `.so`s with no profiler symbol.
15+
with-libfunc-profiling = ["cairo-native/with-libfunc-profiling"]
16+
1017
[dependencies]
1118
anyhow.workspace = true
1219
cairo-lang-sierra.workspace = true

src/executor.rs

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,8 @@
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::libfunc_profile::AotWithProgram;
68
#[cfg(feature = "sierra-emu")]
79
pub use self::emu_contract_executor::EmuContractExecutor;
810
pub use self::{aot::AotNativeExecutor, contract::AotContractExecutor, jit::JitNativeExecutor};
@@ -23,6 +25,11 @@ use crate::{
2325
values::Value,
2426
};
2527
use bumpalo::Bump;
28+
// Profiling and sierra-emu consumers (e.g. blockifier) only ever hold the Sierra
29+
// program as a shared handle. Export the `Arc` alias so they can name it without
30+
// taking a direct `cairo-lang-sierra` dependency.
31+
#[cfg(any(feature = "sierra-emu", feature = "with-libfunc-profiling"))]
32+
pub type ArcProgram = std::sync::Arc<cairo_lang_sierra::program::Program>;
2633
use cairo_lang_sierra::{
2734
extensions::{
2835
circuit::CircuitTypeConcrete,
@@ -44,6 +51,8 @@ mod contract;
4451
#[cfg(feature = "sierra-emu")]
4552
mod emu_contract_executor;
4653
mod jit;
54+
#[cfg(feature = "with-libfunc-profiling")]
55+
mod libfunc_profile;
4756

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

src/executor/libfunc_profile.rs

Lines changed: 188 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,188 @@
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::cell::Cell;
7+
use std::sync::atomic::{AtomicU64, Ordering};
8+
use std::sync::{Arc, Mutex};
9+
10+
use cairo_lang_sierra::program::Program;
11+
use starknet_types_core::felt::Felt;
12+
13+
use crate::error::{Error, Result};
14+
use crate::execution_result::ContractExecutionResult;
15+
use crate::executor::{AotContractExecutor, ArcProgram};
16+
use crate::metadata::profiler::{Profile, ProfilerBinding, ProfilerImpl, LIBFUNC_PROFILE};
17+
use crate::starknet::StarknetSyscallHandler;
18+
use crate::utils::BuiltinCosts;
19+
20+
/// An AOT executor paired with the Sierra program it was compiled from. The program is
21+
/// needed to resolve libfunc profiling samples against the program's declarations, so
22+
/// pairing the two spares callers from keeping their own copy around.
23+
#[derive(Debug)]
24+
pub struct AotWithProgram {
25+
pub executor: AotContractExecutor,
26+
pub program: ArcProgram,
27+
}
28+
29+
/// Process-wide lock that serializes *top-level* profiled runs across threads. The
30+
/// profiler hot-swaps a process-global symbol (`cairo_native__profiler__profile_id`);
31+
/// a concurrent thread would race on that write and on the [`LIBFUNC_PROFILE`] slot
32+
/// bookkeeping. It is acquired only by the outermost profiled frame on each thread
33+
/// (see [`PROFILE_DEPTH`]); nested same-thread calls -- a profiled contract invoking
34+
/// another contract -- re-enter without re-locking, which would otherwise deadlock
35+
/// this non-reentrant mutex.
36+
static PROFILE_LOCK: Mutex<()> = Mutex::new(());
37+
38+
thread_local! {
39+
/// Nesting depth of profiled runs on the current thread. Only the outermost frame
40+
/// (depth 0) takes [`PROFILE_LOCK`]; deeper frames rely on the lock the outer frame
41+
/// already holds. The per-call `old_trace_id` save/restore keeps the global trace-id
42+
/// symbol correct across nesting without any additional locking.
43+
static PROFILE_DEPTH: Cell<usize> = const { Cell::new(0) };
44+
}
45+
46+
impl AotContractExecutor {
47+
/// Run the entrypoint with libfunc-level profiling instrumentation.
48+
///
49+
/// Wraps [`AotContractExecutor::run`] with the bookkeeping the
50+
/// `with-libfunc-profiling` runtime needs:
51+
///
52+
/// 1. Acquires [`PROFILE_LOCK`] so concurrent profile calls serialize on the
53+
/// global trace-id symbol. The lock is recovered if poisoned.
54+
/// 2. Looks up the executor's `cairo_native__profiler__profile_id` symbol. If
55+
/// absent (the .so was compiled without profiling instrumentation) the call
56+
/// returns an error before touching any global state.
57+
/// 3. Allocates a unique trace ID and inserts an empty `ProfilerImpl` slot in
58+
/// [`LIBFUNC_PROFILE`]; points the profile-id symbol at the new ID, saving
59+
/// the previous value.
60+
/// 4. Calls `run`. Per-statement samples accumulate in the slot via the runtime
61+
/// `push_stmt` callback.
62+
/// 5. Drains the slot. On success (and only on success) hands the resulting
63+
/// [`Profile`] to `on_profile`; on failure the callback is not invoked
64+
/// (partial profiles aren't meaningful).
65+
/// 6. A [`ProfilerGuard`] restores the previous trace ID and clears the slot on
66+
/// both the success and unwind paths.
67+
///
68+
/// `program` must be the Sierra program this executor was compiled from; it's used
69+
/// by `get_profile` to map runtime libfunc IDs back to declarations.
70+
#[allow(clippy::too_many_arguments)]
71+
pub fn run_with_libfunc_profile<H, F>(
72+
&self,
73+
program: &Arc<Program>,
74+
selector: Felt,
75+
args: &[Felt],
76+
gas: u64,
77+
builtin_costs: Option<BuiltinCosts>,
78+
syscall_handler: H,
79+
on_profile: F,
80+
) -> Result<ContractExecutionResult>
81+
where
82+
H: StarknetSyscallHandler,
83+
F: FnOnce(Profile),
84+
{
85+
// Acquire the cross-thread lock only at the outermost profiled frame on this
86+
// thread. A profiled contract that calls another contract re-enters this
87+
// function on the same thread; re-locking the non-reentrant `PROFILE_LOCK`
88+
// there would self-deadlock, so nested frames inherit the outer frame's lock.
89+
// Recover from a poisoned lock -- it only gates access to the global trace-id
90+
// symbol, on which we hold no data invariants.
91+
let _profile_lock = PROFILE_DEPTH
92+
.with(|depth| depth.get() == 0)
93+
.then(|| PROFILE_LOCK.lock().unwrap_or_else(|e| e.into_inner()));
94+
PROFILE_DEPTH.with(|depth| depth.set(depth.get() + 1));
95+
let _depth_guard = ProfileDepthGuard;
96+
97+
// Look up the profile-id symbol before touching any global state. If the
98+
// executor wasn't compiled with libfunc-profiling instrumentation, the
99+
// symbol is absent -- return a typed error rather than panicking.
100+
let trace_id_ptr = self
101+
.find_symbol_ptr(ProfilerBinding::ProfileId.symbol())
102+
.ok_or_else(|| {
103+
Error::UnexpectedValue(format!(
104+
"AOT executor missing libfunc-profiling symbol `{}`; \
105+
was the program compiled with libfunc-profiling enabled?",
106+
ProfilerBinding::ProfileId.symbol()
107+
))
108+
})?
109+
.cast::<u64>();
110+
111+
static COUNTER: AtomicU64 = AtomicU64::new(0);
112+
let counter = COUNTER.fetch_add(1, Ordering::Relaxed);
113+
114+
LIBFUNC_PROFILE
115+
.lock()
116+
.unwrap_or_else(|e| e.into_inner())
117+
.insert(counter, ProfilerImpl::new());
118+
119+
// SAFETY: the pointer targets a memref-global emitted into the executor's
120+
// shared library; the executor outlives the call. `PROFILE_LOCK` serializes
121+
// us against any other writer, and the JIT/AOT code reads through the same
122+
// address. Reads/writes are aligned `u64`s.
123+
let old_trace_id = unsafe { *trace_id_ptr };
124+
unsafe {
125+
*trace_id_ptr = counter;
126+
}
127+
128+
let _guard = ProfilerGuard {
129+
trace_id_ptr,
130+
old_trace_id,
131+
counter,
132+
};
133+
134+
let result = self.run(selector, args, gas, builtin_costs, syscall_handler);
135+
136+
// Drain the slot. `ProfilerGuard::drop` would also remove it; doing it here
137+
// means we hold the lock for the shortest time and can hand the profile to
138+
// the callback. Tolerate a poisoned mutex (we'd lose the profile, not state).
139+
let drained = LIBFUNC_PROFILE
140+
.lock()
141+
.unwrap_or_else(|e| e.into_inner())
142+
.remove(&counter);
143+
144+
// Only call the user's callback when `run` succeeded -- a partial profile
145+
// captured against an aborted execution wouldn't be meaningful.
146+
if let (Some(profiler), Ok(_)) = (drained, &result) {
147+
on_profile(profiler.get_profile(program));
148+
}
149+
150+
result
151+
}
152+
}
153+
154+
/// RAII cleanup for the profiler globals. Restores `*trace_id_ptr` on success or
155+
/// unwind. The [`LIBFUNC_PROFILE`] slot at `counter` is normally drained on the
156+
/// success path; this guard removes it if it's still occupied (panic case).
157+
struct ProfilerGuard {
158+
trace_id_ptr: *mut u64,
159+
old_trace_id: u64,
160+
counter: u64,
161+
}
162+
163+
impl Drop for ProfilerGuard {
164+
fn drop(&mut self) {
165+
// SAFETY: same provenance as the construction site. `PROFILE_LOCK` is held
166+
// by the enclosing scope (still in flight while we drop) so no other thread
167+
// races us.
168+
unsafe {
169+
*self.trace_id_ptr = self.old_trace_id;
170+
}
171+
// Tolerate a poisoned mutex silently -- Drop must not panic. Slot leak on
172+
// poison is intentional and matches the behavior of other Drop impls in
173+
// this crate; the alternative (panic in Drop) is worse.
174+
if let Ok(mut profile) = LIBFUNC_PROFILE.lock() {
175+
profile.remove(&self.counter);
176+
}
177+
}
178+
}
179+
180+
/// Decrements [`PROFILE_DEPTH`] on every exit path (including unwind) so the
181+
/// outermost-frame lock bookkeeping stays correct even if the profiled run panics.
182+
struct ProfileDepthGuard;
183+
184+
impl Drop for ProfileDepthGuard {
185+
fn drop(&mut self) {
186+
PROFILE_DEPTH.with(|depth| depth.set(depth.get().saturating_sub(1)));
187+
}
188+
}

src/metadata/profiler.rs

Lines changed: 6 additions & 2 deletions
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)]
@@ -362,7 +362,11 @@ impl ProfilerImpl {
362362

363363
// Push a profiler frame
364364
pub extern "C" fn push_stmt(profile_id: u64, statement_idx: u64, tick_delta: u64) {
365-
let mut profiler = LIBFUNC_PROFILE.lock().unwrap();
365+
// Invoked directly from compiled/JIT'd Cairo code across the C ABI: a panic
366+
// here (e.g. `.unwrap()` on a poisoned mutex) would unwind into machine code
367+
// that cannot handle it. Recover from poison instead of panicking, matching
368+
// the lock handling in `run_with_libfunc_profile`.
369+
let mut profiler = LIBFUNC_PROFILE.lock().unwrap_or_else(|e| e.into_inner());
366370

367371
let Some(profiler) = profiler.get_mut(&profile_id) else {
368372
eprintln!("Could not find libfunc profiler!");

0 commit comments

Comments
 (0)