jit(wasm): CA callee jitframes as nursery GC objects + live-Box virtualizable identity#501
Conversation
WalkthroughThe PR parameterizes wasm frame layouts, makes CALL_ASSEMBLER frames nursery-GC-managed, and adds frame/Ref reload handling across collecting operations. It also replaces virtualizable identity constants with live boxes and supports optional host-stack identity overrides during blackhole resume. ChangesGC-managed CALL_ASSEMBLER frames
Virtualizable resume identity
Estimated code review effort: 5 (Critical) | ~120 minutes Possibly related issues
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 |
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit ea9f654). 1. Regressions to PyPy parity introduced by this patch
2. Other mismatches introduced by this patchNone. 3. Pre-existing mismatches (already present before this patch)
4. Structural adaptations
|
![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
Here are some automated review suggestions for this pull request.
Reviewed commit: a1d16c1626
ℹ️ 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".
| op.opcode, | ||
| OpCode::New | OpCode::NewWithVtable | OpCode::NewArray | OpCode::NewArrayClear | ||
| ) | ||
| (op.opcode.is_call() |
There was a problem hiding this comment.
When residual calls are treated as collecting here, RefHomes::collect still only homes refs that are live after the op via live_across_any; a Ref operand whose last use is this call is deliberately excluded. The direct residual-call path passes those args as raw wasm values, so for a call like f(obj) where obj dies at the call, a minor collection inside f cannot see or update obj and the callee can observe a moved/freed pointer. Please home Ref operands for collecting residual calls as the CA arm does, or keep these direct residual calls on a no-collect path.
Useful? React with 👍 / 👎.
…shots Remove state_field_identity_const and the identity_const parameter chain: build_vable_snapshot_boxes now always emits the identity as SnapshotTagged::Box. The identity is a live input box; rewriting it to TAGCONST excluded it from liveboxes/failarg numbering, baking a trace-time object address into resume data. next_ref() now decodes the identity from the failing frame's failarg slot (resume.py:1566-1578). Assisted-by: Claude
…e pointer after collecting calls wasm_jit_ca_alloc_frame switches from the malloc+LIFO pool to a collecting nursery allocation (rewrite.py gen_malloc_nursery_varsize_frame), retiring the CA_FRAMES pool. Because a minor collection can now move the running frame and the collector cannot update wasm locals, trace bodies reload local 0 from the jitframe shadow-stack top before every post-collection home reload (_reload_frame_if_necessary, assembler.py:1369), and the CA arm reloads the caller frame via the new jf_under_top_ptr() accessor after its callee push. Trampoline fallback paths keep the documented non-moving-frame assumption. fib32 via CA: gc_majors 262 (oldgen variant) -> 4. Synthetic suite 151/151 on wasm, dynasm, cranelift. Assisted-by: Claude
Assisted-by: Claude
Replace the fixed MIN_FRAME_BYTES frame layout for CA-allocated callee frames with a per-token FrameGeometry frozen at the token's first compile: value region sized from the token's actual slot needs, then call area, dispatch key, and Ref homes at per-token offsets (jitframe.py sizes frames by the regalloc frame_depth). Arena execute_token frames keep the fixed floor. compile_bridge declines a chained bridge whose slot needs exceed the frozen geometry. Host-side frame readers and per-bridge gcmap building take the token geometry instead of the global constants. fib32 via CA: gc_minors 5034 -> 1243, wall 1.20s -> 0.66-0.69s (malloc-pool baseline 0.83s). Synthetic suite 151/151 on wasm, dynasm, cranelift. Assisted-by: Claude
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: e80d0bd3bb
ℹ️ 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".
| resumereader.handling_async_forcing(); | ||
| // resume.py:1350 | ||
| resumereader.consume_vref_and_vable(vrefinfo, vinfo, ginfo); | ||
| resumereader.consume_vref_and_vable(vrefinfo, vinfo, ginfo, None); |
There was a problem hiding this comment.
On the async-forcing path (handle_async_forcing -> force_from_resumedata), state-field virtualizables have the same host-stack identity issue that the new blackhole paths handle with vable_identity_override. This call still passes None, so after the snapshot writer stopped encoding the &state identity as a const, a GUARD_NOT_FORCED force can decode the folded-out TAGBOX/null in consume_vable_info before force_all_virtuals, misidentifying the virtualizable. Please pass an equivalent current-call identity here or keep const encoding for this force path.
Useful? React with 👍 / 👎.
| sink.i32_const(ca_reload_fn_ptr as i32); | ||
| sink.call_indirect(0, base); | ||
| sink.i32_wrap_i64(); | ||
| sink.local_set(0); |
There was a problem hiding this comment.
When execute_token takes the legacy Vec-backed frame path (wasm_jitframe_tid() == 0), no jitframe is pushed on the jf shadow stack. Direct helper/alloc paths now emit this reload after collecting calls, so wasm_jit_ca_reload_frame reads an empty jf stack and returns FIRST_ITEM_OFFSET; local 0 is then overwritten with that small offset instead of the Vec frame pointer, and subsequent frame/home accesses go to the wrong address. Please gate this reload to GC-managed JitFrame execution, or keep direct collecting helpers disabled for the legacy frame path.
Useful? React with 👍 / 👎.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 3
🤖 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 `@majit/majit-backend-wasm/src/codegen.rs`:
- Around line 76-117: Update the trampoline-call fallback using
FrameGeometry::compact and jit_call so calls with more than 16 arguments are
rejected before writing call_args entries. Enforce the limit against
FrameGeometry::CALL_AREA_SLOTS’ 16 argument capacity, preserving existing
behavior for valid calls and preventing writes into dispatch_key_ofs or the
Ref-home region.
In `@majit/majit-metainterp/src/jitdriver.rs`:
- Around line 2799-2807: Extract the duplicated virtualizable identity
derivation into a shared `compute_vable_identity_override` method on the
surrounding type, accepting the state and compiled metadata. Replace the inline
`self.meta.virtualizable_info().and_then(...)` blocks in `back_edge_internal`
and `run_back_edge_generic` with calls to this helper, preserving the existing
`Option<i64>` behavior and arguments.
In `@majit/majit-metainterp/src/resume.rs`:
- Around line 6872-6880: Update the identity handling in the resume-data path
around `encoded_identity` and `virtualizable` so `identity_override` consumes
the raw resume item without calling `next_ref()` or decoding it; call
`next_ref()` only when no override is provided, preserving the existing decoded
identity behavior for the non-override path.
🪄 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: Organization UI
Review profile: ASSERTIVE
Plan: Pro
Run ID: b998267b-8ea3-4b2d-a207-a5641dd56509
📒 Files selected for processing (16)
majit/majit-backend-wasm/js/jit_glue.jsmajit/majit-backend-wasm/src/codegen.rsmajit/majit-backend-wasm/src/failguard.rsmajit/majit-backend-wasm/src/lib.rsmajit/majit-backend-wasm/tests/codegen_test.rsmajit/majit-gc/src/shadow_stack.rsmajit/majit-macros/src/jit_interp/codegen_state.rsmajit/majit-metainterp/src/blackhole.rsmajit/majit-metainterp/src/jit_state.rsmajit/majit-metainterp/src/jitdriver.rsmajit/majit-metainterp/src/pyjitpl/dispatch.rsmajit/majit-metainterp/src/resume.rsmajit/majit-metainterp/src/trace_ctx.rspyre/pyre-jit/src/call_jit.rspyre/pyre-wasm-runner/src/main.rspyre/pyre-wasm-runner/src/wasmi_host.rs
| const CALL_AREA_SLOTS: usize = 3 + 16; // result, function, nargs, args | ||
|
|
||
| /// Historical fixed geometry, used by direct codegen tests and by callers | ||
| /// that deliberately need the arena-compatible layout. | ||
| pub const fn fixed() -> Self { | ||
| Self { | ||
| value_slots: MIN_FRAME_BYTES / 8, | ||
| call_result_ofs: CALL_RESULT_OFS, | ||
| call_func_ofs: CALL_FUNC_OFS, | ||
| call_nargs_ofs: CALL_NARGS_OFS, | ||
| call_args_ofs: CALL_ARGS_OFS, | ||
| dispatch_key_ofs: DISPATCH_KEY_OFS, | ||
| home_slot_base: HOME_SLOT_BASE, | ||
| home_slots: 0, | ||
| frame_bytes: (MIN_FRAME_BYTES + SLOT_SIZE as usize) as u32, | ||
| } | ||
| } | ||
|
|
||
| /// Compact geometry for one token. `value_slots` includes frame[0], so | ||
| /// the call area begins immediately after the greatest positional input or | ||
| /// fail-arg slot used by that token. | ||
| pub fn compact(value_slots: usize, home_slots: usize) -> Self { | ||
| let value_slots = value_slots.max(1); | ||
| let call_result_ofs = (value_slots as u64) * SLOT_SIZE; | ||
| let call_func_ofs = call_result_ofs + SLOT_SIZE; | ||
| let call_nargs_ofs = call_func_ofs + SLOT_SIZE; | ||
| let call_args_ofs = call_nargs_ofs + SLOT_SIZE; | ||
| let dispatch_key_ofs = call_result_ofs + Self::CALL_AREA_SLOTS as u64 * SLOT_SIZE; | ||
| let home_slot_base = dispatch_key_ofs + SLOT_SIZE; | ||
| let frame_bytes = home_slot_base + home_slots as u64 * SLOT_SIZE; | ||
| Self { | ||
| value_slots, | ||
| call_result_ofs, | ||
| call_func_ofs, | ||
| call_nargs_ofs, | ||
| call_args_ofs, | ||
| dispatch_key_ofs, | ||
| home_slot_base, | ||
| home_slots, | ||
| frame_bytes: frame_bytes as u32, | ||
| } | ||
| } |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for any arity/args bound check gating call-area residual calls.
rg -nP --type=rust 'call_args_ofs|CALL_ARGS_OFS|CALL_AREA_SLOTS' majit/majit-backend-wasm/src/codegen.rs
# Look for declines related to call arity / too many args.
rg -nP --type=rust 'arity|nargs|too many|args\.len\(\)\s*[<>]=?' majit/majit-backend-wasm/src -C2
# Confirm max_value_slots / home_slots are the only frozen-layout bounds enforced.
rg -nP --type=rust 'exceed frozen frame layout' majit/majit-backend-wasm/src/codegen.rs -C3Repository: youknowone/pyre
Length of output: 30959
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the direct residual-call and trampoline emission paths.
sed -n '2820,2925p' majit/majit-backend-wasm/src/codegen.rs
printf '\n----\n'
sed -n '2925,3090p' majit/majit-backend-wasm/src/codegen.rs
# Look for any explicit caps on call argument counts in codegen.
printf '\n==== ARITY CHECKS ====\n'
rg -nP --type=rust 'saturating_sub\(1\)|call_args\.len\(\)|nargs|arity|too many|Unsupported\(format!\(' majit/majit-backend-wasm/src/codegen.rs -C2Repository: youknowone/pyre
Length of output: 25475
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Search for any global cap or validation on call argument counts.
rg -nP --type=rust 'MAX_.*ARG|max.*arg|call.*limit|too many.*arg|arg count|getarglist\(\)\.len\(\)|saturating_sub\(1\)' majit -C2
# Find any tests or docs describing the maximum supported residual-call arity.
rg -nP --type=rust '16 arg|16 args|CALL_AREA_SLOTS|dispatch_key_ofs|call_args_ofs' majit -C2
# Look for any compile-time or runtime checks in the backend that reject high-arity calls.
rg -nP --type=rust 'Unsupported\(format!\(".*arity|Unsupported\(format!\(".*args|return Err\(BackendError::Unsupported' majit/majit-backend-wasm/src -C3Repository: youknowone/pyre
Length of output: 50371
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Inspect the frame-layout commentary and the actual layout checks in build_wasm_module.
sed -n '35,55p' majit/majit-backend-wasm/src/codegen.rs
printf '\n----\n'
sed -n '1120,1165p' majit/majit-backend-wasm/src/codegen.rsRepository: youknowone/pyre
Length of output: 3404
Reject trampoline calls with >16 args. FrameGeometry::compact() reserves only 16 call-argument slots, but the jit_call fallback still writes call_args.len() entries with no cap. A larger fallback call will overwrite dispatch_key_ofs and the Ref-home region.
🤖 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 `@majit/majit-backend-wasm/src/codegen.rs` around lines 76 - 117, Update the
trampoline-call fallback using FrameGeometry::compact and jit_call so calls with
more than 16 arguments are rejected before writing call_args entries. Enforce
the limit against FrameGeometry::CALL_AREA_SLOTS’ 16 argument capacity,
preserving existing behavior for valid calls and preventing writes into
dispatch_key_ofs or the Ref-home region.
| // The state-field macro's `&state` is host-stack storage, so | ||
| // its identity may be folded out of the failing frame. Ask | ||
| // only an explicit host opt-in for the current call's address; | ||
| // heap virtualizables retain the live resume TAGBOX path. | ||
| let vable_identity_override = self.meta.virtualizable_info().and_then(|info| { | ||
| state | ||
| .blackhole_virtualizable_identity(&compiled_meta, &info.name, info) | ||
| .map(|ptr| ptr as i64) | ||
| }); |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Duplicate vable_identity_override computation across two call sites.
The same override-derivation block (self.meta.virtualizable_info().and_then(|info| state.blackhole_virtualizable_identity(...).map(|ptr| ptr as i64))) is duplicated verbatim in back_edge_internal and run_back_edge_generic. A future change to this logic risks drifting between the two.
♻️ Proposed helper extraction
+ fn compute_vable_identity_override(&self, state: &S, meta: &S::Meta) -> Option<i64> {
+ self.meta.virtualizable_info().and_then(|info| {
+ state
+ .blackhole_virtualizable_identity(meta, &info.name, info)
+ .map(|ptr| ptr as i64)
+ })
+ }Then replace both inline blocks with self.compute_vable_identity_override(state, &compiled_meta) / self.compute_vable_identity_override(state, &meta).
Also applies to: 4896-4902
🤖 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 `@majit/majit-metainterp/src/jitdriver.rs` around lines 2799 - 2807, Extract
the duplicated virtualizable identity derivation into a shared
`compute_vable_identity_override` method on the surrounding type, accepting the
state and compiled metadata. Replace the inline
`self.meta.virtualizable_info().and_then(...)` blocks in `back_edge_internal`
and `run_back_edge_generic` with calls to this helper, preserving the existing
`Option<i64>` behavior and arguments.
| // Consume the encoded identity even when a host supplies an override: | ||
| // it occupies one resume-data item and keeps the reader aligned for | ||
| // the field payload. Heap virtualizables (PyFrame) use this live | ||
| // TAGBOX exactly as RPython does. The state-field macro JIT opts in | ||
| // to `identity_override` because its host-stack `&state` is folded out | ||
| // of backend failargs; at deopt it must use the current call's address, | ||
| // never a trace-time pointer or an unrelated deadframe slot. | ||
| let encoded_identity = self.next_ref(); | ||
| let virtualizable = identity_override.unwrap_or(encoded_identity); |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | ⚡ Quick win
Do not decode the encoded identity when an override is present.
Line 6879 calls next_ref() before applying identity_override. For state-field JITs, the host-stack identity is omitted from backend failargs, so decoding the TAGBOX can index a missing or unrelated deadframe slot. Consume the raw resume item when overridden and decode only when no override is provided.
Proposed fix
- let encoded_identity = self.next_ref();
- let virtualizable = identity_override.unwrap_or(encoded_identity);
+ let virtualizable = match identity_override {
+ Some(identity) => {
+ let _ = self.resumecodereader.next_item();
+ identity
+ }
+ None => self.next_ref(),
+ };📝 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.
| // Consume the encoded identity even when a host supplies an override: | |
| // it occupies one resume-data item and keeps the reader aligned for | |
| // the field payload. Heap virtualizables (PyFrame) use this live | |
| // TAGBOX exactly as RPython does. The state-field macro JIT opts in | |
| // to `identity_override` because its host-stack `&state` is folded out | |
| // of backend failargs; at deopt it must use the current call's address, | |
| // never a trace-time pointer or an unrelated deadframe slot. | |
| let encoded_identity = self.next_ref(); | |
| let virtualizable = identity_override.unwrap_or(encoded_identity); | |
| // Consume the encoded identity even when a host supplies an override: | |
| // it occupies one resume-data item and keeps the reader aligned for | |
| // the field payload. Heap virtualizables (PyFrame) use this live | |
| // TAGBOX exactly as RPython does. The state-field macro JIT opts in | |
| // to `identity_override` because its host-stack `&state` is folded out | |
| // of backend failargs; at deopt it must use the current call's address, | |
| // never a trace-time pointer or an unrelated deadframe slot. | |
| let virtualizable = match identity_override { | |
| Some(identity) => { | |
| let _ = self.resumecodereader.next_item(); | |
| identity | |
| } | |
| None => self.next_ref(), | |
| }; |
🤖 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 `@majit/majit-metainterp/src/resume.rs` around lines 6872 - 6880, Update the
identity handling in the resume-data path around `encoded_identity` and
`virtualizable` so `identity_override` consumes the raw resume item without
calling `next_ref()` or decoding it; call `next_ref()` only when no override is
provided, preserving the existing decoded identity behavior for the non-override
path.
1 participant
Summary
Epic slice A1: wasm CALL_ASSEMBLER callee jitframes become real GC objects, replacing the malloc+LIFO pool, with the two invariants upstream requires for that to be sound.
878e6fc8409majit: encode the virtualizable identity as a live Box in resume snapshots — removesstate_field_identity_constand theidentity_constparameter chain. The virtualizable identity in a resume snapshot is now alwaysSnapshotTagged::Box, i.e. a live failarg decoded from the failing frame's slot at deopt time (pyjitpl.py:3319 appendsvirtualizable_boxes[-1]as a live box; resume.py:1566-1578 decodes it viaget_ref_value). The previous TAGCONST rewrite baked a trace-time object address into resume data — a latent unsoundness once frames can move.a1d16c1626awasm: allocate CA callee jitframes in the nursery and reload the frame pointer after collecting calls —wasm_jit_ca_alloc_frameswitches to a collecting nursery allocation (rewrite.pygen_malloc_nursery_varsize_frame), retiring theCA_FRAMESpool. Since a minor collection can now move the running frame and the collector cannot rewrite wasm locals, trace bodies reload local 0 from the jitframe shadow-stack top before every post-collection home reload — the_reload_frame_if_necessaryport (x86 assembler.py:1369) — and the CA arm reloads the caller frame via the newjf_under_top_ptr()accessor.Verification
Remaining
🤖 Generated with Claude Code
Summary by CodeRabbit