jit: FOR_ITER per-iteration getarrayitem_vable_r iterator reload (#152)#421
Conversation
Reload the FOR_ITER iterator from its value-stack slot each iteration via getarrayitem_vable_r (peekvalue port), giving the loop-carried iterator an in-loop reader so it resolves at the ForIterNext residual instead of aborting ResidualCallArgUnbound. The jit_merge_point reds are [frame, ec] only, so the operand stack is not register-carried across the resident loop; a preamble-only iterator SSA def cannot survive it. Withhold the reload when the loop body holds a CALL result across an in-body POP_JUMP_IF_* branch. That shape reaches a depth>1 kept operand stack at an interior guard whose compiled resume cannot reconstruct the deep slots; enabling the reload lets the walk consume the iterator and run residual calls concretely, then abort mid-body with an unrecoverable half-executed iteration. Withholding aborts at the FOR_ITER header before the consume, so the location interprets, bit-exact. Add kept_stack_deep_var_shortcircuit_mutate.py regression bench: g/h append to a global list, so a torn or doubled iteration is observable. Assisted-by: Claude
WalkthroughAdds a new adversarial Python benchmark script that exercises short-circuit boolean logic with side-effecting mutations, and modifies the JIT codewriter's ForIter instruction handling to detect a call/branch hazard within loop bodies, conditionally skipping the vable-based iterator reload when the hazard is present. ChangesForIter Iterator Hazard 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 |
🤖 Codex parity reviewStatic analysis of this diff vs the local RPython/PyPy sources (commit fbd0b06). 1. Regressions to PyPy parity introduced by this patchNone. 2. Other mismatches introduced by this patch
3. Pre-existing mismatches (already present before this patch)
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: 2
🤖 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 `@pyre/bench/synth/kept_stack_deep_var_shortcircuit_mutate.py`:
- Around line 7-8: The mutation-count invariant in the
kept_stack_deep_var_shortcircuit_mutate benchmark is wrong because the loop in f
appends three times per iteration instead of the intended two, which makes
len(log) undercount against the expected total. Update the logic in f so each
iteration performs exactly two mutations, keeping the existing g and h
references and adjusting the conditional append behavior to avoid the extra log
entry. Ensure the final log length matches the iteration count expectation for
f(40000).
In `@pyre/pyre-jit/src/jit/codewriter.rs`:
- Around line 9546-9576: The FOR_ITER hazard scan in codewriter’s
foriter_deep_kept_call_hazard is too broad because it treats any call-like
opcode plus any conditional branch in the loop body as unsafe. Tighten the scan
so it only flags cases where a branch occurs after a call and the call result is
still retained on the operand stack or otherwise live across that branch, using
the existing scan_state/OpArgState context or equivalent liveness evidence. If
that tracking is not available, keep the current behavior behind a more explicit
conservative check rather than disabling the vable reload for safe loops.
🪄 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: 1858502e-0acd-4ab1-b779-e7e59c9feaf5
📒 Files selected for processing (2)
pyre/bench/synth/kept_stack_deep_var_shortcircuit_mutate.pypyre/pyre-jit/src/jit/codewriter.rs
| # and DOUBLE the mutation. The `log` length must equal the iteration count | ||
| # exactly (2 mutations per iteration): a doubled delivery over-counts, a |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Fix the mutation-count invariant.
Line 28 appends three times per iteration (g, h, then one conditional g/h), but the comment says “2 mutations per iteration.” For f(40000), the expected len(log) is 120000, not 80000.
Suggested correction
-# exactly (2 mutations per iteration): a doubled delivery over-counts, a
+# exactly (3 mutations per iteration): a doubled delivery over-counts, aAlso applies to: 28-28
🤖 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 `@pyre/bench/synth/kept_stack_deep_var_shortcircuit_mutate.py` around lines 7 -
8, The mutation-count invariant in the kept_stack_deep_var_shortcircuit_mutate
benchmark is wrong because the loop in f appends three times per iteration
instead of the intended two, which makes len(log) undercount against the
expected total. Update the logic in f so each iteration performs exactly two
mutations, keeping the existing g and h references and adjusting the conditional
append behavior to avoid the extra log entry. Ensure the final log length
matches the iteration count expectation for f(40000).
| let foriter_deep_kept_call_hazard = { | ||
| let exhaust_target = jump_target_forward( | ||
| code, | ||
| num_instrs, | ||
| py_pc + 1, | ||
| delta.get(op_arg).as_usize(), | ||
| ); | ||
| let mut scan_state = pyre_interpreter::OpArgState::default(); | ||
| let mut has_call = false; | ||
| let mut has_branch = false; | ||
| for scan_pc in 0..num_instrs { | ||
| let (scan_instr, _) = | ||
| scan_state.get(code.instructions[scan_pc]); | ||
| if scan_pc <= py_pc || scan_pc >= exhaust_target { | ||
| continue; | ||
| } | ||
| match scan_instr { | ||
| Instruction::Call { .. } | ||
| | Instruction::CallKw { .. } | ||
| | Instruction::CallFunctionEx | ||
| | Instruction::CallIntrinsic1 { .. } | ||
| | Instruction::CallIntrinsic2 { .. } => has_call = true, | ||
| Instruction::PopJumpIfTrue { .. } | ||
| | Instruction::PopJumpIfFalse { .. } | ||
| | Instruction::PopJumpIfNone { .. } | ||
| | Instruction::PopJumpIfNotNone { .. } => has_branch = true, | ||
| _ => {} | ||
| } | ||
| } | ||
| } | ||
| has_call && has_branch | ||
| }; |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | 🏗️ Heavy lift
Narrow the hazard scan to an actually kept call result.
This flags any loop body containing both a call-like opcode and a conditional jump, even when the call result is not held across that branch (or the branch occurs before the call). That withholds the vable reload for safe FOR_ITER loops and pushes them back to the header-abort path. Please track branch-after-call with stack/liveness evidence that the call result is still on the operand stack, or keep this behind a more explicit conservative gate.
🤖 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 `@pyre/pyre-jit/src/jit/codewriter.rs` around lines 9546 - 9576, The FOR_ITER
hazard scan in codewriter’s foriter_deep_kept_call_hazard is too broad because
it treats any call-like opcode plus any conditional branch in the loop body as
unsafe. Tighten the scan so it only flags cases where a branch occurs after a
call and the call result is still retained on the operand stack or otherwise
live across that branch, using the existing scan_state/OpArgState context or
equivalent liveness evidence. If that tracking is not available, keep the
current behavior behind a more explicit conservative check rather than disabling
the vable reload for safe loops.
![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: 6276359bb5
ℹ️ 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".
| } | ||
| } | ||
| } | ||
| has_call && has_branch |
There was a problem hiding this comment.
This predicate treats any CALL-family opcode anywhere in the loop body plus any POP_JUMP_IF_* anywhere before the exhaust target as the deep-kept-call hazard, so ordinary loops such as for ...: if pred: g(i) or x = g(i); if pred: ... also skip the vable iterator reload even though no call result is held on the operand stack across the branch. The fallback path intentionally keeps the pre-reload behavior, which the added comment says aborts at FOR_ITER and interprets permanently, so this disables JIT residency for a broad class of safe for-loops rather than just the #416/#420 shape. Please base the gate on the call result actually being live deep on the operand stack across the branch, or on a narrower bytecode pattern.
Useful? React with 👍 / 👎.
1 participant
Summary
Restores the per-iteration iterator reload in the
FOR_ITERhandler so thekept_stack_deep_var_{condexpr,nested_call}residency benches JIT-compile instead of abortingResidualCallArgUnbound{pc:230}on theForIterNextresidual. This is a/parityfix — it ports RPython'sw_iterator = self.peekvalue()(pyopcode.py:1303-1304, lowered byjtransform.py:760-767 do_fixed_list_getitemtogetarrayitem_vable_rinside the loop body).The
jit_merge_pointreds are[frame, ec]only (matchinginterp_jit.py:67), so the operand stack lives in the virtualizable image, not in loop-carried registers. An iterator defined once in the loop preamble cannot survive as an SSA register across the compiled resident loop; re-materializing it from the vable each iteration gives it a genuine in-loop reader, socompute_livenesskeeps it live at the-live-marker andbuild_pcdep_color_slotsretains its color→slot entry for the M3 body-guard resume.Withhold gate (deep-kept-call hazard)
The reload is withheld when the loop body holds a CALL result across an in-body conditional branch (both a CALL-family op and a
POP_JUMP_IF_*inside[py_pc+1, exhaust_target)). That shape (e.g.kept_stack_deep_var_shortcircuit) reaches a depth>1 kept operand stack at an interior short-circuit guard whose compiled resume cannot yet reconstruct the deep slots (the #416/#420 deep-kept interior-snapshot gap). With the reload the walk would consume the shared iterator, run the body's residual calls concretely, then abort at the deep guard — an irreversible half-executed iteration that can neither be delivered (doubles a committed effect) nor cleanly dropped (loses the iteration). Withholding aborts at theFOR_ITERheader before the consume, so the location interprets permanently — bit-exact with the pre-reload behaviour. Lifted once the deep-kept interior-guard snapshot lands.Results (true reference =
PYRE_NO_JIT=1)kept_stack_deep_var_condexpr→-266006670— JIT-residentkept_stack_deep_var_nested_call→818440000— JIT-residentkept_stack_deep_var_shortcircuit→840076000— gated, interprets (its f-loop never JIT-compiled anyway)kept_stack_deep_var_shortcircuit_mutate.py→840076000 / 120000— companion whereg/happend to a global list, so a torn or doubled iteration is observable; proves no tear/doubleVerification
check.pydynasm / cranelift / wasm — 184/184 eachcargo test -p pyre-jit297/0,-p pyre-jit-trace12/0PYRE_PCDEP_VALIDATE=1— 0 injection violations on the trio🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Tests