Implement jit_guard_known_class

lib/ruby_vm/mjit/assembler.rb

@@ -153,6 +153,17 @@ module RubyVM::MJIT
           disp: left_disp,
           imm: imm32(right_imm),
         )
+      # CMP r/m64, imm8 (Mod 01: [reg]+disp8)
+      in [[Symbol => left_reg, Integer => left_disp], Integer => right_imm] if r64?(left_reg) && imm8?(left_disp) && imm8?(right_imm)
+        # REX.W + 83 /7 ib
+        # MI: Operand 1: ModRM:r/m (r), Operand 2: imm8/16/32
+        insn(
+          prefix: REX_W,
+          opcode: 0x83,
+          mod_rm: ModRM[mod: Mod01, reg: 7, rm: left_reg],
+          disp: left_disp,
+          imm: imm8(right_imm),
+        )
       # CMP r/m64, imm8 (Mod 11: reg)
       in [Symbol => left_reg, Integer => right_imm] if r64?(left_reg) && imm8?(right_imm)
         # REX.W + 83 /7 ib
@@ -396,6 +407,17 @@ module RubyVM::MJIT
             disp: dst_disp,
             imm: imm32(src_imm),
           )
+        # MOV r/m64, imm32 (Mod 10: [reg]+disp32)
+        in Integer => src_imm if r64?(dst_reg) && imm32?(dst_disp) && imm32?(src_imm)
+          # REX.W + C7 /0 id
+          # MI: Operand 1: ModRM:r/m (w), Operand 2: imm8/16/32/64
+          insn(
+            prefix: REX_W,
+            opcode: 0xc7,
+            mod_rm: ModRM[mod: Mod10, reg: 0, rm: dst_reg],
+            disp: imm32(dst_disp),
+            imm: imm32(src_imm),
+          )
         # MOV r/m64, r64 (Mod 01: [reg]+disp8)
         in Symbol => src_reg if r64?(dst_reg) && imm8?(dst_disp) && r64?(src_reg)
           # REX.W + 89 /r

lib/ruby_vm/mjit/exit_compiler.rb

@@ -1,8 +1,7 @@
 module RubyVM::MJIT
   class ExitCompiler
     def initialize
-      # TODO: Use GC offsets
+      @gc_refs = [] # TODO: GC offsets?
-      @gc_refs = []
     end
 
     # Used for invalidating a block on entry.

lib/ruby_vm/mjit/insn_compiler.rb

@@ -6,6 +6,7 @@ module RubyVM::MJIT
       @ocb = ocb
       @exit_compiler = exit_compiler
       @invariants = Invariants.new(cb, ocb, exit_compiler)
+      @gc_refs = [] # TODO: GC offsets?
       # freeze # workaround a binding.irb issue. TODO: resurrect this
     end
 
@@ -150,7 +151,7 @@ module RubyVM::MJIT
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def getinstancevariable(jit, ctx, asm)
-      # Specialize on compile-time receiver, and split a block for chain guards
+      # Specialize on a compile-time receiver, and split a block for chain guards
       unless jit.at_current_insn?
         defer_compilation(jit, ctx, asm)
         return EndBlock
@@ -578,7 +579,12 @@ module RubyVM::MJIT
     # @param ctx [RubyVM::MJIT::Context]
     # @param asm [RubyVM::MJIT::Assembler]
     def jit_chain_guard(opcode, jit, ctx, asm, side_exit, limit: 10)
-      assert_equal(opcode, :jne) # TODO: support more
+      case opcode
+      when :je, :jne, :jnz
+        # ok
+      else
+        raise ArgumentError, "jit_chain_guard: unexpected opcode #{opcode.inspect}"
+      end
 
       if ctx.chain_depth < limit
         deeper = ctx.dup
@@ -598,13 +604,64 @@ module RubyVM::MJIT
           branch_asm.stub(branch_stub) do
             case branch_stub.shape
             in Default
-              asm.jne(branch_stub.target0.address)
+              asm.public_send(opcode, branch_stub.target0.address)
             end
           end
         end
         branch_stub.compile.call(asm)
       else
