diff --git a/majit/majit-backend-dynasm/src/aarch64/assembler.rs b/majit/majit-backend-dynasm/src/aarch64/assembler.rs index b3ce39da480..5191d09b7ec 100644 --- a/majit/majit-backend-dynasm/src/aarch64/assembler.rs +++ b/majit/majit-backend-dynasm/src/aarch64/assembler.rs @@ -1576,7 +1576,11 @@ impl<'a> AssemblerARM64<'a> { Some(Loc::Frame(f)) => { gcmap_set_bit(gcmap, f.position + JITFRAME_FIXED_SIZE); } - _ => {} + None => {} + Some(other) => panic!( + "guard_gcmap_from_faillocs: a Ref fail argument at {other:?} \ + carries no gcmap bit" + ), } } gcmap @@ -1924,12 +1928,14 @@ impl<'a> AssemblerARM64<'a> { input_slot_depth.max(JITFRAME_FIXED_SIZE + ra.get_final_frame_depth()); self.frame_depth = self.frame_depth.max(frame_slot_depth); - // Sync regalloc frame positions to opref_to_slot for the emitters that - // still read resolve_opref instead of arglocs: emit_call, - // genop_discard_setfield, genop_cond_call_value, genop_alloc_varsize - // and genop_discard_zero_array reach it through load_arg_to_rax / - // load_arg_to_rcx. When regalloc spills a value to a frame slot, that - // slot's position must be visible to resolve_opref. + // Sync regalloc frame positions to opref_to_slot, the map `resolve_opref` + // reads. No live emitter reaches it any more: every remaining consumer + // (`load_arg_to_rax` / `load_arg_to_rcx` / `resolve_const_or` and the + // genops that call them) is `#[allow(dead_code)]`, kept for the upstream + // method boundary. The map cannot express what the regalloc actually + // decides — `before_call` leaves a value bound to a callee-saved member + // of the allocation pool in its register, with no frame slot at all — so + // an emitter that needs an operand location must take it from `arglocs`. // opref_to_slot stores ABSOLUTE jitframe slots (user position + // JITFRAME_FIXED_SIZE) so slot_offset(slot) gives the correct byte // offset without further adjustment. @@ -2654,13 +2660,31 @@ impl<'a> AssemblerARM64<'a> { } // ── Memory stores: opassembler.rs emit_op_setfield_regalloc ── OpCode::SetfieldGc | OpCode::SetfieldRaw => { - if let (Some(Loc::Reg(base)), Some(val_loc)) = (arglocs.first(), arglocs.get(1)) { - let ofs = op.with_field_descr(|fd| fd.offset() as i32).unwrap_or(0); - let field_size = op.with_field_descr(|fd| fd.field_size()).unwrap_or(8); - self.emit_op_setfield_regalloc(base, val_loc, ofs, field_size); - } else { - self.genop_discard_setfield(op); - } + // `rewrite.rs transform_to_gc_load` lowers both opcodes to + // GC_STORE / GC_STORE_INDEXED through `emit_gc_store_or_indexed` + // before the regalloc ever sees them, and `consider_setfield_j2` + // force-allocates a register for every non-constant base, so the + // only shape that could reach the else arm is a constant base on + // an op the rewriter did not consume. Emitting it from + // `resolve_opref` would read whichever slot the lifetime's + // end-of-allocation `current_frame_loc` happens to name and + // clobber x0/x1 behind the regalloc's back; decline instead. + let [base_loc, val_loc] = arglocs else { + panic!( + "{:?} expects two regalloc locations, got {arglocs:?}", + op.opcode, + ); + }; + let Loc::Reg(base) = base_loc else { + panic!( + "{:?} reached the backend with a non-register base {base_loc:?} \ + — `rewrite.rs transform_to_gc_load` must have consumed it", + op.opcode, + ); + }; + let ofs = op.with_field_descr(|fd| fd.offset() as i32).unwrap_or(0); + let field_size = op.with_field_descr(|fd| fd.field_size()).unwrap_or(8); + self.emit_op_setfield_regalloc(base, val_loc, ofs, field_size); } // ── aarch64/opassembler.py _emit_op_gc_load ── // arglocs = [base_loc, ofs_loc, res_loc, imm(nsize)] @@ -2699,7 +2723,11 @@ impl<'a> AssemblerARM64<'a> { }; self.emit_op_gcload_regalloc(&base, ofs_loc, dst, nsize); } - _ => {} + None => {} + Some(other) => panic!( + "GC_LOAD_INDEXED: unhandled base {other:?} — no load is emitted \ + and the destination keeps its previous value" + ), } } } @@ -3127,9 +3155,9 @@ impl<'a> AssemblerARM64<'a> { // ── Allocation (raw, when GC rewriter is not active) ── OpCode::New => self.genop_new(op), OpCode::NewWithVtable => self.genop_new_with_vtable(op), - OpCode::NewArray | OpCode::NewArrayClear => self.genop_new_array(op), - OpCode::Newstr => self.genop_newstr(op), - OpCode::Newunicode => self.genop_newunicode(op), + OpCode::NewArray | OpCode::NewArrayClear => self.genop_new_array(op, arglocs), + OpCode::Newstr => self.genop_newstr(op, arglocs), + OpCode::Newunicode => self.genop_newunicode(op, arglocs), // ── Allocation (rewritten by GC rewriter) ── // aarch64/regalloc.py:958 + assembler.py:682 malloc_cond parity OpCode::CallMallocNursery => { @@ -3677,7 +3705,10 @@ impl<'a> AssemblerARM64<'a> { self.emit_mov_imm64(17, expected.value); dynasm!(self.mc ; .arch aarch64 ; cmp x16, x17); } - _ => {} + other => panic!( + "_cmp_guard_gc_type: unhandled expected typeid {other:?} — no cmp \ + is emitted and the guard branches on stale flags" + ), } } @@ -3763,7 +3794,10 @@ impl<'a> AssemblerARM64<'a> { self.emit_mov_imm64(17, expected.value); dynasm!(self.mc ; .arch aarch64 ; cmp x16, x17); } - _ => {} + other => panic!( + "emit_guard_exception: unhandled expected location {other:?} — no \ + cmp is emitted and the guard branches on stale flags" + ), } } @@ -3784,7 +3818,10 @@ impl<'a> AssemblerARM64<'a> { ); self.emit_str_fp(17, frame.ebp_loc.value); } - _ => {} + other => panic!( + "emit_store_and_reset_exception: unhandled result location \ + {other:?} — the exception value would be dropped" + ), } } dynasm!(self.mc ; .arch aarch64 ; str xzr, [x16]); @@ -5223,43 +5260,6 @@ impl<'a> AssemblerARM64<'a> { // x86/assembler.py:1747 genop_getfield_gc etc. // ---------------------------------------------------------------- - /// Extract the byte offset from an op's FieldDescr. - /// Returns 0 if no field descriptor is present. - fn field_offset_from_descr(op: &Op) -> i32 { - op.with_field_descr(|fd| fd.offset() as i32).unwrap_or(0) - } - - /// Extract the field size from an op's FieldDescr. - /// Returns 8 (WORD) if no field descriptor is present. - fn field_size_from_descr(op: &Op) -> usize { - op.with_field_descr(|fd| fd.field_size()).unwrap_or(8) - } - - /// SETFIELD_GC: [arg0 + offset] = arg1 - fn genop_discard_setfield(&mut self, op: &Op) { - let offset = Self::field_offset_from_descr(op); - let size = Self::field_size_from_descr(op); - - // Load object pointer into rax/x0 and value into rcx/x1. - self.load_arg_to_rax(op.arg(0).to_opref()); - self.load_arg_to_rcx(op.arg(1).to_opref()); - - match size { - 1 => dynasm!(self.mc ; .arch aarch64 - ; strb w1, [x0, offset as u32] - ), - 2 => dynasm!(self.mc ; .arch aarch64 - ; strh w1, [x0, offset as u32] - ), - 4 => dynasm!(self.mc ; .arch aarch64 - ; str w1, [x0, offset as u32] - ), - _ => dynasm!(self.mc ; .arch aarch64 - ; str x1, [x0, offset as u32] - ), - } - } - // ---------------------------------------------------------------- // genop_* — calls // x86/assembler.py _genop_call @@ -5272,42 +5272,6 @@ impl<'a> AssemblerARM64<'a> { } } - /// Emit a function call. `func_arg` is the index of the function - /// pointer arg; call arguments start at `func_arg + 1`. - fn emit_call(&mut self, op: &Op, func_arg: usize) { - let arg_count = op.num_args(); - - dynasm!(self.mc ; .arch aarch64 ; stp x29, x30, [sp, #-16]!); - - for i in (func_arg + 1)..arg_count.min(func_arg + 7) { - let arg = op.arg(i).to_opref(); - let abi_idx = i - func_arg - 1; - match self.resolve_opref(arg) { - ResolvedArg::Slot(offset) => { - let reg = abi_idx as u8; - self.emit_ldr_fp(reg, offset); - } - ResolvedArg::Const(val) => { - let reg = abi_idx as u32; - self.emit_mov_imm64(reg, val); - } - } - } - - match self.resolve_opref(op.arg(func_arg).to_opref()) { - ResolvedArg::Slot(offset) => { - self.emit_ldr_fp(8, offset); - dynasm!(self.mc ; .arch aarch64 ; blr x8); - } - ResolvedArg::Const(val) => { - self.emit_mov_imm64(8, val); - dynasm!(self.mc ; .arch aarch64 ; blr x8); - } - } - - dynasm!(self.mc ; .arch aarch64 ; ldp x29, x30, [sp], #16); - } - /// aarch64/opassembler.py _emit_call + aarch64/callbuilder.py:21-67 /// prepare_arguments. /// @@ -5393,7 +5357,10 @@ impl<'a> AssemblerARM64<'a> { Loc::Immed(im) => { immed_args.push((abi_idx, im.value, im.is_float)); } - _ => {} + // See the x86 twin: Reg/Frame/Immed is the whole range the + // regalloc produces, and silently skipping anything else + // passes a stale register as that argument. + other => panic!("call argument {abi_idx} has unsupported location {other:?}"), } } @@ -5449,11 +5416,15 @@ impl<'a> AssemblerARM64<'a> { if fnloc_in_ip1 { dynasm!(self.mc ; .arch aarch64 ; blr x17); } else { - if let Some(Loc::Immed(i)) = fnloc { - let val = i.value; - self.emit_mov_imm64(8, val); - dynasm!(self.mc ; .arch aarch64 ; blr x8); - } + // Unlike x86, which always ends in `call rax`, this arm is the + // only place a `blr` is emitted — an unhandled fnloc spelling + // would emit no call at all and fall through with a stale x0. + let Some(Loc::Immed(i)) = fnloc else { + panic!("unsupported AArch64 call target {fnloc:?}"); + }; + let val = i.value; + self.emit_mov_imm64(8, val); + dynasm!(self.mc ; .arch aarch64 ; blr x8); } if stack_bytes != 0 { @@ -5495,11 +5466,6 @@ impl<'a> AssemblerARM64<'a> { } } - /// assembler.py _genop_call — internal call implementation. - fn _genop_call(&mut self, op: &Op) { - self.emit_call(op, 0); - } - fn _genop_call_with_arglocs(&mut self, op: &Op, arglocs: &[Loc]) { let func_index = 3 + usize::from(op.opcode.is_call_release_gil()); self.emit_call_from_arglocs(arglocs, func_index); @@ -6381,11 +6347,11 @@ impl<'a> AssemblerARM64<'a> { } /// NEW_ARRAY / NEW_ARRAY_CLEAR: allocate an array. - fn genop_new_array(&mut self, op: &Op) { + fn genop_new_array(&mut self, op: &Op, arglocs: &[Loc]) { let (base_size, item_size) = op .with_array_descr(|ad| (ad.base_size() as i64, ad.item_size() as i64)) .unwrap_or((8, 8)); - self.genop_alloc_varsize(op, base_size, item_size); + self.genop_alloc_varsize(op, arglocs, base_size, item_size); } // ---------------------------------------------------------------- @@ -6409,7 +6375,10 @@ impl<'a> AssemblerARM64<'a> { dynasm!(self.mc ; .arch aarch64 ; ldr x17, [x16]); self.emit_str_fp(17, frame.ebp_loc.value); } - _ => {} + None => {} + Some(other) => { + panic!("genop_save_exc_class: unhandled result location {other:?}") + } } } @@ -6428,7 +6397,10 @@ impl<'a> AssemblerARM64<'a> { } Loc::Frame(frame) => this.emit_ldr_fp(17, frame.ebp_loc.value), Loc::Immed(imm) => this.emit_mov_imm64(17, imm.value), - _ => {} + other => panic!( + "genop_restore_exception: unhandled operand {other:?} — x17 still \ + holds the previously loaded operand, which the store below writes" + ), }; load_to_x17(self, &arglocs[1]); // value self.emit_mov_imm64(16, crate::jit_exc_value_addr() as i64); @@ -6652,17 +6624,21 @@ impl<'a> AssemblerARM64<'a> { /// COND_CALL_VALUE_I/R: if arg(0) == 0, call function; else result = arg(0). /// - /// The predicate comes from its regalloc location, not `resolve_opref`, - /// which only recognises constants and frame slots — a predicate left - /// register-resident has no slot mapping. It is loaded into x0 rather - /// than the ip0 scratch because on the not-taken path the predicate IS the - /// result, and `store_rax_to_result` reads it from there. + /// Every operand comes from its regalloc location, not `resolve_opref`, + /// which only recognises constants and frame slots. `before_call` spills + /// only `CALLER_RESP` (x0..x13), so a value bound to x19/x20 stays + /// register-resident across it and has no slot mapping at all. + /// + /// The predicate is loaded into x0 rather than the ip0 scratch because on + /// the not-taken path the predicate IS the result, and + /// `store_rax_to_result` reads it from there. x0 is caller-saved, so + /// `before_call` guarantees it is not itself one of the arglocs. fn genop_cond_call_value(&mut self, op: &Op, arglocs: &[Loc]) { self.emit_load_to_rax(arglocs[0]); let skip_label = self.mc.new_dynamic_label(); dynasm!(self.mc ; .arch aarch64 ; cbnz x0, =>skip_label); - self.emit_call(op, 1); + self.emit_call_from_arglocs(arglocs, 1); dynasm!(self.mc ; .arch aarch64 ; =>skip_label); @@ -6680,16 +6656,16 @@ impl<'a> AssemblerARM64<'a> { /// (`builtin_string_array_descr` in `runner.rs`), which encodes /// `get_array_token(rstr.STR, ...)` — basesize includes the +1 /// extra_item_after_alloc null terminator. - fn genop_newstr(&mut self, op: &Op) { + fn genop_newstr(&mut self, op: &Op, arglocs: &[Loc]) { let (base_size, item_size, type_id) = Self::array_token_from_descr(op, 16, 1); - self.genop_alloc_lowlevel_string(op, type_id, base_size, item_size); + self.genop_alloc_lowlevel_string(op, arglocs, type_id, base_size, item_size); } /// NEWUNICODE: allocate a unicode string (4-byte chars). /// Basesize = 16 (no extra_item_after_alloc), itemsize = 4. - fn genop_newunicode(&mut self, op: &Op) { + fn genop_newunicode(&mut self, op: &Op, arglocs: &[Loc]) { let (base_size, item_size, type_id) = Self::array_token_from_descr(op, 16, 4); - self.genop_alloc_lowlevel_string(op, type_id, base_size, item_size); + self.genop_alloc_lowlevel_string(op, arglocs, type_id, base_size, item_size); } /// Read `(base_size, item_size)` from the injected ArrayDescr. @@ -6709,11 +6685,20 @@ impl<'a> AssemblerARM64<'a> { fn genop_alloc_lowlevel_string( &mut self, op: &Op, + arglocs: &[Loc], type_id: i64, base_size: i64, item_size: i64, ) { - self.load_arg_to_rax(op.arg(0).to_opref()); + // The length comes from its regalloc location, for the reason + // `genop_alloc_varsize` states below: `consider_raw_call_like` plans + // these opcodes, and `before_call` leaves a value bound to a + // callee-saved member of the allocation pool in its register, where + // `resolve_opref` cannot see it. + let [len_loc, ..] = arglocs else { + panic!("lowlevel string allocation expects a length location, got {arglocs:?}"); + }; + self.emit_load_to_rax(*len_loc); dynasm!(self.mc ; .arch aarch64 ; mov x3, x0); self.emit_mov_imm64(0, type_id); self.emit_mov_imm64(1, base_size); @@ -6732,12 +6717,19 @@ impl<'a> AssemblerARM64<'a> { } } - /// Shared implementation for NEWSTR / NEWUNICODE / NEW_ARRAY. + /// Shared implementation for NEW_ARRAY. /// Allocates base_size + length * item_size bytes, zero-fills, /// and writes length to the header. - fn genop_alloc_varsize(&mut self, op: &Op, base_size: i64, item_size: i64) { - // arg(0) = length - self.load_arg_to_rax(op.arg(0).to_opref()); + fn genop_alloc_varsize(&mut self, op: &Op, arglocs: &[Loc], base_size: i64, item_size: i64) { + // The length comes from its regalloc location. `consider_raw_call_like_j2` + // plans these opcodes, and `before_call` leaves a value bound to a + // callee-saved member of the allocation pool in its register, where + // `resolve_opref` cannot see it. The load is consumed by the very next + // instruction, so rax/x0 doubling as ABI arg0 is safe here. + let [len_loc, ..] = arglocs else { + panic!("varsize allocation expects a length location, got {arglocs:?}"); + }; + self.emit_load_to_rax(*len_loc); let malloc_ptr = libc::malloc as *const () as i64; let memset_ptr = libc::memset as *const () as i64; @@ -6779,7 +6771,14 @@ impl<'a> AssemblerARM64<'a> { /// ZERO_ARRAY: zero a range in an array. /// arg(0)=base, arg(1)=start, arg(2)=size, arg(3)=scale_start, arg(4)=scale_size. fn genop_discard_zero_array(&mut self, op: &Op, arglocs: &[Loc]) { - let [base_loc, start_loc, size_loc, ..] = arglocs else { + let [ + base_loc, + start_loc, + size_loc, + scale_start_loc, + scale_size_loc, + ] = arglocs + else { panic!("ZERO_ARRAY expects five regalloc locations, got {arglocs:?}"); }; if matches!(size_loc, Loc::Immed(i) if i.value == 0) { @@ -6789,8 +6788,20 @@ impl<'a> AssemblerARM64<'a> { .with_array_descr(|ad| (ad.base_size() as i64, ad.item_size() as i64)) .unwrap_or((8, 8)); - let scale_start = self.resolve_const_or(op.arg(3).to_opref(), 1); - let scale_size = self.resolve_const_or(op.arg(4).to_opref(), 1); + // The scale operands are the `st.const_int(scale)` pair that + // `rewrite.rs` emits for every ZERO_ARRAY, so `make_sure_var_in_reg` + // hands them back as bare immediates (`return_constant` with no + // selected_reg). Read them from there rather than re-resolving the + // op: a non-constant scale is an invariant break the emitter cannot + // encode, and failing loud declines the trace instead of silently + // scaling by one. + let Loc::Immed(scale_start) = scale_start_loc else { + panic!("ZERO_ARRAY scale_start must be an immediate, got {scale_start_loc:?}"); + }; + let Loc::Immed(scale_size) = scale_size_loc else { + panic!("ZERO_ARRAY scale_size must be an immediate, got {scale_size_loc:?}"); + }; + let (scale_start, scale_size) = (scale_start.value, scale_size.value); // aarch64/opassembler.py:755-839: first compute the byte destination // in ip0/x16. ip0/ip1 are never managed by regalloc. diff --git a/majit/majit-backend-dynasm/src/regalloc.rs b/majit/majit-backend-dynasm/src/regalloc.rs index c77b5aa744d..8308bf2c855 100644 --- a/majit/majit-backend-dynasm/src/regalloc.rs +++ b/majit/majit-backend-dynasm/src/regalloc.rs @@ -5279,6 +5279,24 @@ impl<'a> RegAlloc<'a> { self.perform_with_gcmap_ptr(i, arglocs, result_loc, gcmap, output); } + /// `emit_call_from_arglocs` silently falls back to all-Int argument types + /// when the descr's arity disagrees with the call, which would place a + /// float argument in a GPR. Every cond-call-value producer builds the arg + /// list and `arg_types` from the same source, so check it here rather than + /// letting the emitter guess — `consider_call_j2` carries the same check + /// for plain calls. `num_args` counts the leading value and function + /// operands, which the descr does not describe. + fn check_cond_call_value_descr_arity(op: &Op, num_args: usize) { + if !op.opcode.is_cond_call_value() { + return; + } + let descr_arc = op.getdescr().expect("cond_call_value without CallDescr"); + let calldescr = descr_arc + .as_call_descr() + .expect("cond_call_value without CallDescr"); + assert_eq!(calldescr.arg_types().len(), num_args - 2); + } + fn consider_raw_call_like( &mut self, op: &Op, @@ -5286,6 +5304,7 @@ impl<'a> RegAlloc<'a> { output: &mut Vec, save_regs: u8, ) { + Self::check_cond_call_value_descr_arity(op, op.num_args()); let type_index = OpTypeIndex::from_parts( self.inputargs, self.operations, @@ -5339,6 +5358,7 @@ impl<'a> RegAlloc<'a> { output: &mut Vec, save_regs: u8, ) { + Self::check_cond_call_value_descr_arity(op, args.len()); let type_index = OpTypeIndex::from_parts( self.inputargs, self.operations, diff --git a/majit/majit-backend-dynasm/src/runner.rs b/majit/majit-backend-dynasm/src/runner.rs index db93d131fa6..26d9a24f2ee 100644 --- a/majit/majit-backend-dynasm/src/runner.rs +++ b/majit/majit-backend-dynasm/src/runner.rs @@ -4770,6 +4770,80 @@ mod tests { assert_eq!(backend.get_ref_value(&frame, 0), second_payload); } + extern "C" fn return_int_passthrough(arg: i64) -> i64 { + arg + } + + /// A `COND_CALL_VALUE_I` argument that is an op result — not an inputarg — + /// and is still live after the call. + /// + /// `consider_raw_call_like_j2` runs `before_call` before it reads the + /// locations, and `spill_or_move_registers_before_call` prefers *moving* a + /// survivor into a free callee-saved register over spilling it. The + /// argument therefore reaches the emitter as a bare `Loc::Reg` (measured: + /// x19 on AArch64) whose lifetime has no frame slot at that point, which is + /// the one shape `resolve_opref` cannot represent. + /// + /// The callee is the identity, so a wrong argument location surfaces + /// directly in the result. `i2` is deliberately NOT equal to `i1`: reading + /// the argument from the slot that the recycled frame position happens to + /// name would otherwise return the right answer by accident. + /// + /// The two older cond-call fixtures pass zero call arguments, so their + /// argument loop is empty and neither can see this. + #[test] + fn test_cond_call_value_passes_a_register_resident_op_result_argument() { + let mut backend = DynasmBackend::new(); + backend.attach_default_test_descrs(); + + // arg 0 is the cond-call predicate, passed 0 so the call is taken. + let inputargs = vec![InputArg::new_int(0), InputArg::new_int(1)]; + let mut constants: indexmap::IndexMap = indexmap::IndexMap::new(); + constants.insert(200, return_int_passthrough as *const () as usize as i64); + backend.set_constants(constants); + + let cond_call = mk_op( + OpCode::CondCallValueI, + &[ + OpRef::input_arg_int(0), + OpRef::int_op(200), + OpRef::int_op(2), + ], + 3, + ); + cond_call.setdescr(make_plain_call_descr(vec![Type::Int], Type::Int)); + + let ops = vec![ + mk_op( + OpCode::Label, + &[OpRef::input_arg_int(0), OpRef::input_arg_int(1)], + OpRef::NONE.raw(), + ), + // i2 = i1 + i1, an op result distinct from every inputarg value. + mk_op( + OpCode::IntAdd, + &[OpRef::input_arg_int(1), OpRef::input_arg_int(1)], + 2, + ), + cond_call, + // Naming i2 again keeps it live across the cond-call, which is what + // sends it down `before_call`'s move-to-callee-saved arm instead of + // letting it die there. + mk_op( + OpCode::Finish, + &[OpRef::int_op(3), OpRef::int_op(2)], + OpRef::NONE.raw(), + ), + ]; + + let token = JitCellToken::new(1617); + backend.compile_loop(&inputargs, &ops, &token).unwrap(); + + let frame = backend.execute_token(&token, &[Value::Int(0), Value::Int(7)]); + assert!(backend.get_latest_descr(&frame).is_finish()); + assert_eq!(backend.get_int_value(&frame, 0), 14); + } + #[test] fn test_label_uses_absolute_jitframe_input_slots_for_resolve_opref_ops() { let mut gc = MiniMarkGC::with_config(GcConfig { diff --git a/majit/majit-backend-dynasm/src/x86/assembler.rs b/majit/majit-backend-dynasm/src/x86/assembler.rs index 26d786eaa1c..cc20bb5d2c4 100644 --- a/majit/majit-backend-dynasm/src/x86/assembler.rs +++ b/majit/majit-backend-dynasm/src/x86/assembler.rs @@ -1035,6 +1035,10 @@ impl<'a> Assembler386<'a> { /// rpython/jit/metainterp/history.py:220 `box.type` parity. /// Single source of truth: `op.type_` for ops, `inputarg.tp` for /// inputargs, the `Const` variant tag for constants. + /// + /// Unreached since the call emitters moved onto regalloc arglocs, which + /// carry the type in the `Loc` itself; kept per the note above. + #[allow(dead_code)] #[inline] fn opref_type(&self, opref: OpRef) -> Option { self.opref_type_at(opref, None) @@ -1332,7 +1336,10 @@ impl<'a> Assembler386<'a> { _ => {} } } - _ => {} + other => panic!( + "emit_binop_reg_loc: unhandled source {other:?} — no arithmetic is \ + emitted and the destination keeps its previous value" + ), } } @@ -1382,7 +1389,10 @@ impl<'a> Assembler386<'a> { let scratch = crate::regloc::X86_64_SCRATCH_REG.value; dynasm!(self.mc ; .arch x64 ; mov Rq(scratch), QWORD i.value ; test Rq(scratch), Rq(scratch)); } - _ => {} + other => panic!( + "emit_test_loc: unhandled operand {other:?} — no test is emitted \ + and the following branch reads stale flags" + ), } } @@ -2329,7 +2339,11 @@ impl<'a> Assembler386<'a> { Some(Loc::Frame(f)) => { gcmap_set_bit(gcmap, f.position + JITFRAME_FIXED_SIZE); } - _ => {} + None => {} + Some(other) => panic!( + "guard_gcmap_from_faillocs: a Ref fail argument at {other:?} \ + carries no gcmap bit" + ), } } gcmap @@ -2656,12 +2670,14 @@ impl<'a> Assembler386<'a> { input_slot_depth.max(JITFRAME_FIXED_SIZE + ra.get_final_frame_depth()); self.frame_depth = self.frame_depth.max(frame_slot_depth); - // Sync regalloc frame positions to opref_to_slot for the emitters that - // still read resolve_opref instead of arglocs: emit_call, - // genop_discard_setfield, genop_cond_call_value, genop_alloc_varsize - // and genop_discard_zero_array reach it through load_arg_to_rax / - // load_arg_to_rcx. When regalloc spills a value to a frame slot, that - // slot's position must be visible to resolve_opref. + // Sync regalloc frame positions to opref_to_slot, the map `resolve_opref` + // reads. No live emitter reaches it any more: every remaining consumer + // (`load_arg_to_rax` / `load_arg_to_rcx` / `resolve_const_or` and the + // genops that call them) is `#[allow(dead_code)]`, kept for the upstream + // method boundary. The map cannot express what the regalloc actually + // decides — `before_call` leaves a value bound to a callee-saved member + // of the allocation pool in its register, with no frame slot at all — so + // an emitter that needs an operand location must take it from `arglocs`. // opref_to_slot stores ABSOLUTE jitframe slots (user position + // JITFRAME_FIXED_SIZE) so slot_offset(slot) gives the correct byte // offset without further adjustment. @@ -3061,7 +3077,10 @@ impl<'a> Assembler386<'a> { ; mul Rq(scratch) ); } - _ => {} + other => panic!( + "UINT_MUL_HIGH: unhandled source {other:?} — no mul is \ + emitted and edx:eax keep their previous values" + ), } if dst.value != crate::regloc::EDX.value { dynasm!(self.mc ; .arch x64 ; mov Rq(dst.value), rdx); @@ -3427,13 +3446,31 @@ impl<'a> Assembler386<'a> { } // ── Memory stores: opassembler.rs emit_op_setfield_regalloc ── OpCode::SetfieldGc | OpCode::SetfieldRaw => { - if let (Some(Loc::Reg(base)), Some(val_loc)) = (arglocs.first(), arglocs.get(1)) { - let ofs = op.with_field_descr(|fd| fd.offset() as i32).unwrap_or(0); - let field_size = op.with_field_descr(|fd| fd.field_size()).unwrap_or(8); - self.emit_op_setfield_regalloc(base, val_loc, ofs, field_size); - } else { - self.genop_discard_setfield(op); - } + // `rewrite.rs transform_to_gc_load` lowers both opcodes to + // GC_STORE / GC_STORE_INDEXED through `emit_gc_store_or_indexed` + // before the regalloc ever sees them, and `consider_setfield_j2` + // force-allocates a register for every non-constant base, so the + // only shape that could reach the else arm is a constant base on + // an op the rewriter did not consume. Emitting it from + // `resolve_opref` would read whichever slot the lifetime's + // end-of-allocation `current_frame_loc` happens to name and + // clobber rax/rcx behind the regalloc's back; decline instead. + let [base_loc, val_loc] = arglocs else { + panic!( + "{:?} expects two regalloc locations, got {arglocs:?}", + op.opcode, + ); + }; + let Loc::Reg(base) = base_loc else { + panic!( + "{:?} reached the backend with a non-register base {base_loc:?} \ + — `rewrite.rs transform_to_gc_load` must have consumed it", + op.opcode, + ); + }; + let ofs = op.with_field_descr(|fd| fd.offset() as i32).unwrap_or(0); + let field_size = op.with_field_descr(|fd| fd.field_size()).unwrap_or(8); + self.emit_op_setfield_regalloc(base, val_loc, ofs, field_size); } // arglocs = [base_loc, ofs_loc, res_loc, imm(nsize)]. // `base_loc` may be Loc::Immed when the load is from a @@ -3607,7 +3644,10 @@ impl<'a> Assembler386<'a> { ; mov [Rq(base.value) + Rq(ofs_reg.value) * $scale + offset], Rq(scratch)), } } - _ => {} + other => panic!( + "emit_store_scaled: unhandled value location {other:?} — no store \ + is emitted" + ), } }}; } @@ -3643,7 +3683,10 @@ impl<'a> Assembler386<'a> { ; mov [Rq(base.value) + Rq(ofs_reg.value) + offset], Rq(scratch)), } } - _ => {} + other => panic!( + "emit_store_unscaled: unhandled value location {other:?} — no store \ + is emitted" + ), } }}; } @@ -4252,9 +4295,9 @@ impl<'a> Assembler386<'a> { // ── Allocation (raw, when GC rewriter is not active) ── OpCode::New => self.genop_new(op), OpCode::NewWithVtable => self.genop_new_with_vtable(op), - OpCode::NewArray | OpCode::NewArrayClear => self.genop_new_array(op), - OpCode::Newstr => self.genop_newstr(op), - OpCode::Newunicode => self.genop_newunicode(op), + OpCode::NewArray | OpCode::NewArrayClear => self.genop_new_array(op, arglocs), + OpCode::Newstr => self.genop_newstr(op, arglocs), + OpCode::Newunicode => self.genop_newunicode(op, arglocs), // ── Allocation (rewritten by GC rewriter) ── OpCode::CallMallocNursery => { self.genop_call_malloc_nursery(op, result_loc); @@ -4966,7 +5009,10 @@ impl<'a> Assembler386<'a> { dynasm!(self.mc ; .arch x64 ; cmp Rq(reg), [rbp + ofs]); } Loc::Immed(value) => self.emit_cmp_imm64(reg, value.value), - _ => {} + other => panic!( + "emit_cmp_reg_loc_i64: unhandled operand {other:?} — no cmp is \ + emitted and the following branch reads stale flags" + ), } } @@ -5043,7 +5089,10 @@ impl<'a> Assembler386<'a> { ; mov [rbp + ofs], Rq(scratch) ); } - _ => {} + other => panic!( + "emit_store_and_reset_exception: unhandled result location \ + {other:?} — the exception value would be dropped" + ), } } dynasm!(self.mc ; .arch x64 @@ -6361,18 +6410,6 @@ impl<'a> Assembler386<'a> { // genop_* — memory operations // x86/assembler.py:1747 genop_getfield_gc etc. - /// Extract the byte offset from an op's FieldDescr. - /// Returns 0 if no field descriptor is present. - fn field_offset_from_descr(op: &Op) -> i32 { - op.with_field_descr(|fd| fd.offset() as i32).unwrap_or(0) - } - - /// Extract the field size from an op's FieldDescr. - /// Returns 8 (WORD) if no field descriptor is present. - fn field_size_from_descr(op: &Op) -> usize { - op.with_field_descr(|fd| fd.field_size()).unwrap_or(8) - } - /// x86/assembler.py:1746 genop_discard_setfield — sized store via regalloc. /// Stage non-register values through X86_64_SCRATCH_REG (r11), mirroring /// the aarch64 path that uses x16. @@ -6424,7 +6461,10 @@ impl<'a> Assembler386<'a> { Loc::Reg(ofs_r) => { self.emit_gcload_sized(base, 0, Some(ofs_r), dst, abs_size, signed); } - _ => {} + other => panic!( + "emit_op_gcload_regalloc: unhandled offset {other:?} — no load is \ + emitted and the destination keeps its previous value" + ), } } @@ -6658,35 +6698,6 @@ impl<'a> Assembler386<'a> { } } - /// SETFIELD_GC: [arg0 + offset] = arg1 - fn genop_discard_setfield(&mut self, op: &Op) { - let offset = Self::field_offset_from_descr(op); - let size = Self::field_size_from_descr(op); - - // Load object pointer into rax/x0 and value into rcx/x1. - self.load_arg_to_rax(op.arg(0).to_opref()); - self.load_arg_to_rcx(op.arg(1).to_opref()); - - match size { - 1 => dynasm!(self.mc - ; .arch x64 - ; mov [rax + offset], cl - ), - 2 => dynasm!(self.mc - ; .arch x64 - ; mov [rax + offset], cx - ), - 4 => dynasm!(self.mc - ; .arch x64 - ; mov [rax + offset], ecx - ), - _ => dynasm!(self.mc - ; .arch x64 - ; mov [rax + offset], rcx - ), - } - } - // genop_* — calls // x86/assembler.py _genop_call @@ -6697,61 +6708,6 @@ impl<'a> Assembler386<'a> { } } - /// Emit a function call. `func_arg` is the index of the function - /// pointer arg; call arguments start at `func_arg + 1`. - fn emit_call(&mut self, op: &Op, func_arg: usize) { - let arg_count = op.num_args(); - let call_arg_count = arg_count.saturating_sub(func_arg + 1); - let descr_arc = op.getdescr(); - let arg_types = descr_arc - .as_ref() - .and_then(|descr| descr.as_call_descr()) - .map(|descr| descr.arg_types().to_vec()) - .filter(|types| types.len() == call_arg_count) - .unwrap_or_else(|| { - (func_arg + 1..arg_count) - .map(|i| self.opref_type(op.arg(i).to_opref()).unwrap_or(Type::Int)) - .collect() - }); - let (placements, stack_slots) = Self::build_abi_arg_placements(&arg_types); - - dynasm!(self.mc ; .arch x64 ; push rbp); - let call_area_adjust = self.emit_reserve_abi_call_area(1, stack_slots); - - for i in (func_arg + 1)..arg_count { - let arg = op.arg(i).to_opref(); - let abi_idx = i - func_arg - 1; - let placement = placements[abi_idx]; - let arg_type = arg_types[abi_idx]; - match self.resolve_opref(arg) { - ResolvedArg::Slot(offset) => { - self.emit_abi_arg_from_mem(placement, offset, arg_type) - } - ResolvedArg::Const(val) => { - self.emit_abi_arg_from_imm(placement, val as i64, arg_type) - } - } - } - - match self.resolve_opref(op.arg(func_arg).to_opref()) { - ResolvedArg::Slot(offset) => { - dynasm!(self.mc ; .arch x64 - ; mov rax, [rbp + offset] - ; call rax - ); - } - ResolvedArg::Const(val) => { - dynasm!(self.mc ; .arch x64 - ; mov rax, QWORD val as i64 - ; call rax - ); - } - } - - self.emit_release_abi_call_area(call_area_adjust); - dynasm!(self.mc ; .arch x64 ; pop rbp); - } - /// aarch64/opassembler.py _emit_call. /// arglocs = [resloc, size, sign, func, args...] for normal CALLs and /// [resloc, size, sign, saveerr, func, args...] for CALL_RELEASE_GIL. @@ -6795,7 +6751,11 @@ impl<'a> Assembler386<'a> { Loc::Frame(f) => self.emit_abi_arg_from_mem(placement, f.ebp_loc.value, arg_type), Loc::Reg(r) => self.emit_abi_arg_from_reg(placement, *r, arg_type), Loc::Immed(i) => self.emit_abi_arg_from_imm(placement, i.value, arg_type), - _ => {} + // `RegisterManager::loc` and `make_sure_var_in_reg` only ever + // yield Reg/Frame/Immed, so any other spelling is a regalloc + // change this emitter has not been taught. Falling through + // would drop the argument and call with a stale register. + other => panic!("call argument {abi_idx} has unsupported location {other:?}"), } } @@ -6850,7 +6810,9 @@ impl<'a> Assembler386<'a> { let val = i.value; dynasm!(self.mc ; .arch x64 ; mov rax, QWORD val); } - _ => {} + // `call rax` is emitted unconditionally below, so leaving rax + // unwritten here would call whatever it happened to hold. + other => panic!("unsupported x86-64 call target {other:?}"), } } dynasm!(self.mc ; .arch x64 ; call rax); @@ -6892,11 +6854,6 @@ impl<'a> Assembler386<'a> { } } - /// assembler.py _genop_call — internal call implementation. - fn _genop_call(&mut self, op: &Op) { - self.emit_call(op, 0); - } - fn _genop_call_with_arglocs(&mut self, op: &Op, arglocs: &[Loc]) { let func_index = 3 + usize::from(op.opcode.is_call_release_gil()); self.emit_call_from_arglocs(op, arglocs, func_index); @@ -7754,11 +7711,11 @@ impl<'a> Assembler386<'a> { } /// NEW_ARRAY / NEW_ARRAY_CLEAR: allocate an array. - fn genop_new_array(&mut self, op: &Op) { + fn genop_new_array(&mut self, op: &Op, arglocs: &[Loc]) { let (base_size, item_size) = op .with_array_descr(|ad| (ad.base_size() as i64, ad.item_size() as i64)) .unwrap_or((8, 8)); - self.genop_alloc_varsize(op, base_size, item_size); + self.genop_alloc_varsize(op, arglocs, base_size, item_size); } // genop_* — misc @@ -7783,7 +7740,10 @@ impl<'a> Assembler386<'a> { ; mov [rbp + ofs], Rq(scratch) ); } - _ => {} + None => {} + Some(other) => { + panic!("genop_save_exc_class: unhandled result location {other:?}") + } } } @@ -7820,7 +7780,10 @@ impl<'a> Assembler386<'a> { ; pop rax ); } - _ => {} + other => panic!( + "genop_restore_exception: unhandled operand {other:?} — the store is \ + skipped and the exception cell keeps its previous contents" + ), } }; store_loc_to(self, crate::jit_exc_value_addr() as i64, &arglocs[1]); @@ -8020,18 +7983,23 @@ impl<'a> Assembler386<'a> { /// COND_CALL_VALUE_I/R: if arg(0) == 0, call function; else result = arg(0). /// - /// The predicate comes from its regalloc location for the same reason as - /// `genop_discard_cond_call`. It is loaded into rax rather than the - /// scratch because on the not-taken path the predicate IS the result, and - /// `store_rax_to_result` reads it from there. `consider_raw_call_like_j2` - /// runs `before_call` before computing arglocs, so no argloc is a - /// caller-saved register and this load cannot clobber one. + /// Every operand comes from its regalloc location. `before_call` only + /// spills the registers in `SAVE_AROUND_CALL_CORE_REGS`; a value bound to + /// one of the callee-saved members of `ALL_CORE_REGS` (ebx/r12..r15, plus + /// esi/edi on Win64) stays register-resident across it, and such a value + /// has no `opref_to_slot` entry — `resolve_opref` would panic on it, or + /// read whatever unrelated slot the lifetime's `current_frame_loc` names. + /// + /// The predicate is loaded into rax rather than the scratch because on + /// the not-taken path the predicate IS the result, and + /// `store_rax_to_result` reads it from there. rax is caller-saved, so + /// `before_call` guarantees it is not itself one of the arglocs. fn genop_cond_call_value(&mut self, op: &Op, arglocs: &[Loc]) { self.emit_load_to_rax(arglocs[0]); let skip_label = self.mc.new_dynamic_label(); dynasm!(self.mc ; .arch x64 ; test rax, rax ; jnz =>skip_label); - self.emit_call(op, 1); + self.emit_call_from_arglocs(op, arglocs, 1); dynasm!(self.mc ; .arch x64 ; =>skip_label); @@ -8047,16 +8015,16 @@ impl<'a> Assembler386<'a> { /// (`builtin_string_array_descr` in `runner.rs`), which encodes /// `get_array_token(rstr.STR, ...)` — basesize includes the +1 /// extra_item_after_alloc null terminator. - fn genop_newstr(&mut self, op: &Op) { + fn genop_newstr(&mut self, op: &Op, arglocs: &[Loc]) { let (base_size, item_size, type_id) = Self::array_token_from_descr(op, 16, 1); - self.genop_alloc_lowlevel_string(op, type_id, base_size, item_size); + self.genop_alloc_lowlevel_string(op, arglocs, type_id, base_size, item_size); } /// NEWUNICODE: allocate a unicode string (4-byte chars). /// Basesize = 16 (no extra_item_after_alloc), itemsize = 4. - fn genop_newunicode(&mut self, op: &Op) { + fn genop_newunicode(&mut self, op: &Op, arglocs: &[Loc]) { let (base_size, item_size, type_id) = Self::array_token_from_descr(op, 16, 4); - self.genop_alloc_lowlevel_string(op, type_id, base_size, item_size); + self.genop_alloc_lowlevel_string(op, arglocs, type_id, base_size, item_size); } /// Read `(base_size, item_size)` from the injected ArrayDescr. @@ -8076,11 +8044,20 @@ impl<'a> Assembler386<'a> { fn genop_alloc_lowlevel_string( &mut self, op: &Op, + arglocs: &[Loc], type_id: i64, base_size: i64, item_size: i64, ) { - self.load_arg_to_rax(op.arg(0).to_opref()); + // The length comes from its regalloc location, for the reason + // `genop_alloc_varsize` states below: `consider_raw_call_like` plans + // these opcodes, and `before_call` leaves a value bound to a + // callee-saved member of the allocation pool in its register, where + // `resolve_opref` cannot see it. + let [len_loc, ..] = arglocs else { + panic!("lowlevel string allocation expects a length location, got {arglocs:?}"); + }; + self.emit_load_to_rax(*len_loc); self.emit_abi_int_arg_from_reg(3, 0); self.emit_abi_int_arg_from_imm(0, type_id); self.emit_abi_int_arg_from_imm(1, base_size); @@ -8094,12 +8071,19 @@ impl<'a> Assembler386<'a> { } } - /// Shared implementation for NEWSTR / NEWUNICODE / NEW_ARRAY. + /// Shared implementation for NEW_ARRAY. /// Allocates base_size + length * item_size bytes, zero-fills, /// and writes length to the header. - fn genop_alloc_varsize(&mut self, op: &Op, base_size: i64, item_size: i64) { - // arg(0) = length - self.load_arg_to_rax(op.arg(0).to_opref()); + fn genop_alloc_varsize(&mut self, op: &Op, arglocs: &[Loc], base_size: i64, item_size: i64) { + // The length comes from its regalloc location. `consider_raw_call_like_j2` + // plans these opcodes, and `before_call` leaves a value bound to a + // callee-saved member of the allocation pool in its register, where + // `resolve_opref` cannot see it. The load is consumed by the very next + // instruction, so rax/x0 doubling as ABI arg0 is safe here. + let [len_loc, ..] = arglocs else { + panic!("varsize allocation expects a length location, got {arglocs:?}"); + }; + self.emit_load_to_rax(*len_loc); let malloc_ptr = libc::malloc as *const () as i64; let memset_ptr = libc::memset as *const () as i64; @@ -8137,7 +8121,14 @@ impl<'a> Assembler386<'a> { /// ZERO_ARRAY: zero a range in an array. /// arg(0)=base, arg(1)=start, arg(2)=size, arg(3)=scale_start, arg(4)=scale_size. fn genop_discard_zero_array(&mut self, op: &Op, arglocs: &[Loc]) { - let [base_loc, start_loc, size_loc, ..] = arglocs else { + let [ + base_loc, + start_loc, + size_loc, + scale_start_loc, + scale_size_loc, + ] = arglocs + else { panic!("ZERO_ARRAY expects five regalloc locations, got {arglocs:?}"); }; if matches!(size_loc, Loc::Immed(i) if i.value == 0) { @@ -8147,8 +8138,20 @@ impl<'a> Assembler386<'a> { .with_array_descr(|ad| (ad.base_size() as i64, ad.item_size() as i64)) .unwrap_or((8, 8)); - let scale_start = self.resolve_const_or(op.arg(3).to_opref(), 1); - let scale_size = self.resolve_const_or(op.arg(4).to_opref(), 1); + // The scale operands are the `st.const_int(scale)` pair that + // `rewrite.rs` emits for every ZERO_ARRAY, so `make_sure_var_in_reg` + // hands them back as bare immediates (`return_constant` with no + // selected_reg). Read them from there rather than re-resolving the + // op: a non-constant scale is an invariant break the emitter cannot + // encode, and failing loud declines the trace instead of silently + // scaling by one. + let Loc::Immed(scale_start) = scale_start_loc else { + panic!("ZERO_ARRAY scale_start must be an immediate, got {scale_start_loc:?}"); + }; + let Loc::Immed(scale_size) = scale_size_loc else { + panic!("ZERO_ARRAY scale_size must be an immediate, got {scale_size_loc:?}"); + }; + let (scale_start, scale_size) = (scale_start.value, scale_size.value); // x86/regalloc.py:1436-1503 + assembler.py:2694-2725. Materialize // the effective address in r11, PyPy's reserved x86-64 scratch GPR. diff --git a/pyre/pyre-jit-trace/build.rs b/pyre/pyre-jit-trace/build.rs index b46ee57cd3f..6bcc7308332 100644 --- a/pyre/pyre-jit-trace/build.rs +++ b/pyre/pyre-jit-trace/build.rs @@ -67,7 +67,7 @@ fn emit_llbc_extraction_placeholders() { std::fs::write(format!("{out_dir}/jitcodes.bin"), b"").unwrap(); std::fs::write( format!("{out_dir}/jitcodes_index.bin"), - bincode::serialize(&(Vec::::new(), vec![0_u32])).unwrap(), + bincode::serialize(&(Vec::::new(), Vec::::new(), vec![0_u32])).unwrap(), ) .unwrap(); std::fs::write( diff --git a/pyre/pyre-jit-trace/build/prepass.rs b/pyre/pyre-jit-trace/build/prepass.rs index 7bfa593b4bf..ec3026da97f 100644 --- a/pyre/pyre-jit-trace/build/prepass.rs +++ b/pyre/pyre-jit-trace/build/prepass.rs @@ -1217,16 +1217,47 @@ fn real_main() { let mut jitcodes_bin = Vec::new(); let mut jitcode_names = Vec::with_capacity(frozen_jitcodes.len()); let mut jitcode_offsets = Vec::with_capacity(frozen_jitcodes.len() + 1); + // Graph identity per jitcode, index-aligned with `jitcode_names`. + // `CallControl::get_jitcode` keys its map by `CallPath` exactly as + // `call.py get_jitcode` keys by graph, but the name it stores beside + // it is the path's last segment, chosen to stay readable in dumps. + // That display name is not unique — this pipeline has 238 colliding + // groups, `from_obj` 52 ways — so a runtime consumer resolving a + // descent target by name cannot address the collided bodies at all. + // Carrying the allocation key lets it name one unambiguously. + let mut jitcode_paths: Vec = Vec::with_capacity(frozen_jitcodes.len()); + let path_by_ptr: std::collections::HashMap = pipeline + .jitcodes_by_path + .iter() + .map(|(path, jitcode)| { + ( + std::sync::Arc::as_ptr(jitcode) as usize, + path.canonical_key(), + ) + }) + .collect(); jitcode_offsets.push(0_u32); - for jitcode in &frozen_jitcodes { + for (index, jitcode) in frozen_jitcodes.iter().enumerate() { jitcode_names.push(jitcode.name.clone()); + // `frozen_jitcodes` is a clone-per-entry of `pipeline.jitcodes` + // built in place, so the two share an index but not an identity; + // key the lookup off the original Arc. A jitcode the codewriter + // minted without a graph key (synthetic shells) carries the empty + // string, which `compute_pathed_jitcode_index` refuses to match. + jitcode_paths.push( + path_by_ptr + .get(&(std::sync::Arc::as_ptr(&pipeline.jitcodes[index]) as usize)) + .cloned() + .unwrap_or_default(), + ); jitcodes_bin.extend(bincode::serialize(jitcode).unwrap()); jitcode_offsets.push( u32::try_from(jitcodes_bin.len()) .expect("serialized jitcodes.bin exceeds the u32 offset range"), ); } - let jitcodes_index_bin = bincode::serialize(&(jitcode_names, jitcode_offsets)).unwrap(); + let jitcodes_index_bin = + bincode::serialize(&(jitcode_names, jitcode_paths, jitcode_offsets)).unwrap(); std::fs::write(format!("{out_dir}/jitcodes.bin"), &jitcodes_bin).unwrap(); std::fs::write(format!("{out_dir}/jitcodes_index.bin"), &jitcodes_index_bin).unwrap(); // Keep the shell fnaddr beside each dense index. A translated PyPy diff --git a/pyre/pyre-jit-trace/src/jitcode_runtime.rs b/pyre/pyre-jit-trace/src/jitcode_runtime.rs index 732c7b5c668..ca5d07f4f9e 100644 --- a/pyre/pyre-jit-trace/src/jitcode_runtime.rs +++ b/pyre/pyre-jit-trace/src/jitcode_runtime.rs @@ -27,6 +27,10 @@ use majit_translate::jitcode::{BhDescr, DescrTable, JitCode}; struct JitCodeIndex { names: Vec, + /// The `CallPath` each jitcode was allocated under, `::`-joined by + /// `CallPath::canonical_key`, index-aligned with `names`. Empty for a + /// jitcode the codewriter minted without a graph key. + paths: Vec, offsets: Vec, } @@ -69,8 +73,8 @@ thread_local! { fn load_jitcode_index() -> &'static JitCodeIndex { const INDEX_BYTES: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/jitcodes_index.bin")); const BODY_BYTES: &[u8] = include_bytes!(concat!(env!("OUT_DIR"), "/jitcodes.bin")); - let (names, offsets): (Vec, Vec) = bincode::deserialize(INDEX_BYTES) - .unwrap_or_else(|e| { + let (names, paths, offsets): (Vec, Vec, Vec) = + bincode::deserialize(INDEX_BYTES).unwrap_or_else(|e| { panic!( "pyre-jit-trace: failed to deserialize jitcodes_index.bin \ ({} bytes): {e}", @@ -84,10 +88,21 @@ fn load_jitcode_index() -> &'static JitCodeIndex { names.len(), offsets.len(), ); + assert_eq!( + paths.len(), + names.len(), + "pyre-jit-trace: jitcode index has {} names but {} graph keys", + names.len(), + paths.len(), + ); assert_eq!(offsets.first().copied(), Some(0)); assert_eq!(offsets.last().copied(), Some(BODY_BYTES.len() as u32)); assert!(offsets.windows(2).all(|pair| pair[0] <= pair[1])); - Box::leak(Box::new(JitCodeIndex { names, offsets })) + Box::leak(Box::new(JitCodeIndex { + names, + paths, + offsets, + })) } fn jitcode_index() -> &'static JitCodeIndex { @@ -407,6 +422,37 @@ fn compute_named_jitcode_index(name: &str) -> Option { first } +/// Resolve a JitCode by the graph key the codewriter allocated it under, +/// spelled as `CallPath::canonical_key` — the `::`-joined segments, e.g. +/// `pyre_interpreter::opcode_ops::compare_value`. +/// +/// `CallControl::get_jitcode` keys `self.jitcodes` by `CallPath`, mirroring +/// `call.py get_jitcode`'s graph-keyed dict, and takes the JitCode's `name` +/// from the path's last segment only so dumps stay readable. Two graphs may +/// therefore share a name — 238 groups do — and the by-name resolver has no +/// way to tell them apart, so it panics rather than pick a body. Resolving by +/// the allocation key has no such ambiguity: `IndexMap` makes it +/// unique by construction. +/// +/// `None` when the pipeline holds no such graph, so a caller can decline +/// rather than assume the helper was compiled. +pub(crate) fn compute_pathed_jitcode_index(canonical_path: &str) -> Option { + if canonical_path.is_empty() { + return None; + } + jitcode_index() + .paths + .iter() + .position(|path| path == canonical_path) +} + +/// The JitCode allocated under `canonical_path`, or `None` when the build-time +/// pipeline does not contain that graph. See [`compute_pathed_jitcode_index`]. +#[allow(dead_code)] +pub(crate) fn pathed_jitcode(canonical_path: &str) -> Option> { + get_jitcode_by_index(compute_pathed_jitcode_index(canonical_path)?) +} + /// Resolve an ordinary portal-closure JitCode by its unique graph leaf name. /// Prefer stable graph paths at build time; this runtime helper exists for /// diagnostics and tests whose serialized artifact stores names only.