YJIT: Use dynamic dispatch for megamorphic send (ruby#8125)

yjit.rb

@@ -270,11 +270,14 @@ def _print_stats(out: $stderr) # :nodoc:
  out.puts "num_send: " + format_number(13, stats[:num_send])
  out.puts "num_send_known_class: " + format_number_pct(13, stats[:num_send_known_class], stats[:num_send])
  out.puts "num_send_polymorphic: " + format_number_pct(13, stats[:num_send_polymorphic], stats[:num_send])
+ out.puts "num_send_megamorphic: " + format_number_pct(13, stats[:num_send_megamorphic], stats[:num_send])
  out.puts "num_send_dynamic: " + format_number_pct(13, stats[:num_send_dynamic], stats[:num_send])
  if stats[:num_send_x86_rel32] != 0 || stats[:num_send_x86_reg] != 0
  out.puts "num_send_x86_rel32: " + format_number(13, stats[:num_send_x86_rel32])
  out.puts "num_send_x86_reg: " + format_number(13, stats[:num_send_x86_reg])
  end
+ out.puts "num_getivar_megamorphic: " + format_number(13, stats[:num_getivar_megamorphic])
+ out.puts "num_setivar_megamorphic: " + format_number(13, stats[:num_setivar_megamorphic])
 
  out.puts "iseq_stack_too_large: " + format_number(13, stats[:iseq_stack_too_large])
  out.puts "iseq_too_long: " + format_number(13, stats[:iseq_too_long])
@@ -302,7 +305,6 @@ def _print_stats(out: $stderr) # :nodoc:
  out.puts "freed_iseq_count: " + format_number(13, stats[:freed_iseq_count])
  out.puts "invalidation_count: " + format_number(13, stats[:invalidation_count])
  out.puts "constant_state_bumps: " + format_number(13, stats[:constant_state_bumps])
- out.puts "get_ivar_max_depth: " + format_number(13, stats[:get_ivar_max_depth])
  out.puts "inline_code_size: " + format_number(13, stats[:inline_code_size])
  out.puts "outlined_code_size: " + format_number(13, stats[:outlined_code_size])
  out.puts "code_region_size: " + format_number(13, stats[:code_region_size])

yjit/src/codegen.rs

@@ -1971,17 +1971,17 @@ fn gen_get_ivar(
  // Check if the comptime receiver is a T_OBJECT
  let receiver_t_object = unsafe { RB_TYPE_P(comptime_receiver, RUBY_T_OBJECT) };
  // Use a general C call at the last chain to avoid exits on megamorphic shapes
- let last_chain = asm.ctx.get_chain_depth() as i32 == max_chain_depth - 1;
- if last_chain {
- gen_counter_incr(asm, Counter::get_ivar_max_depth);
+ let megamorphic = asm.ctx.get_chain_depth() as i32 >= max_chain_depth;
+ if megamorphic {
+ gen_counter_incr(asm, Counter::num_getivar_megamorphic);
  }
 
  // If the class uses the default allocator, instances should all be T_OBJECT
  // NOTE: This assumes nobody changes the allocator of the class after allocation.
  // Eventually, we can encode whether an object is T_OBJECT or not
  // inside object shapes.
  // too-complex shapes can't use index access, so we use rb_ivar_get for them too.
- if !receiver_t_object || uses_custom_allocator || comptime_receiver.shape_too_complex() || last_chain {
+ if !receiver_t_object || uses_custom_allocator || comptime_receiver.shape_too_complex() || megamorphic {
  // General case. Call rb_ivar_get().
  // VALUE rb_ivar_get(VALUE obj, ID id)
  asm.comment("call rb_ivar_get()");
@@ -2182,11 +2182,16 @@ fn gen_setinstancevariable(
 
  // Check if the comptime receiver is a T_OBJECT
  let receiver_t_object = unsafe { RB_TYPE_P(comptime_receiver, RUBY_T_OBJECT) };
+ // Use a general C call at the last chain to avoid exits on megamorphic shapes
+ let megamorphic = asm.ctx.get_chain_depth() as i32 >= SET_IVAR_MAX_DEPTH;
+ if megamorphic {
+ gen_counter_incr(asm, Counter::num_setivar_megamorphic);
+ }
 
  // If the receiver isn't a T_OBJECT, or uses a custom allocator,
  // then just write out the IV write as a function call.
  // too-complex shapes can't use index access, so we use rb_ivar_get for them too.
- if !receiver_t_object || uses_custom_allocator || comptime_receiver.shape_too_complex() || (asm.ctx.get_chain_depth() as i32) >= SET_IVAR_MAX_DEPTH {
+ if !receiver_t_object || uses_custom_allocator || comptime_receiver.shape_too_complex() || megamorphic {
  asm.comment("call rb_vm_setinstancevariable()");
 
  let ic = jit.get_arg(1).as_u64(); // type IVC
@@ -6555,6 +6560,11 @@ fn gen_send_general(
  if asm.ctx.get_chain_depth() > 1 {
  gen_counter_incr(asm, Counter::num_send_polymorphic);
  }
+ // If megamorphic, let the caller fallback to dynamic dispatch
+ if asm.ctx.get_chain_depth() as i32 >= SEND_MAX_DEPTH {
+ gen_counter_incr(asm, Counter::num_send_megamorphic);
+ return None;
+ }
 
  jit_guard_known_klass(
  jit,

yjit/src/stats.rs

@@ -396,9 +396,6 @@ make_counters! {
 
  constant_state_bumps,
 
- // Not using "getivar_" to exclude this from exit reasons
- get_ivar_max_depth,
-
  // Currently, it's out of the ordinary (might be impossible) for YJIT to leave gaps in
  // executable memory, so this should be 0.
  exec_mem_non_bump_alloc,
@@ -407,11 +404,15 @@ make_counters! {
 
  num_send,
  num_send_known_class,
+ num_send_megamorphic,
  num_send_polymorphic,
  num_send_x86_rel32,
  num_send_x86_reg,
  num_send_dynamic,
 
+ num_getivar_megamorphic,
+ num_setivar_megamorphic,
+
  iseq_stack_too_large,
  iseq_too_long,