Add fast-path overflow operations with batched profiler counting

Author: 3tty0n

rpython/jit/codewriter/genextension.py

@@ -157,9 +157,10 @@ def generate(self):
         self.jitcode._genext_source = "\n".join(allcode)
         # Import rop for opnum constants used in type-specialized recording
         from rpython.jit.metainterp.resoperation import rop
+        from rpython.rlib.rarithmetic import ovfcheck
         d = {"ConstInt": ConstInt, "ConstPtr": ConstPtr, "ConstFloat": ConstFloat, "JitCode": JitCode, "ChangeFrame": ChangeFrame,
              "lltype": lltype, "rstr": rstr, 'llmemory': llmemory, 'OBJECTPTR': OBJECTPTR, 'support': support,
-             'rop': rop}
+             'rop': rop, 'ovfcheck': ovfcheck}
         d.update(self.globals)
         source = py.code.Source(self.jitcode._genext_source)
         exec source.compile() in d
@@ -2137,6 +2138,62 @@ def emit_unspecialized_goto_if_not_int_ne(self):
     def emit_unspecialized_goto_if_not_int_eq(self):
         return self._emit_goto_if_not_int_comparison_fast("INT_EQ", "==")
 
+    def _emit_int_ovf_fast(self, rop_name, py_op):
+        lines = []
+        _, label, arg0, arg1, _, result = self.insn
+
+        target_pc = self.get_target_pc(label)
+
+        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 with overflow check, 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("self.metainterp.ovf_flag = False")
+        lines.append("try:")
+        lines.append("    _res = ovfcheck(_v0 %s _v1)" % py_op)
+        lines.append("except OverflowError:")
+        lines.append("    self.metainterp.ovf_flag = True")
+        lines.append("    _res = 0")
+
+        lines.append("# fast-path: record overflow op directly, skip heapcache")
+        lines.append("_op = self.metainterp.history.record2_int(rop.%s, %s, %s, _res)" % (
+            rop_name, box0, box1))
+        lines.append("self.registers_i[%d] = _op" % result.index)
+        lines.append("i%d = _res" % result.index)
+
+        lines.append("self.handle_possible_overflow_error(%d, %d, _op)" % (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 - {result}, 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 - {result}, 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_int_add_jump_if_ovf(self):
+        return self._emit_int_ovf_fast("INT_ADD_OVF", "+")
+
+    def emit_unspecialized_int_sub_jump_if_ovf(self):
+        return self._emit_int_ovf_fast("INT_SUB_OVF", "-")
+
+    def emit_unspecialized_int_mul_jump_if_ovf(self):
+        return self._emit_int_ovf_fast("INT_MUL_OVF", "*")
+
     def emit_unspecialized_switch(self):
         lines = []
         arg0, descr = self._get_args()

rpython/jit/metainterp/pyjitpl.py

@@ -2545,6 +2545,12 @@ def __init__(self, staticdata, jitdriver_sd, force_finish_trace=False):
 
         self.box_names_memo = {}
 
+        # Batched profiler counting for fast-path operations
+        # Instead of calling profiler.count_ops() for every operation,
+        # we batch counts and flush them at strategic points (guards, end of trace)
+        self._batched_ops_count = 0
+        self._batched_recorded_ops_count = 0
+
         self.aborted_tracing_jitdriver = None
         self.aborted_tracing_greenkey = None
 
@@ -2554,6 +2560,30 @@ def __init__(self, staticdata, jitdriver_sd, force_finish_trace=False):
         self.force_finish_trace = force_finish_trace
         self.trace_length_at_last_tco = -1
 
+    @always_inline
+    def batch_op_count(self):
+        """Batch an operation count instead of calling profiler directly.
+
+        This is used by fast-path recording methods to avoid the overhead
+        of calling profiler.count_ops() for every operation. Counts are
+        flushed at strategic points (guards, end of trace).
+        """
+        self._batched_ops_count += 1
+        self._batched_recorded_ops_count += 1
+
+    def flush_batched_counts(self):
+        """Flush batched operation counts to the profiler.
+
+        Called before guards and at end of trace to ensure accurate counting.
+        """
+        profiler = self.staticdata.profiler
+        if self._batched_ops_count > 0:
+            profiler.count(Counters.OPS, self._batched_ops_count)
+            self._batched_ops_count = 0
+        if self._batched_recorded_ops_count > 0:
+            profiler.count(Counters.RECORDED_OPS, self._batched_recorded_ops_count)
+            self._batched_recorded_ops_count = 0
+
     def retrace_needed(self, trace, exported_state):
         self.partial_trace = trace
         self.retracing_from = self.potential_retrace_position
@@ -2699,6 +2729,8 @@ def check_recursion_invariant(self):
             raise AssertionError
 
     def generate_guard(self, opnum, box=None, extraarg=None, resumepc=-1):
+        # Flush batched counts before guard to ensure accurate profiling
+        self.flush_batched_counts()
         if isinstance(box, Const):    # no need for a guard
             return
         if opnum == rop.GUARD_EXCEPTION:
@@ -2848,8 +2880,8 @@ def _record_helper(self, opnum, resvalue, descr, *argboxes):
     def _record_int_binop(self, opnum, resvalue, b1, b2):
         if not we_are_translated():
             PyjitplCounters._record_int_binop_calls += 1
-        profiler = self.staticdata.profiler
-        profiler.count_ops(opnum, Counters.RECORDED_OPS)
+        # Use batched counting instead of calling profiler directly
+        self.batch_op_count()
         if self.framestack:
             self.framestack[-1].jitcode.traced_operations += 1
         op = self.history.record2_int(opnum, b1, b2, resvalue)
@@ -2859,8 +2891,8 @@ def _record_int_binop(self, opnum, resvalue, b1, b2):
     def _record_int_unop(self, opnum, resvalue, b1):
         if not we_are_translated():
             PyjitplCounters._record_int_binop_calls += 1
-        profiler = self.staticdata.profiler
-        profiler.count_ops(opnum, Counters.RECORDED_OPS)
+        # Use batched counting instead of calling profiler directly
+        self.batch_op_count()
         if self.framestack:
             self.framestack[-1].jitcode.traced_operations += 1
         op = self.history.record1_int(opnum, b1, resvalue)
@@ -2870,8 +2902,8 @@ def _record_int_unop(self, opnum, resvalue, b1):
     def _record_float_binop(self, opnum, resvalue, b1, b2):
         if not we_are_translated():
             PyjitplCounters._record_int_binop_calls += 1
-        profiler = self.staticdata.profiler
-        profiler.count_ops(opnum, Counters.RECORDED_OPS)
+        # Use batched counting instead of calling profiler directly
+        self.batch_op_count()
         if self.framestack:
             self.framestack[-1].jitcode.traced_operations += 1
         op = self.history.record2_float(opnum, b1, b2, resvalue)
@@ -3060,6 +3092,7 @@ def compile_and_run_once(self, jitdriver_sd, *args):
             original_boxes = self.initialize_original_boxes(jitdriver_sd, *args)
             return self._compile_and_run_once(original_boxes)
         finally:
+            self.flush_batched_counts()  # Flush before end_tracing
             self.staticdata.profiler.end_tracing()
             debug_stop('jit-tracing')
 
@@ -3101,6 +3134,7 @@ def handle_guard_failure(self, resumedescr, deadframe):
             self.run_blackhole_interp_to_cancel_tracing(stb)
         finally:
             self.resumekey_original_loop_token = None
+            self.flush_batched_counts()  # Flush before end_tracing
             self.staticdata.profiler.end_tracing()
             debug_stop('jit-tracing')