-        asm.jne(side_exit)
+        asm.public_send(opcode, side_exit)
+      end
+    end
+
+    # @param jit [RubyVM::MJIT::JITState]
+    # @param ctx [RubyVM::MJIT::Context]
+    # @param asm [RubyVM::MJIT::Assembler]
+    def jit_guard_known_class(jit, ctx, asm, known_klass, obj_opnd, comptime_obj, side_exit, limit: 5)
+      if known_klass == NilClass
+        asm.incr_counter(:send_guard_nil)
+        return CantCompile
+      elsif known_klass == TrueClass
+        asm.incr_counter(:send_guard_true)
+        return CantCompile
+      elsif known_klass == FalseClass
+        asm.incr_counter(:send_guard_false)
+        return CantCompile
+      elsif known_klass == Integer
+        asm.incr_counter(:send_guard_integer)
+        return CantCompile
+      elsif known_klass == Symbol
+        asm.incr_counter(:send_guard_symbol)
+        return CantCompile
+      elsif known_klass == Float
+        asm.incr_counter(:send_guard_float)
+        return CantCompile
+      elsif known_klass.singleton_class?
+        asm.comment('guard known object with singleton class')
+        asm.mov(:rax, C.to_value(comptime_obj))
+        asm.cmp(obj_opnd, :rax)
+        jit_chain_guard(:jne, jit, ctx, asm, side_exit, limit:)
+      else
+        # If obj_opnd isn't already a register, load it.
+        if obj_opnd.is_a?(Array)
+          asm.mov(:rax, obj_opnd)
+          obj_opnd = :rax
+        end
+
+        # Check that the receiver is a heap object
+        # Note: if we get here, the class doesn't have immediate instances.
+        asm.comment('guard not immediate')
+        asm.test(obj_opnd, C.RUBY_IMMEDIATE_MASK)
+        jit_chain_guard(:jnz, jit, ctx, asm, side_exit, limit:)
+        asm.cmp(obj_opnd, Qfalse)
+        jit_chain_guard(:je, jit, ctx, asm, side_exit, limit:)
+
+        # Bail if receiver class is different from known_klass
+        klass_opnd = [obj_opnd, C.RBasic.offsetof(:klass)]
+        asm.comment('guard known class')
+        asm.mov(:rcx, to_value(known_klass))
+        asm.cmp(klass_opnd, :rcx)
+        jit_chain_guard(:jne, jit, ctx, asm, side_exit, limit:)
       end
     end
 
@@ -714,11 +771,16 @@ module RubyVM::MJIT
       mid = C.vm_ci_mid(ci)
       flags = C.vm_ci_flag(ci)
 
+      # Specialize on a compile-time receiver, and split a block for chain guards
       unless jit.at_current_insn?
         defer_compilation(jit, ctx, asm)
         return EndBlock
       end
 
+      # Generate a side exit
+      side_exit = side_exit(jit, ctx)
+
+      # Calculate a receiver index
       if flags & C.VM_CALL_KW_SPLAT != 0
         # recv_index calculation may not work for this
         asm.incr_counter(:send_kw_splat)
@@ -726,18 +788,14 @@ module RubyVM::MJIT
       end
       recv_index = argc + ((flags & C.VM_CALL_ARGS_BLOCKARG == 0) ? 0 : 1)
 
+      # Get a compile-time receiver and its class
       comptime_recv = jit.peek_at_stack(recv_index)
       comptime_recv_klass = C.rb_class_of(comptime_recv)
 
       # Guard the receiver class (part of vm_search_method_fastpath)
-      if comptime_recv_klass.singleton_class?
+      recv_opnd = [SP, C.VALUE.size * (ctx.sp_offset - 1 - recv_index)]
-        asm.comment('guard known object with singleton class')
+      megamorphic_exit = counted_exit(side_exit, :send_klass_megamorphic)
-        asm.mov(:rax, C.to_value(comptime_recv))
+      if jit_guard_known_class(jit, ctx, asm, comptime_recv_klass, recv_opnd, comptime_recv, megamorphic_exit) == CantCompile
-        asm.cmp([SP, C.VALUE.size * (ctx.sp_offset - 1 - recv_index)], :rax)
-        asm.jne(side_exit(jit, ctx))
-      else
-        # TODO: support more classes
-        asm.incr_counter(:send_guard_known_object)
         return CantCompile
       end
 
