diff --git a/test/ruby/test_zjit.rb b/test/ruby/test_zjit.rb index 9b4271ffc098a5..8f83e858a659c6 100644 --- a/test/ruby/test_zjit.rb +++ b/test/ruby/test_zjit.rb @@ -18,6 +18,15 @@ def test_enabled RUBY end + def test_stats_enabled + assert_runs 'false', <<~RUBY, stats: false + RubyVM::ZJIT.stats_enabled? + RUBY + assert_runs 'true', <<~RUBY, stats: true + RubyVM::ZJIT.stats_enabled? + RUBY + end + def test_enable_through_env child_env = {'RUBY_YJIT_ENABLE' => nil, 'RUBY_ZJIT_ENABLE' => '1'} assert_in_out_err([child_env, '-v'], '') do |stdout, stderr| @@ -1486,10 +1495,13 @@ def test_require_rubygems_with_auto_compact end def test_stats - assert_runs 'true', %q{ + assert_runs '[true, true]', %q{ def test = 1 test - RubyVM::ZJIT.stats[:zjit_insns_count] > 0 + [ + RubyVM::ZJIT.stats[:zjit_insns_count] > 0, + RubyVM::ZJIT.stats(:zjit_insns_count) > 0, + ] }, stats: true end @@ -1890,6 +1902,17 @@ def make_range_then_exit(v) }, call_threshold: 2 end + def test_raise_in_second_argument + assert_compiles '{ok: true}', %q{ + def write(hash, key) + hash[key] = raise rescue true + hash + end + + write({}, :ok) + } + end + private # Assert that every method call in `test_script` can be compiled by ZJIT diff --git a/zjit.c b/zjit.c index 54d9f7ed86ad2f..9cc8b514232c03 100644 --- a/zjit.c +++ b/zjit.c @@ -358,7 +358,7 @@ rb_zjit_singleton_class_p(VALUE klass) // Primitives used by zjit.rb. Don't put other functions below, which wouldn't use them. VALUE rb_zjit_assert_compiles(rb_execution_context_t *ec, VALUE self); -VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self); +VALUE rb_zjit_stats(rb_execution_context_t *ec, VALUE self, VALUE target_key); VALUE rb_zjit_stats_enabled_p(rb_execution_context_t *ec, VALUE self); // Preprocessed zjit.rb generated during build diff --git a/zjit.rb b/zjit.rb index 7f98c5adc7b97c..871b1a91ca8a1f 100644 --- a/zjit.rb +++ b/zjit.rb @@ -19,10 +19,15 @@ def enabled? Primitive.cexpr! 'RBOOL(rb_zjit_enabled_p)' end + # Check if `--zjit-stats` is used + def stats_enabled? + Primitive.rb_zjit_stats_enabled_p + end + # Return ZJIT statistics as a Hash - def stats - stats = Primitive.rb_zjit_stats - return nil if stats.nil? + def stats(key = nil) + stats = Primitive.rb_zjit_stats(key) + return stats if stats.nil? || !key.nil? if stats.key?(:vm_insns_count) && stats.key?(:zjit_insns_count) stats[:total_insns_count] = stats[:vm_insns_count] + stats[:zjit_insns_count] diff --git a/zjit/src/codegen.rs b/zjit/src/codegen.rs index 16d44c74b0e37e..7ff5564d220304 100644 --- a/zjit/src/codegen.rs +++ b/zjit/src/codegen.rs @@ -353,10 +353,10 @@ fn gen_insn(cb: &mut CodeBlock, jit: &mut JITState, asm: &mut Assembler, functio Insn::Jump(branch) => no_output!(gen_jump(jit, asm, branch)), Insn::IfTrue { val, target } => no_output!(gen_if_true(jit, asm, opnd!(val), target)), Insn::IfFalse { val, target } => no_output!(gen_if_false(jit, asm, opnd!(val), target)), - Insn::SendWithoutBlock { cd, state, self_val, args, .. } => gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args)), + Insn::SendWithoutBlock { cd, state, .. } => gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)), // Give up SendWithoutBlockDirect for 6+ args since asm.ccall() doesn't support it. - Insn::SendWithoutBlockDirect { cd, state, self_val, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self - gen_send_without_block(jit, asm, *cd, &function.frame_state(*state), opnd!(self_val), opnds!(args)), + Insn::SendWithoutBlockDirect { cd, state, args, .. } if args.len() + 1 > C_ARG_OPNDS.len() => // +1 for self + gen_send_without_block(jit, asm, *cd, &function.frame_state(*state)), Insn::SendWithoutBlockDirect { cme, iseq, self_val, args, state, .. } => gen_send_without_block_direct(cb, jit, asm, *cme, *iseq, opnd!(self_val), opnds!(args), &function.frame_state(*state)), // Ensure we have enough room fit ec, self, and arguments // TODO remove this check when we have stack args (we can use Time.new to test it) @@ -859,26 +859,19 @@ fn gen_send_without_block( asm: &mut Assembler, cd: *const rb_call_data, state: &FrameState, - self_val: Opnd, - args: Vec, ) -> lir::Opnd { - gen_spill_locals(jit, asm, state); - // Spill the receiver and the arguments onto the stack. - // They need to be on the interpreter stack to let the interpreter access them. - // TODO: Avoid spilling operands that have been spilled before. - // TODO: Despite https://github.com/ruby/ruby/pull/13468, Kokubun thinks this should - // spill the whole stack in case it raises an exception. The HIR might need to change - // for opt_aref_with, which pushes to the stack in the middle of the instruction. - asm_comment!(asm, "spill receiver and arguments"); - for (idx, &val) in [self_val].iter().chain(args.iter()).enumerate() { - // Currently, we don't move the SP register. So it's equal to the base pointer. - let stack_opnd = Opnd::mem(64, SP, idx as i32 * SIZEOF_VALUE_I32); - asm.mov(stack_opnd, val); - } + // Note that it's incorrect to use this frame state to side exit because + // the state might not be on the boundary of an interpreter instruction. + // For example, `opt_aref_with` pushes to the stack and then sends. + asm_comment!(asm, "spill frame state"); // Save PC and SP gen_save_pc(asm, state); - gen_save_sp(asm, 1 + args.len()); // +1 for receiver + gen_save_sp(asm, state.stack().len()); + + // Spill locals and stack + gen_spill_locals(jit, asm, state); + gen_spill_stack(jit, asm, state); asm_comment!(asm, "call #{} with dynamic dispatch", ruby_call_method_name(cd)); unsafe extern "C" { diff --git a/zjit/src/hir.rs b/zjit/src/hir.rs index 1fdca5a688ac63..1cd31497d8e0eb 100644 --- a/zjit/src/hir.rs +++ b/zjit/src/hir.rs @@ -818,7 +818,7 @@ impl<'a> std::fmt::Display for InsnPrinter<'a> { } Ok(()) }, - Insn::Snapshot { state } => write!(f, "Snapshot {}", state), + Insn::Snapshot { state } => write!(f, "Snapshot {}", state.print(self.ptr_map)), Insn::Defined { op_type, v, .. } => { // op_type (enum defined_type) printing logic from iseq.c. // Not sure why rb_iseq_defined_string() isn't exhaustive. @@ -2513,11 +2513,10 @@ pub struct FrameState { locals: Vec, } -impl FrameState { - /// Get the opcode for the current instruction - pub fn get_opcode(&self) -> i32 { - unsafe { rb_iseq_opcode_at_pc(self.iseq, self.pc) } - } +/// Print adaptor for [`FrameState`]. See [`PtrPrintMap`]. +pub struct FrameStatePrinter<'a> { + inner: &'a FrameState, + ptr_map: &'a PtrPrintMap, } /// Compute the index of a local variable from its slot index @@ -2624,14 +2623,24 @@ impl FrameState { args.extend(self.locals.iter().chain(self.stack.iter()).map(|op| *op)); args } + + /// Get the opcode for the current instruction + pub fn get_opcode(&self) -> i32 { + unsafe { rb_iseq_opcode_at_pc(self.iseq, self.pc) } + } + + pub fn print<'a>(&'a self, ptr_map: &'a PtrPrintMap) -> FrameStatePrinter<'a> { + FrameStatePrinter { inner: self, ptr_map } + } } -impl std::fmt::Display for FrameState { +impl Display for FrameStatePrinter<'_> { fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result { - write!(f, "FrameState {{ pc: {:?}, stack: ", self.pc)?; - write_vec(f, &self.stack)?; + let inner = self.inner; + write!(f, "FrameState {{ pc: {:?}, stack: ", self.ptr_map.map_ptr(inner.pc))?; + write_vec(f, &inner.stack)?; write!(f, ", locals: ")?; - write_vec(f, &self.locals)?; + write_vec(f, &inner.locals)?; write!(f, " }}") } } @@ -3195,9 +3204,11 @@ pub fn iseq_to_hir(iseq: *const rb_iseq_t) -> Result { let aref_arg = fun.push_insn(block, Insn::Const { val: Const::Value(get_arg(pc, 0)) }); let args = vec![aref_arg]; + let mut send_state = state.clone(); + send_state.stack_push(aref_arg); + let send_state = fun.push_insn(block, Insn::Snapshot { state: send_state }); let recv = state.stack_pop()?; - let exit_id = fun.push_insn(block, Insn::Snapshot { state: exit_state }); - let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: exit_id }); + let send = fun.push_insn(block, Insn::SendWithoutBlock { self_val: recv, cd, args, state: send_state }); state.stack_push(send); } YARVINSN_opt_neq => { @@ -3920,6 +3931,12 @@ mod tests { expected_hir.assert_eq(&actual_hir); } + #[track_caller] + pub fn assert_function_hir_with_frame_state(function: Function, expected_hir: Expect) { + let actual_hir = format!("{}", FunctionPrinter::with_snapshot(&function)); + expected_hir.assert_eq(&actual_hir); + } + #[track_caller] fn assert_compile_fails(method: &str, reason: ParseError) { let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", method)); @@ -5154,11 +5171,18 @@ mod tests { eval(" def test(a) = a['string lit triggers aref_with'] "); - assert_method_hir("test", expect![[r#" + + let iseq = crate::cruby::with_rubyvm(|| get_method_iseq("self", "test")); + assert!(iseq_contains_opcode(iseq, YARVINSN_opt_aref_with)); + let function = iseq_to_hir(iseq).unwrap(); + assert_function_hir_with_frame_state(function, expect![[r#" fn test@:2: bb0(v0:BasicObject, v1:BasicObject): - v3:StringExact[VALUE(0x1000)] = Const Value(VALUE(0x1000)) + v2:Any = Snapshot FrameState { pc: 0x1000, stack: [], locals: [v1] } + v3:StringExact[VALUE(0x1008)] = Const Value(VALUE(0x1008)) + v4:Any = Snapshot FrameState { pc: 0x1010, stack: [v1, v3], locals: [v1] } v5:BasicObject = SendWithoutBlock v1, :[], v3 + v6:Any = Snapshot FrameState { pc: 0x1018, stack: [v5], locals: [v1] } CheckInterrupts Return v5 "#]]); diff --git a/zjit/src/stats.rs b/zjit/src/stats.rs index fa8b741eea913e..ce185597c402d0 100644 --- a/zjit/src/stats.rs +++ b/zjit/src/stats.rs @@ -72,29 +72,41 @@ fn incr_counter(counter: Counter, amount: u64) { /// Return a Hash object that contains ZJIT statistics #[unsafe(no_mangle)] -pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE) -> VALUE { +pub extern "C" fn rb_zjit_stats(_ec: EcPtr, _self: VALUE, target_key: VALUE) -> VALUE { if !zjit_enabled_p() { return Qnil; } - fn set_stat(hash: VALUE, key: &str, value: u64) { - unsafe { rb_hash_aset(hash, rust_str_to_sym(key), VALUE::fixnum_from_usize(value as usize)); } + macro_rules! set_stat { + ($hash:ident, $key:expr, $value:expr) => { + let key = rust_str_to_sym($key); + // Evaluate $value only when it's needed + if key == target_key { + return VALUE::fixnum_from_usize($value as usize); + } else if $hash != Qnil { + #[allow(unused_unsafe)] + unsafe { rb_hash_aset($hash, key, VALUE::fixnum_from_usize($value as usize)); } + } + } } - let hash = unsafe { rb_hash_new() }; + let hash = if target_key.nil_p() { + unsafe { rb_hash_new() } + } else { + Qnil + }; let counters = ZJITState::get_counters(); for &counter in DEFAULT_COUNTERS { - let counter_val = unsafe { *counter_ptr(counter) }; - set_stat(hash, &counter.name(), counter_val); + set_stat!(hash, &counter.name(), unsafe { *counter_ptr(counter) }); } // Set counters that are enabled when --zjit-stats is enabled if get_option!(stats) { - set_stat(hash, "zjit_insns_count", counters.zjit_insns_count); + set_stat!(hash, "zjit_insns_count", counters.zjit_insns_count); if unsafe { rb_vm_insns_count } > 0 { - set_stat(hash, "vm_insns_count", unsafe { rb_vm_insns_count }); + set_stat!(hash, "vm_insns_count", unsafe { rb_vm_insns_count }); } }