jit: journal the portal coordinate and snapshot the locals the f_locals fold writes - #1434
Conversation
|
No actionable comments were generated in the recent review. 🎉 ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Plus Run ID: 📒 Files selected for processing (9)
Included review availability: Your plan provides up to 2 included reviews per hour; 0 remain after this review. WalkthroughThe JIT now uses per-frame journals to preserve ChangesFrame locals mirror lifecycle
Estimated code review effort: 4 (Complex) | ~45 minutes Merge Risk: ⚪ Minimal · up to The change preserves frame coordinates and mirrored locals across aborted walks and residual calls, with targeted coverage across all supported backends; no actionable merge-blocking risk remains after normal checks and review. Poem
🚥 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 |
There was a problem hiding this comment.
💡 Codex Review
https://github.com/youknowone/pyre/blob/c0689c199e3a2f88732e10cee660e8ec5446a546/pyre-jit-trace/src/jitcode_dispatch/specialize.rs#L2622
Separate fold snapshots from residual-force undo
When a standard-frame f_locals fold is followed by any non-escaping residual, this capture remains in the force-specific ESCAPE_FLUSH_UNDO slot, and execute_residual_call unconditionally calls restore_escape_flush_undo() after that residual. Consequently, code such as two successive proxy["x"] reads sees the current shadow value on the first call, then restores the pre-fold locals array and sees an old value or KeyError on the second, corrupting the recording-time result and potentially the compiled trace. The fold snapshot needs a lifecycle distinct from the temporary force capture so ordinary residual completion cannot consume it.
AGENTS.md reference: AGENTS.md:L26-L33
https://github.com/youknowone/pyre/blob/c0689c199e3a2f88732e10cee660e8ec5446a546/pyre-jit-trace/src/jitcode_dispatch/residual_call.rs#L1607-L1608
Clear the locals capture when a walk commits
When the folded f_locals operation is near the end of a successfully committed walk and no later residual consumes this capture, the normal commit epilogue never calls discard_escape_flush_undo; it is only cleared when another trace starts. Because the permanently registered extra-root walker visits every saved slot, objects that were present in the frame before the fold remain GC-reachable for the lifetime of the thread if no later trace occurs, potentially retaining an arbitrarily large locals graph. The committed walk path should explicitly discard this snapshot.
AGENTS.md reference: AGENTS.md:L98-L104
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit dc52193). Files in the reviewed diff1. Regressions to PyPy parity introduced by this patchNone. 2. Other mismatches introduced by this patchNone. 3. Pre-existing mismatches (already present before this patch)
4. Structural adaptations
|
c0689c1 to
364dffd
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 364dffd33e
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| /// One frame's pre-fold `locals_cells_stack_w` image. See | ||
| /// [`FBW_LOCALS_MIRROR_UNDO`]. | ||
| pub(crate) struct FbwLocalsMirrorUndo { | ||
| frame: usize, |
There was a problem hiding this comment.
Root the frame held by the locals undo journal
When f_locals is folded for a nursery-allocated standard virtualizable, such as a bridge or recursive JIT frame, a subsequent proxy or residual allocation can relocate that frame before the non-commit epilogue. The extra-root walker forwards entry.slots but never this raw usize, so fbw_locals_mirror_rollback later dereferences the abandoned nursery address instead of the live red frame, failing to restore its locals and potentially corrupting reused nursery memory. Store and visit this frame identity as a forwardable GC root rather than a raw address.
AGENTS.md reference: AGENTS.md:L26-L33
Useful? React with 👍 / 👎.
364dffd to
1ecfab2
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1ecfab2e38
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| // Journaled like the per-opcode publication: this store lands whether | ||
| // or not the walk commits, and a walk that does not commit replays the | ||
| // frame from its pre-walk coordinate. | ||
| crate::jitcode_dispatch::fbw_note_last_instr_undo(cur_ptr as usize); |
There was a problem hiding this comment.
Forward the frame stored in the coordinate undo
When this positive-depth _getframe fold lands on a nursery-allocated portal frame and a later residual allocation triggers a minor collection before a non-committing walk exit, CURRENT_FRAME is forwarded but the raw usize saved by fbw_note_last_instr_undo is not. fbw_exit_last_instr_rollback then restores last_instr through the abandoned address, leaving the live frame at the walk-time coordinate and potentially writing into reused nursery memory; store this identity in a forwardable root slot and visit it alongside the other journal roots.
AGENTS.md reference: AGENTS.md:L26-L33
Useful? React with 👍 / 👎.
…irrors `try_walker_specialize_sys_getframe` wrote the portal frame's `last_instr` without recording the value it displaced, and `walker_write_back_standard_frame_locals` mirrored the virtualizable shadow into the live locals array without recording the values it overwrote. Both land whether or not the walk commits, and a walk that does not commit replays the frame from its pre-walk state. `FBW_EXIT_LAST_INSTR_UNDO` becomes one entry per frame rather than one entry overall: a walk can displace both the recording frame's coordinate (the per-opcode publication) and the portal frame's, and a single cell kept only whichever came first. `fbw_note_last_instr_undo` is the shared recorder; the rollback walks the entries newest first. The locals get their own journal, `FBW_LOCALS_MIRROR_UNDO`, rather than the residual force's `ESCAPE_FLUSH_UNDO`. That capture is consumed by the tail `restore_escape_flush_undo` in `try_execute_residual_call_via_executor`, which runs after every non-forcing residual, so a mirror recorded there is reverted by the next call in the same walk: `f_locals['x']` then answers with the previous iteration's value on the trace-recording iteration, measured on both dynasm and cranelift. The new journal is consumed only by the walk-end legs, beside the coordinate undo, and runs after `restore_escape_flush_undo` so the older pre-fold image stands where both are armed for one frame. It snapshots only the `nlocals` region the mirror writes, leaving the operand stack alone, and the store-journal root area forwards its entries. `synth/getframe_flocals_mirror_survives_residual` reads the proxy after a `str(i)` call in the same body. It reports `mismatches: 1` when the mirror is recorded in the escape-flush capture and `mismatches: 0` with the journal; check.py calls the former `wrong output`. Both restore paths re-arm the frame/array write barrier after writing the saved image back. The forward writes arm it per store because boxing allocates between them; a restore writes values that are already boxed, so one arming after the loop covers the image. `restore_escape_flush_undo` was missing the same arming before this change. Both journals name their frame by raw address, and a JIT-created frame can be nursery-resident, so a minor collection between the eager write and the restore drags it out and leaves a forwarding stub at the recorded address. All three restores now resolve the address through `gc_current_object_address` — the same reload `PyFrame::live_mut` performs for the interpreter's own field writes — at record time so entries key on one identity and at restore time so an entry recorded before the move still lands. One application suffices: the drag-out promotes the frame out of the nursery, and nothing outside it moves. Assisted-by: Claude
1ecfab2 to
dc52193
Compare
Follow-up to #1401. Two eager walk-time writes had no undo, so a walk that does
not commit — which replays the frame from its pre-walk state — replayed on top
of them.
1. The portal frame's
last_instrtry_walker_specialize_sys_getframepublishes the portal frame's coordinatewhen a positive-depth
_getframefold lands on it. The per-opcode publishernext to it already journals, but through a single
Cell, so whichever writecame first in a walk was the only one recorded — and a walk can displace both
frames.
FBW_EXIT_LAST_INSTR_UNDObecomes one entry per frame,fbw_note_last_instr_undois the shared recorder, and the rollback walks entries newest first.
2. The locals the
f_localsfold mirrorswalker_write_back_standard_frame_localsmirrors the virtualizable shadow intothe live
locals_cells_stack_warray, because pyre answers the attribute with a3.14
FrameLocalsProxythat reads the array lazily instead of copying out of itat the call.
The first version of this fix recorded that image in the residual force's
ESCAPE_FLUSH_UNDO, and that was wrong. The tail oftry_execute_residual_call_via_executorcallsrestore_escape_flush_undoafter every non-forcing residual, so the next call in the same walk reverted
the mirror and the proxy answered from before the fold. Measured with an
env-gated A/B on one binary:
one wrong answer on the trace-recording iteration, identical on dynasm and
cranelift;
check.pyreports it aswrong output.So the locals get their own journal,
FBW_LOCALS_MIRROR_UNDO:restore_escape_flush_undo, so where a force flushed thesame frame the older pre-fold image is the one that stands;
nlocalsregion the mirror writes, leaving the operandstack to whoever owns it;
can be the only reference left once the live slot holds the mirrored one.
mark_escape_flush_undo_pendingis still deliberately not armed anywhere here:that flag means "this walk withdrew its commit", which a fold cannot know.
Regression fixture
synth/getframe_flocals_mirror_survives_residualis the snippet above, and itdiscriminates:
mismatches: 1/FAIL … wrong outputwith the escape-flushcapture,
mismatches: 0with the journal. Baselines recorded for all threebackends (
loops_compiled=1,loops_aborted=0).GC write barrier (Codex parity review)
Both restore paths now re-arm
frame_array_write_barrierafter writing thesaved image back. A restored value can be nursery-young while the frame and its
array are old-gen, and nothing re-traces it without a remembered-set entry. The
forward writes arm it per store because boxing allocates between them; a restore
writes values that are already boxed, so nothing can collect mid-loop and one
arming after it covers the image.
restore_escape_flush_undowas missing the same arming before this change —Codex flagged it as pre-existing. It is fixed here too: arming a barrier is
monotonic, so the only cost is an entry the collector would otherwise have
skipped, and the mechanism is the one this PR is about.
Frame relocation (Codex P1)
Both journals name their frame by raw address, and a JIT-created frame can be
nursery-resident: a minor collection between the eager write and the walk-end
restore drags it out and leaves a forwarding stub at the recorded address, so
the restore would write the abandoned copy and leave the live frame at its
walk-time state.
PyFrame::live_mutdocuments the same hazard for theinterpreter's own field writes.
All three restores now resolve the address through
gc_current_object_address,at record time so entries key on one identity and at restore time so an entry
recorded before the move still lands. One application is enough: the drag-out
promotes the frame out of the nursery, and nothing outside the nursery moves.
Resolving beats rooting the address here — the query answers
addrunchangedfor anything the GC does not own, so it needs no ownership gate, and it does not
turn a journal entry into a strong root for a frame that is already reachable.
FBW_EXIT_LAST_INSTR_UNDOandESCAPE_FLUSH_UNDOboth carried this onmain;they are fixed here alongside the new journal.
Not fixed
The
FrameLocalsProxywrite-through vs. virtualizable-shadow coherence gap isreal but pre-existing on
main— gating the fold to the inline receiver alonestill reproduces it — so it stays out of this PR.
FLAG_ESCAPEDis still not rolled back, on risk asymmetry: a stale flag costsan unnecessary force, while clearing a flag another unjournaled site set costs a
skipped force and a dangling vref.
Verification
All three backends built from this head:
getframe_**inlined_callee**locals*gc_*traceback_*Every pre-existing jit-stat unmoved.
cargo test --all --no-default-features --features dynasm,cpyextclean;cargo fmt --all --checkclean;check-new-line-citations.py --base origin/mainreports no new line-numbercitations across 168 added Rust lines.