RJIT: Optimize Kernel#is_a?

ruby · Mar 19, 2023 · 824cf88 · 824cf88
1 parent cc9330f
commit 824cf88
Show file tree

Hide file tree

Showing 4 changed files with 59 additions and 4 deletions.
diff --git a/lib/ruby_vm/rjit/insn_compiler.rb b/lib/ruby_vm/rjit/insn_compiler.rb
@@ -1376,7 +1376,7 @@ def objtostring(jit, ctx, asm)
       recv = ctx.stack_opnd(0)
       comptime_recv = jit.peek_at_stack(0)
 
-      if C::RB_TYPE_P(comptime_recv, C::RUBY_T_STRING)
+      if C.RB_TYPE_P(comptime_recv, C::RUBY_T_STRING)
         side_exit = side_exit(jit, ctx)
 
         jit_guard_known_klass(jit, ctx, asm, C.rb_class_of(comptime_recv), recv, comptime_recv, side_exit)
@@ -2461,6 +2461,53 @@ def jit_rb_false(jit, ctx, asm, argc, _known_recv_class)
       true
     end
 
+    # @param jit [RubyVM::RJIT::JITState]
+    # @param ctx [RubyVM::RJIT::Context]
+    # @param asm [RubyVM::RJIT::Assembler]
+    def jit_rb_kernel_is_a(jit, ctx, asm, argc, known_recv_class)
+      if argc != 1
+        return false
+      end
+
+      # If this is a super call we might not know the class
+      if known_recv_class.nil?
+        return false
+      end
+
+      # Important note: The output code will simply `return true/false`.
+      # Correctness follows from:
+      #  - `known_recv_class` implies there is a guard scheduled before here
+      #    for a particular `CLASS_OF(lhs)`.
+      #  - We guard that rhs is identical to the compile-time sample
+      #  - In general, for any two Class instances A, B, `A < B` does not change at runtime.
+      #    Class#superclass is stable.
+
+      sample_rhs = jit.peek_at_stack(0)
+      sample_lhs = jit.peek_at_stack(1)
+
+      # We are not allowing module here because the module hierachy can change at runtime.
+      if C.RB_TYPE_P(sample_rhs, C::RUBY_T_CLASS)
+        return false
+      end
+      sample_is_a = C.obj_is_kind_of(sample_lhs, sample_rhs)
+
+      side_exit = side_exit(jit, ctx)
+      asm.comment('Kernel#is_a?')
+      asm.mov(:rax, to_value(sample_rhs))
+      asm.cmp(ctx.stack_opnd(0), :rax)
+      asm.jne(counted_exit(side_exit, :send_is_a_class_mismatch))
+
+      ctx.stack_pop(2)
+
+      stack_ret = ctx.stack_push
+      if sample_is_a
+        asm.mov(stack_ret, Qtrue)
+      else
+        asm.mov(stack_ret, Qfalse)
+      end
+      return true
+    end
+
     # @param jit [RubyVM::RJIT::JITState]
     # @param ctx [RubyVM::RJIT::Context]
     # @param asm [RubyVM::RJIT::Assembler]
@@ -2710,8 +2757,8 @@ def register_cfunc_codegen_funcs
 
       register_cfunc_method(NilClass, :nil?, :jit_rb_true)
       register_cfunc_method(Kernel, :nil?, :jit_rb_false)
-      #register_cfunc_method(Kernel, :is_a?, :jit_rb_kernel_is_a)
-      #register_cfunc_method(Kernel, :kind_of?, :jit_rb_kernel_is_a)
+      register_cfunc_method(Kernel, :is_a?, :jit_rb_kernel_is_a)
+      register_cfunc_method(Kernel, :kind_of?, :jit_rb_kernel_is_a)
       #register_cfunc_method(Kernel, :instance_of?, :jit_rb_kernel_instance_of)
 
       register_cfunc_method(BasicObject, :==, :jit_rb_obj_equal)
@@ -4274,7 +4321,7 @@ def static_symbol?(obj)
 
     def dynamic_symbol?(obj)
       return false if C::SPECIAL_CONST_P(obj)
-      C::RB_TYPE_P(obj, C::RUBY_T_SYMBOL)
+      C.RB_TYPE_P(obj, C::RUBY_T_SYMBOL)
     end
 
     def shape_too_complex?(obj)

diff --git a/rjit_c.h b/rjit_c.h
@@ -40,6 +40,7 @@ RJIT_RUNTIME_COUNTERS(
     send_stackoverflow,
     send_arity,
     send_c_tracing,
+    send_is_a_class_mismatch,
 
     send_blockarg_not_nil_or_proxy,
     send_blockiseq,

diff --git a/rjit_c.rb b/rjit_c.rb
@@ -306,6 +306,10 @@ def rb_hash_stlike_lookup(hash, key)
         return result;
       }
     end
+
+    def obj_is_kind_of(obj, c)
+      Primitive.cexpr! 'rb_obj_is_kind_of(obj, c)'
+    end
   end
 
   #
@@ -371,6 +375,7 @@ def rb_iseqw_to_iseq(iseqw)
   C::RUBY_SPECIAL_SHIFT = Primitive.cexpr! %q{ SIZET2NUM(RUBY_SPECIAL_SHIFT) }
   C::RUBY_SYMBOL_FLAG = Primitive.cexpr! %q{ SIZET2NUM(RUBY_SYMBOL_FLAG) }
   C::RUBY_T_ARRAY = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_ARRAY) }
+  C::RUBY_T_CLASS = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_CLASS) }
   C::RUBY_T_ICLASS = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_ICLASS) }
   C::RUBY_T_MASK = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_MASK) }
   C::RUBY_T_MODULE = Primitive.cexpr! %q{ SIZET2NUM(RUBY_T_MODULE) }
@@ -1214,6 +1219,7 @@ def C.rb_rjit_runtime_counters
       send_stackoverflow: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_stackoverflow)")],
       send_arity: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_arity)")],
       send_c_tracing: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_c_tracing)")],
+      send_is_a_class_mismatch: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_is_a_class_mismatch)")],
       send_blockarg_not_nil_or_proxy: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_blockarg_not_nil_or_proxy)")],
       send_blockiseq: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_blockiseq)")],
       send_block_handler: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_rjit_runtime_counters *)NULL)), send_block_handler)")],

diff --git a/tool/rjit/bindgen.rb b/tool/rjit/bindgen.rb
@@ -424,6 +424,7 @@ def push_target(target)
       RUBY_SPECIAL_SHIFT
       RUBY_SYMBOL_FLAG
       RUBY_T_ARRAY
+      RUBY_T_CLASS
       RUBY_T_ICLASS
       RUBY_T_MASK
       RUBY_T_MODULE