jit: admit direct heap stores in abort-free FOR_ITER bodies (#57)#448
Conversation
|
Caution Review failedThe pull request is closed. ℹ️ Recent review info⚙️ Run configurationConfiguration used: Organization UI Review profile: ASSERTIVE Plan: Pro Run ID: 📒 Files selected for processing (1)
WalkthroughExtends the FOR_ITER JIT eligibility gate in eval.rs by adding an abort-free detection helper for loop bodies, then widens the allow-list to permit direct heap-store opcodes and LOAD_NAME reads for such bodies. Corresponding unit tests are updated to reflect the new permissive behavior and its boundaries. ChangesFOR_ITER abort-free gating
Estimated code review effort: 3 (Moderate) | ~25 minutes Possibly related PRs
Poem
✨ 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 e628269). Codex did not produce a report (exit 1). Last log lines: |
![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: dc87f83d1d
ℹ️ 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".
| | I::StoreSubscr | ||
| | I::StoreAttr { .. } | ||
| | I::StoreName { .. } | ||
| | I::StoreGlobal { .. } |
There was a problem hiding this comment.
Admitting these store opcodes makes the frame JIT-eligible even though the current recovery path explicitly refuses to deliver the consumed FOR_ITER item once a store has committed. If the trace walk aborts after the store on a later allowed opcode such as a CALL, LOAD_ATTR, or branch guard, fbw_foriter_inflight_take* returns None; the iterator has already advanced, so the rest of that iteration is skipped rather than resumed. That changes program results for common bodies like for x in xs: obj.x = x; total += g(x) where the store has run but the later local update/call is dropped; previously these frames stayed interpreted and preserved Python semantics.
Useful? React with 👍 / 👎.
youknowone
changed the title
|
P1 confirmed — reverted. This is a real correctness regression, and the reproductions show it is not specific to the store opcodes:
The FBW refuse-drop drops the rest of an iteration once any body effect has committed and a later op aborts — a partial-iteration wrong answer ( Remediation: reverted the — commented by Claude |
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 82054dcb9d
ℹ️ 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".
| || (abort_free | ||
| && matches!( | ||
| body_instr, | ||
| I::StoreSubscr | ||
| | I::StoreAttr { .. } |
There was a problem hiding this comment.
The earlier comment named CALL/LOAD_ATTR/branch; in the current diff those are excluded, but BinaryOp and other base-listed guard/residual opcodes remain allowed after this store widening. For a syntactically call/branch-free body such as for x in data: buf[0] = x; total = total + x with an int-to-float transition like the existing foriter57/f1_polyonly pattern, this admits the frame: the store residual marks a body effect, then a later arithmetic guard abort makes Option-C refuse re-delivery, so the consumed item never reaches the remaining StoreFast and total silently drops that element. The abort-free test needs to reject stores before any opcode that can still abort, not just explicit calls and branches.
Useful? React with 👍 / 👎.
A FOR_ITER frame ran entirely in the interpreter when any loop body contained STORE_SUBSCR/STORE_ATTR/STORE_NAME/STORE_GLOBAL, since those opcodes were absent from the for_iter_body_op_is_jit_safe allow-list. Add for_iter_body_is_abort_free() and permit those store opcodes plus LOAD_NAME when the body contains no call, no nested FOR_ITER, and no conditional/forward branch. Such a body has no reachable mid-loop walk abort, so the eager store is committed exactly once and cannot be dropped or doubled (#57). LOAD_ATTR stays excluded; the call-bearing path is unchanged. Assisted-by: Claude
Rewrite the `for_iter_bodies_all_jit_safe` doc comment to state that the whole gate compensates for the FBW single-pass walker's Option-C refuse-drop, which drops the rest of an iteration on a mid-body abort, whereas the register-machine metainterp resumes exactly at the guard's `(jitcode, pc)` (`blackhole_from_resumedata`, resume.py:1312) and never drops. Records the convergence path (#215 P2 / #73 / task#50) and states why abort-free bodies are the observable-deviation-free subset and why LOAD_ATTR stays excluded. Assisted-by: Claude
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 2daab3dde4
ℹ️ 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".
| | I::StoreName { .. } | ||
| | I::StoreGlobal { .. } | ||
| | I::LoadName { .. } |
There was a problem hiding this comment.
This admits STORE_NAME/LOAD_NAME in any abort-free FOR_ITER body, including class bodies and exec(code, g, l) frames where w_locals != w_globals. The trace-side implementation does not preserve that namespace: MIFrame::store_name_value records the write to trace_globals_ptr() (pyre-jit-trace/src/opcode_handler_impls.rs:887-889), and load_name_value first tries a module-global cell path from concrete_namespace. A hot class loop such as class C: ...; for i in range(n): x = i can therefore start writing i/x to module globals, or read a global shadow, once this gate lets the loop JIT; before this change these FOR_ITER bodies stayed interpreted because the NAME opcodes were not allow-listed.
Useful? React with 👍 / 👎.
1 participant
What
Widen the JIT
FOR_ITERloop-body allow-list (lever A / #57) so that astraight-line
FOR_ITERframe whose body writes the heap is JIT-compiledinstead of running the whole frame in the interpreter. Admits the direct
heap-store opcodes
STORE_SUBSCR/STORE_ATTR/STORE_NAME/STORE_GLOBALand the
LOAD_NAMEthat reads module globals — but only in an abort-freebody (no call, branch, or nested loop).
unsupported_jit_shape→for_iter_bodies_all_jit_safepreviously rejected anyFOR_ITERframe whose loop body contained a heap store, sending the entireframe to the interpreter. That frame-level rejection — not call-inlining
coverage — is the dominant JIT perf gap:
whileloops hit ~10.7×, function-localforloops ~14.75×, but anyforloop that writes the heap (d[i]=,o.x=,module global) got 0× (never JIT).
Why abort-free only
This gate is a conservative adaptation, not an upstream mechanism. The FBW
single-pass walker recovers from a mid-body walk abort by refusing to
re-deliver the in-flight
FOR_ITERitem (Option-C deliver/refuse), which dropsthe rest of that iteration. The register-machine metainterp instead resumes
exactly: it reads the precise
(jitcode, pc)from the guard's resume data,setpositions a blackhole interpreter there and runs the body to its end(
blackhole_from_resumedata,resume.py:1312), so it never drops or doubles.An abort-free body cannot reach the deliver/refuse path at all: with no call,
branch, or nested loop there is no abort after the store, so the store commits
exactly once and its result coincides with the exact-resume result. That makes
abort-free stores the observable-deviation-free subset.
Admitting store-bearing bodies that can abort (calls, branches, method loads)
was explored and reverted: when a body effect commits and a later op aborts,
the refuse-drop skips the rest of that iteration — a partial-iteration wrong
answer (
for x: d[k]=v; l.append(x)→lshort by the number of aborts), not aclean iteration drop. Converging those bodies requires porting the exact-resume
path (#215 P2 / #73 / task#50); until then they stay interpreted and correct.
Validation
check.pydynasm 159/159 + cranelift 159/159 (2/2 backend runs).foriter57run_matrix(JIT vs interp oracle) MATRIX OK on both backends.for_itergate unit tests 9/9: straight-lineSTORE_SUBSCR/STORE_ATTRbodies admitted; append / store+branch / store+call bodies rejected.
d[k]=v; l.append,s.add; l.append,h(x); l.append) now match the interpreter exactly under JIT.Perf (all previously 0× / never-JIT)
g = g + ibuf[i] = i(subscript store)d[k] = i🤖 Generated with Claude Code
Summary by CodeRabbit
New Features
for-loop bodies that are simple and branch-free.Bug Fixes