majit: resolve a boxing cluster's header pointers across the block boundary - #1141
Conversation
…undary `fuse_boxing_alloc` read `ob_header.ob_type` through `const_ref_addr`, which gave up when no operation produced the value. A call ends a block, so a cluster whose header stores sit before one carries the type pointer into the next block as a link argument while its `ConstRefAddr` producer stays in the predecessor; the lookup found no producer and reported the vtable unresolved. The comment in `intobject.rs` recorded that outcome for the whole tree — the pass fired nowhere — and attributed it to an empty `HostStaticAddrs.pytypes`. That attribution is wrong: the table is populated in the build-script pipeline and the addresses resolve once the walk crosses the block boundary. Follow the links when the value is a `Block.inputargs` phi, require every predecessor to agree, and treat a link argument that is not a variable as a disagreement. Seven sites fuse: `w_int_box_slow`, `w_int_new_unique`, `w_complex_new`, `w_object_mutable_cell_new`, `w_int_mutable_cell_new`, `w_dict_new_unmanaged_side_table_value`, `w_module_new_aliasing_dict`. Firing the pass exposes what it drops. The fusion discards the whole `ob_header` subtree and the runtime re-synthesises both header words from the kept vtable (`materialize_virtual_object`, `w_class = get_instantiate(vtable)`). That stands in for the dropped stores only where the cluster's own `w_class` is `get_instantiate(&T)` for the very `&T` its `ob_type` names. `w_bytes_subclass_from_bytes` stores a base `ob_type` beside a `w_class` read back off the shadow stack, so re-synthesising it would answer the base type for `type(B(b'x'))`. Require the two to agree and decline otherwise. Dropping the whole header subtree is itself the deviation. `jtransform.py` suppresses one field, by the literal name `typeptr` (908-911, 952-954), and `heaptracker.py:66-69` skips that name while recursing into every other field of a nested header struct — upstream a second header word is an ordinary traced field, as the Python-level class is for user subclasses (`mapdict.py:751-752`). The faithful shape is to keep `ob_type` vtable-derived and re-emit `w_class` as an ordinary payload store, which would also fuse the subclass constructors this guard declines. It needs per-group header `w_class` fielddescrs (`function_header_w_class_descr` is the pattern) because the shared `w_class_descr()` has `index_in_parent == 0` and would be dropped silently by the materialize replay, and `virtualize.rs:944-958` names the read/write descr-identity split as the prerequisite. Guarding first. No runtime effect is claimed. `bench/synth/list_pop_append` reads the same with and without this change, and a negative control that also reverts the fusion — the shape whose regression motivated the boundary — does not reproduce it, so that bench no longer discriminates the mechanism. The cast walk and the phi walk are shared by both resolvers, and the `get_instantiate` match pins the full `pyre_object::pyobject` owner path as `prune_dead_boxing_remnants` does. Four single-block fixtures modelled `ob_type` alone or fed `w_class` the type pointer instead of the instantiate slot; they now build the header the constructors build. A new test pins the decline, with a control row that fuses. Assisted-by: Claude
WalkthroughThe boxing fusion pass now resolves type pointers across additional IR relationships and fuses allocations only when ChangesBoxing fusion validation
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 5a89db8). 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)
4. Structural adaptations
|
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 5a89db8de4
ℹ️ 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".
| let addr = link | ||
| .args | ||
| .get(slot) | ||
| .and_then(LinkArg::as_variable) | ||
| .and_then(|arg| resolve_addr(graph, arg, depth - 1, terminal))?; |
There was a problem hiding this comment.
Preserve off-GC mutable cells when resolving the phi
When a traced module-dict update reaches w_object_mutable_cell_new, this new cross-block resolution makes its legacy malloc_typed cluster fuse into NewWithVtable, whose backends allocate a collector-owned object. However, walk_module_value_slot explicitly assumes these cells are off-GC and therefore forwards only ObjectMutableCell.w_value, not the module-dict slot containing the cell (pyre/pyre-object/src/celldict.rs:223-255). At the next collection the managed cell can consequently move or be reclaimed while the module dict retains its old pointer, causing a stale-pointer/UAF failure. Keep these legacy off-GC allocations unfused, or update their ownership and root-walking contract together.
AGENTS.md reference: AGENTS.md:L14-L18
Useful? React with 👍 / 👎.
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-translate/src/model.rs`:
- Around line 3007-3039: Update resolve_addr so an OpKind::Input producer is
skipped instead of being resolved through terminal, allowing resolution to
continue through the consuming block’s inputargs and predecessor links. Preserve
existing behavior for other operation producers. Add a diamond regression
covering matching predecessor addresses (fuses) and mismatching addresses
(declines), then run cargo check --features dynasm, cargo test --features
dynasm, and record the eight JIT benchmark results before merging.
🪄 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: 5282e3e8-fb1b-442c-a1dd-f72c2a214db3
📒 Files selected for processing (3)
majit/majit-translate/src/model.rspyre/pyre-jit-trace/src/state.rspyre/pyre-object/src/intobject.rs
| let producer = graph | ||
| .blocks | ||
| .iter() | ||
| .flat_map(|b| &b.operations) | ||
| .find(|o| o.result.as_ref() == Some(var))?; | ||
| match &producer.kind { | ||
| .find(|o| o.result.as_ref() == Some(var)); | ||
| let Some(producer) = producer else { | ||
| let (target_id, slot) = graph.blocks.iter().find_map(|b| { | ||
| b.inputargs | ||
| .iter() | ||
| .position(|arg| arg == var) | ||
| .map(|slot| (b.id, slot)) | ||
| })?; | ||
| let mut resolved: Option<i64> = None; | ||
| for block in &graph.blocks { | ||
| for link in &block.exits { | ||
| if link.target != target_id { | ||
| continue; | ||
| } | ||
| // A link whose arity does not cover the slot is a malformed | ||
| // edge; treat it as a disagreement rather than reading past it. | ||
| let addr = link | ||
| .args | ||
| .get(slot) | ||
| .and_then(LinkArg::as_variable) | ||
| .and_then(|arg| resolve_addr(graph, arg, depth - 1, terminal))?; | ||
| match resolved { | ||
| None => resolved = Some(addr), | ||
| Some(seen) if seen == addr => {} | ||
| Some(_) => return None, | ||
| } | ||
| } | ||
| } | ||
| return resolved; |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | 🏗️ Heavy lift
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== locate file =="
git ls-files | grep -F 'majit/majit-translate/src/model.rs' || true
echo "== outline around relevant symbols =="
wc -l majit/majit-translate/src/model.rs
ast-grep outline majit/majit-translate/src/model.rs --match install_phi --view expanded || true
ast-grep outline majit/majit-translate/src/model.rs --match resolve_addr --view expanded || true
echo "== relevant source snippets =="
sed -n '2980,3045p' majit/majit-translate/src/model.rs
sed -n '3045,3115p' majit/majit-translate/src/model.rs
sed -n '6340,6390p' majit/majit-translate/src/model.rs
echo "== OpKind and resolver terminal definitions/usages =="
rg -n "enum OpKind|Input|resolve_terminal|resolve_addr|InputArg|LinkArg|install_phi" majit/majit-translate -g '*.rs' -g '*.py'Repository: youknowone/pyre
Length of output: 50372
🏁 Script executed:
#!/bin/bash
set -euo pipefail
# Behavioral probe: model the resolver decisions from the source shown in the review.
python3 - <<'PY'
from dataclasses import dataclass
from typing import Optional, Generic, TypeVar
T = TypeVar('T')
class OpKind:
Input = "Input"
# placeholder other ops omitted
`@dataclass`(frozen=True)
class Variable:
name: str
`@dataclass`
class LinkArg:
value: Optional[Variable]
def as_variable(self):
return self.value
`@dataclass`
class Link:
target: str
args: list
`@dataclass`
class Operation:
op_kind: str
result: Optional[Variable]
def kind(self):
return self.op_kind
`@dataclass`
class Block:
id: str
inputargs: list
operations: list
exits: list
`@dataclass`
class Graph:
blocks: list
def resolve_terminal(v, depth):
# This is the current terminal behavior from the source: only constant values are resolved.
return None
def resolve_addr(graph, var, depth, terminal, use_block=None):
# Original producer-first implementation from the review's concern.
producer = None
for b in graph.blocks:
for o in b.operations:
if o.result is not None and o.result == var:
producer = o
break
if producer:
break
if producer:
if producer.kind() == "Input":
if terminal.resolve_terminal(var, depth) is not None:
return terminal.resolve_terminal(var, depth)
return None
# Other producers omitted; return address from other ops
return resolve_addr(graph, producer, depth - 1, terminal)
# fallback to block inputarg search as in the original code
target_id = slot = None
for b in graph.blocks:
for pos,arg in enumerate(b.inputargs):
if arg == var:
target_id = b.id
slot = pos
break
if target_id:
break
if target_id is None:
return None
resolved = None
for block in graph.blocks:
for link in block.exits:
if link.target != target_id:
continue
addr = None
try:
arg = link.args[slot]
except IndexError:
continue
if arg is not None:
av = arg.as_variable()
if av is not None:
addr = resolve_addr(graph, av, depth - 1, terminal)
if resolved is None:
resolved = addr
elif addr != resolved:
return None
return resolved
phi = Variable("phi1")
g = Graph([
Block("A", [], [Operation(OpKind.Input, phi)], [Link("B", [Variable("a")])]),
Block("B", [phi], [], []),
])
terminal = type("Term", (), {"resolve_terminal": staticmethod(resolve_terminal)})()
print("current_resolve_returns:", resolve_addr(g, phi, 20, terminal))
print("current_input_producers_selected:", any(o.kind() == "Input" and o.result == phi for b in g.blocks for o in b.operations))
print("fallback_never_runs:", resolve_addr(g, phi, 20, terminal) is None)
# The current producer-first path does not trace the predecessor link value from A.
# It also cannot distinguish a diamond with matching predecessor arguments from mismatching ones
# unless producers and terminals both resolve, so phi values coming from inputargs are lost.
PYRepository: youknowone/pyre
Length of output: 248
🏁 Script executed:
#!/bin/bash
set -euo pipefail
echo "== OpKind definition =="
sed -n '70,105p' majit/majit-translate/src/flowspace/operation.rs
echo "== resolve_addr context with surrounding match arms =="
sed -n '3005,3095p' majit/majit-translate/src/model.rs
echo "== find const_ref_addr call sites =="
rg -n "const_ref_addr\(|get_instantiate_arg_addr\(|resolve_vtable_addr" majit/majit-translate/src/model.rs
echo "== install_phi call sites/tests =="
rg -n "install_phi\(" majit/majit-translate/src/model.rs -A 8 -B 8
echo "== tests near resolving headers / const address =="
rg -n "const_ref_addr|resolve_addr|header stores|fused|must keep" majit/majit-translate/src/model.rs -A 20 -B 20Repository: youknowone/pyre
Length of output: 50371
Track the consuming inputarg slot when resolving phis.
resolve_addr picks the first matching operation result, and install_phi always creates a corresponding OpKind::Input producer. Inputs do not reach terminal, so valid cross-block phis return None and fuse_boxing_alloc can decline without matching predecessor addresses.
For OpKind::Input, skip the producer result and continue into the Block.inputargs path that the resolver already documents. Add a diamond regression where matching predecessor addresses fuse and mismatching addresses decline. Include cargo check --features dynasm, cargo test --features dynasm, and the eight JIT benchmark results before merge.
🤖 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-translate/src/model.rs` around lines 3007 - 3039, Update
resolve_addr so an OpKind::Input producer is skipped instead of being resolved
through terminal, allowing resolution to continue through the consuming block’s
inputargs and predecessor links. Preserve existing behavior for other operation
producers. Add a diamond regression covering matching predecessor addresses
(fuses) and mismatching addresses (declines), then run cargo check --features
dynasm, cargo test --features dynasm, and record the eight JIT benchmark results
before merging.
Source: Coding guidelines
|
Both reds on this PR — The identical failure occurs on It is also JIT-independent — — commented by Claude |
…ss-block walk (#1159) * jit: print the orthodox list sub-walk decline pc under PYRE_FBW_DEBUG_ABORT `OrthodoxSubWalkTraceUnsupported` carries the pc but the three decline arms dropped it. The identifier the decline names elsewhere is a symbolic fnaddr minted in `pyre-jit-trace/build.rs`, so a runtime reverse-name registry has nothing to read; the pc is what pairs the decline with a bytecode offset. Assisted-by: Claude * intobject: trace w_int_new's allocation again, as wrapint does #1131 put a `dont_look_inside` boundary around the int box's allocating tail (`w_int_box_slow`) because the sub-jitcode walk declined on the `SyntheticTransparentCtor` the `malloc_typed` arm lowered to, and named its own exit condition: "Drop the boundary once the fusion resolves a vtable there." #1141 landed that resolution -- `fuse_boxing_alloc` now follows the header pointers across the block boundary each call ends -- and its comment records the condition met, "it does now fire here", while keeping the boundary. `wrapint` (`objspace/std/intobject.py:903-921`) carries no `@dont_look_inside`; its comment reads "this whole function is getting inlined into every caller", and it allocates with `instantiate(W_IntObject)` then `w_res.intval = x`, the alloc-then-init pair the rtyper lowers to `new_with_vtable` + `setfield_gc`. Put the `malloc_typed` arm back in `w_int_new` where the fusion rewrites it into that pair, and keep the collector-heap arm as `w_int_gc_alloc` behind its own boundary: that arm carries a blocker that is still real, the wasm backend not lowering the offset-0 `ob_type` store faithfully. `bench/synth/list_pop_append` does not decide this. It reads the same either way (2.6/2.3/2.6 against 2.5/2.5/2.3, three runs each), and #1141 records a negative control -- boundary removed with the fusion reverted -- that failed to reproduce the regression the boundary was added for, so the bench is uninformative in both directions. The trace shape is what differs and it is observable: with this change the compiled loop carries a `NewWithVtable` whose descr is `W_IntObject.intval` and no residual call to any `w_int_*` boxing symbol. A residual call can never be virtualized. check.py ALL PASSED on all three backends -- dynasm 417/417, cranelift 416/416, wasm 412/412 -- and cargo test --workspace green. The box is still not virtualized away, for a reason outside this change: `orthodox_list_append_commit` deliberately forces the value with a ptr->int->ptr identity pair so the descended sub-walk reads the current iteration's payload, and that force lands before the class guard that would otherwise fold. Assisted-by: Claude * majit: test fuse_boxing_alloc across the links a split cluster crosses `resolve_addr` steps through `Block.inputargs` and requires every predecessor to agree, but every `fuse_boxing_alloc` case built its cluster in a single block, so neither behaviour was reached by a test: the only case added with the walk is a decline that resolves inside one block. Four rows over the same one-payload `W_FloatObject` cluster, differing only in where the header values come from: one relay block between the producer and the ctor, two relay blocks, two predecessors of a merge block naming one type, and two naming different types. Each asserts the fused count, the address stamped on the `NewWithVtable`, and whether the `malloc_typed` survives as a residual. Measured by ablation on this tree: returning `None` for a phi fails the "one link crossing" row, and dropping the disagreement arm fails the "predecessors naming two types" row. The six pre-existing `fuse_boxing_alloc` tests pass under both ablations. Assisted-by: Claude
fuse_boxing_alloc— pyre's port ofrewrite_op_malloc— never fired anywherein the tree. This makes it fire at seven sites, and guards the miscompile that
firing it exposes.
Why it never fired
The pass resolves
ob_header.ob_typethroughconst_ref_addr, which gave upwhenever no operation produced the value. A call ends a block, so a boxing
cluster whose header stores sit before one carries the type pointer into the
next block as a
Block.inputargsphi while itsConstRefAddrproducer staysin the predecessor. The lookup found no producer and reported the vtable
unresolved.
The in-tree comment recorded that outcome and attributed it to an empty
HostStaticAddrs.pytypes. That attribution is wrong — the table is populatedin the build-script pipeline, and the addresses resolve as soon as the walk
crosses the block boundary. The comment is corrected here.
The fix follows the links when the value is a phi, requires every
predecessor to agree, and treats a non-variable link argument as a
disagreement. The cast walk and the phi walk are now shared by both resolvers.
Census: 0 → 7 fusions —
w_int_box_slow,w_int_new_unique,w_complex_new,w_object_mutable_cell_new,w_int_mutable_cell_new,w_dict_new_unmanaged_side_table_value,w_module_new_aliasing_dict.What firing it exposes
The fusion discards the whole
ob_headersubtree and the runtimere-synthesises both header words from the kept vtable
(
materialize_virtual_object,w_class = get_instantiate(vtable)). Thatsubstitution is only valid where the cluster's own
w_classisget_instantiate(&T)for the very&Titsob_typenames.w_bytes_subclass_from_bytesstores a baseob_typebeside aw_classreadback off the shadow stack. Re-synthesising it would answer the base type
for
type(B(b'x')). This requires the two to agree and declines otherwise,pinned by a test with a control row that fuses.
The deviation, named
Dropping the whole header subtree is itself the deviation from upstream.
jtransform.pysuppresses exactly one field, by the literal nametypeptr(908-911, 952-954), and
heaptracker.py:66-69skips that name whilerecursing into every other field of a nested header struct. Upstream, a
second header word is an ordinary traced field — as the Python-level class is
for user subclasses (
mapdict.py:751-752).The faithful shape is to keep
ob_typevtable-derived and re-emitw_classasan ordinary payload store, which would also fuse the subclass constructors this
guard declines. That needs per-group header
w_classfielddescrs(
function_header_w_class_descris the pattern) — the sharedw_class_descr()has
index_in_parent == 0and would be dropped silently by the materializereplay — and
virtualize.rs:944-958names the read/write descr-identity splitas the prerequisite. Guarding first; the path is recorded in the commit
message.
No performance claim
bench/synth/list_pop_appendreads the same with and without this change, anda pre-registered negative control — which also reverts the fusion, i.e. the
exact shape whose regression motivated the
dont_look_insideboundary — doesnot reproduce that regression. A negative control that fails to differ
voids every arm, so the bench no longer discriminates this mechanism and
nothing was decided from it. In particular the peer-merged boundary was left
untouched. This is stated in the commit message and in the
intobject.rscomment.
Verification
pyre/check.py(full): dynasm 415/415, cranelift 414/414, wasm410/410, exit 0, no jitstats drift, CPython gate passed
cargo test -p majit-translate -p pyre-jit-trace: 21 suites okFour single-block fixtures modelled
ob_typealone, or fedw_classthe typepointer instead of the instantiate slot; they now build the header the
constructors actually build.
🤖 Generated with Claude Code
Summary by CodeRabbit
Bug Fixes
Documentation
Tests