@@ -1029,5 +1087,10 @@ module RubyVM::MJIT
     def def_iseq_ptr(cme_def)
       C.rb_iseq_check(cme_def.body.iseq.iseqptr)
     end
+
+    def to_value(obj)
+      @gc_refs << obj
+      C.to_value(obj)
+    end
   end
 end

mjit_c.h

@@ -109,16 +109,23 @@ MJIT_RUNTIME_COUNTERS(
     vm_insns_count,
     mjit_insns_count,
 
+    send_args_splat,
+    send_klass_megamorphic,
     send_kw_splat,
-    send_guard_known_object,
+    send_kwarg,
     send_missing_cme,
+    send_not_iseq,
     send_private,
     send_protected,
-    send_not_iseq,
-    send_args_splat,
-    send_kwarg,
     send_tailcall,
 
+    send_guard_nil,
+    send_guard_true,
+    send_guard_false,
+    send_guard_integer,
+    send_guard_symbol,
+    send_guard_float,
+
     getivar_megamorphic,
     getivar_not_heap,
     getivar_not_t_object,

mjit_c.rb

@@ -467,6 +467,14 @@ module RubyVM::MJIT # :nodoc: all
     @RB_BUILTIN ||= self.rb_builtin_function
   end
 
+  def C.RBasic
+    @RBasic ||= CType::Struct.new(
+      "RBasic", Primitive.cexpr!("SIZEOF(struct RBasic)"),
+      flags: [self.VALUE, Primitive.cexpr!("OFFSETOF((*((struct RBasic *)NULL)), flags)")],
+      klass: [self.VALUE, Primitive.cexpr!("OFFSETOF((*((struct RBasic *)NULL)), klass)")],
+    )
+  end
+
   def C.RObject
     @RObject ||= CType::Struct.new(
       "RObject", Primitive.cexpr!("SIZEOF(struct RObject)"),
@@ -859,15 +867,21 @@ module RubyVM::MJIT # :nodoc: all
       "rb_mjit_runtime_counters", Primitive.cexpr!("SIZEOF(struct rb_mjit_runtime_counters)"),
       vm_insns_count: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), vm_insns_count)")],
       mjit_insns_count: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), mjit_insns_count)")],
+      send_args_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_args_splat)")],
+      send_klass_megamorphic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_klass_megamorphic)")],
       send_kw_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_kw_splat)")],
-      send_guard_known_object: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_known_object)")],
+      send_kwarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_kwarg)")],
       send_missing_cme: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_missing_cme)")],
+      send_not_iseq: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_not_iseq)")],
       send_private: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_private)")],
       send_protected: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_protected)")],
-      send_not_iseq: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_not_iseq)")],
-      send_args_splat: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_args_splat)")],
-      send_kwarg: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_kwarg)")],
       send_tailcall: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_tailcall)")],
+      send_guard_nil: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_nil)")],
+      send_guard_true: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_true)")],
+      send_guard_false: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_false)")],
+      send_guard_integer: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_integer)")],
+      send_guard_symbol: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_symbol)")],
+      send_guard_float: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), send_guard_float)")],
       getivar_megamorphic: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_megamorphic)")],
       getivar_not_heap: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_not_heap)")],
       getivar_not_t_object: [CType::Immediate.parse("size_t"), Primitive.cexpr!("OFFSETOF((*((struct rb_mjit_runtime_counters *)NULL)), getivar_not_t_object)")],
@@ -922,10 +936,6 @@ module RubyVM::MJIT # :nodoc: all
     @shape_id_t ||= CType::Immediate.find(Primitive.cexpr!("SIZEOF(shape_id_t)"), Primitive.cexpr!("SIGNED_TYPE_P(shape_id_t)"))
   end
 
-  def C.RBasic
-    CType::Stub.new(:RBasic)
-  end
-
   def C.rb_id_table
     CType::Stub.new(:rb_id_table)
   end

tool/mjit/bindgen.rb

@@ -400,6 +400,7 @@ generator = BindingGenerator.new(
     IC
     IVC
     RB_BUILTIN
+    RBasic
     RObject
     attr_index_t
     compile_branch