jit: gh#394 warmup follow-ons — cache FrameRoot shadow-stack cell + correct walk-end handback vsd - #698
Merged
Conversation
FrameRoot::frame() read the current frame root through shadow_stack::get(depth), which resolves the SHADOW_STACK thread-local on every call (a _tlv_get_addr per access on macOS). eval_loop_jit re-reads the root through frame() several times per opcode step, so the thread-local resolution recurs across the whole warmup loop. Resolve the thread-local cell once in FrameRoot::new and hold the pointer in the root (ShadowStackSlot); frame() now re-reads entries[depth] through the cached cell via slot_get. The cell address is stable for the owning thread's life (the invariant MutatorEntry already relies on for STW walks), and FrameRoot is a same-thread stack local, so the cached pointer cannot outlive or cross its thread. slot_get takes the same transient borrow the old get() did; a re-read still indexes the live entries Vec, so a Vec realloc between reads is unaffected. Also const-initialize the SHADOW_STACK thread-local (ShadowStack::new is now const with an empty Vec instead of Vec::with_capacity(64)), dropping the per-access lazy-init guard. Assisted-by: Claude
WalkthroughThe change adds a cached per-thread shadow-stack slot API, uses it for JIT frame-root retrieval, simplifies post-collection dispatch reseeding, and synchronizes restart instruction and operand-stack state. ChangesJIT and GC integration
Estimated code review effort: 3 (Moderate) | ~20 minutes Sequence Diagram(s)sequenceDiagram
participant JIT as JIT eval
participant Root as FrameRoot
participant Stack as Shadow stack
JIT->>Root: Create frame root
Root->>Stack: Resolve and cache shadow-stack slot
JIT->>Root: Retrieve frame
Root->>Stack: Read entry through cached slot
JIT->>JIT: Reseed dispatch state after collection
JIT->>JIT: Correct restart instruction and operand depth
Possibly related PRs
Suggested reviewers: 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 |
Wrap the slot_get expression to the form rustfmt 1.9 produces. Whitespace only. Assisted-by: Claude
The non-flush walk-end handback set the resume pc but not valuestackdepth: when a loop-header marker sits inside a super-instruction, the walk closes the loop at loop_header_pc + 1 and leaves valuestackdepth advanced through the super-instruction, while the handback wrote only last_instr/next_instr. The frame was handed back with pc and operand depth inconsistent, and eval_loop_jit patched it downstream by recomputing instruction.stack_effect every opcode step and, when the projected push overflowed the frame array and the depth table agreed, advancing the pc instead of dispatching. Make the handback symmetric with the flush leg and the blackhole legs (apply_blackhole_crn_handoff): after set_last_instr_from_next_instr, call correct_resume_vsd(restart_pc) to re-derive the operand depth from the resume pc. The handed-back frame is then self-consistent, so the per-step reconciliation and its stack_effect recompute are removed. valuestackdepth is now only ever compared against the precomputed depth table (depth_based_vsd_for_wcode), never re-derived from a per-step stack effect -- the dispatch loop no longer computes a stack effect at all, and the depth table is the single source of truth. rustpython's stack_effect_info is no longer consulted on the hot path. Assisted-by: Claude
youknowone
changed the title
Owner
Author
|
Added commit A/B for this commit (cand-1 baseline vs this commit, interleaved best-of-N, 21 rounds): +1.39% best-of-N / +0.17% median, bands overlap — perf-neutral-to-slightly-positive. It lands as a correctness/orthodoxy cleanup that removes machinery; the A/B confirms no regression. Gates green: check.py 241/241 both backends byte-identical, gc_stress 12/12, pyre-jit 320+12, golden — commented by Claude |
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit 370e908). 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)None. 4. Structural adaptations
|
Owner
Author
|
CI status: all green except one pre-existing failure.
— commented by Claude |
youknowone
added a commit
that referenced
this pull request
Aug 13, 2026
`run_perfn_walk` set `WALK_END_RESTART_PC` for every `CloseLoop`. The portal consumes that cell only on the leg where the walk-end flush declined, and a decline keeps the legacy replay, whose contract is that the frame still holds pre-walk state. Applying the resume pc there moved `last_instr` and `valuestackdepth` to the loop header while the locals stayed at the trace entry, so the frame carried values from two points at once: `_process_class` resumed with `cmp_fields` unbound and `field` holding a bound method pushed ~14 lines later, raising `TypeError: 'str' object is not an iterator` at dataclasses.py:1170. Set the cell only when the resume pc differs from the loop header — the marker legs it was added for (#698), where a loop-header marker inside a super-instruction leaves the frame advanced past the header. Reproduced on darwin-aarch64 and linux-aarch64 with a 15-line dataclass loop under `PYRE_JIT="threshold=3,function_threshold=3"`; green after the change at thresholds 3, 4, 5, 6, 8, 12, 20, 40, 80 and 200. Assisted-by: Claude
youknowone
added a commit
that referenced
this pull request
Aug 13, 2026
`run_perfn_walk` set `WALK_END_RESTART_PC` for every `CloseLoop`. The portal consumes that cell only on the leg where the walk-end flush declined, and a decline keeps the legacy replay, whose contract is that the frame still holds pre-walk state. Applying the resume pc there moved `last_instr` and `valuestackdepth` to the loop header while the locals stayed at the trace entry, so the frame carried values from two points at once: `_process_class` resumed with `cmp_fields` unbound and `field` holding a bound method pushed ~14 lines later, raising `TypeError: 'str' object is not an iterator` at dataclasses.py:1170. Set the cell only when the resume pc differs from the loop header — the marker legs it was added for (#698), where a loop-header marker inside a super-instruction leaves the frame advanced past the header. Reproduced on darwin-aarch64 and linux-aarch64 with a 15-line dataclass loop under `PYRE_JIT="threshold=3,function_threshold=3"`; green after the change at thresholds 3, 4, 5, 6, 8, 12, 20, 40, 80 and 200. Assisted-by: Claude
youknowone
added a commit
that referenced
this pull request
Aug 14, 2026
`run_perfn_walk` set `WALK_END_RESTART_PC` for every `CloseLoop`. The portal consumes that cell only on the leg where the walk-end flush declined, and a decline keeps the legacy replay, whose contract is that the frame still holds pre-walk state. Applying the resume pc there moved `last_instr` and `valuestackdepth` to the loop header while the locals stayed at the trace entry, so the frame carried values from two points at once: `_process_class` resumed with `cmp_fields` unbound and `field` holding a bound method pushed ~14 lines later, raising `TypeError: 'str' object is not an iterator` at dataclasses.py:1170. Set the cell only when the resume pc differs from the loop header — the marker legs it was added for (#698), where a loop-header marker inside a super-instruction leaves the frame advanced past the header. Reproduced on darwin-aarch64 and linux-aarch64 with a 15-line dataclass loop under `PYRE_JIT="threshold=3,function_threshold=3"`; green after the change at thresholds 3, 4, 5, 6, 8, 12, 20, 40, 80 and 200. Assisted-by: Claude
youknowone
added a commit
that referenced
this pull request
Aug 14, 2026
`run_perfn_walk` set `WALK_END_RESTART_PC` for every `CloseLoop`. The portal consumes that cell only on the leg where the walk-end flush declined, and a decline keeps the legacy replay, whose contract is that the frame still holds pre-walk state. Applying the resume pc there moved `last_instr` and `valuestackdepth` to the loop header while the locals stayed at the trace entry, so the frame carried values from two points at once: `_process_class` resumed with `cmp_fields` unbound and `field` holding a bound method pushed ~14 lines later, raising `TypeError: 'str' object is not an iterator` at dataclasses.py:1170. Set the cell only when the resume pc differs from the loop header — the marker legs it was added for (#698), where a loop-header marker inside a super-instruction leaves the frame advanced past the header. Reproduced on darwin-aarch64 and linux-aarch64 with a 15-line dataclass loop under `PYRE_JIT="threshold=3,function_threshold=3"`; green after the change at thresholds 3, 4, 5, 6, 8, 12, 20, 40, 80 and 200. Assisted-by: Claude
youknowone
added a commit
that referenced
this pull request
Aug 14, 2026
…ection (#1198) * jit: do not hand back a walk-end resume pc the flush declined `run_perfn_walk` set `WALK_END_RESTART_PC` for every `CloseLoop`. The portal consumes that cell only on the leg where the walk-end flush declined, and a decline keeps the legacy replay, whose contract is that the frame still holds pre-walk state. Applying the resume pc there moved `last_instr` and `valuestackdepth` to the loop header while the locals stayed at the trace entry, so the frame carried values from two points at once: `_process_class` resumed with `cmp_fields` unbound and `field` holding a bound method pushed ~14 lines later, raising `TypeError: 'str' object is not an iterator` at dataclasses.py:1170. Set the cell only when the resume pc differs from the loop header — the marker legs it was added for (#698), where a loop-header marker inside a super-instruction leaves the frame advanced past the header. Reproduced on darwin-aarch64 and linux-aarch64 with a 15-line dataclass loop under `PYRE_JIT="threshold=3,function_threshold=3"`; green after the change at thresholds 3, 4, 5, 6, 8, 12, 20, 40, 80 and 200. Assisted-by: Claude * cpyext: address review findings on the extension slice Replace the three cpyext tables' `LazyLock<Mutex<..>>` with a `ForkMutex<T>` whose lock word `after_fork_child` rebuilds in place, keeping the payload: the child inherits the parent's mappings, so the loaded libraries and the raw-mirror census must survive while only the stale lock word is replaced. Seed `__spec__`/`__loader__`/`__package__`/`__file__` on a natively resolved extension module from `_bootstrap_external.spec_from_file_location`, the way the source and builtin branches of the same `load_part` already do. Accept `create_dynamic(spec, file)`: the fixed arity of 1 rejected the two- argument call with a TypeError before the loader ran. Accept `m_size == 0` in `PyModule_Create2` alongside `-1`; neither allocates per-module state. Set `PY_VERSION_HEX` to `0x030E06F0`, matching `sys.hexversion`. The previous `0x030E0000` sorts below `0x030E00F0` (3.14.0 final). Root `path_list` before the `__path__` store, which allocates. Assert `PyObjectRef` and `majit_ir::GcRef` have the same size and alignment at the root-forwarding cast. Assisted-by: Claude * call: reload constructor arguments from the shadow stack `type_descr_call_with_mode` pinned neither the type nor its arguments, and `call_with_kwargs_in_ctx` pinned them at entry but then forwarded the incoming slices raw. Both build the `__init__` argument list after `__new__` has run Python code, so a minor collection during `__new__` left the forwarded slice holding pre-move addresses; `__init__` stored one of those into an instance attribute, and the next collection tripped over it through the remembered set. Both paths now read the type and every argument back through `pyre_object::gc_roots`, which is the shape `type_descr_call_impl` already used. Assisted-by: Claude * _ctypes: reject a symbol that resolves to address 0 `dlsym` reports a miss by returning NULL, and a resolver that itself returns NULL leaves `dlerror` unset, so `lookup_function_symbol_addr` reported success with address 0. The unix `lookup_symbol` now rejects that, matching `rdynload.dlsym`, and `_ctypes.dlsym` goes through `lookup_symbol` instead of calling the host lookup directly. Assisted-by: Claude * cpyext: reject an init symbol that resolves to address 0 `load_extension_module` transmuted the looked-up address to the init signature without checking it, so a symbol resolving to NULL became a call through a null pointer. Assisted-by: Claude * gc: describe the mapdict storage trace as it is implemented `mapdict_storage_custom_trace`'s doc called `storage` an off-GC `Box<Vec<PyObjectRef>>` and said `instance_walk_boxed_storage` consults the map to skip erased unboxed slots. Neither holds: `storage` is a GC-managed leaf block allocated stable and non-moving by `alloc_mapdict_storage_block`, and the walk iterates `0..capacity` unconditionally, which `erase_unboxed` licenses by storing an ordinary `GC_INT_ARRAY` reference in the slot. Assisted-by: Claude * jit: publish the gcmap around the varsize nursery slowpath call `CallMallocNurseryVarsize` stored a null gcmap into the jitframe before calling `dynasm_nursery_slowpath_varsize`, which can collect. A null gcmap tells the collector the frame holds no references, so the slots the register allocator spilled into (it uses `SAVE_ALL_REGS` here) are not traced and the values they hold are not forwarded. The fixed-size siblings (`CallMallocNursery` / `CallMallocNurseryHeaderless` / `CallMallocNurseryVarsizeFrame`) already push `pending_malloc_nursery_gcmap` and spill the registers into the jitframe slots the gcmap's bits name; this brings the varsize path to the same shape on both backends. The register allocator already attaches a gcmap to this op through `perform_with_gcmap`, so the value was available and unused. This does not change the `-m test.test_unittest` GC crash: 5/5 runs before and after the change abort with the same `GC BUG: invalid type_id`. A constant-length array lowers through `gen_malloc_nursery` (the `total_size >= 0` arm of `handle_new_array`), not this op. Assisted-by: Claude * list: root the subscript operands across the replacement materialization `STORE_SUBSCR` pops the container, key and value off the value stack before dispatching, so the frame no longer roots any of them. `setitem_list_slice` then held all three as bare addresses across `slice_unpack` (which honors `__index__`), `collect_iterable` (which runs the iterable's own Python code) and the two `w_list_new` allocations. Publish the operands on the shadow stack at entry and re-read each after every step that can collect; the extended-slice loop brackets its item root per iteration rather than pushing one root per element. Measured with `MAJIT_GC_NURSERY_POISON=1` on `-m test.test_unittest`, aarch64 Linux: before the change 3/3 runs abort in `switch_to_object_strategy` on a receiver whose header and body both read the poison fill; after it, 0/5 and the suite reaches `Ran 1090 tests ... OK`. `PYRE_NO_JIT=1` reproduces the same abort at the same site, so the window is in the interpreter. This does not change the `GC BUG: invalid type_id` abort on the same fixture: 5/5 runs with poison off before and after. Assisted-by: Claude
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Follow-on to #692, carrying two independent gh#394 warmup cleanups that a re-profile of the #692 result surfaced. Each is separately gated and A/B-measured.
Commit 1 — cache the shadow-stack cell in
FrameRoot; const-initSHADOW_STACK#692 cut the per-opcode-step
FrameRoot::frame()reload count ineval_loop_jit(gh#394 driver residual); a re-profile showedmajit_gc::shadow_stack::getstill the #1 non-idle warmup self-time frame — the remaining ~3 seeds/step each pay a thread-local resolution (_tlv_get_addron macOS) becauseSHADOW_STACK.with()runs on every access.ShadowStackSlot(a cached*const RefCell<ShadowStack>) plusshadow_stack_slot()/unsafe slot_get(). The cell address is stable for the owning thread's life — the same invariantMutatorEntryalready relies on for STW walks.FrameRoot::newresolves the thread-local cell once and stores the slot;frame()re-readsentries[depth]through the cached slot instead of resolving the thread-local on every call.FrameRootis a same-thread stack local, so the cached pointer can neither outlive nor cross its thread.slot_gettakes the same transient borrow the oldget()did, and re-reads still index the liveentriesVec, so a Vec realloc between reads is unaffected.const-initialize theSHADOW_STACKthread-local (ShadowStack::newis nowconstwith an empty Vec instead ofVec::with_capacity(64)), dropping the per-access lazy-init guard.Orthodoxy. RPython does not re-resolve the root-stack base per access: it lives in a plain global (
gcdatasingleton,framework.py), the C backend resolves it once per function (gc_enter_roots_frame,funcgen.py), the x86 JIT caches the root-stack top in a register across call boundaries (assembler.py:_load_shadowstack_top_in_ebx), andshadowcolor.pyeven hoists root saves out of loops. pyre's per-accessSHADOW_STACK.with()is the divergence; caching the resolution once per activation and reusing it between GC safepoints brings the access cost back in line with upstream. The cache never spans a thread-switch/GIL boundary (single stack-localFrameRootper activation).Perf (A/B, interleaved best-of-N, EPOCHREALTIME): run 1 (15 rounds) +3.66% best-of-N / +3.85% median, bands disjoint; run 2 (12 rounds) +3.71% / +3.52%, disjoint. Every round faster.
Commit 2 — correct
valuestackdepthon the super-instruction walk-end handbackThe non-flush walk-end handback in
compile_and_run_onceset the resume pc but notvaluestackdepth: when a loop-header marker sits inside a super-instruction, the walk closes the loop atloop_header_pc + 1and leavesvaluestackdepthadvanced through the super-instruction, while the handback wrote onlylast_instr/next_instr. The frame was handed back with pc and operand depth inconsistent, andeval_loop_jitpatched it downstream by recomputinginstruction.stack_effectevery opcode step and, when the projected push overflowed the frame array and the depth table agreed, advancing the pc instead of dispatching. That per-stepstack_effectwas the #2 non-idle warmup self-time frame.restore_resume_state_from) and the blackhole legs (apply_blackhole_crn_handoff): afterset_last_instr_from_next_instr(restart_pc), callcorrect_resume_vsd(frame, restart_pc)to re-derive the operand depth from the resume pc via the depth table.stack_effectrecompute are removed.valuestackdepthis now only ever compared against the precomputed depth table (depth_based_vsd_for_wcode/ liveness), never re-derived per dispatch step — the depth table is the single source of truth, and rustpython'sstack_effect_infoleaves the hot path.Orthodoxy. PyPy's dispatch loop never computes a per-step stack effect:
valuestackdepthis an incrementally-maintained counter (pyframe.pypush/popvalue),stack_effectis compile-time only (co_stacksize,astcompiler/assemble.py), and the JIT frame is a virtualizable restored positionally from resume data in one shot (rpython/jit/metainterp/virtualizable.py), so pc and vsd are never mutually inconsistent at a merge point. Fixing the seam so the frame is self-consistent — rather than patching the inconsistency downstream — is the orthodox shape.Perf (A/B, interleaved best-of-N, 21 rounds): +1.39% best-of-N / +0.17% median, bands overlap. Perf-neutral-to-slightly-positive; this commit lands as a correctness/orthodoxy cleanup that removes machinery (the per-step reconciliation), with the A/B confirming no regression.
Correctness (whole branch)
check.py --backend dynasm,cranelift: 241/241 dynasm + 241/241 cranelift, byte-identical both backends (measured on both the pre- and post-commit-2 base).MAJIT_GC_STRESS=1gc_stress 12/12 — the UAF gate: a moving collection at every allocation window.cargo test -p majit-gc191,cargo test -p pyre-jit320+12 (both eval-loop twins).3004320000) for both commits.— opened by Claude
Summary by CodeRabbit
Bug Fixes
Performance