Skip to content

Commit 1c7759c

Browse files
youknowoneyouknowone
authored
Flatten Python eval loop with trampoline for reduced call overhead (#8431)
* Factor with_iframe into enter_iframe/exit_iframe helpers Extract the frame entry (recursion check, TLS link, exception save) and exit (materialization sync, TLS restore, GC tracking) logic from with_iframe into standalone enter_iframe/exit_iframe methods. with_iframe now calls them, with no behavioral change. This prepares for the trampoline loop where enter/exit are called individually rather than wrapped around a closure. Assisted-by: Claude * Allocate InterpreterFrame and LocalsPlus together on the datastack Add InterpreterFrame::new_on_datastack() that bump-allocates both the InterpreterFrame struct and its LocalsPlus data array in a single datastack push, eliminating one allocation per function call. Update datastack_frame_size_bytes_for_code() to include InterpreterFrame size. Convert invoke_prepared_exact_args() and the invoke() fast path to use the combined allocation. Add release_datastack_frame() method on InterpreterFrame that drops all localsplus values, runs field destructors (trace, temporary_refs, retained_back, etc.), and returns the datastack base pointer for pop. Assisted-by: Claude * Implement trampoline loop for Python-to-Python calls Add ExecutionResult::TailCall variant and a trampoline in run_frame_fast that flattens Python-to-Python calls into a single Rust stack frame instead of recursing through the eval loop. CallPyExactArgs now prepares the callee frame on the datastack and returns TailCall when tailcall_enabled is set (run_iframe path only). The trampoline dispatches via a state machine (EnterCallee / ReturnValue / Unwind) in a single loop, avoiding mutual recursion between helper functions that would exhaust the C stack. Exception propagation through suspended frames uses trampoline_handle_exception which adds traceback entries and calls unwind_blocks on each caller. Assisted-by: Claude * Optimize trampoline: skip recursion check, avoid temporary_refs mutex, add bound method TailCall - Add enter_iframe_unchecked for trampoline callee entry (recursion already checked by specialization_call_recursion_guard) - Move callable ownership from per-frame temporary_refs mutex to trampoline-local SuspendedFrame.owned_refs via VM side channel - Add TailCall support for CallBoundMethodExactArgs - Move args directly from caller stack to callee fastlocals - Read materialized pointer once in exit_iframe Incremental call overhead: ~55 ns -> ~35 ns Assisted-by: Claude * Use UnsafeCell for pending_tailcall_refs The VM is per-thread so RefCell's runtime borrow checking is unnecessary overhead. Replace with UnsafeCell for direct access. Assisted-by: Claude * Remove .claude/settings.json from tracking Assisted-by: Claude * Replace SendPtr with Option<SendNonNull> for type-safe pending tailcall Use NonNull + Option instead of raw *mut T with manual null checks. The compiler enforces non-null via the type system, and Option<NonNull> has the same size as a raw pointer thanks to niche optimization. Also extract take_pending_tailcall helper to deduplicate the pattern. Assisted-by: Claude * Restrict PendingFrame to private, expose only set/take methods Rename SendNonNull to PendingFrame and make it fully private: the struct, its field, and the pending_tailcall_frame Cell are all non-pub. External code accesses the side channel only through set_pending_tailcall (pub(crate)) and take_pending_tailcall (private). This ensures the unsafe Send+Sync impl cannot be reused elsewhere without justifying a new safety argument. Assisted-by: Claude * Restore C-stack overflow check in trampoline enter_iframe_unchecked enter_iframe_unchecked was skipping C-stack checks under the assumption that the trampoline stays in one Rust stack frame. But each run_iframe call still consumes Rust stack, so deep Python recursion through the trampoline can exhaust the C stack (observed as STATUS_STACK_OVERFLOW on Windows CI). Keep the C-stack check (every 8th call) while still skipping the Python recursion depth check (already done by specialization_call_recursion_guard). Assisted-by: Claude * Address code review: entry frame double-free, panic safety, cleanup - Fix double-free: mark entry frame with is_entry flag in SuspendedFrame so the trampoline skips its datastack release (the caller owns that cleanup) - Clear iframe.previous in exit_iframe before unlinking the chain, matching with_frame and resume_gen_frame behavior - Add scopeguard in with_iframe for panic safety - Use saturating_sub(1) for lasti in trampoline_handle_exception - Extract datastack_iframe_localsplus_offset helper to avoid duplicated alignment computation Assisted-by: Claude
1 parent 1fa76a3 commit 1c7759c

6 files changed

Lines changed: 783 additions & 79 deletions

File tree

‎.claude/settings.json‎

Lines changed: 0 additions & 15 deletions
This file was deleted.

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

Lines changed: 16 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -64,9 +64,9 @@ fn format_missing_args(
6464
#[pyclass(module = false, name = "function", traverse = "manual")]
6565
#[derive(Debug)]
6666
pub struct PyFunction {
67-
code: PyAtomicRef<PyCode>,
68-
globals: PyDictRef,
69-
builtins: PyObjectRef,
67+
pub(crate) code: PyAtomicRef<PyCode>,
68+
pub(crate) globals: PyDictRef,
69+
pub(crate) builtins: PyObjectRef,
7070
pub(crate) closure: Option<PyRef<PyTuple<PyCellRef>>>,
7171
defaults_and_kwdefaults: PyMutex<(Option<PyTupleRef>, Option<PyDictRef>)>,
7272
name: PyMutex<PyStrRef>,
@@ -617,11 +617,6 @@ impl Py<PyFunction> {
617617

618618
// Fast path: stack-allocated InterpreterFrame, no FrameObject.
619619
// No refcount inc for code — it's alive via self.code for the call duration.
620-
let nlocalsplus = code.localspluskinds.len();
621-
let max_stackdepth = code.max_stackdepth as usize;
622-
let localsplus =
623-
crate::frame::LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm);
624-
625620
let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) {
626621
crate::frame::FrameLocals::lazy()
627622
} else if let Some(locals) = locals {
@@ -634,22 +629,21 @@ impl Py<PyFunction> {
634629

635630
// Use self.as_object() as raw pointer — no refcount inc/dec.
636631
// The function is alive on the caller's stack for the call duration.
637-
let mut iframe = crate::frame::InterpreterFrame::new(
632+
let iframe = crate::frame::InterpreterFrame::new_on_datastack(
638633
&self.code,
639634
&self.globals,
640635
&self.builtins,
641636
Some(self.as_object()),
642-
localsplus,
643637
locals,
644638
self.closure.as_ref().map_or(&[], |c| c.as_slice()),
645-
crate::frame::FrameOwner::Thread,
639+
vm,
646640
);
647641
let result = self
648-
.fill_locals_from_args_iframe(&mut iframe, func_args, vm)
649-
.and_then(|()| vm.run_frame_fast(&mut iframe));
642+
.fill_locals_from_args_iframe(iframe, func_args, vm)
643+
.and_then(|()| vm.run_frame_fast(iframe));
650644
// Release data stack memory — must happen on both success and error.
651645
unsafe {
652-
if let Some(base) = iframe.localsplus.release_datastack() {
646+
if let Some(base) = iframe.release_datastack_frame() {
653647
vm.datastack_pop(base);
654648
}
655649
}
@@ -797,10 +791,6 @@ impl Py<PyFunction> {
797791
vm: &VirtualMachine,
798792
) -> PyResult {
799793
let code = &*self.code;
800-
let nlocalsplus = code.localspluskinds.len();
801-
let max_stackdepth = code.max_stackdepth as usize;
802-
let localsplus =
803-
crate::frame::LocalsPlus::new_on_datastack(nlocalsplus, max_stackdepth, vm);
804794

805795
let locals = if code.flags.contains(bytecode::CodeFlags::NEWLOCALS) {
806796
crate::frame::FrameLocals::lazy()
@@ -810,15 +800,14 @@ impl Py<PyFunction> {
810800
))
811801
};
812802

813-
let mut iframe = crate::frame::InterpreterFrame::new(
803+
let iframe = crate::frame::InterpreterFrame::new_on_datastack(
814804
code,
815805
&self.globals,
816806
&self.builtins,
817807
Some(self.as_object()),
818-
localsplus,
819808
locals,
820809
self.closure.as_ref().map_or(&[], |c| c.as_slice()),
821-
crate::frame::FrameOwner::Thread,
810+
vm,
822811
);
823812

824813
// Fill arguments directly into fastlocals
@@ -829,9 +818,9 @@ impl Py<PyFunction> {
829818
}
830819
}
831820

832-
let result = vm.run_frame_fast(&mut iframe);
821+
let result = vm.run_frame_fast(iframe);
833822
unsafe {
834-
if let Some(base) = iframe.localsplus.release_datastack() {
823+
if let Some(base) = iframe.release_datastack_frame() {
835824
vm.datastack_pop(base);
836825
}
837826
}
@@ -887,8 +876,10 @@ pub(crate) fn datastack_frame_size_bytes_for_code(code: &Py<PyCode>) -> Option<u
887876
return None;
888877
}
889878
let nlocalsplus = code.localspluskinds.len();
890-
let capacity = nlocalsplus.checked_add(code.max_stackdepth as usize)?;
891-
capacity.checked_mul(core::mem::size_of::<usize>())
879+
Some(crate::frame::datastack_iframe_total_bytes(
880+
nlocalsplus,
881+
code.max_stackdepth as usize,
882+
))
892883
}
893884

894885
impl PyPayload for PyFunction {

‎crates/vm/src/coroutine.rs‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@ impl ExecutionResult {
2323
};
2424
PyIterReturn::StopIteration(arg)
2525
}
26+
Self::TailCall => unreachable!("TailCall in generator/coroutine"),
2627
}
2728
}
2829
}
@@ -104,6 +105,7 @@ impl Coro {
104105
self.clear_frame_locals_on_close();
105106
}
106107
Ok(ExecutionResult::Yield(_)) => {}
108+
Ok(ExecutionResult::TailCall) => unreachable!("TailCall in generator/coroutine"),
107109
}
108110
}
109111

0 commit comments

Comments
 (0)