Add fast path for goto_if

Author: 3tty0n

rpython/jit/codewriter/genextension.py

@@ -1127,6 +1127,22 @@ def emit_specialized_int_is_true(self):
         self._emit_jump(lines)
         return lines
 
+    def emit_specialized_int_is_zero(self):
+        arg, = self._get_args()
+        result = self.insn[self.resindex]
+        # Constant folding: if arg is a constant, evaluate at compile time
+        val = self._get_constant_int_value(arg)
+        if val is not None:
+            lines = ["i%s = %d" % (result.index, int(val == 0))]
+            self._emit_jump(lines)
+            return lines
+        lines = ["i%s = int(%s == 0)" % (
+            result.index,
+            self._get_as_unboxed(arg)
+        )]
+        self._emit_jump(lines)
+        return lines
+
     def emit_specialized_guard_class(self):
         lines = ['# guard_class, argument is already constant']
         arg, = self._get_args()
@@ -1146,6 +1162,8 @@ def emit_specialized_getfield_raw_i(self):
             self._emit_jump(lines, constant_registers=self.constant_registers.union({res}))
             return lines
         raise Unsupported
+    emit_specialized_getfield_raw_r = emit_specialized_getfield_raw_i
+    emit_specialized_getfield_raw_f = emit_specialized_getfield_raw_i
 
     def emit_specialized_getfield_gc_i_pure(self):
         if self.insn[2].is_always_pure():
@@ -1175,6 +1193,19 @@ def emit_specialized_getarrayitem_gc_i_pure(self):
     emit_specialized_getarrayitem_gc_r_pure = emit_specialized_getarrayitem_gc_i_pure
     emit_specialized_getarrayitem_gc_f_pure = emit_specialized_getarrayitem_gc_i_pure
 
+    def emit_specialized_getarrayitem_raw_i_pure(self):
+        lines = []
+        arg, index, descr, res = self._get_args_and_res()
+        TYPE, ITEM = _get_ptrtype_itemtype_from_arraydescr(descr)
+        resultcast = _find_result_cast(ITEM)
+        lines.append('%s = %sllmemory.cast_adr_to_ptr(support.int2adr(i%s), %s)[%s])' % (
+            self._get_as_unboxed(res), resultcast, arg.index, self._add_global(TYPE),
+            self._get_as_unboxed(index)))
+        self._emit_jump(lines, constant_registers=self.constant_registers.union({res}))
+        return lines
+    emit_specialized_getarrayitem_raw_r_pure = emit_specialized_getarrayitem_raw_i_pure
+    emit_specialized_getarrayitem_raw_f_pure = emit_specialized_getarrayitem_raw_i_pure
+
     def emit_specialized_getfield_gc_i(self):
         arg, descr = self._get_args()
         res = self.insn[self.resindex]
@@ -1430,6 +1461,14 @@ def emit_specialized_unicodelen(self):
         self._emit_jump(lines)
         return lines
 
+    def emit_specialized_unicodegetitem(self):
+        arg0, arg1 = self._get_args()
+        result = self.insn[self.resindex]
+        lines = ["i%s = ord(lltype.cast_opaque_ptr(lltype.Ptr(rstr.UNICODE), r%d).chars[%s])" % (
+            result.index, arg0.index, self._get_as_unboxed(arg1))]
+        self._emit_jump(lines)
+        return lines
+
     def _get_type_prefix(self, arg):
         if isinstance(arg, Constant) or isinstance(arg, Register):
             # TODO: this logic also works for the 'else' case. probably.
@@ -2039,23 +2078,64 @@ def emit_unspecialized_goto_if_not_comparison(self, name, symbol):
         lines.append("continue")
         return lines
 
