Skip to content

Optimize fast_locals and atomic ordering #7289

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
youknowone merged 5 commits into from
Mar 1, 2026
Merged

Optimize fast_locals and atomic ordering #7289

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
6895ddb
Relax RefCount atomic ordering from SeqCst to Arc pattern
youknowone Mar 1, 2026
89703b6
Reuse existing Vec via prepend_arg in execute_call
youknowone Mar 1, 2026
6a9be9f
Skip downcast_ref checks in invoke when tracing is disabled
youknowone Mar 1, 2026
958853f
Replace fastlocals PyMutex with UnsafeCell-based FastLocals
youknowone Mar 1, 2026
516ff0f
Auto-format: cargo fmt --all
github-actions[bot] Mar 1, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 18 additions & 14 deletions crates/common/src/refcount.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -92,26 +92,26 @@ impl RefCount {
/// Get current strong count
#[inline]
pub fn get(&self) -> usize {
State::from_raw(self.state.load(Ordering::SeqCst)).strong() as usize
State::from_raw(self.state.load(Ordering::Relaxed)).strong() as usize
}

/// Increment strong count
#[inline]
pub fn inc(&self) {
let val = State::from_raw(self.state.fetch_add(COUNT, Ordering::SeqCst));
let val = State::from_raw(self.state.fetch_add(COUNT, Ordering::Relaxed));
if val.destructed() || (val.strong() as usize) > STRONG - 1 {
refcount_overflow();
}
if val.strong() == 0 {
// The previous fetch_add created a permission to run decrement again
self.state.fetch_add(COUNT, Ordering::SeqCst);
self.state.fetch_add(COUNT, Ordering::Relaxed);
}
}

#[inline]
pub fn inc_by(&self, n: usize) {
debug_assert!(n <= STRONG);
let val = State::from_raw(self.state.fetch_add(n * COUNT, Ordering::SeqCst));
let val = State::from_raw(self.state.fetch_add(n * COUNT, Ordering::Relaxed));
if val.destructed() || (val.strong() as usize) > STRONG - n {
refcount_overflow();
}
Expand All @@ -121,7 +121,7 @@ impl RefCount {
#[inline]
#[must_use]
pub fn safe_inc(&self) -> bool {
let mut old = State::from_raw(self.state.load(Ordering::SeqCst));
let mut old = State::from_raw(self.state.load(Ordering::Relaxed));
loop {
if old.destructed() {
return false;
Expand All @@ -130,11 +130,11 @@ impl RefCount {
refcount_overflow();
}
let new_state = old.add_strong(1);
match self.state.compare_exchange(
match self.state.compare_exchange_weak(
old.as_raw(),
new_state.as_raw(),
Ordering::SeqCst,
Ordering::SeqCst,
Ordering::Relaxed,
Ordering::Relaxed,
) {
Ok(_) => return true,
Err(curr) => old = State::from_raw(curr),
Expand All @@ -146,27 +146,31 @@ impl RefCount {
#[inline]
#[must_use]
pub fn dec(&self) -> bool {
let old = State::from_raw(self.state.fetch_sub(COUNT, Ordering::SeqCst));
let old = State::from_raw(self.state.fetch_sub(COUNT, Ordering::Release));

// LEAKED objects never reach 0
if old.leaked() {
return false;
}

old.strong() == 1
if old.strong() == 1 {
core::sync::atomic::fence(Ordering::Acquire);
return true;
}
false
}

/// Mark this object as leaked (interned). It will never be deallocated.
pub fn leak(&self) {
debug_assert!(!self.is_leaked());
let mut old = State::from_raw(self.state.load(Ordering::SeqCst));
let mut old = State::from_raw(self.state.load(Ordering::Relaxed));
loop {
let new_state = old.with_leaked(true);
match self.state.compare_exchange(
match self.state.compare_exchange_weak(
old.as_raw(),
new_state.as_raw(),
Ordering::SeqCst,
Ordering::SeqCst,
Ordering::AcqRel,
Ordering::Relaxed,
) {
Ok(_) => return,
Err(curr) => old = State::from_raw(curr),
Expand Down
3 changes: 2 additions & 1 deletion crates/vm/src/builtins/frame.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -690,8 +690,9 @@ impl Py<Frame> {
}

// Clear fastlocals
// SAFETY: Frame is not executing (detached or stopped).
{
let mut fastlocals = self.fastlocals.lock();
let fastlocals = unsafe { self.fastlocals.borrow_mut() };
for slot in fastlocals.iter_mut() {
*slot = None;
}
Expand Down
3 changes: 2 additions & 1 deletion crates/vm/src/builtins/function.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -224,7 +224,8 @@ impl PyFunction {
// See also: PyEval_EvalCodeWithName in cpython:
// https://github.com/python/cpython/blob/main/Python/ceval.c#L3681

let mut fastlocals = frame.fastlocals.lock();
// SAFETY: Frame was just created and not yet executing.
let fastlocals = unsafe { frame.fastlocals.borrow_mut() };

let mut args_iter = func_args.args.into_iter();

Expand Down
7 changes: 3 additions & 4 deletions crates/vm/src/builtins/super.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -85,10 +85,9 @@ impl Initializer for PySuper {
if frame.code.arg_count == 0 {
return Err(vm.new_runtime_error("super(): no arguments"));
}
// Lock fastlocals briefly to get arg[0], then drop the guard
// before calling get_cell_contents (which also locks fastlocals).
let first_arg = frame.fastlocals.lock()[0].clone();
let obj = first_arg
// SAFETY: Frame is current and not concurrently mutated.
let obj = unsafe { frame.fastlocals.borrow() }[0]
.clone()
.or_else(|| {
if let Some(cell2arg) = frame.code.cell2arg.as_deref() {
cell2arg[..frame.code.cellvars.len()]
Expand Down
Loading
Loading