ZJIT: Count unoptimized Send (ruby#14801)

* ZJIT: Count unoptimized `Send`

This includes `Send` in `send fallback reasons` to guide future
optimizations.

* ZJIT: Create dedicated def_type counter for Send

Author: st0012

zjit.rb

@@ -159,7 +159,8 @@ def stats_string
     print_counters_with_prefix(prefix: 'not_annotated_cfuncs_', prompt: 'not annotated C methods', buf:, stats:, limit: 20)
 
     # Show fallback counters, ordered by the typical amount of fallbacks for the prefix at the time
-    print_counters_with_prefix(prefix: 'unspecialized_def_type_', prompt: 'not optimized method types', buf:, stats:, limit: 20)
+    print_counters_with_prefix(prefix: 'unspecialized_send_def_type_', prompt: 'not optimized method types for send', buf:, stats:, limit: 20)
+    print_counters_with_prefix(prefix: 'unspecialized_send_without_block_def_type_', prompt: 'not optimized method types for send_without_block', buf:, stats:, limit: 20)
     print_counters_with_prefix(prefix: 'not_optimized_yarv_insn_', prompt: 'not optimized instructions', buf:, stats:, limit: 20)
     print_counters_with_prefix(prefix: 'send_fallback_', prompt: 'send fallback reasons', buf:, stats:, limit: 20)
 

zjit/src/codegen.rs

@@ -15,7 +15,7 @@ use crate::invariants::{
 };
 use crate::gc::{append_gc_offsets, get_or_create_iseq_payload, get_or_create_iseq_payload_ptr, IseqCodePtrs, IseqPayload, IseqStatus};
 use crate::state::ZJITState;
-use crate::stats::{send_fallback_counter, exit_counter_for_compile_error, incr_counter, incr_counter_by, send_fallback_counter_for_method_type, send_fallback_counter_ptr_for_opcode, CompileError};
+use crate::stats::{send_fallback_counter, exit_counter_for_compile_error, incr_counter, incr_counter_by, send_fallback_counter_for_method_type, send_without_block_fallback_counter_for_method_type, send_fallback_counter_ptr_for_opcode, CompileError};
 use crate::stats::{counter_ptr, with_time_stat, Counter, Counter::{compile_time_ns, exit_compile_error}};
 use crate::{asm::CodeBlock, cruby::*, options::debug, virtualmem::CodePtr};
 use crate::backend::lir::{self, asm_comment, asm_ccall, Assembler, Opnd, Target, CFP, C_ARG_OPNDS, C_RET_OPND, EC, NATIVE_STACK_PTR, NATIVE_BASE_PTR, SCRATCH_OPND, SP};
@@ -1617,6 +1617,9 @@ fn gen_incr_send_fallback_counter(asm: &mut Assembler, reason: SendFallbackReaso
             gen_incr_counter_ptr(asm, send_fallback_counter_ptr_for_opcode(opcode));
         }
         SendWithoutBlockNotOptimizedMethodType(method_type) => {
+            gen_incr_counter(asm, send_without_block_fallback_counter_for_method_type(method_type));
+        }
+        SendNotOptimizedMethodType(method_type) => {
             gen_incr_counter(asm, send_fallback_counter_for_method_type(method_type));
         }
         _ => {}

zjit/src/hir.rs

@@ -535,6 +535,9 @@ pub enum SendFallbackReason {
     SendWithoutBlockCfuncArrayVariadic,
     SendWithoutBlockNotOptimizedMethodType(MethodType),
     SendWithoutBlockDirectTooManyArgs,
+    SendPolymorphic,
+    SendNoProfiles,
+    SendNotOptimizedMethodType(MethodType),
     CCallWithFrameTooManyArgs,
     ObjToStringNotString,
     /// Initial fallback reason for every instruction, which should be mutated to
@@ -2123,6 +2126,42 @@ impl Function {
                             self.push_insn_id(block, insn_id); continue;
                         }
                     }
+                    // This doesn't actually optimize Send yet, just replaces the fallback reason to be more precise.
+                    // TODO: Optimize Send
+                    Insn::Send { recv, cd, state, .. } => {
+                        let frame_state = self.frame_state(state);
+                        let klass = if let Some(klass) = self.type_of(recv).runtime_exact_ruby_class() {
+                            // If we know the class statically, use it to fold the lookup at compile-time.
+                            klass
+                        } else {
+                            let Some(recv_type) = self.profiled_type_of_at(recv, frame_state.insn_idx) else {
+                                if get_option!(stats) {
+                                    match self.is_polymorphic_at(recv, frame_state.insn_idx) {
+                                        Some(true) => self.set_dynamic_send_reason(insn_id, SendPolymorphic),
+                                        // If the class isn't known statically, then it should not also be monomorphic
+                                        Some(false) => panic!("Should not have monomorphic profile at this point in this branch"),
+                                        None => self.set_dynamic_send_reason(insn_id, SendNoProfiles),
+                                    }
+                                }
+                                self.push_insn_id(block, insn_id); continue;
+                            };
+                            recv_type.class()
+                        };
+                        let ci = unsafe { get_call_data_ci(cd) }; // info about the call site
+                        let mid = unsafe { vm_ci_mid(ci) };
+                        // Do method lookup
+                        let mut cme = unsafe { rb_callable_method_entry(klass, mid) };
+                        if cme.is_null() {
+                            self.set_dynamic_send_reason(insn_id, SendNotOptimizedMethodType(MethodType::Null));
+                            self.push_insn_id(block, insn_id); continue;
+                        }
+                        // Load an overloaded cme if applicable. See vm_search_cc().
+                        // It allows you to use a faster ISEQ if possible.
+                        cme = unsafe { rb_check_overloaded_cme(cme, ci) };
+                        let def_type = unsafe { get_cme_def_type(cme) };
+                        self.set_dynamic_send_reason(insn_id, SendNotOptimizedMethodType(MethodType::from(def_type)));
+                        self.push_insn_id(block, insn_id); continue;
+                    }
                     Insn::GetConstantPath { ic, state, .. } => {
                         let idlist: *const ID = unsafe { (*ic).segments };
                         let ice = unsafe { (*ic).entry };

zjit/src/stats.rs

@@ -148,6 +148,9 @@ make_counters! {
         send_fallback_send_without_block_cfunc_array_variadic,
         send_fallback_send_without_block_not_optimized_method_type,
         send_fallback_send_without_block_direct_too_many_args,
+        send_fallback_send_polymorphic,
+        send_fallback_send_no_profiles,
+        send_fallback_send_not_optimized_method_type,
         send_fallback_ccall_with_frame_too_many_args,
         send_fallback_obj_to_string_not_string,
         send_fallback_not_optimized_instruction,
@@ -185,20 +188,35 @@ make_counters! {
     dynamic_getivar_count,
     dynamic_setivar_count,
 
-    // Method call def_type related to fallback to dynamic dispatch
-    unspecialized_def_type_iseq,
-    unspecialized_def_type_cfunc,
-    unspecialized_def_type_attrset,
-    unspecialized_def_type_ivar,
-    unspecialized_def_type_bmethod,
-    unspecialized_def_type_zsuper,
-    unspecialized_def_type_alias,
-    unspecialized_def_type_undef,
-    unspecialized_def_type_not_implemented,
-    unspecialized_def_type_optimized,
-    unspecialized_def_type_missing,
-    unspecialized_def_type_refined,
-    unspecialized_def_type_null,
+    // Method call def_type related to send without block fallback to dynamic dispatch
+    unspecialized_send_without_block_def_type_iseq,
+    unspecialized_send_without_block_def_type_cfunc,
+    unspecialized_send_without_block_def_type_attrset,
+    unspecialized_send_without_block_def_type_ivar,
+    unspecialized_send_without_block_def_type_bmethod,
+    unspecialized_send_without_block_def_type_zsuper,
+    unspecialized_send_without_block_def_type_alias,
+    unspecialized_send_without_block_def_type_undef,
+    unspecialized_send_without_block_def_type_not_implemented,
+    unspecialized_send_without_block_def_type_optimized,
+    unspecialized_send_without_block_def_type_missing,
+    unspecialized_send_without_block_def_type_refined,
+    unspecialized_send_without_block_def_type_null,
+
+    // Method call def_type related to send fallback to dynamic dispatch
+    unspecialized_send_def_type_iseq,
+    unspecialized_send_def_type_cfunc,
+    unspecialized_send_def_type_attrset,
+    unspecialized_send_def_type_ivar,
+    unspecialized_send_def_type_bmethod,
+    unspecialized_send_def_type_zsuper,
+    unspecialized_send_def_type_alias,
+    unspecialized_send_def_type_undef,
+    unspecialized_send_def_type_not_implemented,
+    unspecialized_send_def_type_optimized,
+    unspecialized_send_def_type_missing,
+    unspecialized_send_def_type_refined,
+    unspecialized_send_def_type_null,
 
     // Writes to the VM frame
     vm_write_pc_count,
@@ -320,30 +338,54 @@ pub fn send_fallback_counter(reason: crate::hir::SendFallbackReason) -> Counter
         SendWithoutBlockCfuncArrayVariadic        => send_fallback_send_without_block_cfunc_array_variadic,
         SendWithoutBlockNotOptimizedMethodType(_) => send_fallback_send_without_block_not_optimized_method_type,
         SendWithoutBlockDirectTooManyArgs         => send_fallback_send_without_block_direct_too_many_args,
+        SendPolymorphic                           => send_fallback_send_polymorphic,
+        SendNoProfiles                            => send_fallback_send_no_profiles,
+        SendNotOptimizedMethodType(_)             => send_fallback_send_not_optimized_method_type,
         CCallWithFrameTooManyArgs                 => send_fallback_ccall_with_frame_too_many_args,
         ObjToStringNotString                      => send_fallback_obj_to_string_not_string,
         NotOptimizedInstruction(_)                => send_fallback_not_optimized_instruction,
     }
 }
 
+pub fn send_without_block_fallback_counter_for_method_type(method_type: crate::hir::MethodType) -> Counter {
+    use crate::hir::MethodType::*;
+    use crate::stats::Counter::*;
+
+    match method_type {
+        Iseq => unspecialized_send_without_block_def_type_iseq,
+        Cfunc => unspecialized_send_without_block_def_type_cfunc,
+        Attrset => unspecialized_send_without_block_def_type_attrset,
+        Ivar => unspecialized_send_without_block_def_type_ivar,
+        Bmethod => unspecialized_send_without_block_def_type_bmethod,
+        Zsuper => unspecialized_send_without_block_def_type_zsuper,
+        Alias => unspecialized_send_without_block_def_type_alias,
+        Undefined => unspecialized_send_without_block_def_type_undef,
+        NotImplemented => unspecialized_send_without_block_def_type_not_implemented,
+        Optimized => unspecialized_send_without_block_def_type_optimized,
+        Missing => unspecialized_send_without_block_def_type_missing,
+        Refined => unspecialized_send_without_block_def_type_refined,
+        Null => unspecialized_send_without_block_def_type_null,
+    }
+}
+
 pub fn send_fallback_counter_for_method_type(method_type: crate::hir::MethodType) -> Counter {
     use crate::hir::MethodType::*;
     use crate::stats::Counter::*;
 
     match method_type {
-        Iseq => unspecialized_def_type_iseq,
-        Cfunc => unspecialized_def_type_cfunc,
-        Attrset => unspecialized_def_type_attrset,
-        Ivar => unspecialized_def_type_ivar,
-        Bmethod => unspecialized_def_type_bmethod,
-        Zsuper => unspecialized_def_type_zsuper,
-        Alias => unspecialized_def_type_alias,
-        Undefined => unspecialized_def_type_undef,
-        NotImplemented => unspecialized_def_type_not_implemented,
-        Optimized => unspecialized_def_type_optimized,
-        Missing => unspecialized_def_type_missing,
-        Refined => unspecialized_def_type_refined,
-        Null => unspecialized_def_type_null,
+        Iseq => unspecialized_send_def_type_iseq,
+        Cfunc => unspecialized_send_def_type_cfunc,
+        Attrset => unspecialized_send_def_type_attrset,
+        Ivar => unspecialized_send_def_type_ivar,
+        Bmethod => unspecialized_send_def_type_bmethod,
+        Zsuper => unspecialized_send_def_type_zsuper,
+        Alias => unspecialized_send_def_type_alias,
+        Undefined => unspecialized_send_def_type_undef,
+        NotImplemented => unspecialized_send_def_type_not_implemented,
+        Optimized => unspecialized_send_def_type_optimized,
+        Missing => unspecialized_send_def_type_missing,
+        Refined => unspecialized_send_def_type_refined,
+        Null => unspecialized_send_def_type_null,
     }
 }