+    def _emit_goto_if_not_int_comparison_fast(self, rop_name, py_op):
+        """Generate fast-path code for goto_if_not_int_* that skips heapcache."""
+        lines = []
+        _, arg0, arg1, arg2 = self.insn  # left, right, label
+
+        target_pc = self.get_target_pc(arg2)
+
+        # Handle constant folding case
+        self._emit_n_ary_if([arg0, arg1], lines)
+        specializer = self.work_list.specialize_insn(
+            self.insn, self.constant_registers.union({arg0, arg1}), self.orig_pc)
+        lines.append("    pc = %d" % (specializer.get_pc(),))
+        lines.append("    continue")
+
+        # Fast-path: compute comparison and record directly, skip heapcache
+        self._emit_sync_registers(lines)
+        box0 = self._get_as_box_after_sync(arg0)
+        box1 = self._get_as_box_after_sync(arg1)
+        lines.append("_v0 = %s" % self._get_as_unboxed_after_sync(arg0))
+        lines.append("_v1 = %s" % self._get_as_unboxed_after_sync(arg1))
+        lines.append("_cond = int(_v0 %s _v1)" % py_op)
+        lines.append("# fast-path: record comparison directly, skip heapcache")
+        lines.append("condbox = self.metainterp.history.record2_int(rop.%s, %s, %s, _cond)" % (
+            rop_name, box0, box1))
+
+        # Call opimpl_goto_if_not for guard generation
+        lines.append("self.opimpl_goto_if_not(condbox, %d, %d, replace=False)" % (target_pc, self.orig_pc))
+        lines.append("pc = self.pc")
+        lines.append("if pc == %s:" % (target_pc,))
+        specializer = self.work_list.specialize_pc(
+            self.constant_registers, target_pc)
+        lines.append("    pc = %s" % (specializer.spec_pc,))
+        lines.append("else:")
+        next_pc = self.work_list.pc_to_nextpc[self.orig_pc]
+        specializer = self.work_list.specialize_pc(
+            self.constant_registers, next_pc)
+        lines.append("    assert self.pc == %s" % (specializer.orig_pc,))
+        lines.append("    pc = %s" % (specializer.spec_pc,))
+        lines.append("continue")
+        return lines
+
     def emit_unspecialized_goto_if_not_int_lt(self):
-        return self.emit_unspecialized_goto_if_not_comparison("int_lt", "<")
+        return self._emit_goto_if_not_int_comparison_fast("INT_LT", "<")
 
     def emit_unspecialized_goto_if_not_int_gt(self):
-        return self.emit_unspecialized_goto_if_not_comparison("int_gt", ">")
+        return self._emit_goto_if_not_int_comparison_fast("INT_GT", ">")
 
     def emit_unspecialized_goto_if_not_int_le(self):
-        return self.emit_unspecialized_goto_if_not_comparison("int_le", "<=")
+        return self._emit_goto_if_not_int_comparison_fast("INT_LE", "<=")
 
     def emit_unspecialized_goto_if_not_int_ge(self):
-        return self.emit_unspecialized_goto_if_not_comparison("int_ge", ">=")
+        return self._emit_goto_if_not_int_comparison_fast("INT_GE", ">=")
 
     def emit_unspecialized_goto_if_not_int_ne(self):
-        return self.emit_unspecialized_goto_if_not_comparison("int_ne", "!=")
+        return self._emit_goto_if_not_int_comparison_fast("INT_NE", "!=")
 
     def emit_unspecialized_goto_if_not_int_eq(self):
-        return self.emit_unspecialized_goto_if_not_comparison("int_eq", "==")
+        return self._emit_goto_if_not_int_comparison_fast("INT_EQ", "==")
 
     def emit_unspecialized_switch(self):
         lines = []

rpython/jit/codewriter/test/test_genextension.py

@@ -56,8 +56,12 @@ def jit_shortcut(self): # test
                 i22 = ri22.getint()
                 pc = 116
                 continue
-            condbox = self.opimpl_int_gt(ri22, ConstInt(4))
-            self.opimpl_goto_if_not(condbox, 16, 0)
+            _v0 = self.registers_i[22].getint()
+            _v1 = 4
+            _cond = int(_v0 > _v1)
+            # fast-path: record comparison directly, skip heapcache
+            condbox = self.metainterp.history.record2_int(rop.INT_GT, self.registers_i[22], ConstInt(4), _cond)
+            self.opimpl_goto_if_not(condbox, 16, 0, replace=False)
             pc = self.pc
             if pc == 16:
                 pc = 16
@@ -627,16 +631,15 @@ def test_int_invert():
     newpc = insn_specializer.get_pc()
     assert newpc == 5
     s = insn_specializer.make_code()
