Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
153 changes: 99 additions & 54 deletions majit/majit-backend-dynasm/src/x86/assembler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -101,6 +101,46 @@ mod tests {

use crate::runner::DynasmBackend;

/// x86/regalloc.py `consider_call_malloc_nursery` binds the result with
/// `force_allocate_reg(op, selected_reg=ecx)`; only FrameManager may spill
/// it later. Storing it into a slot here as well grew every recursive
/// loop's JitFrame by one slot per allocation, which changes how much the
/// process allocates and so shifts the minor-collection schedule the
/// jitcounter decay is driven by.
#[test]
fn malloc_nursery_result_does_not_grow_frame_depth() {
let mut backend = DynasmBackend::new();
backend.attach_default_test_descrs();

let malloc = Rc::new(Op::new(
OpCode::CallMallocNursery,
&[Operand::from_opref(OpRef::const_int(32))],
));
malloc.pos.set(OpRef::ref_op(0));

let finish = Op::new(OpCode::Finish, &[Operand::from_bound_op(&malloc)]);
finish.pos.set(OpRef::void_op(1));
finish.set_fail_arg_types(vec![Type::Ref]);
finish.setfailargs(vec![].into());

let token = JitCellToken::new(517);
backend
.compile_loop(&[], &[malloc, Rc::new(finish)], &token)
.expect("compile register-resident nursery result trace");

let compiled = token
.compiled
.get()
.expect("compiled code")
.downcast_ref::<super::CompiledCode>()
.expect("dynasm compiled code");
assert_eq!(
compiled.frame_depth.load(Ordering::Acquire),
super::JITFRAME_FIXED_SIZE,
"a register-resident nursery result must not allocate a shadow frame slot",
);
}

fn compile_eval_breaker_poll_trace(
trace_id: u64,
word_addr: usize,
Expand Down Expand Up @@ -4175,7 +4215,7 @@ impl<'a> Assembler386<'a> {
// store next GUARD_NOT_FORCED's descr ptr to jf_force_descr
// BEFORE the call, so forcing code knows which guard to resume.
self._store_force_index_if_next_guard(ops, op_index, fail_index);
self.genop_call_assembler(op, arglocs);
self.genop_call_assembler(op, arglocs, result_loc);
}
OpCode::CondCallN => self.genop_discard_cond_call(op, arglocs),
OpCode::CondCallValueI | OpCode::CondCallValueR => {
Expand Down Expand Up @@ -4300,20 +4340,10 @@ impl<'a> Assembler386<'a> {
dynasm!(self.mc ; .arch x64 ; mov Rq(rv), rcx);
}
dynasm!(self.mc ; .arch x64 ; =>done);
// Spill the result to the regalloc-assigned jitframe slot
// directly from `result_reg`. The malloc-nursery clobber
// set is only ECX/EDX (SAVE_DEFAULT_REGS), so RAX may hold
// a live Box the regalloc bound to it across this op;
// routing the spill through RAX (`mov rax, result_reg`)
// silently destroyed that value — this mirrors the
// fixed-size `genop_call_malloc_nursery` final block, which
// was already corrected for the same reason.
let pos = op.pos.get();
if !pos.is_none() {
let slot = self.allocate_slot(pos);
let offset = Self::slot_offset(slot);
dynasm!(self.mc ; .arch x64 ; mov [rbp + offset], Rq(rv));
}
// The payload already sits in `result_reg` on both paths, and
// that register is the delivery contract (see
// `genop_call_malloc_nursery`); an extra store grew
// `frame_depth` by one slot per allocation.
}
// x86/assembler.py malloc_cond_varsize parity
// arglocs = [lengthloc, imm(itemsize), imm(kind)]
Expand Down Expand Up @@ -4377,8 +4407,20 @@ impl<'a> Assembler386<'a> {
// result slot.
self.emit_propagate_exception_if_zero(0);
dynasm!(self.mc ; .arch x64 ; mov QWORD [rbp + gcmap_ofs], 0);
if !op.pos.get().is_none() {
self.store_rax_to_result(op.pos.get());
// Unlike the other three nursery paths, this one leaves the
// helper's return in RAX rather than in the regalloc result
// register, so the move is emitted here instead of the store
// that used to grow `frame_depth` by a slot per allocation.
// `consider_call_malloc_nursery_varsize` forces the result to
// `MALLOC_NURSERY_RESULT`, so this is never a no-op.
let Some(Loc::Reg(r)) = result_loc else {
panic!(
"CallMallocNurseryVarsize result_loc must be a register; got {result_loc:?}"
);
};
if r.value != crate::regloc::EAX.value {
let rv = r.value;
dynasm!(self.mc ; .arch x64 ; mov Rq(rv), rax);
}
}
// x86/assembler.py `genop_discard_check_memory_error`
Expand Down Expand Up @@ -6984,7 +7026,7 @@ impl<'a> Assembler386<'a> {
/// `reload_frame_if_necessary` because a minor GC during the callee
/// may have moved the caller jitframe; the popped rbp is the
/// pre-GC address while the shadow stack carries the updated one.
fn genop_call_assembler(&mut self, op: &Op, arglocs: &[Loc]) {
fn genop_call_assembler(&mut self, op: &Op, arglocs: &[Loc], result_loc: Option<&Loc>) {
// handle_call_assembler (rewrite.py) always pre-builds the
// callee jitframe — storing every inputarg, and for a virtualizable
// passing the forced vable object as the second arg — so the backend
Expand Down Expand Up @@ -7027,9 +7069,7 @@ impl<'a> Assembler386<'a> {
} else {
dynasm!(self.mc ; .arch x64 ; xor eax, eax);
}
if !op.pos.get().is_none() {
self.store_rax_to_result(op.pos.get());
}
self.move_call_assembler_result(result_type, result_loc);
return;
}

Expand Down Expand Up @@ -7104,8 +7144,36 @@ impl<'a> Assembler386<'a> {
; =>merge
);
}
if !op.pos.get().is_none() {
self.store_rax_to_result(op.pos.get());
self.move_call_assembler_result(result_type, result_loc);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Keep CALL_ASSEMBLER results visible to legacy consumers

When an x86 dynasm trace uses a CALL_ASSEMBLER result as the predicate of a following COND_CALL_N or COND_CALL_VALUE_*, this now leaves the value only in result_loc. Both conditional-call emitters still ignore that predicate's regalloc location and call load_arg_to_rax, whose resolve_opref only recognizes constants and frame slots; because the removed result spill also supplied the slot mapping, compiling this valid trace can now panic with “unmapped non-constant OpRef.” Pass the predicate argloc into those emitters, or retain materialization until every legacy consumer uses regalloc locations.

AGENTS.md reference: AGENTS.md:L184-L185

Useful? React with 👍 / 👎.

}

/// Materialize a CALL_ASSEMBLER result from the raw bits both paths leave
/// in RAX into the regalloc-assigned location.
///
/// The previous shape spilled RAX to a fresh JitFrame slot, which grew
/// `frame_depth` by one slot per call and left a `Float` result taking the
/// helper path in RAX while the regalloc expected XMM0. Only the fast path
/// happened to leave it in XMM0 as a side effect of its `movq rax, xmm0`
/// normalisation.
///
/// x86/regalloc.py `_consider_call_assembler` binds the result through
/// `after_call`, so it is `eax` for an int or ref and `xmm0` for a float;
/// the move is elided when the value already sits there.
fn move_call_assembler_result(&mut self, result_type: Type, result_loc: Option<&Loc>) {
match (result_type, result_loc) {
(Type::Void, None) => {}
(Type::Float, Some(Loc::Reg(r))) if r.is_xmm => {
dynasm!(self.mc ; .arch x64 ; movq Rx(r.value), rax);
}
(_, Some(Loc::Reg(r))) if !r.is_xmm => {
if r.value != crate::regloc::EAX.value {
dynasm!(self.mc ; .arch x64 ; mov Rq(r.value), rax);
}
}
_ => panic!(
"CALL_ASSEMBLER result must use its regalloc result register: \
type={result_type:?} loc={result_loc:?}"
),
}
}

Expand Down Expand Up @@ -7480,25 +7548,11 @@ impl<'a> Assembler386<'a> {
let _ = gcmap_ofs;

dynasm!(self.mc ; .arch x64 ; =>done);
// Spill the result to the regalloc-assigned jitframe slot. Stage
// it directly from `result_reg`; routing through RAX (the previous
// shape) silently clobbered any live Box the regalloc bound to
// RAX across this op, since the malloc-nursery clobber set is
// only ECX/EDX.
let pos = op.pos.get();
if !pos.is_none() {
let slot = self.allocate_slot(pos);
let offset = Self::slot_offset(slot);
// malloc_cond / malloc_cond_varsize (assembler.py:2556,2604) keep
// the allocated pointer in ecx and route it through the regalloc-
// assigned `result_reg`. Anything else here would spill a stale
// RAX (now caller-live) instead of the object pointer.
let Some(Loc::Reg(r)) = result_loc else {
panic!("CallMallocNursery result_loc must be a register; got {result_loc:?}");
};
let rv = r.value;
dynasm!(self.mc ; .arch x64 ; mov [rbp + offset], Rq(rv));
}
// x86/regalloc.py `consider_call_malloc_nursery` binds the result with
// `force_allocate_reg(op, selected_reg=ecx)`, so the register IS the
// delivery contract; FrameManager emits a spill only where a later
// lifetime boundary needs one. Storing it here as well grew
// `frame_depth` by one slot for every allocation.
}

/// Headerless fixed-size nursery allocation. `size` excludes any GC
Expand Down Expand Up @@ -7584,18 +7638,9 @@ impl<'a> Assembler386<'a> {
let _ = gcmap_ofs;

dynasm!(self.mc ; .arch x64 ; =>done);
let pos = op.pos.get();
if !pos.is_none() {
let slot = self.allocate_slot(pos);
let offset = Self::slot_offset(slot);
let Some(Loc::Reg(r)) = result_loc else {
panic!(
"CallMallocNurseryHeaderless result_loc must be a register; got {result_loc:?}"
);
};
let rv = r.value;
dynasm!(self.mc ; .arch x64 ; mov [rbp + offset], Rq(rv));
}
// The result register is the delivery contract (see
// `genop_call_malloc_nursery`); an extra store here grew `frame_depth`
// by one slot per allocation.
}

/// NEW: allocate a fixed-size object. Requires GC runtime.
Expand Down
29 changes: 17 additions & 12 deletions majit/majit-translate/src/codewriter/jitcode.rs
Original file line number Diff line number Diff line change
Expand Up @@ -205,10 +205,11 @@ pub struct JitCode {
/// so an answer already computed for the original holds for it too.
#[derive(Debug, Default, Clone)]
pub struct DerivedBodyFacts {
/// Whether descending into this body can reach a residual call whose
/// funcbox is an un-lowered helper's symbolic hash. Read through
/// [`JitCode::descent_reaches_unlowered_helper_call`].
descent_reaches_unlowered_helper_call: OnceLock<bool>,
/// The un-lowered helper a descent into this body can reach, named by
/// the symbolic hash standing in for its funcbox, or `None` when the
/// descent reaches no such call. Read through
/// [`JitCode::descent_unlowered_helper_blocker`].
descent_unlowered_helper_blocker: OnceLock<Option<i64>>,
}

mod oncelock_usize_serde {
Expand Down Expand Up @@ -312,16 +313,20 @@ impl JitCode {
.expect("JitCode body not yet set — call set_body() before body_mut()")
}

/// Whether descending into this body can reach a residual call whose
/// funcbox is an un-lowered helper's symbolic hash, computing the answer
/// with `compute` the first time it is asked. The property is fixed by the
/// assembled body, so the first answer is the only one this instance gives:
/// `body_mut` needs `&mut self`, which `runtime_fnaddr_patch` can only take
/// before the jitcode is published behind an `Arc`.
pub fn descent_reaches_unlowered_helper_call(&self, compute: impl FnOnce() -> bool) -> bool {
/// The un-lowered helper a descent into this body can reach, named by
/// the symbolic hash standing in for its funcbox, computing the answer
/// with `compute` the first time it is asked. `None` means the descent
/// reaches no such call. The property is fixed by the assembled body, so
/// the first answer is the only one this instance gives: `body_mut` needs
/// `&mut self`, which `runtime_fnaddr_patch` can only take before the
/// jitcode is published behind an `Arc`.
pub fn descent_unlowered_helper_blocker(
&self,
compute: impl FnOnce() -> Option<i64>,
) -> Option<i64> {
*self
.derived
.descent_reaches_unlowered_helper_call
.descent_unlowered_helper_blocker
.get_or_init(compute)
}

Expand Down
4 changes: 4 additions & 0 deletions majit/majit-translate/src/front/mir.rs
Original file line number Diff line number Diff line change
Expand Up @@ -2213,6 +2213,10 @@ fn lower_unstructured_with_static_addrs_and_attrs(
tail_forwarded_returns,
)
.map_err(LowerError::Unsupported)?;
// Fold each raise site's `PyError` constructor into its
// materialisation call, so the transparent constructor — which has
// no host symbol and therefore no address — leaves this graph.
crate::front::result_exc::fuse_kind_ctor_raise(&mut lo.graph);

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Badge Move raise fusion out of the translator special case

Remove this bespoke front-end fusion and express the opaque raise path in the interpreter source, or fix constructor lowering generically. This call makes the generated JIT recognize one exact PyError::type_error/literal-message CFG and substitute a helper that the interpreter never calls; consequently formatted messages, additional constructors, or harmless CFG reshaping silently bypass the fix and remain inline blockers. That is precisely the source/JIT divergence the repository requires generation fixes to avoid.

AGENTS.md reference: AGENTS.md:L12-L15

Useful? React with 👍 / 👎.

Copy link
Copy Markdown
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Taking this one seriously rather than waving it off, because one half of it is
literally accurate and I want to separate that half from the part I think does
not hold.

Accurate: pyerror_type_error_to_exc_object has no interpreter-side
caller. Its only references are the fnaddr registration, the fusion table in
front/result_exc.rs, and a test. If the standard for "the JIT is generated
from the interpreter source" is "every residual the JIT calls is also on a live
interpreter path", this helper fails it, and I am not going to argue otherwise.

Where I think the divergence claim does not hold: the helper is interpreter
source, and its body is the exact sequential composition of the pair it
replaces — not a re-implementation of it:

pub unsafe fn pyerror_type_error_to_exc_object(w_msg: *mut PyObject) -> *mut PyObject {
    let msg = unsafe { w_str_get_wtf8(w_msg) }.to_owned();
    PyError::type_error(msg).to_exc_object()
}

There is no behaviour the generated JIT can observe here that the unfused pair
would not produce. The previous spelling already passed the same receiver
address to the same body via CallTarget::Method{to_exc_object, receiver PyError}, so this does not introduce a new ABI or a new GC exposure either.

On "formatted messages and other constructors silently bypass the fix":
that is correct, and for formatted messages it is required, not incidental. The
helper does an unchecked w_str_get_wtf8 deref, and a runtime-formatted
message need not be a box_str_constant object — so the declining path
(producer alloc::fmt::format) is the safe direction. Declining falls back to
exactly today's lowering; nothing regresses, it simply is not improved.

On "additional constructors": I measured this before writing the table rather
than assuming it. type_error is the only PyError constructor that reaches a
raise site in the __pyre_wrap_* family — the rule fires at 582 of 681
constructors across 301 distinct wrapper graphs. A nine-constructor table would
have been speculative; one entry is the corpus.

On "fix constructor lowering generically" — this is the real point, and I
concede it is the better fix.
I could not reach it. The generic fix is making
PyError::new lower, and it is blocked on two independent things: it is generic
over impl Into<Wtf8Buf> and returns an aggregate, and its struct construction
falls back to a symbolic transparent ctor because of the layout —
message: Wtf8Buf is 24 bytes against one 8-byte descr row, with kind at
offset 54 of 56. Closing that means changing PyError's layout, which an
adversarial review already refused for the trace-New variant.

So my honest summary: this is a mitigation with a measured ceiling (union
blocker census: CLEAR 0 → 73 of 561; the PyError-ctor bucket 560 → 376), not
the generic generation fix, and it does not pretend to be one. Whether a
measured mitigation is worth carrying while the generic fix is blocked is a
call about this repository's policy, not something I should decide unilaterally
@youknowone, if you would rather this come out until constructor lowering
can be fixed properly, say so and I will drop the front/result_exc.rs rule
and its table entry; the rest of the commit (the published helper and its
fnaddr test) stands on its own.

commented by Claude

if result_exc_ok_is_unit {
// Stamp `FUNC.RESULT = void`. The exception-link lowering
// already returns the unit `()` (the callee no longer
Expand Down
Loading
Loading