jit: admit a retrace's close onto its attach token, align the retrace path with unroll.py, and fix the wasm livelock it exposed - #1134
Conversation
…hether the trace carries a LABEL `has_cross_loop_terminal_jump` answered `has_jump && !has_label`, while the code generator answers the same question by descr identity in `find_loop_label_index` — `x86/assembler.py:2463`'s `target_token in self.target_tokens_currently_compiling`. The two disagree on a trace that defines LABELs of its own and whose terminal JUMP names none of them: codegen takes the external arm and emits `return_call_indirect(external_jump_slot)`, but the predicate answered false, so neither `compile_loop` nor `compile_bridge` called `resolve_cross_loop_jump_target` and the arm ran with `external_jump_key = 0` and `external_jump_slot = source_func_handle`. Key 0 re-enters the target at its function entry, so a peeled target re-runs its preamble against mid-loop state and the induction variable never advances — the livelock `compile_bridge` already documents above `bridge_is_loop_closing`, reached through a shape the predicate did not recognise. A tail call keeps the stack and heap flat, so it presents as a hang, not a crash. Both callers keep the one predicate. Widening is safe at each: they either resolve a real target through `resolve_cross_loop_jump_target` — adopting its frozen frame geometry, which a tail call requires anyway — or decline with `BackendError::Unsupported`, dropping the trace back to the interpreter. Assisted-by: Claude
…en this compilation attaches to `_jump_to_existing_trace` walks `jitcelltoken.target_tokens` (unroll.py:171), so every candidate upstream offers belongs to one JitCellToken. pyre seeds the candidates from `compiled_loops[green_key].front_target_tokens`, a green-key side table that survives a recompile, so a candidate can name a token of an earlier, retired compilation. The unroll pass therefore discarded every close whose JUMP did not name the body token this compilation had just pushed. Carry the token the artifact will be installed under and admit a close onto any of its target tokens. `compile_retrace` resolves it from the source guard's own loop token, which is where `compile.py:797-811` attaches the result, so such a close stays inside one code buffer. A `compile_loop` has no such token yet — `compile.py:287-289` expresses the same thing by resetting `jitcell_token.target_tokens` to `[start_descr]`, leaving nothing matchable — so it admits nothing beyond its own body. `synth/retrace_outer_loop_type_flip` goes from `loops_aborted=2 retraces_compiled=0 guard_failures=590` to `loops_aborted=0 retraces_compiled=1 bridges_compiled=1 guard_failures=201` on all three backends. Assisted-by: Claude
WalkthroughThe change fixes cross-loop jump classification and unroll loop-close validation by using local label descriptors and JitCellToken ownership. Benchmark statistics now record successful bridge and retrace compilation. ChangesCross-loop target handling
Estimated code review effort: 4 (Complex) | ~45 minutes 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
Here are some automated review suggestions for this pull request.
Reviewed commit: b87ad9e8aa
ℹ️ 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".
| /// `compile.py:355` — the `JitCellToken` this compilation's artifact will be | ||
| /// installed under, when there is one. `compile_retrace` resolves it from | ||
| /// `get_procedure_token(greenkey)` and the result is attached as a bridge to | ||
| /// that same token (`compile.py:797-811`), so a close onto one of its | ||
| /// `target_tokens` stays inside one code buffer. | ||
| /// | ||
| /// `None` for a fresh `compile_loop`, which mints its own token afterwards. | ||
| /// `compile.py:287-289` expresses the same thing by *resetting* | ||
| /// `jitcell_token.target_tokens` to `[start_descr]`, leaving nothing for | ||
| /// `_jump_to_existing_trace` to match; pyre seeds the previous | ||
| /// compilation's tokens instead, so the distinction has to be carried here. | ||
| pub attach_jitcell_token_number: Option<u64>, |
There was a problem hiding this comment.
This adds a second, numeric ownership channel solely to compensate for sourcing target_tokens from the green-key front_target_tokens side table, whereas upstream passes the exact JitCellToken into UnrolledLoopData and scans that object's target_tokens directly. The change therefore preserves the structural divergence and makes close admission depend on separately maintained and later restamped token numbers; restore the upstream JCT-owned target-token list rather than layering this special case over the side table.
AGENTS.md reference: AGENTS.md:L231-L233
Useful? React with 👍 / 👎.
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit 66199e0). 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
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🤖 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-metainterp/src/optimizeopt/unroll.rs`:
- Around line 254-265: Correct the documentation for attach_jitcell_token_number
to state that its setter derives the token from retrace_resumekey.source_descr
via majit_backend::descr_owning_jct, representing the guard’s owning loop token.
Remove the claim that compile_retrace resolves this field through
get_procedure_token(greenkey), while distinguishing that live procedure-token
lookup as the separate loop_jitcell_token value.
🪄 Autofix
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 Plus
Run ID: 27410962-789f-4110-912c-cb0b1137da73
📒 Files selected for processing (7)
majit/majit-backend-wasm/src/codegen.rsmajit/majit-backend-wasm/src/lib.rsmajit/majit-metainterp/src/optimizeopt/unroll.rsmajit/majit-metainterp/src/pyjitpl.rspyre/bench/synth/retrace_outer_loop_type_flip.cranelift.jitstatspyre/bench/synth/retrace_outer_loop_type_flip.dynasm.jitstatspyre/bench/synth/retrace_outer_loop_type_flip.wasm.jitstats
| /// `compile.py:355` — the `JitCellToken` this compilation's artifact will be | ||
| /// installed under, when there is one. `compile_retrace` resolves it from | ||
| /// `get_procedure_token(greenkey)` and the result is attached as a bridge to | ||
| /// that same token (`compile.py:797-811`), so a close onto one of its | ||
| /// `target_tokens` stays inside one code buffer. | ||
| /// | ||
| /// `None` for a fresh `compile_loop`, which mints its own token afterwards. | ||
| /// `compile.py:287-289` expresses the same thing by *resetting* | ||
| /// `jitcell_token.target_tokens` to `[start_descr]`, leaving nothing for | ||
| /// `_jump_to_existing_trace` to match; pyre seeds the previous | ||
| /// compilation's tokens instead, so the distinction has to be carried here. | ||
| pub attach_jitcell_token_number: Option<u64>, |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win
Fix the field doc: attach_jitcell_token_number is not resolved via get_procedure_token(greenkey).
The comment states that compile_retrace resolves this token from get_procedure_token(greenkey). The actual (and only) setter, in majit/majit-metainterp/src/pyjitpl.rs (lines 7679-7690), derives it from retrace_resumekey.source_descr's owning JitCellToken through majit_backend::descr_owning_jct. That is the guard's own loop token, which can differ from the loop's current live procedure token returned by get_procedure_token(greenkey) if the loop was recompiled since the guard was created. compile_retrace already resolves the live procedure token separately, into the unrelated loop_jitcell_token variable.
Update the doc to describe the actual derivation path so future readers auditing RPython parity do not rely on the wrong resolution source.
📝 Proposed doc fix
/// `compile.py:355` — the `JitCellToken` this compilation's artifact will be
- /// installed under, when there is one. `compile_retrace` resolves it from
- /// `get_procedure_token(greenkey)` and the result is attached as a bridge to
- /// that same token (`compile.py:797-811`), so a close onto one of its
+ /// installed under, when there is one. For a guard-originated retrace,
+ /// `compile_retrace` resolves it from the source guard's own descr-owning
+ /// `JitCellToken` (`descr_owning_jct`), and the result is attached as a
+ /// bridge to that same token (`compile.py:797-811`), so a close onto one of its
/// `target_tokens` stays inside one code buffer.📝 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.
| /// `compile.py:355` — the `JitCellToken` this compilation's artifact will be | |
| /// installed under, when there is one. `compile_retrace` resolves it from | |
| /// `get_procedure_token(greenkey)` and the result is attached as a bridge to | |
| /// that same token (`compile.py:797-811`), so a close onto one of its | |
| /// `target_tokens` stays inside one code buffer. | |
| /// | |
| /// `None` for a fresh `compile_loop`, which mints its own token afterwards. | |
| /// `compile.py:287-289` expresses the same thing by *resetting* | |
| /// `jitcell_token.target_tokens` to `[start_descr]`, leaving nothing for | |
| /// `_jump_to_existing_trace` to match; pyre seeds the previous | |
| /// compilation's tokens instead, so the distinction has to be carried here. | |
| pub attach_jitcell_token_number: Option<u64>, | |
| /// `compile.py:355` — the `JitCellToken` this compilation's artifact will be | |
| /// installed under, when there is one. For a guard-originated retrace, | |
| /// `compile_retrace` resolves it from the source guard's own descr-owning | |
| /// `JitCellToken` (`descr_owning_jct`), and the result is attached as a | |
| /// bridge to that same token (`compile.py:797-811`), so a close onto one of its | |
| /// `target_tokens` stays inside one code buffer. | |
| /// | |
| /// `None` for a fresh `compile_loop`, which mints its own token afterwards. | |
| /// `compile.py:287-289` expresses the same thing by *resetting* | |
| /// `jitcell_token.target_tokens` to `[start_descr]`, leaving nothing for | |
| /// `_jump_to_existing_trace` to match; pyre seeds the previous | |
| /// compilation's tokens instead, so the distinction has to be carried here. | |
| pub attach_jitcell_token_number: Option<u64>, |
🤖 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/optimizeopt/unroll.rs` around lines 254 - 265,
Correct the documentation for attach_jitcell_token_number to state that its
setter derives the token from retrace_resumekey.source_descr via
majit_backend::descr_owning_jct, representing the guard’s owning loop token.
Remove the claim that compile_retrace resolves this field through
get_procedure_token(greenkey), while distinguishing that live procedure-token
lookup as the separate loop_jitcell_token value.
…time boxes, and neither the budget nor the disable sentinel from the loop path Three divergences from `optimize_peeled_loop` (unroll.py:112-180), which contains no `retraced_count` bookkeeping at all and calls `disable_retracing_if_max_retrace_guards` exactly once. The `force_boxes=true` retry at unroll.py:161-168 is unconditional. pyre had copied `optimize_bridge`'s accounting (unroll.py:213-226) onto the loop path, so every loop compile whose first match missed spent one unit of the per-JitCellToken retrace budget that upstream reserves for bridges. The budget's own increment lives on the bridge path (`tok.set_retraced_count(tok.get_retraced_count() + 1)`, unroll.py:213-215) and is untouched. The two arms differed only in that bookkeeping, so they collapse into one call. `disable_retracing_if_max_retrace_guards` ran before the close ladder, and a second time over the combined preamble+body list. Upstream runs it after both `jump_to_preamble` early returns — only for a loop that closed — over `self._newoperations`, the peeled body. Either write sets `retraced_count = u32::MAX`, which the early check in `compile_loop` reads as a permanent "skipping recompile" for that green key. `unroll.py:231-233` threads the bridge's own `runtime_boxes` into `ExportedState`, and unroll.py:153/166 pass `state.runtime_boxes` to `jump_to_existing_trace`. The only assignment in pyre was on the `compile_loop` branch, so `generate_guards` saw an empty list on every retrace and each runtime-guided arm — GUARD_VALUE, GUARD_NONNULL, GUARD_CLASS/GUARD_NONNULL_CLASS, the `IntBound::make_guards` fallback — declined for want of a value to read. The length-mismatch fallback in unroll.rs is left alone; this makes its "a trace with no recorded JUMP" comment true again. No jit-stats and no wallclock movement across the corpus: all three backends gate clean, and an in-place revert A/B on dynasm with the arms interleaved reads 0.31-0.33s reverted against 0.32-0.34s applied. `retrace_limit` defaults to 0 (`rpython/rlib/jit.py:595`), so the retrace legs stay dormant for code that does not raise it. Assisted-by: Claude
`guard_failures` 13 -> 2 with `loops_compiled` unchanged at 2. Not this branch's doing: reverting this branch's three metainterp files and rebuilding the wasm module reads the same 2, so the move belongs to the base this branch was rebased onto — `jit: residualise the whole int-box tail, and decline a kept-stack branch holding a NULL ConstPtr` (#1131). `check.py` reports the drop as IMPROVED and refuses to pass until the baseline is recorded. Assisted-by: Claude
youknowone
changed the title
1 participant
Four commits. The first two are the measured change; the third is parity alignment with no
measured movement; the fourth is a baseline that belongs to
main.1.
_jump_to_existing_tracewas offered candidates it could never take_jump_to_existing_tracewalksjitcelltoken.target_tokens(unroll.py:171), so everycandidate upstream offers belongs to one
JitCellToken. pyre seeds the candidates fromcompiled_loops[green_key].front_target_tokens, a green-key side table that survives arecompile, so a candidate can name a token of an earlier, retired compilation. The unroll
pass defended against that by discarding every close whose JUMP did not name the body
token this compilation had just pushed — including the legitimate ones.
This carries the token the artifact will be installed under and admits a close onto any of
its target tokens.
compile_retraceresolves it from the source guard's own loop token,which is where
compile.py:797-811attaches the result, so such a close stays inside onecode buffer. A
compile_loophas no such token yet —compile.py:287-289expresses thesame thing by resetting
jitcell_token.target_tokensto[start_descr], leaving nothingmatchable — so it admits nothing beyond its own body.
synth/retrace_outer_loop_type_flip, the acceptance fixture landed with #1127, on allthree backends:
loops_compiledbridges_compiledretraces_compiledloops_abortedguard_failures2. The wasm livelock the admission exposed
Admitting the close made that fixture hang forever on wasm — flat 89 MB RSS, so a
constant-stack-depth livelock rather than a leak.
PYRE_WASM_DUMP_ALL_TRACESshowed theretrace module tail-calling the loop with dispatch key 0, while the loop's entry
dispatch is
br_table 0 (;@6;) 1 (;@5;) 3 (;@3;) 0 (;@6;)— key 0 lands on the functionentry, so the peeled target re-ran its preamble against mid-loop state and the induction
variable never advanced.
has_cross_loop_terminal_jumpdecided "does this JUMP leave the trace?" by the structuralproxy
has_jump && !has_label, while the code generator decides the same question by descridentity in
find_loop_label_index—x86/assembler.py:2463'starget_token in self.target_tokens_currently_compiling. The two disagree on exactly oneshape: a trace that defines LABELs of its own and whose terminal JUMP names none of them,
which is precisely a retrace attached as a bridge. Codegen took the external arm; the
predicate had answered false, so neither caller ever called
resolve_cross_loop_jump_targetand the arm ran withexternal_jump_key = 0andexternal_jump_slot = source_func_handle.compile_bridgealready documents this livelock abovebridge_is_loop_closing("the wasmchaining hang on nbody / fannkuch") — it was reached through a shape its guard did not
recognise. Both callers now share the one token-based predicate.
3. Three divergences from
optimize_peeled_loop, with no measured movementoptimize_peeled_loop(unroll.py:112-180) contains noretraced_countbookkeeping atall and calls
disable_retracing_if_max_retrace_guardsexactly once.force_boxes=trueretry (unroll.py:161-168) is unconditional. pyre had copiedoptimize_bridge's accounting (unroll.py:213-226) onto the loop path, so every loopcompile whose first match missed spent one unit of the per-
JitCellTokenretrace budgetupstream reserves for bridges. The budget's own increment lives on the bridge path
(
unroll.py:213-215) and is untouched.disable_retracing_if_max_retrace_guardsran before the close ladder, and a second timeover the combined preamble+body list. Upstream runs it after both
jump_to_preambleearly returns — only for a loop that closed — over
self._newoperations, the peeledbody. Either write sets
retraced_count = u32::MAX, whichcompile_loop's early checkreads as a permanent "skipping recompile" for that green key.
unroll.py:231-233threads the bridge's ownruntime_boxesintoExportedState, andunroll.py:153/:166passstate.runtime_boxestojump_to_existing_trace. The onlyassignment in pyre was on the
compile_loopbranch, sogenerate_guardssaw an emptylist on every retrace and each runtime-guided arm — GUARD_VALUE, GUARD_NONNULL,
GUARD_CLASS/GUARD_NONNULL_CLASS, the
IntBound::make_guardsfallback — declined for wantof a value to read.
virtualstate.py:550-555calls those runtime values the "educatedguess" that picks guard-vs-retrace.
This commit moves nothing measurable. All three backends gate clean with no jit-stats
change, and an in-place revert A/B on dynasm with the arms interleaved reads 0.31-0.33s
reverted against 0.32-0.34s applied.
retrace_limitdefaults to 0(
rpython/rlib/jit.py:595), so the retrace legs stay dormant for code that does not raiseit. It is parity work, not a speedup.
4.
synth/unary_negative's wasm baselineguard_failures13 -> 2. Not this branch's doing — reverting this branch's three metainterpfiles and rebuilding the wasm module reads the same 2, so it belongs to #1131, which this
branch was rebased onto mid-session.
Gate
Full
check.pyon all three backends, on a freshly re-extracted LLBC, at the tip of thisbranch: dynasm 415/415, cranelift 414/414, wasm 409/409 + the one IMPROVED above. No
SNAPDIFFin any log.An earlier run of the same gate, at load 12-36, reported
synth/list_pop_appendandcpython-suite test_pickletoolsfailing; both went away at load 2.4 on identical code, andmain's own run at this branch's original parent failslist_pop_appendon ubuntu andwindows.
— authored by Claude