gh-149335: Avoid JIT trace buffer asserts with overhead above FITNESS_INITIAL (GH-149633)

(cherry picked from commit 441af3a)

Co-authored-by: Hai Zhu <haiizhu@outlook.com>

Author: cocolato

Include/internal/pycore_optimizer.h

@@ -31,9 +31,8 @@ extern "C" {
  * 4. A push followed by a matching return is net-zero on frame-specific
  *    fitness, excluding per-slot costs.
  */
-#define MAX_TARGET_LENGTH          (UOP_MAX_TRACE_LENGTH / 2)
 #define OPTIMIZER_EFFECTIVENESS    2
-#define FITNESS_INITIAL            (MAX_TARGET_LENGTH * OPTIMIZER_EFFECTIVENESS)
+#define MAX_TARGET_LENGTH          (FITNESS_INITIAL / OPTIMIZER_EFFECTIVENESS)
 
 /* Exit quality thresholds: trace stops when fitness < exit_quality.
  * Higher = trace is more willing to stop here. */

Include/internal/pycore_uop.h

@@ -36,14 +36,18 @@ typedef struct _PyUOpInstruction{
 #endif
 } _PyUOpInstruction;
 
-// This is the length of the trace we translate initially.
+// Fitness is the target length of the trace we translate initially. The uop
+// buffer has a small amount of extra space for entry/loop-closing overhead.
 #if defined(Py_DEBUG) && defined(_Py_JIT)
     // With asserts, the stencils are a lot larger
-#define UOP_MAX_TRACE_LENGTH 1000
+#define FITNESS_INITIAL 1000
 #else
-#define UOP_MAX_TRACE_LENGTH 2500
+#define FITNESS_INITIAL 2500
 #endif
 
+#define UOP_TRACE_BUFFER_OVERHEAD 10
+#define UOP_MAX_TRACE_LENGTH (FITNESS_INITIAL + UOP_TRACE_BUFFER_OVERHEAD)
+
 /* Bloom filter with m = 256
  * https://en.wikipedia.org/wiki/Bloom_filter */
 #ifdef HAVE_GCC_UINT128_T

Lib/test/test_capi/test_opt.py

@@ -5922,6 +5922,44 @@ def __next__(self):
         """), PYTHON_JIT="1", PYTHON_JIT_STRESS="1")
         self.assertEqual(result[0].rc, 0, result)
 
+    def test_149335_trace_buffer_guard(self):
+        # https://github.com/python/cpython/issues/149335
+
+        result = script_helper.run_python_until_end('-c', textwrap.dedent("""
+        import sys
+
+        def f1():
+            for i_3178 in 0, 2, 10:
+                mv162 = 162
+
+            mv3 = mv1 = mv_165 = mv16 = \
+            mv167 = mv168 = \
+            mv169 = \
+                mv_1403_170 = \
+                169
+
+            mv_1403_170
+
+            mv_172 = mv_3 = mv_4 = mv175 = mv176 = mv17 = mv178 = mv179 = mv0 = mv1 = mv182 = (
+            mv3
+            ) = mv4 = mv185 = mv186 = mv187 = mv18 = mv189 = mv0 = mv1 = mv192 = mv3 = mv4 = (
+            mv195
+            ) = mv196 = mv197 = mv_198 = mv19 = mv0 = mv1 = mv2 = mv3 = mv4 = mv05 = mv06 = (
+            mv07
+            ) = mv08 = mv09 = mv0 = mv1 = mv2 = mv3 = mv4 = mv15 = mv16 = mv17 = mv18 = mv19 = (
+            mv0
+            ) = mv1 = mv_2 = mv3 = mv4 = mv_25 = mv_26 = mv_27 = mv_28 = mv_29 = mv0 = mv1 = (
+            mv2
+            ) = mv_1403 = mv4 = mv35 = mv36 = mv37 = mv38 = mv39 = mv0 = -sys.maxsize / 3
+
+            mv1 = mv_12 = mv3 = mv_14 = mv45 = sys.float_info.epsilon
+            mv46 = sys.float_info.epsilon
+
+        for i in range(15000):
+            f1()
+        """), PYTHON_JIT="1")
+        self.assertEqual(result[0].rc, 0, result)
+
     def test_144068_daemon_thread_jit_cleanup(self):
         result = script_helper.run_python_until_end('-c', textwrap.dedent("""
         import threading

Python/pystate.c

@@ -634,7 +634,7 @@ init_interpreter(PyInterpreterState *interp,
     // Trace fitness configuration
     init_policy(&interp->opt_config.fitness_initial,
                 "PYTHON_JIT_FITNESS_INITIAL",
-                FITNESS_INITIAL, EXIT_QUALITY_CLOSE_LOOP, UOP_MAX_TRACE_LENGTH - 1);
+                FITNESS_INITIAL, EXIT_QUALITY_CLOSE_LOOP, FITNESS_INITIAL);
 
     interp->opt_config.specialization_enabled = !is_env_enabled("PYTHON_SPECIALIZATION_OFF");
     interp->opt_config.uops_optimize_enabled = !is_env_disabled("PYTHON_UOPS_OPTIMIZE");