+    # fast-path: directly compute and record, skip heapcache
     assert s == """\
-ri0 = self.registers_i[0]
-if isinstance(ri0, ConstInt):
-    i0 = ri0.getint()
-    pc = %d
-    continue
-else:
-    self.registers_i[1] = self.opimpl_int_invert(ri0)
+_v0 = self.registers_i[0].getint()
+_res = ~_v0
+_op = self.metainterp.history.record1_int(rop.INT_INVERT, self.registers_i[0], _res)
+self.registers_i[1] = _op
+i1 = _res
 pc = 7
-continue""" % (work_list.OFFSET + 7)
+continue"""
     next_constant_registers = insn_specializer.get_next_constant_registers()
     assert next_constant_registers == set()
 
@@ -684,6 +687,7 @@ def test_goto_if_not_int_lt():
     newpc = insn_specializer.get_pc()
     assert newpc == 5
     s = insn_specializer.make_code()
+    # fast-path: directly compute and record comparison, skip heapcache
     assert s == """\
 ri0 = self.registers_i[0]
 ri1 = self.registers_i[1]
@@ -692,8 +696,12 @@ def test_goto_if_not_int_lt():
     i1 = ri1.getint()
     pc = 117
     continue
-condbox = self.opimpl_int_lt(ri0, ri1)
-self.opimpl_goto_if_not(condbox, 17, 5)
+_v0 = self.registers_i[0].getint()
+_v1 = self.registers_i[1].getint()
+_cond = int(_v0 < _v1)
+# fast-path: record comparison directly, skip heapcache
+condbox = self.metainterp.history.record2_int(rop.INT_LT, self.registers_i[0], self.registers_i[1], _cond)
+self.opimpl_goto_if_not(condbox, 17, 5, replace=False)
 pc = self.pc
 if pc == 17:
     pc = 17
@@ -702,9 +710,10 @@ def test_goto_if_not_int_lt():
     pc = 6
 continue"""
 
-    # unspecialized case
+    # unspecialized case with constant register
     insn_specializer = work_list.specialize_pc({i2}, 5)
     s = insn_specializer.make_code()
+    # fast-path with register sync
     assert s == """\
 ri0 = self.registers_i[0]
 ri1 = self.registers_i[1]
@@ -713,9 +722,13 @@ def test_goto_if_not_int_lt():
     i1 = ri1.getint()
     pc = 119
     continue
-condbox = self.opimpl_int_lt(ri0, ri1)
 glob0(self, i2) # jit_sync_regs_i2
-self.opimpl_goto_if_not(condbox, 17, 5)
+_v0 = self.registers_i[0].getint()
+_v1 = self.registers_i[1].getint()
+_cond = int(_v0 < _v1)
+# fast-path: record comparison directly, skip heapcache
+condbox = self.metainterp.history.record2_int(rop.INT_LT, self.registers_i[0], self.registers_i[1], _cond)
+self.opimpl_goto_if_not(condbox, 17, 5, replace=False)
 pc = self.pc
 if pc == 17:
     pc = 120
@@ -1443,15 +1456,15 @@ def test_int_is_true():
     newpc = insn_specializer.get_pc()
     assert newpc == 5
     s = insn_specializer.make_code()
+    # fast-path: directly compute and record, skip heapcache
     assert s == """\
-ri0 = self.registers_i[0]
-if isinstance(ri0, ConstInt):
-    i0 = ri0.getint()
-    pc = %d
-    continue
-self.registers_i[1] = self.opimpl_int_is_true(ri0)
+_v0 = self.registers_i[0].getint()
+_res = int(bool(_v0))
+_op = self.metainterp.history.record1_int(rop.INT_IS_TRUE, self.registers_i[0], _res)
+self.registers_i[1] = _op
+i1 = _res
 pc = 7
-continue""" % (work_list.OFFSET + 7)
+continue"""
     next_constant_registers = insn_specializer.get_next_constant_registers()
     assert next_constant_registers == set()
 
@@ -1552,7 +1565,7 @@ def get_extra_info(self):
     continue
 else:
     ri0 = self.registers_i[0]
-    v0 = self.do_residual_call(ConstPtr(lltype.cast_opaque_ptr(llmemory.GCREF, glob3)), [ri0], glob2, 5)
+    v0 = self.do_residual_call(ConstPtr(lltype.cast_opaque_ptr(llmemory.GCREF, glob3)), [self.registers_i[0]], glob2, 5)
     i1 = v0.getint()
     self.registers_i[1] = v0
     pc = 7