jit(aarch64): close the loop gap vs PyPy in the backend; fix NaN float compares on both dynasm arches

majit/majit-backend-dynasm/src/aarch64/opassembler.rs

@@ -188,9 +188,11 @@ impl<'a> AssemblerARM64<'a> {
             // (regalloc.py:161) admits 0..4095 for ADD_ri; anything else
             // falls back to `mov x16, #ofs; add x16, x16, index`.
             if (0..4096).contains(&ofs) {
-                dynasm!(self.mc ; .arch aarch64
-                    ; mov x16, X(index.value)
-                    ; add x16, x16, ofs as u32);
+                // `opassembler.py:403 ADD_ri(ip0, index, ofs)` is a single
+                // instruction; dynasm's `add Xd|SP, Xn|SP, #uimm` form rejects a
+                // dynamic register operand, so encode the word directly
+                // (`codebuilder.py:113 ADD_ri`).
+                self.emit_add_ri(16, index.value, ofs as u32);
             } else {
                 self.emit_mov_imm64(16, ofs);
                 dynasm!(self.mc ; .arch aarch64

majit/majit-backend-dynasm/src/aarch64/regalloc.rs

@@ -267,14 +267,29 @@ impl<'a> RegAlloc<'a> {
     /// aarch64 `int_is_true` / `int_is_zero`: shares the 3-op `prepare_unary`
     /// shape from regalloc.py:456 since `cmp Xn, #0 ; cset Wd, ne` keeps
     /// the input register live while writing a fresh destination.
+    ///
+    /// When the next op consumes the flags, `regalloc.py:469
+    /// prepare_comp_unary` allocates no destination at all and
+    /// `opassembler.py:210 emit_comp_op_int_is_true` emits the `cmp` alone,
+    /// returning the condition.  `force_allocate_reg_or_cc` spells that as
+    /// the frame-register sentinel, which the `IntIsTrue` / `IntIsZero` emit
+    /// arms recognise through `flush_cc`.
     pub(crate) fn consider_int_is_true_j2(
         &mut self,
         dst: OpRef,
         arg: OpRef,
         i: usize,
         output: &mut Vec<RegAllocOp>,
     ) {
-        self.consider_unary_int_j2(dst, arg, i, output);
+        assert!(
+            !arg.is_constant(),
+            "prepare_comp_unary expects a non-const arg; got constant OpRef {arg:?} (should have been folded earlier)"
+        );
+        let arg_loc = self.make_sure_var_in_reg(arg, Type::Int, &[], None, false);
+        self.possibly_free_var(arg, Type::Int);
+        let ops_ref: &[majit_ir::Op] = self.operations;
+        let res = self.force_allocate_reg_or_cc(dst, ops_ref, i);
+        self.perform(i, vec![arg_loc], Some(res), output);
     }
 
     /// aarch64/regalloc.py:397 `prepare_op_uint_mul_high = prepare_op_int_mul`.

majit/majit-backend-dynasm/src/regalloc.rs

@@ -3603,7 +3603,6 @@ impl<'a> RegAlloc<'a> {
     /// - `CondCallN` via `guard_success_cc` (see
     ///   `genop_discard_cond_call`, mirrors `x86/assembler.py:2526
     ///   cond_call`).
-    #[cfg(target_arch = "x86_64")]
     fn next_op_can_accept_cc(&self, ops: &[Op], i: usize, result: OpRef) -> bool {
         if i + 1 >= ops.len() {
             return false;
@@ -3666,19 +3665,15 @@ impl<'a> RegAlloc<'a> {
     /// following guard. Otherwise force-allocate a general-purpose
     /// register (with `need_lower_byte` so `SETcc r8b` is encodable).
     ///
-    /// The CC-sentinel path is x86-only: the aarch64 CompOp emit at
-    /// `aarch64/assembler.rs` unconditionally emits `setcc result_loc`
-    /// and has no `flush_cc` equivalent, so receiving `frame_reg` as
-    /// the result there would clobber x29 (the frame pointer). On
-    /// non-x86 architectures, fall through to a plain force-allocate.
+    /// Both backends recognise the sentinel: `x86/assembler.rs flush_cc`
+    /// and `aarch64/assembler.rs flush_cc` publish `guard_success_cc` and
+    /// emit nothing when `result_loc` is the frame register, so neither
+    /// clobbers rbp / x29.
     pub(crate) fn force_allocate_reg_or_cc(&mut self, result: OpRef, ops: &[Op], i: usize) -> Loc {
-        #[cfg(target_arch = "x86_64")]
         if self.next_op_can_accept_cc(ops, i, result) {
             self.rm.force_allocate_frame_reg(result);
             return Loc::Reg(arch_regalloc::frame_reg());
         }
-        #[cfg(not(target_arch = "x86_64"))]
-        let _ = (ops, i);
         Loc::Reg(self.force_allocate_reg(result, Type::Int, &[], None, true))
     }
 
@@ -4137,8 +4132,10 @@ impl<'a> RegAlloc<'a> {
         if !vx_in_reg && !vy_in_reg && !vx.is_constant() {
             arglocs[0] = self.make_sure_var_in_reg(vx, Type::Float, &[], None, false);
         }
-        let result_loc =
-            Loc::Reg(self.force_allocate_reg(op.pos.get(), Type::Int, &[], None, false));
+        // x86/regalloc.py:682 — a float comparison whose only consumer is the
+        // next guard leaves its answer in the flags, like the integer one.
+        let ops_ref: &[Op] = self.operations;
+        let result_loc = self.force_allocate_reg_or_cc(op.pos.get(), ops_ref, i);
         self.perform(i, arglocs, Some(result_loc), output);
     }
 
@@ -4156,7 +4153,8 @@ impl<'a> RegAlloc<'a> {
         if !lhs_in_reg && !rhs_in_reg && !lhs.is_constant() {
             arglocs[0] = self.make_sure_var_in_reg(lhs, Type::Float, &[], None, false);
         }
-        let result_loc = Loc::Reg(self.force_allocate_reg(dst, Type::Int, &[], None, false));
+        let ops_ref: &[Op] = self.operations;
+        let result_loc = self.force_allocate_reg_or_cc(dst, ops_ref, i);
         self.perform(i, arglocs, Some(result_loc), output);
     }
 

majit/majit-backend-dynasm/src/x86/assembler.rs

@@ -3314,11 +3314,29 @@ impl<'a> Assembler386<'a> {
                         self.regalloc_mov(b_loc, &Loc::Reg(scratch));
                         scratch
                     };
-                    dynasm!(self.mc ; .arch x64 ; ucomisd Rx(a.value), Rx(b.value));
-                    if let Some(Loc::Reg(r)) = result_loc {
-                        let cc = Self::float_opcode_to_cc(op.opcode);
-                        self.emit_setcc(cc, r.value);
+                    // `assembler.py:1322 _cmpop_float`: UCOMISD sets
+                    // ZF = PF = CF = 1 when either operand is NaN, so only the
+                    // `A` / `AE` forms are already false on an unordered
+                    // compare.  FLOAT_LT / FLOAT_LE reach them by comparing in
+                    // the reverse order (`rev_cond`); FLOAT_EQ / FLOAT_NE have
+                    // no such form and take the parity fixup instead.
+                    let (lhs, rhs, cc, need_parity) = match op.opcode {
+                        OpCode::FloatLt => (b, a, CC_A, false),
+                        OpCode::FloatLe => (b, a, CC_AE, false),
+                        OpCode::FloatGt => (a, b, CC_A, false),
+                        OpCode::FloatGe => (a, b, CC_AE, false),
+                        OpCode::FloatEq => (a, b, CC_E, true),
+                        _ => (a, b, CC_NE, true),
+                    };
+                    dynasm!(self.mc ; .arch x64 ; ucomisd Rx(lhs.value), Rx(rhs.value));
+                    if need_parity {
+                        self.emit_if_parity_clear_zero_and_carry();
                     }
+                    // `assembler.py:1345 genop_cmp_float` ends in `flush_cc`, so
+                    // a comparison whose only consumer is the next guard keeps
+                    // its answer in the flags instead of materialising a
+                    // boolean the guard would immediately re-test.
+                    self.flush_cc(cc, result_loc);
                 }
             }
             // ── Casts ──
@@ -5069,6 +5087,22 @@ impl<'a> Assembler386<'a> {
         dynasm!(self.mc ; .arch x64 ; movzx Rd(dst_reg), Rb(dst_reg));
     }
 
+    /// `assembler.py:1314 _if_parity_clear_zero_and_carry`.
+    ///
+    /// UCOMISD sets PF on an unordered compare, together with ZF and CF, so
+    /// `sete` / `setb` / `setbe` would report NaN as equal / less-than and
+    /// `setne` would report it as not-not-equal.  `cmp rbp, 0` on the frame
+    /// pointer — never null inside compiled code — clears ZF and CF, and is
+    /// jumped over when PF is clear.
+    fn emit_if_parity_clear_zero_and_carry(&mut self) {
+        let ordered = self.mc.new_dynamic_label();
+        dynasm!(self.mc ; .arch x64
+            ; jnp =>ordered
+            ; cmp rbp, 0
+            ; =>ordered
+        );
+    }
+
     /// x86/assembler.py:1286 `flush_cc` parity.
     ///
     /// After emitting a CMP/TEST that leaves a boolean in the
@@ -5081,18 +5115,26 @@ impl<'a> Assembler386<'a> {
     /// value for non-guard consumers (e.g. boolean stored into a
     /// frame slot).
     fn flush_cc(&mut self, cond: u8, result_loc: Option<&Loc>) {
+        // `assembler.py:1293 flush_cc` opens with
+        // `assert self.guard_success_cc == rx86.cond_none` — a condition
+        // still pending here was published by an earlier op and never
+        // consumed, which would make the following guard branch on it.
+        debug_assert!(
+            self.guard_success_cc.is_none(),
+            "flush_cc: guard_success_cc already set",
+        );
         let frame_reg_value = crate::x86::regalloc::frame_reg().value;
         if let Some(Loc::Reg(r)) = result_loc {
             if r.value == frame_reg_value {
                 // Sentinel: the next op accepts cc.
-                debug_assert!(
-                    self.guard_success_cc.is_none(),
-                    "flush_cc: guard_success_cc already set",
-                );
                 self.guard_success_cc = Some(cond);
                 return;
             }
-            dynasm!(self.mc ; .arch x64 ; mov Rq(r.value), 0);
+            // `assembler.py:1300` clears the destination with `MOV imm0`
+            // before `SET_ir` because `SETcc` writes only the low byte.
+            // `emit_setcc` ends in `movzx r32, r8`, which zeroes bits 8..63
+            // on its own, so the same two-instruction sequence is spelled
+            // without the leading MOV.
             self.emit_setcc(cond, r.value);
         }
     }
@@ -5128,19 +5170,6 @@ impl<'a> Assembler386<'a> {
         }
     }
 
-    /// Map a float comparison OpCode to a condition code (after ucomisd).
-    fn float_opcode_to_cc(opcode: OpCode) -> u8 {
-        match opcode {
-            OpCode::FloatLt => CC_B,  // ucomisd: below = less than
-            OpCode::FloatLe => CC_BE, // below or equal
-            OpCode::FloatGt => CC_A,  // above
-            OpCode::FloatGe => CC_AE, // above or equal
-            OpCode::FloatEq => CC_E,  // equal
-            OpCode::FloatNe => CC_NE, // not equal
-            _ => CC_E,
-        }
-    }
-
     /// Guard with faillocs — emit conditional jump and store faillocs on descr.
     fn implement_guard_with_faillocs(
         &mut self,

majit/majit-translate/src/front/mir.rs

@@ -4628,12 +4628,15 @@ impl<'a> Lowering<'a> {
             //
             // Upstream materialises the address as `Constant(funcptr)` of
             // lltype `Ptr(FuncType)` (`rtyper.getcallable`, and
-            // `sub_helper_funcptr_constant` for the sub-helper twins), whose
-            // `getkind` is `r`.  The slot here stays `Int` because majit
-            // materialises a funcptr as its integer address everywhere else
-            // (`jtransform.rs direct_funcptr_value` emits `ConstInt(fnaddr)`,
-            // which the assembler encodes through the `'i'` argcode), and the
-            // flowspace fold gives the define a `Signed` legacy slot to match.
+            // `sub_helper_funcptr_constant` for the sub-helper twins).
+            // `FuncType._gckind` is `raw`, so `getkind` maps that pointer to
+            // `int`, not to `ref` — the `Int` slot here IS that mapping rather
+            // than a departure from it, and re-stamping it `Ref` would be the
+            // deviation.  majit materialises a funcptr as its integer address
+            // everywhere else as well (`jtransform.rs direct_funcptr_value`
+            // emits `ConstInt(fnaddr)`, which the assembler encodes through the
+            // `'i'` argcode), and the flowspace fold gives the define a
+            // `Signed` legacy slot to match.
             DecodedConst::FnPath(segments) => {
                 let mut synthetic = Vec::with_capacity(segments.len() + 1);
                 synthetic.push(crate::model::FN_CONST_HEAD.to_string());

pyre/pyre-interpreter/src/eval.rs

@@ -607,17 +607,11 @@ pub fn capture_pyframe_root_area() -> *const () {
 /// slot instead holds a `JitVirtualRef`, and reading that as a `PyFrame` would
 /// interpret its `virtual_token` word as frame fields.  Hop through the vref
 /// instead.  A still-virtual vref ends the walk: the frames it stands for have
-/// no heap image to visit, and `virtualref.py:157 force_virtual_if_necessary`
+/// no heap image to visit, and `virtualref.py force_virtual_if_necessary`
 /// cannot run here because materializing one allocates.
 #[inline]
 unsafe fn chain_next_frame(f_backref: *mut PyFrame) -> *mut PyFrame {
-    unsafe {
-        if majit_metainterp::virtualref::ptr_is_virtual_ref(f_backref as *const u8) {
-            majit_metainterp::virtualref::vref_forced(f_backref as *const u8) as *mut PyFrame
-        } else {
-            f_backref
-        }
-    }
+    crate::executioncontext::vref_referent(f_backref)
 }
 
 /// Walk one captured thread's active frame and interpreter root state.

pyre/pyre-interpreter/src/executioncontext.rs

@@ -49,9 +49,34 @@ pub fn register_force_vref_hook(f: ForceVRefFn) {
     let _ = FORCE_VREF_HOOK.set(f);
 }
 
+/// The frame a chain slot NAMES, read WITHOUT forcing —
+/// `virtualref.py force_virtual`'s trailing `return vref.forced`.
+///
+/// Exact for the whole recording walk: `virtual_ref_during_tracing` writes
+/// `forced = real_object` at allocation and only `continue_tracing` ever
+/// rewrites it.  Null when the vref is still virtual with nothing
+/// materialized — a live compiled frame — so callers must read that as
+/// "names no reachable frame", never as a match.
+///
+/// For identity tests only, never for handing a frame to application code.
+/// Forcing would be wrong here, not merely expensive: a live vref carries
+/// `TOKEN_TRACING_RESCALL` across a residual, `force_virtual` clears it, and
+/// that cleared token is the one marker `tracing_after_residual_call` reads as
+/// "the callee forced this vref".  A reader that forced would report its own
+/// read as a callee escape.
+#[inline]
+pub fn vref_referent(ptr: *mut PyFrame) -> *mut PyFrame {
+    if unsafe { majit_metainterp::virtualref::ptr_is_virtual_ref(ptr as *const u8) } {
+        unsafe { majit_metainterp::virtualref::vref_forced(ptr as *const u8) as *mut PyFrame }
+    } else {
+        ptr
+    }
+}
+
 /// Force a vref stored in the frame chain (`topframeref` / `f_backref`).
-/// `virtualref.py:135`: `if inst.typeptr != jit_virtual_ref_vtable: return inst`
-/// (the pointer already *is* the frame) else materialize via `force_virtual`.
+/// `virtualref.py force_virtual_if_necessary`: `if inst.typeptr !=
+/// jit_virtual_ref_vtable: return inst` (the pointer already *is* the frame)
+/// else materialize via `force_virtual`.
 #[inline]
 pub(crate) fn force_vref(ptr: *mut PyFrame) -> *mut PyFrame {
     if unsafe { majit_metainterp::virtualref::ptr_is_virtual_ref(ptr as *const u8) } {
@@ -907,7 +932,7 @@ impl ExecutionContext {
             self.w_tracefunc = pyre_object::PY_NULL;
         } else {
             self.force_all_frames(false);
-            // executioncontext.py:296-298 — increase the JIT's
+            // executioncontext.py settrace — increase the JIT's
             // trace_limit when a tracefunc is installed; tracing
             // generates a ton of extra ops per bytecode.
             crate::call::set_jit_param("trace_limit", 10000);
@@ -918,7 +943,7 @@ impl ExecutionContext {
         self.w_tracefunc
     }
 
-    /// pypy/interpreter/executioncontext.py:303-310 setprofile.
+    /// `executioncontext.py setprofile`.
     pub fn setprofile(&mut self, w_func: PyObjectRef) -> Result<(), crate::PyError> {
         if w_func.is_null() || w_func == pyre_object::w_none() {
             self.profilefunc = None;
@@ -929,19 +954,19 @@ impl ExecutionContext {
         }
     }
 
-    /// pypy/interpreter/executioncontext.py:312-313 getprofile.
+    /// `executioncontext.py getprofile`.
     pub fn getprofile(&self) -> PyObjectRef {
         self.w_profilefuncarg
     }
 
-    /// pypy/interpreter/executioncontext.py:315-321 setllprofile.
+    /// `executioncontext.py setllprofile`.
     pub fn setllprofile(
         &mut self,
         func: Option<ProfileFunc>,
         w_arg: PyObjectRef,
     ) -> Result<(), crate::PyError> {
         if func.is_some() {
-            // executioncontext.py:317-318 `if w_arg is None: raise
+            // executioncontext.py setllprofile: `if w_arg is None: raise
             // ValueError("Cannot call setllprofile with real None")`.
             // The check is against RPython-level None (== null in pyre);
             // Python-level `w_none()` (`space.w_None`) is a valid user
@@ -958,9 +983,24 @@ impl ExecutionContext {
         Ok(())
     }
 
+    /// `executioncontext.py force_all_frames` — "force" every frame in the
+    /// sense of the JIT, so one that is running in assembler fails its next
+    /// `GUARD_NOT_FORCED` and falls back to interpreted execution, where the
+    /// freshly installed trace / profile callback is honoured.
+    ///
+    /// Upstream gets that effect from the walk itself: `f_backref` holds a
+    /// `jit.virtual_ref`, so `getnextframe_nohidden`'s `frame.f_backref()` is a
+    /// `jit_force_virtual`, and `virtualref.force_virtual` runs
+    /// `ResumeGuardForcedDescr.force_now` on a token that still names a live
+    /// JIT frame.  Pyre's walk calls [`force_vref`] at the same points, but
+    /// nothing stores a `JitVirtualRef` in the chain yet, so it is the identity
+    /// and the walk forces nothing.  Until the tracer emits `VIRTUAL_REF` at
+    /// the inline push, this consumer — whose whole purpose is the force —
+    /// states it directly.
     pub fn force_all_frames(&mut self, is_being_profiled: bool) {
         let mut frame = self.gettopframe_nohidden();
         while !frame.is_null() {
+            force_frame(frame);
             if is_being_profiled {
                 unsafe {
                     (*frame).getorcreatedebug(-1).is_being_profiled = true;

pyre/pyre-interpreter/src/jit_fnaddr.rs

@@ -1479,12 +1479,10 @@ pub fn jit_trace_fnaddrs() -> Vec<(&'static str, i64)> {
     // store leaves to native ops, leaving `list_write_barrier(l)` as a
     // residual call. Register it so the codewriter resolves the residual to a
     // runtime-patchable address instead of a `symbolic_fnaddr_for_path` hash
-    // the inline sub-walk must decline. The address is also what the walker
-    // matches on to drop the residual entirely when the backend GC rewrite
-    // already covers the store (`FbwWalkMode::append_inplace_wb_covered`);
-    // with an off-GC ItemsBlock the residual stays, because there the
-    // collector reaches the block's slots only through the remembered
-    // `W_ListObject`.
+    // the inline sub-walk must decline. The residual barrier remembers the
+    // enclosing `W_ListObject`, whose trace reaches every item slot, and is
+    // the only thing keeping an appended `old -> young` element reachable
+    // across a minor collection.
     push_alias_pair(
         &mut entries,
         "pyre_object::listobject::list_write_barrier",