Skip to content

Commit d68d02b

Browse files
youknowoneyouknowone
committed
vm: reuse cleared frame blocks and shorten interpreter hot paths
- `datastack` remembers the most recently popped frame block. `push_frame` reports an exact LIFO reuse, and `setup_datastack_frame` then skips zero-filling localsplus. - Small-int loads push the context's cached int as a borrowed stack reference instead of taking a new one. - Calls to jitted functions go straight to `execute_call_vectorcall`. - Binary-op specialization reads ints through the new `PyInt::try_to_i64_fast` instead of the generic primitive conversion, and `try_to_bool` moves its non-bool path into a `#[cold]` helper. - Dict caches read an entry through a keys-version stamp (`get_index_if_keys_version`) rather than an entry-index hint. - Frame publishing caches a pointer to `ThreadSlot::top_iframe` in a thread-local `Cell` instead of borrowing `CURRENT_THREAD_SLOT`.
1 parent 82a7be9 commit d68d02b

10 files changed

Lines changed: 371 additions & 206 deletions

File tree

‎crates/vm/src/builtins/bool.rs‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -34,13 +34,20 @@ impl<'a> TryFromBorrowedObject<'a> for bool {
3434

3535
impl PyObjectRef {
3636
/// Convert Python bool into Rust bool.
37+
#[inline(always)]
3738
pub fn try_to_bool(self, vm: &VirtualMachine) -> PyResult<bool> {
3839
if self.is(&vm.ctx.true_value) {
3940
return Ok(true);
4041
} else if self.is(&vm.ctx.false_value) {
4142
return Ok(false);
4243
}
4344

45+
self.try_to_bool_slow(vm)
46+
}
47+
48+
#[cold]
49+
#[inline(never)]
50+
fn try_to_bool_slow(self, vm: &VirtualMachine) -> PyResult<bool> {
4451
let slots = &self.class().slots;
4552

4653
// 1. Try nb_bool slot first

‎crates/vm/src/builtins/dict.rs‎

Lines changed: 8 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -114,11 +114,6 @@ impl PyDict {
114114
&self.entries
115115
}
116116

117-
/// Monotonically increasing version for mutation tracking.
118-
pub(crate) fn version(&self) -> u64 {
119-
self.entries.version()
120-
}
121-
122117
/// Returns all keys as a Vec, atomically under a single read lock.
123118
/// Thread-safe: prevents "dictionary changed size during iteration" errors.
124119
pub fn keys_vec(&self) -> Vec<PyObjectRef> {
@@ -817,18 +812,15 @@ impl Py<PyDict> {
817812
}
818813
}
819814

820-
/// Fast lookup using a cached entry index hint.
821-
pub(crate) fn get_item_opt_hint<K: DictKey + ?Sized>(
815+
/// Read a cached exact-dict entry after validating its key-layout stamp.
816+
#[inline]
817+
pub(crate) fn get_item_by_index_and_keys_version(
822818
&self,
823-
key: &K,
824-
hint: u16,
825-
vm: &VirtualMachine,
826-
) -> PyResult<Option<PyObjectRef>> {
827-
if self.exact_dict(vm) {
828-
self.entries.get_hint(vm, key, usize::from(hint))
829-
} else {
830-
self.get_item_opt(key, vm)
831-
}
819+
version: u16,
820+
index: u16,
821+
) -> Option<PyObjectRef> {
822+
self.entries
823+
.get_index_if_keys_version(u32::from(version), usize::from(index))
832824
}
833825

834826
/// Lookup trying a cached entry index hint first.

‎crates/vm/src/builtins/function.rs‎

Lines changed: 28 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -550,6 +550,20 @@ impl Py<PyFunction> {
550550
self.code.flags.contains(bytecode::CodeFlags::OPTIMIZED)
551551
}
552552

553+
/// Whether this function currently has native JIT code. Adaptive Python
554+
/// call specializations must yield to that entry point.
555+
#[inline]
556+
pub(crate) fn is_jitted(&self) -> bool {
557+
#[cfg(feature = "jit")]
558+
{
559+
self.jitted_code.lock().is_some()
560+
}
561+
#[cfg(not(feature = "jit"))]
562+
{
563+
false
564+
}
565+
}
566+
553567
pub fn invoke_with_locals(
554568
&self,
555569
func_args: FuncArgs,
@@ -643,8 +657,8 @@ impl Py<PyFunction> {
643657
.and_then(|()| vm.run_frame_fast(iframe));
644658
// Release data stack memory — must happen on both success and error.
645659
unsafe {
646-
if let Some(base) = iframe.release_datastack_frame() {
647-
vm.datastack_pop(base);
660+
if let Some((base, size)) = iframe.release_datastack_frame() {
661+
vm.datastack_pop_frame(base, size);
648662
}
649663
}
650664
result
@@ -823,8 +837,8 @@ impl Py<PyFunction> {
823837

824838
let result = vm.run_frame_fast(iframe);
825839
unsafe {
826-
if let Some(base) = iframe.release_datastack_frame() {
827-
vm.datastack_pop(base);
840+
if let Some((base, size)) = iframe.release_datastack_frame() {
841+
vm.datastack_pop_frame(base, size);
828842
}
829843
}
830844
result
@@ -1619,6 +1633,16 @@ pub(crate) fn vectorcall_function(
16191633
let code: &Py<PyCode> = &zelf.code;
16201634

16211635
let has_kwargs = kwnames.is_some_and(|kw| !kw.is_empty());
1636+
if zelf.is_jitted() {
1637+
let func_args = if has_kwargs {
1638+
FuncArgs::from_vectorcall(&args, nargs, kwnames)
1639+
} else {
1640+
args.truncate(nargs);
1641+
FuncArgs::from(args)
1642+
};
1643+
return zelf.invoke(func_args, vm);
1644+
}
1645+
16221646
let is_simple = !has_kwargs
16231647
&& code.flags.contains(bytecode::CodeFlags::OPTIMIZED)
16241648
&& !code.flags.contains(bytecode::CodeFlags::VARARGS)

‎crates/vm/src/builtins/int.rs‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -305,6 +305,22 @@ impl PyInt {
305305
&self.value
306306
}
307307

308+
/// Extract the inline magnitude without the generic primitive-conversion path.
309+
#[inline(always)]
310+
pub(crate) fn try_to_i64_fast(&self) -> Option<i64> {
311+
let bits = self.value.bits();
312+
if bits > i64::BITS as u64 {
313+
return None;
314+
}
315+
let magnitude = self.value.iter_u64_digits().next().unwrap_or(0);
316+
let signed_magnitude = i64::try_from(magnitude).ok();
317+
match self.value.sign() {
318+
Sign::Minus if magnitude == 1u64 << 63 => Some(i64::MIN),
319+
Sign::Minus => signed_magnitude.map(|value| -value),
320+
Sign::NoSign | Sign::Plus => signed_magnitude,
321+
}
322+
}
323+
308324
/// Fast decimal string conversion, using i64 path when possible.
309325
#[inline]
310326
#[must_use]

‎crates/vm/src/datastack.rs‎

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -61,6 +61,9 @@ pub struct DataStack {
6161
top: *mut u8,
6262
/// End of usable space in the current chunk.
6363
limit: *mut u8,
64+
/// Most recently popped full-frame allocation whose localsplus slots were
65+
/// cleared before the pop. An exact LIFO reuse can skip zero-filling them.
66+
reusable_frame: Option<(*mut u8, usize)>,
6467
}
6568

6669
impl DataStack {
@@ -73,7 +76,12 @@ impl DataStack {
7376
// Skip one ALIGN-sized slot in the root chunk so that `pop()` never
7477
// frees it (`push_chunk` convention).
7578
let top = unsafe { top.add(ALIGN) };
76-
Self { chunk, top, limit }
79+
Self {
80+
chunk,
81+
top,
82+
limit,
83+
reusable_frame: None,
84+
}
7785
}
7886

7987
/// Check if the current chunk has at least `size` bytes available.
@@ -91,6 +99,24 @@ impl DataStack {
9199
/// (LIFO order).
92100
#[inline(always)]
93101
pub fn push(&mut self, size: usize) -> *mut u8 {
102+
self.reusable_frame = None;
103+
self.push_inner(size)
104+
}
105+
106+
/// Allocate a full interpreter frame and report whether it exactly reuses
107+
/// a just-cleared frame block.
108+
#[inline(always)]
109+
pub fn push_frame(&mut self, size: usize) -> (*mut u8, bool) {
110+
let aligned_size = (size + ALIGN - 1) & !(ALIGN - 1);
111+
let reusable_frame = self.reusable_frame.take();
112+
let ptr = self.push_inner(size);
113+
let reused =
114+
reusable_frame.is_some_and(|(base, old_size)| base == ptr && old_size == aligned_size);
115+
(ptr, reused)
116+
}
117+
118+
#[inline(always)]
119+
fn push_inner(&mut self, size: usize) -> *mut u8 {
94120
let aligned_size = (size + ALIGN - 1) & !(ALIGN - 1);
95121
unsafe {
96122
if self.top.add(aligned_size) <= self.limit {
@@ -138,6 +164,25 @@ impl DataStack {
138164
/// and all allocations made after it must already have been popped.
139165
#[inline(always)]
140166
pub unsafe fn pop(&mut self, base: *mut u8) {
167+
self.reusable_frame = None;
168+
unsafe { self.pop_inner(base) };
169+
}
170+
171+
/// Pop a full frame whose localsplus slots have already been cleared.
172+
///
173+
/// # Safety
174+
/// `base` and `size` must describe the most recent allocation returned by
175+
/// `push_frame`, every later allocation must already be popped, and all
176+
/// localsplus slots in the frame must have been cleared.
177+
#[inline(always)]
178+
pub unsafe fn pop_frame(&mut self, base: *mut u8, size: usize) {
179+
unsafe { self.pop_inner(base) };
180+
let aligned_size = (size + ALIGN - 1) & !(ALIGN - 1);
181+
self.reusable_frame = Some((base, aligned_size));
182+
}
183+
184+
#[inline(always)]
185+
unsafe fn pop_inner(&mut self, base: *mut u8) {
141186
debug_assert!(!base.is_null());
142187
if self.is_in_current_chunk(base) {
143188
// Common case: base is within the current chunk.

‎crates/vm/src/dict_inner.rs‎

Lines changed: 17 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,7 @@ use alloc::fmt;
2020
use core::mem::size_of;
2121
use core::ops::ControlFlow;
2222
use core::sync::atomic::{
23-
AtomicU32, AtomicU64,
23+
AtomicU32,
2424
Ordering::{AcqRel, Acquire, Relaxed, Release},
2525
};
2626
use num_traits::ToPrimitive;
@@ -39,7 +39,6 @@ type EntryIndex = usize;
3939

4040
pub(crate) struct Dict<T = PyObjectRef> {
4141
inner: PyRwLock<DictInner<T>>,
42-
version: AtomicU64,
4342
/// Keys-version stamp, assigned lazily by `assign_keys_version` and
4443
/// reset to 0 whenever the key set changes. Value-only updates keep it.
4544
///
@@ -202,7 +201,6 @@ impl<T: Clone> Clone for Dict<T> {
202201
fn clone(&self) -> Self {
203202
Self {
204203
inner: PyRwLock::new(self.inner.read().clone()),
205-
version: AtomicU64::new(0),
206204
keys_version: AtomicU32::new(0),
207205
}
208206
}
@@ -217,7 +215,6 @@ impl<T> Default for Dict<T> {
217215
indices: vec![IndexEntry::FREE; 8],
218216
entries: Vec::new(),
219217
}),
220-
version: AtomicU64::new(0),
221218
keys_version: AtomicU32::new(0),
222219
}
223220
}
@@ -362,16 +359,6 @@ impl<T> DictInner<T> {
362359
type PopInnerResult<T> = ControlFlow<Option<DictEntry<T>>>;
363360

364361
impl<T: Clone> Dict<T> {
365-
/// Monotonically increasing version counter for mutation tracking.
366-
pub(crate) fn version(&self) -> u64 {
367-
self.version.load(Acquire)
368-
}
369-
370-
/// Bump the version counter after any mutation.
371-
fn bump_version(&self) {
372-
self.version.fetch_add(1, Release);
373-
}
374-
375362
/// Current keys-version stamp, or 0 if none has been assigned since the
376363
/// last key-set change. Equal nonzero stamps guarantee an unchanged key
377364
/// set (values may differ).
@@ -500,7 +487,6 @@ impl<T: Clone> Dict<T> {
500487
)]
501488
if entry.index == index_index {
502489
let removed = core::mem::replace(&mut entry.value, value);
503-
self.bump_version();
504490
// defer dec RC
505491
break Some(removed);
506492
} else {
@@ -517,7 +503,6 @@ impl<T: Clone> Dict<T> {
517503
}
518504
self.invalidate_keys_version();
519505
inner.unchecked_push(index_index, hash, key.to_pyobject(vm), value, entry_index);
520-
self.bump_version();
521506
break None;
522507
}
523508
};
@@ -616,7 +601,6 @@ impl<T: Clone> Dict<T> {
616601
match inner.entries.get_mut(hint) {
617602
Some(Some(entry)) if key.key_is(&entry.key) => {
618603
let removed = core::mem::replace(&mut entry.value, value);
619-
self.bump_version();
620604
drop(inner);
621605
// defer dec RC until after the lock is released
622606
drop(removed);
@@ -656,6 +640,22 @@ impl<T: Clone> Dict<T> {
656640
}
657641
}
658642

643+
/// Read an entry directly when a cached keys-version still describes the
644+
/// dictionary layout. The version is rechecked while holding the read lock
645+
/// so the entry index and value are observed from the same key-set state.
646+
#[inline]
647+
pub(crate) fn get_index_if_keys_version(&self, version: u32, index: usize) -> Option<T> {
648+
let inner = self.read();
649+
if self.keys_version.load(Acquire) != version {
650+
return None;
651+
}
652+
inner
653+
.entries
654+
.get(index)
655+
.and_then(Option::as_ref)
656+
.map(|entry| entry.value.clone())
657+
}
658+
659659
fn _get_inner<K: DictKey + ?Sized>(
660660
&self,
661661
vm: &VirtualMachine,
@@ -701,7 +701,6 @@ impl<T: Clone> Dict<T> {
701701
inner.indices.resize(8, IndexEntry::FREE);
702702
inner.used = 0;
703703
inner.filled = 0;
704-
self.bump_version();
705704
// defer dec rc
706705
core::mem::take(&mut inner.entries)
707706
};
@@ -830,7 +829,6 @@ impl<T: Clone> Dict<T> {
830829
}
831830
self.invalidate_keys_version();
832831
inner.unchecked_push(index_index, hash, key.to_owned(), value, entry);
833-
self.bump_version();
834832
break None;
835833
};
836834
Ok(())
@@ -867,7 +865,6 @@ impl<T: Clone> Dict<T> {
867865
value.clone(),
868866
index_entry,
869867
);
870-
self.bump_version();
871868
return Ok(value);
872869
}
873870
}
@@ -905,7 +902,6 @@ impl<T: Clone> Dict<T> {
905902
let ret = (key_obj.clone(), value.clone());
906903
self.invalidate_keys_version();
907904
inner.unchecked_push(index_index, hash, key_obj, value, index_entry);
908-
self.bump_version();
909905
return Ok(ret);
910906
}
911907
}
@@ -1117,7 +1113,6 @@ impl<T: Clone> Dict<T> {
11171113
} = IndexEntry::DUMMY;
11181114
inner.used -= 1;
11191115
let removed = slot.take();
1120-
self.bump_version();
11211116
Ok(ControlFlow::Break(removed))
11221117
}
11231118

@@ -1152,7 +1147,6 @@ impl<T: Clone> Dict<T> {
11521147
// entry.index always refers valid index
11531148
inner.indices.get_unchecked_mut(entry.index)
11541149
} = IndexEntry::DUMMY;
1155-
self.bump_version();
11561150
Some((entry.key, entry.value))
11571151
}
11581152

0 commit comments

Comments
 (0)