Flatten Python eval loop with trampoline for reduced call overhead - #51
Flatten Python eval loop with trampoline for reduced call overhead#51youknowone wants to merge 8 commits into
Conversation
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
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
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
…, 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
The VM is per-thread so RefCell's runtime borrow checking is unnecessary overhead. Replace with UnsafeCell for direct access. Assisted-by: Claude
Assisted-by: Claude
|
Warning Review limit reached
Next review available in: 40 minutes Enable usage-based reviews in Billing to review now. Otherwise, wait until the next included review is available. How can I continue?After more reviews become available, a review can be triggered using the To avoid repeated limits, reduce automatic review volume by pausing incremental auto-reviews earlier, using label-based review opt-in, excluding WIP or generated PR titles, or requesting reviews manually when the PR is ready. If your team needs uninterrupted high-volume reviews, an organization admin can enable usage-based reviews. How do review limits work?CodeRabbit enforces per-developer PR review limits for each organization. Most developers receive the normal plan review availability. For paid Pro and Pro+ PR reviews, CodeRabbit uses adaptive limits for sustained high-volume activity. When a developer's recent PR review activity reaches the 95th percentile or higher among CodeRabbit users, additional reviews become available more gradually as earlier reviews age out of the rolling window. Please refer docs for additional details. Review details⚙️ Run configurationConfiguration used: Path: .coderabbit.yml Review profile: CHILL Plan: Pro Plus Run ID: 📒 Files selected for processing (3)
📝 WalkthroughWalkthroughThe VM now uses datastack-backed interpreter frames and an iterative tail-call trampoline for exact Python-function and bound-method calls. Frame entry, exit, exception handling, reference retention, and coroutine result handling were updated. A session-start setup hook was removed. ChangesDatastack tail-call execution
Estimated code review effort: 4 (Complex) | ~45 minutes Sequence Diagram(s)sequenceDiagram
participant PythonCall
participant ExecutingFrame
participant VirtualMachine
participant InterpreterFrame
PythonCall->>ExecutingFrame: invoke exact Python function or bound method
ExecutingFrame->>VirtualMachine: prepare tail call
VirtualMachine->>InterpreterFrame: allocate callee on datastack
VirtualMachine->>VirtualMachine: publish pending frame and retained references
VirtualMachine->>InterpreterFrame: enter callee
InterpreterFrame-->>VirtualMachine: return value or exception
VirtualMachine->>ExecutingFrame: resume suspended caller
Possibly related issues
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
| /// A `*mut T` wrapper that implements `Send`. | ||
| /// Only valid when the pointer is exclusively used by one thread at a time. | ||
| #[repr(transparent)] | ||
| pub(crate) struct SendPtr<T>(pub(crate) *mut T); |
There was a problem hiding this comment.
Is this better than
pub(crate) struct SendPtr<T>(pub(crate) NonNull<T>);
and Option<SendPtr<T>>?
There was a problem hiding this comment.
Actionable comments posted: 2
🧹 Nitpick comments (4)
crates/vm/src/frame.rs (2)
2292-2308: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winKeep the size calculation and the allocation layout in sync.
datastack_iframe_total_bytesduplicates the padding formula used innew_on_datastack(Line 997-999). If one formula changes, the allocation and the space check silently disagree, and the space guard inspecialization_has_datastack_space_for_funcbecomes wrong.Extract one shared helper that returns both the aligned localsplus offset and the total size, and call it from both places.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/vm/src/frame.rs` around lines 2292 - 2308, Extract a shared helper near datastack_iframe_total_bytes that computes the aligned LocalsPlus offset and total allocation size, returning both values. Update datastack_iframe_total_bytes and new_on_datastack to use this helper instead of maintaining separate padding calculations, keeping specialization_has_datastack_space_for_func consistent with the allocation layout.
10613-10646: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the shared callee-frame construction.
tailcall_prepare_frameandtailcall_prepare_bound_method_framerepeat the same logic: derivecode, buildFrameLocalsfromCodeFlags::NEWLOCALS, and callInterpreterFrame::new_on_datastackwith the same six arguments. Only the receiver handling differs.Extract a helper such as
tailcall_new_callee_frame(func, vm) -> &mut InterpreterFrameand call it from both places. This keeps the two paths from drifting apart.Also applies to: 10679-10710
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/vm/src/frame.rs` around lines 10613 - 10646, Extract the shared callee-frame construction from tailcall_prepare_frame and tailcall_prepare_bound_method_frame into a helper such as tailcall_new_callee_frame, using the function’s code, globals, builtins, receiver, locals, and closure to create the InterpreterFrame. Replace the duplicated code in both methods with calls to the helper while preserving their distinct receiver handling.crates/vm/src/vm/mod.rs (2)
2226-2235: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winAdd a panic guard for the recursion depth.
with_iframecallsenter_iframe, thenf, thenexit_iframe. Iffpanics,exit_iframedoes not run. The recursion depth stays incremented and the TLS frame chain keeps pointing at the dead iframe.
with_framehandles this with ascopeguard::guardthat decrements the depth on unwind (Lines 2031-2033). Apply the same protection here. The same gap exists inrun_frame_fastand in the trampoline, where a panic leaves everySuspendedFrameunreleased.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/vm/src/vm/mod.rs` around lines 2226 - 2235, Add panic-safe cleanup for iframe execution in with_iframe by using the existing scopeguard pattern from with_frame, ensuring exit_iframe runs during unwinding as well as normal returns. Apply equivalent unwind guards to run_frame_fast and the trampoline so recursion depth, TLS frame links, and all SuspendedFrame resources are restored or released when callbacks panic.
1463-1497: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winExtract the repeated result-dispatch logic.
The block that drains
pending_tailcall_refs, pushes aSuspendedFrame, and readspending_tailcall_framerepeats four times. The block that callsexit_iframe, thenrelease_datastack_frame, thendatastack_poprepeats four times. This is lifecycle-critical code, so any future divergence between the copies causes a leak or a double free.Extract two helpers, for example
fn take_pending_tailcall(&self) -> (*mut InterpreterFrame, Vec<PyObjectRef>)andunsafe fn finish_iframe(&self, iframe: &mut InterpreterFrame, entry: IframeEntryState), then call them from each arm.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@crates/vm/src/vm/mod.rs` around lines 1463 - 1497, Extract the repeated tail-call and iframe cleanup logic from the result-dispatch code into two helpers: a take-pending-tailcall helper that drains pending_tailcall_refs, reads and clears pending_tailcall_frame, and returns the frame with owned references; and an unsafe finish-iframe helper that calls exit_iframe, releases the datastack frame, and pops it when present. Replace all four duplicated blocks with these helpers, preserving the existing SuspendedFrame construction and action behavior.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@crates/vm/src/frame.rs`:
- Around line 976-984: Make Frame::new_on_datastack safe only with a
type-enforced lifetime, or mark it unsafe/return a raw pointer or NonNull<Self>;
update its call sites in builtins/function.rs and tailcall_prepare_frame to
handle the revised contract and require explicit unsafe dereferencing where
applicable.
In `@crates/vm/src/vm/mod.rs`:
- Around line 2208-2212: Update exit_iframe to clear iframe.previous to zero
before restoring the old frame chain via set_current_frame, matching the
ordering used by with_frame. Keep the existing exception restoration and
recursion-depth cleanup unchanged, and ensure the raw previous pointer is
invalidated before the caller’s stack allocation can be released.
---
Nitpick comments:
In `@crates/vm/src/frame.rs`:
- Around line 2292-2308: Extract a shared helper near
datastack_iframe_total_bytes that computes the aligned LocalsPlus offset and
total allocation size, returning both values. Update
datastack_iframe_total_bytes and new_on_datastack to use this helper instead of
maintaining separate padding calculations, keeping
specialization_has_datastack_space_for_func consistent with the allocation
layout.
- Around line 10613-10646: Extract the shared callee-frame construction from
tailcall_prepare_frame and tailcall_prepare_bound_method_frame into a helper
such as tailcall_new_callee_frame, using the function’s code, globals, builtins,
receiver, locals, and closure to create the InterpreterFrame. Replace the
duplicated code in both methods with calls to the helper while preserving their
distinct receiver handling.
In `@crates/vm/src/vm/mod.rs`:
- Around line 2226-2235: Add panic-safe cleanup for iframe execution in
with_iframe by using the existing scopeguard pattern from with_frame, ensuring
exit_iframe runs during unwinding as well as normal returns. Apply equivalent
unwind guards to run_frame_fast and the trampoline so recursion depth, TLS frame
links, and all SuspendedFrame resources are restored or released when callbacks
panic.
- Around line 1463-1497: Extract the repeated tail-call and iframe cleanup logic
from the result-dispatch code into two helpers: a take-pending-tailcall helper
that drains pending_tailcall_refs, reads and clears pending_tailcall_frame, and
returns the frame with owned references; and an unsafe finish-iframe helper that
calls exit_iframe, releases the datastack frame, and pops it when present.
Replace all four duplicated blocks with these helpers, preserving the existing
SuspendedFrame construction and action behavior.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Path: .coderabbit.yml
Review profile: CHILL
Plan: Pro Plus
Run ID: c68c8fc7-c5cb-4d65-aad2-f285d7a7536d
📒 Files selected for processing (6)
.claude/settings.jsoncrates/vm/src/builtins/function.rscrates/vm/src/coroutine.rscrates/vm/src/frame.rscrates/vm/src/vm/mod.rscrates/vm/src/vm/thread.rs
💤 Files with no reviewable changes (1)
- .claude/settings.json
| pub(crate) fn new_on_datastack<'a>( | ||
| code: &Py<PyCode>, | ||
| globals: &Py<PyDict>, | ||
| builtins: &PyObject, | ||
| func_obj: Option<&PyObject>, | ||
| locals: FrameLocals, | ||
| closure: &[PyCellRef], | ||
| vm: &VirtualMachine, | ||
| ) -> &'a mut Self { |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟠 Major | 🏗️ Heavy lift
Bound the returned reference or make the constructor unsafe.
new_on_datastack is a safe function that returns &'a mut Self with a caller-chosen lifetime. A safe caller can therefore pick any lifetime, including one that outlives datastack_pop, and use the reference after the memory is reclaimed. The doc comment states the contract, but the type system does not enforce it.
Mark the function unsafe, or return *mut Self (or NonNull<Self>) so that every dereference is an explicit unsafe block at the call sites in crates/vm/src/builtins/function.rs and tailcall_prepare_frame.
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/vm/src/frame.rs` around lines 976 - 984, Make Frame::new_on_datastack
safe only with a type-enforced lifetime, or mark it unsafe/return a raw pointer
or NonNull<Self>; update its call sites in builtins/function.rs and
tailcall_prepare_frame to handle the revised contract and require explicit
unsafe dereferencing where applicable.
| if save_exc { | ||
| self.restore_exception(saved_exc); | ||
| } | ||
| // Restore the frame chain BEFORE clearing temporary_refs, so | ||
| // top_frame no longer points at the materialized FrameObject | ||
| // when its last strong reference is released. | ||
| let _ = crate::vm::thread::set_current_frame(old_chain); | ||
| self.recursion_depth.update(|d| d - 1); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Clear iframe.previous before you unlink the frame from the chain.
with_frame stores 0 into iframe.previous before it restores the old chain (Lines 2088-2097). The comment there states the reason: previous can point to a stack-allocated iframe that is freed when the caller exits. exit_iframe omits that store.
For a materialized datastack frame, the FrameObject outlives exit_iframe and keeps previous pointing at the caller iframe. After the caller's datastack allocation is popped, any consumer that walks previous reads freed memory. crates/vm/src/vm/thread.rs walks previous in reinit_frame_slot_after_fork, and set_current_frame dereferences the published iframe pointer.
retained_back covers f_back, but it does not cover the raw previous walkers.
🐛 Proposed fix
if save_exc {
self.restore_exception(saved_exc);
}
+ // Clear previous before unlinking — it may point to a datastack
+ // iframe that is freed when the caller's frame is released.
+ {
+ #[allow(unused_imports)]
+ use rustpython_common::atomic::Radium;
+ let live_iframe = unsafe { &*iframe_ptr };
+ live_iframe
+ .previous
+ .store(0, core::sync::atomic::Ordering::Relaxed);
+ }
let _ = crate::vm::thread::set_current_frame(old_chain);
self.recursion_depth.update(|d| d - 1);📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if save_exc { | |
| self.restore_exception(saved_exc); | |
| } | |
| // Restore the frame chain BEFORE clearing temporary_refs, so | |
| // top_frame no longer points at the materialized FrameObject | |
| // when its last strong reference is released. | |
| let _ = crate::vm::thread::set_current_frame(old_chain); | |
| self.recursion_depth.update(|d| d - 1); | |
| if save_exc { | |
| self.restore_exception(saved_exc); | |
| } | |
| // Clear previous before unlinking — it may point to a datastack | |
| // iframe that is freed when the caller's frame is released. | |
| { | |
| #[allow(unused_imports)] | |
| use rustpython_common::atomic::Radium; | |
| let live_iframe = unsafe { &*iframe_ptr }; | |
| live_iframe | |
| .previous | |
| .store(0, core::sync::atomic::Ordering::Relaxed); | |
| } | |
| let _ = crate::vm::thread::set_current_frame(old_chain); | |
| self.recursion_depth.update(|d| d - 1); |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@crates/vm/src/vm/mod.rs` around lines 2208 - 2212, Update exit_iframe to
clear iframe.previous to zero before restoring the old frame chain via
set_current_frame, matching the ordering used by with_frame. Keep the existing
exception restoration and recursion-depth cleanup unchanged, and ensure the raw
previous pointer is invalidated before the caller’s stack allocation can be
released.
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
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
Summary
Flatten the Python-to-Python call path so that
CallPyExactArgsandCallBoundMethodExactArgsno longer recurse through new Rust stack frames. Instead, the bytecode loop returns aTailCallsignal and a trampoline swaps frames in a single loop — matching CPython 3.12+'s approach.Changes
Phase 0: Factor
with_iframeenter_iframe/exit_iframefromwith_iframewithIframeEntryStatestructPhase 1: Datastack-allocated InterpreterFrame
InterpreterFrame::new_on_datastack()bump-allocates both the InterpreterFrame and its LocalsPlus in a single datastack pushrelease_datastack_frame()drops values and returns the base pointer fordatastack_popPhase 2–4: Trampoline loop + TailCall + exception propagation
ExecutionResult::TailCallvariantpending_tailcall_frameside channel onVirtualMachinerun_frame_fast_trampoline()withVec<SuspendedFrame>stacktailcall_prepare_frame()builds callee frame on datastack and stores pointertrampoline_handle_exception()— unwinds through suspended callers, attaches traceback entriesPhase 5: Optimizations
enter_iframe_uncheckedskips recursion/C-stack checks in trampoline (already verified byspecialization_call_recursion_guard)temporary_refsmutex to trampoline-localSuspendedFrame.owned_refsvia VM side channel — eliminates mutex lock + Vec allocation per callCallBoundMethodExactArgsmaterializedpointer once inexit_iframeUnsafeCellforpending_tailcall_refs(VM is per-thread)Performance
Measured on Apple Silicon. The remaining gap to the ≤30 ns target is addressable by InterpreterFrame hot/cold field splitting (separate work).
Test coverage
All existing tests pass:
test_frame,test_traceback,test_generators,test_sys,test_pdb,test_exceptions,test_call,test_funcattrs,test_descr,test_faulthandler.Deep recursion (
RecursionError), exception propagation through trampoline,sys._getframe()chain, andtry/exceptacross call boundaries all verified.Refs: #40
🤖 Generated with Claude Code
Summary by CodeRabbit
Performance
Reliability
Maintenance