jit: hand the CALL residual an execution context, not a caller frame - #1046
Conversation
WalkthroughThis PR changes interpreter and JIT call dispatch to pass an explicit ChangesExecution-context dispatch and abort replay handling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Residual as Residual JIT call
participant CallJit as pyre-jit/call_jit.rs
participant Ctx as PyExecutionContext
participant Call as pyre-interpreter/call.rs
participant Eval as evaluator/JIT
Residual->>CallJit: invoke residual or portal call
CallJit->>Ctx: read/install pinned execution context
CallJit->>Call: call_user_function_with_ctx / call_kw_in_ctx / call_function_ex_in_ctx
Call->>Ctx: resolve active profiling frame
Call->>Eval: execute user function, builtin, or type call
Eval-->>Call: PyResult
Call-->>CallJit: PyResult
CallJit->>Ctx: restore previous TLS context
CallJit-->>Residual: return result
Possibly related PRs
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 |
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
https://github.com/youknowone/pyre/blob/df3c1e863c560cdbd80be18a41e4c1de191a166a/pyre-jit/src/call_jit.rs#L4836
When a compiled CALL_KW residual reaches this path, the copied ABI values (callable, kwnames, and the positional slice) are no longer protected by the old forced parent frame and call_kw(frame, ...) wrapper. call_kw_in_ctx can allocate/collect during keyword binding or callee-frame allocation before those operands are reloaded from any shadow-stack roots; as noted by the neighboring bh_call_fn_impl pattern, the backend gcmap keeps the original objects alive but cannot rewrite these copied native parameters after a moving GC. Please pin and reload these residual operands before the frameless dispatch.
https://github.com/youknowone/pyre/blob/df3c1e863c560cdbd80be18a41e4c1de191a166a/pyre-jit/src/call_jit.rs#L5113-L5117
For CALL_FUNCTION_EX residuals, replacing call_function_ex(parent_frame, ...) with this frameless entry removes the FrameLocalsRoot that used to protect the caller stack, but the raw callable, self_or_null, starargs, and kwargs_or_null parameters are not pinned before call_function_ex_in_ctx starts unpacking. In particular, non-tuple/list *args or mapping **kwargs can execute Python and trigger a moving collection before call_with_kwargs_in_ctx installs roots, leaving these copied native parameters stale. Root/reload the residual operands here the same way bh_call_fn_impl does before any Python-level unpacking.
ℹ️ 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 4377be8). Files in the reviewed diff1. Regressions to PyPy parity introduced by this patch
2. Other mismatches introduced by this patchNone. 3. Pre-existing mismatches (already present before this patch)None identified within the reviewed behavior. 4. Structural adaptations
|
…_use `classify_compare_box_use` decoded `compare_pc` against the snapshot root's JitCode and bailed out with `Other` on every inlined sub-walk, since the callee's offset does not index the root's bytes. Take the bytes from the frame the walk is in (`framestack.last().w_code` -> `sub_jitcode_body_for_code`). The liveness tables condition 4 reads are keyed on the root jitcode index and do not describe a callee, so the sub-walk path reports the shape with `arms_dead: false`: `walker_newbool_guarded` uses it, and the box elision in `compare_box_provably_dead` still requires condition 4 and keeps declining inside a sub-walk. On pyre/bench/fib_recursive.py this drops 3 `CallR(newbool)` residuals whose results are read by nothing (2 of them in the hottest loop). Recorded jitstats are unchanged (loops_compiled=1 bridges_compiled=3 loops_aborted=0 guard_failures=406). Assisted-by: Claude
…f the caller frame `bh_call_fn_impl` resolved a caller frame with `getexecutioncontext().gettopframe_raw()`, which is `force_vref`. While the trace records, that chain slot holds a vref stamped `TOKEN_TRACING_RESCALL` by `tracing_before_residual_call`; `force_virtual` clears the token (`virtualref.py:161-167`) and `tracing_after_residual_call` reads the cleared token as "the callee forced this vref", so `vrefs_after_residual_call` -> `stop_tracking_virtualref` recorded the non-NULL `VIRTUAL_REF_FINISH` ahead of the call. The helper's own read was the escape. `bhimpl_residual_call_r_r` (blackhole.py:1227) is `cpu.bh_call_r(func, None, args_r, ...)` and carries no frame; the callee frame takes its execution context from the space. Pass the execution context: * `prepare_user_call` takes `execution_context`, the only field it read off the caller frame. * `call_user_function_residual` is replaced by `call_user_function_residual_with_ctx` — `call_user_function_plain_with_ctx` with the JIT-aware eval function. The caller-side `FrameLocalsRoot` goes with the caller frame. * `create_callee_frame_in_ctx` splits out of `create_callee_frame_impl`, which read the caller only for `execution_context`. * the self-recursion probe reads `vref_referent`, the non-forcing chain read documented for identity tests, and declines on a null referent. * `set_last_exec_ctx` takes the `ec` already in hand on the user-function and cold paths. * the `PYRE_BH_NULL_ARG` diagnostic resolves the frame itself. `bh_call_kw_impl` and `bh_call_function_ex_fn` still call `gettopframe_raw()`: `call_kw` / `call_function_ex` reach `call_function_carrier_with_mode`, which reads `frame.get_is_being_profiled()` and hands the frame to `call_args_and_c_profile` (baseobjspace.py:1243). Measured on dynasm, one binary pair, output byte-identical on every bench. Minor collections: fib_recursive 5267 -> 4645, linrec 483 -> 405, treerec 1710 -> 1498, nbody 414 -> 414. fib_recursive child CPU 0.868x and 0.886x across the two arm orders (min of 9, interleaved). The compiled traces are unchanged: the `after opt` op counts match the baseline binary on `inline_gate_operand_provenance` (61/61/61/49) and `kept_stack_deep_var_shortcircuit` (105/131). `cargo test -p pyre-jit -p pyre-interpreter --features dynasm`: 842 passed. `check.py --backend dynasm,cranelift,wasm`: no correctness failure; dynasm 375 passed / cranelift 376 / wasm 371, with 5 jit-stats rows moving. No `.jitstats` re-recorded here. Assisted-by: Claude
`call_valuestack` (baseobjspace.py:1243) is the only call dispatcher upstream hands a frame. `call_args` (descroperation.py:189) and `Function.call_args` (function.py:79) take none, and the C-level profile check sits at the opcode level, where `self` is the frame (pyopcode.py:1402 CALL_FUNCTION_KW, :1429 CALL_FUNCTION_EX). pyre had that profile check inside the generic dispatcher, so a caller frame had to be threaded through the whole chain. Convert `call_user_function_resolved`, `call_callable_with_mode`, `call_non_function_callable_with_mode`, `call_kw_in_ctx`, `call_function_ex_in_ctx`, `call_with_kwargs_in_ctx` and `type_descr_call_with_mode` to take `*const PyExecutionContext`, and add `c_profile_frame`, which resolves a caller frame only when a profiler is installed — `executioncontext.py:147-149 call_trace` sets `is_being_profiled` only while `profilefunc is not None`, and `:121-123 _c_call_return_trace` clears it, so testing the execution context first preserves the implication. `call_callable`, `call_kw`, `call_function_ex` and `call_with_kwargs` keep their frame-taking signatures as wrappers that install `FrameLocalsRoot` and delegate to the `_in_ctx` bodies, so no external caller changes. `bh_call_kw_impl` and `bh_call_function_ex_fn` no longer resolve a caller frame, matching `bh_call_fn_impl`. Assisted-by: Claude
`latch_abort_call_resume` and `run_perfn_walk`'s `call_forward_abort` each
listed the same two DispatchError variants, so a walk that aborted with
`BranchGuardUnrestorableKeptStackPermanent` or `BranchGuardKeptStackUnsupported`
inside a top-level inline sub-walk had no gh#467 carrier. Its remaining leg,
the kept-stack branch flush, declines on `abort_in_subwalk`, and the legacy drop
behind that decline re-enters the outer frame at its entry — re-running the
residual calls the authoritative walk had already executed concretely.
`threading.Thread.start` then calls `_start_joinable_thread` a second time on a
handle that is already started.
Both allowlists now carry the two kept-stack variants. The soundness gate is
unchanged and central: `try_commit_entry_carrier_call` commits
`WalkEndResume::Rewind { effects_at_resume_point }`, which `commit_walk_end`
declines when the executed-effect odometer moved since the outer CALL, and
`latch_abort_call_resume` latches only for a top-level inline whose callee
executed no concrete effect. A third field on `call_forward_abort` keeps the
new variants off the `MidBody` carrier, which only the original two latch. The
kept-stack flush leg now checks `WALK_END_FLUSH_COMMITTED` first so it cannot
reposition a frame the carrier leg already flushed.
`[fbw-abort]` additionally names the walked code object and the executed-effect
count under `PYRE_FBW_DEBUG_ABORT`.
Measured, dynasm and cranelift: the new parity test fails 4/4 on the parent
commit (deterministically, `RuntimeError: thread already started`) and passes
4/4 with the change; `check.py --backend dynasm,cranelift` reports no functional
failures and no jit-stats delta attributable to the change (the two `improved:
guard_failures` rows reproduce identically on a parent-commit binary);
`cargo test --workspace --features dynasm` and the parity suite are green.
Assisted-by: Claude
…lowers `exception_subclass_attrs` dynasm and cranelift 3 -> 2, and `inline_gate_operand_provenance` dynasm 5 -> 4. All three OS jobs on PR #1046 reported the identical value (macos-latest, ubuntu-24.04, windows-latest), so one number satisfies every host. `inline_gate_operand_provenance`'s optimized-trace op counts are unchanged across the change (61/61/61/49 before and after), so the move is deopt accounting rather than codegen. The two wasm rows check.py reports as regressed are not recorded. Assisted-by: Claude
df3c1e8 to
d49bef2
Compare
There was a problem hiding this comment.
💡 Codex Review
https://github.com/youknowone/pyre/blob/d49bef22a13e5a0c0de22bfc18528cfbd58a515c/pyre-jit/src/call_jit.rs#L4800-L4802
When this residual takes the new frameless path, the incoming callable, null_or_self, kwnames, and positional values are still copied native parameters; unlike bh_call_fn_impl, they are not pinned/reloaded before call_kw_in_ctx can allocate or enter descriptor/keyword handling. In a JIT CALL_KW residual where a minor collection happens before the eventual callee frame owns those objects, the moving GC can forward the real roots while these copied pointers remain stale, leading to wrong arguments or a crash.
https://github.com/youknowone/pyre/blob/d49bef22a13e5a0c0de22bfc18528cfbd58a515c/pyre-jit/src/call_jit.rs#L5079-L5085
This frameless residual passes copied native PyObjectRef parameters straight into call_function_ex_in_ctx, but that helper may run arbitrary Python while unpacking a non-tuple *args or resolving a non-dict **kwargs before any callee frame roots these values. As bh_call_fn_impl notes for the same ABI boundary, the backend gcmap keeps objects alive but cannot rewrite these copied parameters after a moving collection, so CALL_FUNCTION_EX residuals can continue with stale callable/starargs/kwargs pointers unless they are pinned and reloaded first.
ℹ️ 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".
`foriter_call_resume_drops_iteration` loops_aborted 4 -> 2, guard_failures 5182 -> 5165, bridges_compiled 23 -> 26: two loops that aborted now compile, and all three fields move toward the dynasm and cranelift baselines (1 / 5150 / 27). `kept_stack_deep_var_shortcircuit` guard_failures 823 -> 824, with loops_compiled, bridges_compiled and loops_aborted unchanged. dynasm is bit-identical across the change (6 / 4 / 0 / 819 both before and after), so the count is wasm-side. The fixture has no self-recursive call, so `bh_call_self_recursive_portal` declines both before and after; what the change removes ahead of that decline is the `gettopframe_raw()` that forced and materialized the caller frame at every `g(i)` / `h(i)` residual. Both values are what ubuntu-24.04 measured on PR #1046 — the only host that runs the wasm backend — and reproduce three runs out of three locally. Restoring call.rs and call_jit.rs to their origin/main contents returns both rows to their previously recorded baselines. `fbw_rolled_back_with_effects` enters both files because 64880ab added the counter. Assisted-by: Claude
1 participant
The residual-call helpers a compiled trace reaches (
bh_call_fn_impl,bh_call_kw_impl,bh_call_function_ex_fn) resolved their caller frame withgetexecutioncontext().gettopframe_raw(). That call isforce_vref, andforcing the frame-chain vref while its token is
TOKEN_TRACING_RESCALLisprecisely how the protocol reports "the callee escaped this frame":
virtualref.py:161-167—force_virtualon aTRACING_RESCALLtoken clearsthe token, "reset token as the marker that this 'virtual' escapes".
tracing_after_residual_callthen reads the cleared token, andvrefs_after_residual_call→stop_tracking_virtualrefrecords a non-NULLVIRTUAL_REF_FINISH(vrefbox, virtualbox)ahead of theCALL.So the trace materialised the caller frame on every execution of the loop —
the reader's own read was recorded as the callee's escape. Upstream never hits
this:
bhimpl_residual_call_r_r(blackhole.py:1227) passescpu.bh_call_r(func, None, args_r, …), no frame at all.jit: thread the execution context through the CALL residual instead of the caller frameprepare_user_callonly ever readframe.execution_context, so it takes theexecution context directly and
call_user_function_residual_with_ctxskips thecaller-side
FrameLocalsRoot.bh_call_fn_implasserts a pinned executioncontext instead of resolving a frame, and
bh_call_self_recursive_portal's"is the caller running this same code object?" test uses the non-forcing
vref_referent— which exists for exactly this (executioncontext.rs:71-85:"For identity tests only, never for handing a frame to application code… A
reader that forced would report its own read as a callee escape.").
interp: separate the frame-taking call dispatcher from the frameless oneUpstream has exactly one frame-taking dispatcher,
call_valuestack(
baseobjspace.py:1243).call_args(descroperation.py:189) andFunction.call_args(function.py:79) take no frame, and the C-level profilecheck lives at the opcode level, where
selfis the frame(
pyopcode.py:1402 CALL_FUNCTION_KW,:1429 CALL_FUNCTION_EX). pyre had thatcheck inside the generic dispatcher, which is what forced a frame through the
whole chain.
Seven functions now take
*const PyExecutionContext(
call_user_function_resolved,call_callable_with_mode,call_non_function_callable_with_mode,call_kw_in_ctx,call_function_ex_in_ctx,call_with_kwargs_in_ctx,type_descr_call_with_mode). The newc_profile_frameresolves a frame onlywhen a profiler is actually installed —
executioncontext.py:147-149 call_tracesets
is_being_profiledonly whileprofilefunc is not Noneand:121-123 _c_call_return_traceclears it, so testing the execution context firstpreserves the implication.
call_callable,call_kw,call_function_exandcall_with_kwargskeep their frame-taking signatures as wrappers that installFrameLocalsRootand delegate, so no external caller changes.bh_call_kw_implandbh_call_function_ex_fnare frameless as of this commit.The only
gettopframe_rawleft incall_jit.rsis the env-gatedPYRE_BH_NULL_ARGdiagnostic.jit(fbw): scan the walked frame's own jitcode in classify_compare_box_useSeparate, self-contained:
classify_compare_box_usedecodedcompare_pcagainst the snapshot root's JitCode and so bailed out with
Otheron everyinlined sub-walk, because a callee's offset does not index the root's bytes.
Take the bytes from the frame the walk is actually in
(
framestack.last().w_code→sub_jitcode_body_for_code). The liveness tablescondition 4 reads are keyed on the root jitcode index and do not describe a
callee, so the sub-walk path reports the shape with
arms_dead: false; the boxelision in
compare_box_provably_deadstill requires condition 4 and keepsdeclining inside a sub-walk. On
bench/fib_recursive.pythis drops threeCallR(newbool)residuals whose results are read by nothing (two of them in thehottest loop), with recorded jit-stats unchanged.
Measurements
Measured against a base that already contains the commit above, so the numbers
below belong to the two frameless commits. Minor collections (
MAJIT_LOG=1,count of
[gc][minor] start— deterministic and load-independent), plusinterleaved min-of-9 child CPU:
fib_recursivelinrectreerecnbody— control, no self-recursive residualfib_recursivechild CPU 0.868x (base,fix order) and 0.886x (fix,baseorder). All program outputs byte-identical.
jit-stats
Five rows move, and all five are recorded.
exception_subclass_attrsguard_failuresinline_gate_operand_provenanceguard_failuresforiter_call_resume_drops_iterationkept_stack_deep_var_shortcircuitguard_failuresThe three native rows came back with the identical number from all three
OS jobs on this PR's first CI run (macos-latest, ubuntu-24.04, windows-latest),
so one recorded value satisfies every host — the condition that blocks the
usual macOS-only re-record. Their optimized-trace op counts are unchanged
across the change (
inline_gate_operand_provenance61/61/61/49,kept_stack_deep_var_shortcircuit105/131), so the moves are deopt accounting.The two wasm rows were attributed and root-caused before being recorded:
call.rsandcall_jit.rsto theirorigin/maincontents —leaving everything else on the branch in place — returns both rows to
their previous baselines. So they belong to the frameless commit, not to the
classify_compare_box_usecommit, not to the kept-stack commit, and not tothe newer base.
foriter_call_resume_drops_iterationis a net improvement with a mechanicalreading:
loops_aborted 4 → 2andguard_failures 5182 → 5165, andbridges_compiled 23 → 26is two aborted loops becoming compiled. All threefields move toward the dynasm/cranelift committed baselines (1 / 5150 / 27).
kept_stack_deep_var_shortcircuitisguard_failures +1withloops_compiled,bridges_compiledandloops_abortedall flat. dynasm isbit-identical across the change (6 / 4 / 0 / 819 before and after), so the
count is wasm-side, and it reproduces 3 runs out of 3. The fixture has no
self-recursive call, so
bh_call_self_recursive_portaldeclines both beforeand after; what the change removes ahead of that decline is the
gettopframe_raw()that forced and materialised the caller frame at everyg(i)/h(i)residual. The fixture's own assertion —len(log)exactly3N, which a doubled delivery over-counts and a dropped iterationunder-counts — passes.
The
call_argscommit moves no counter at all: its gate output isrow-for-row identical to the frameless commit's, on two different bases.
fbw_rolled_back_with_effectsenters the two wasm files because 64880ab(#1056) added that counter.
jit: route kept-stack branch aborts through the CALL-forward flushAlso on this branch (authored in a parallel session).
latch_abort_call_resumeand
run_perfn_walk'scall_forward_aborteach listed the same twoDispatchErrorvariants, so a walk aborting withBranchGuardUnrestorableKeptStackPermanentorBranchGuardKeptStackUnsupportedinside a top-level inline sub-walk had no gh#467 carrier; the legacy drop behind
the declining kept-stack flush re-entered the outer frame at its entry and
re-ran residual calls the authoritative walk had already executed concretely —
threading.Thread.startcalling_start_joinable_threadtwice. Both allowlistsnow carry the two kept-stack variants, and the flush leg checks
WALK_END_FLUSH_COMMITTEDfirst. Its parity test fails 4/4 on the parent commitand passes 4/4 with the change.
Left for a follow-up
Three sites in
call.rs(build_class's non-type metaclass,__prepare__, and__init_subclass__) still resolve a frame throughgettopframe_raw()purely tosatisfy the old
call_with_kwargs(frame, …)signature — and silently drop thekeywords when that frame is null, which
compiling.py:213-219does not do(
space.call_args(w_meta, Arguments(…))passes them unconditionally).call_with_kwargs_in_ctxmakes the frame unnecessary there now, but deliveringkeywords where they were previously dropped is a behaviour change orthogonal to
this branch and wants its own gate.
Verification
cargo test -p pyre-jit -p pyre-interpreter --features dynasm— 490 + 329 + 23passed.
pyre/check.py --backend dynasm,cranelift,wasm— zero correctness failuresthroughout; before the rows were recorded the only failures were the five
jit-stats snapshots above. With the three native rows recorded: dynasm
379/379 and cranelift 379/379 ALL PASSED.
🤖 Generated with Claude Code