gh-146073: Add fitness/exit quality mechanism for JIT trace frontend

Include/cpython/pystats.h

@@ -144,6 +144,8 @@ typedef struct _optimization_stats {
     uint64_t unknown_callee;
     uint64_t trace_immediately_deopts;
     uint64_t executors_invalidated;
+    uint64_t fitness_terminated_traces;
+    uint64_t best_exit_fallback;
     UOpStats opcode[PYSTATS_MAX_UOP_ID + 1];
     uint64_t unsupported_opcode[256];
     uint64_t trace_length_hist[_Py_UOP_HIST_SIZE];

Include/internal/pycore_interp_structs.h

@@ -449,6 +449,20 @@ typedef struct _PyOptimizationConfig {
     uint16_t side_exit_initial_value;
     uint16_t side_exit_initial_backoff;
 
+    // Trace fitness thresholds
+    uint16_t fitness_initial;
+    uint16_t fitness_initial_side;
+    uint16_t fitness_per_instruction;
+    uint16_t fitness_branch_biased;
+    uint16_t fitness_branch_unbiased;
+    uint16_t fitness_backward_edge;
+    uint16_t fitness_frame_entry;
+
+    // Exit quality thresholds for fitness-based trace termination
+    uint16_t exit_quality_enter_executor;
+    uint16_t exit_quality_default;
+    uint16_t exit_quality_specializable;
+
     // Optimization flags
     bool specialization_enabled;
     bool uops_optimize_enabled;

Include/internal/pycore_optimizer.h

@@ -15,6 +15,23 @@ extern "C" {
 #include "pycore_optimizer_types.h"
 #include <stdbool.h>
 
+/* Default fitness configuration values for trace quality control.
+ * These can be overridden via PYTHON_JIT_FITNESS_* environment variables. */
+#define FITNESS_INITIAL             1000
+#define FITNESS_INITIAL_SIDE         800
+#define FITNESS_PER_INSTRUCTION        2
+#define FITNESS_BRANCH_BIASED          5
+#define FITNESS_BRANCH_UNBIASED       25
+#define FITNESS_BACKWARD_EDGE         80
+#define FITNESS_FRAME_ENTRY           10
+
+/* Default exit quality constants for fitness-based trace termination.
+ * Higher values mean better places to stop the trace.
+ * These can be overridden via PYTHON_JIT_EXIT_QUALITY_* environment variables. */
+#define EXIT_QUALITY_ENTER_EXECUTOR  500
+#define EXIT_QUALITY_DEFAULT         200
+#define EXIT_QUALITY_SPECIALIZABLE    50
+
 
 typedef struct _PyJitUopBuffer {
     _PyUOpInstruction *start;
@@ -101,7 +118,11 @@ typedef struct _PyJitTracerPreviousState {
 } _PyJitTracerPreviousState;
 
 typedef struct _PyJitTracerTranslatorState {
-    int jump_backward_seen;
+    int32_t fitness;              // Current trace fitness, starts high, decrements
+    int32_t best_exit_quality;    // Best exit quality seen so far
+    int best_exit_buffer_pos;     // Position in code_buffer of best exit (-1=none)
+    uint32_t best_exit_target;    // Bytecode target of best exit point
+    int frame_depth;              // Current inline depth (0 = root frame)
 } _PyJitTracerTranslatorState;
 
 typedef struct _PyJitTracerState {

Python/optimizer.c

@@ -549,8 +549,6 @@ dynamic_exit_uop[MAX_UOP_ID + 1] = {
 };
 
 
-#define CONFIDENCE_RANGE 1000
-#define CONFIDENCE_CUTOFF 333
 
 #ifdef Py_DEBUG
 #define DPRINTF(level, ...) \
@@ -598,6 +596,86 @@ add_to_trace(
     ((uint32_t)((INSTR) - ((_Py_CODEUNIT *)(CODE)->co_code_adaptive)))
 
 
+/* Compute branch bias from the 16-bit branch history register.
+ * Returns 0 (completely unpredictable, 50/50) to 8 (fully biased). */
+static inline int
+compute_branch_bias(uint16_t history)
+{
+    int ones = _Py_popcount32((uint32_t)history);
+    return abs(ones - 8);
+}
+
+/* Compute exit quality for the current trace position.
+ * Higher values mean it's a better place to stop the trace. */
+static inline int32_t
+compute_exit_quality(_Py_CODEUNIT *target_instr, int opcode,
+                     const _PyOptimizationConfig *cfg)
+{
+    if (target_instr->op.code == ENTER_EXECUTOR) {
+        return (int32_t)cfg->exit_quality_enter_executor;
+    }
+    if (_PyOpcode_Caches[_PyOpcode_Deopt[opcode]] > 0) {
+        return (int32_t)cfg->exit_quality_specializable;
+    }
+    return (int32_t)cfg->exit_quality_default;
+}
+
+/* Try to truncate the trace to the best recorded exit point.
+ * Returns 1 if successful, 0 if no valid best exit exists.
+ * Enforces progress constraints: the fallback position must satisfy
+ * the minimum trace length requirements. */
+static inline int
+try_best_exit_fallback(
+    _PyJitUopBuffer *trace,
+    _PyJitTracerTranslatorState *ts,
+    bool progress_needed)
+{
+    int best_pos = ts->best_exit_buffer_pos;
+    if (best_pos <= 0) {
+        return 0;
+    } else if (progress_needed && best_pos <= CODE_SIZE_NO_PROGRESS) {
+        return 0;
+    } else if (!progress_needed && best_pos <= CODE_SIZE_EMPTY) {
+        return 0;
+    }
+    trace->next = trace->start + best_pos;
+    /* Caller must add terminator (_EXIT_TRACE) after this */
+    return 1;
+}
+
+/* Update trace fitness after translating one bytecode instruction. */
+static inline void
+update_trace_fitness(
+    _PyJitTracerTranslatorState *ts,
+    int opcode,
+    _Py_CODEUNIT *target_instr,
+    const _PyOptimizationConfig *cfg)
+{
+    ts->fitness -= cfg->fitness_per_instruction;
+
+    switch (opcode) {
+        case POP_JUMP_IF_FALSE:
+        case POP_JUMP_IF_TRUE:
+        case POP_JUMP_IF_NONE:
+        case POP_JUMP_IF_NOT_NONE: {
+            int bias = compute_branch_bias(target_instr[1].cache);
+            /* Linear interpolation: bias 0 → unbiased penalty, bias 8 → biased penalty */
+            int penalty = cfg->fitness_branch_unbiased
+                - (bias * (cfg->fitness_branch_unbiased - cfg->fitness_branch_biased)) / 8;
+            ts->fitness -= penalty;
+            break;
+        }
+        case JUMP_BACKWARD:
+        case JUMP_BACKWARD_JIT:
+        case JUMP_BACKWARD_NO_JIT:
+            ts->fitness -= cfg->fitness_backward_edge;
+            break;
+        default:
+            break;
+    }
+}
+
+
 static int
 is_terminator(const _PyUOpInstruction *uop)
 {
@@ -730,17 +808,44 @@ _PyJit_translate_single_bytecode_to_trace(
         goto unsupported;
     }
 
+    // Track frame depth changes for fitness (only for supported frame transitions)
+    if (frame != tracer->prev_state.instr_frame) {
+        _PyJitTracerTranslatorState *ts_depth = &tracer->translator_state;
+        if (frame->previous == tracer->prev_state.instr_frame) {
+            ts_depth->frame_depth++;
+            // Penalty scales with depth: shallow inlining is cheap,
+            // deep inlining gets progressively more expensive.
+            int32_t penalty = (int32_t)tstate->interp->opt_config.fitness_frame_entry
+                            * ts_depth->frame_depth;
+            ts_depth->fitness -= penalty;
+        } else if (ts_depth->frame_depth > 0) {
+            ts_depth->frame_depth--;
+        }
+    }
+
     if (oparg > 0xFFFF) {
         DPRINTF(2, "Unsupported: oparg too large\n");
         unsupported:
         {
-            // Rewind to previous instruction and replace with _EXIT_TRACE.
+            // If we have a high-quality best_exit (enter_executor, etc.),
+            // prefer it over rewinding to last _SET_IP — this covers the
+            // main unsupported path, not just the edge case.
+            _PyJitTracerTranslatorState *ts_unsup = &tracer->translator_state;
+            if (ts_unsup->best_exit_quality > (int32_t)tstate->interp->opt_config.exit_quality_default &&
+                try_best_exit_fallback(trace, ts_unsup, progress_needed)) {
+                ADD_TO_TRACE(_EXIT_TRACE, 0, 0, ts_unsup->best_exit_target);
+                uop_buffer_last(trace)->operand1 = true; // is_control_flow
+                OPT_STAT_INC(best_exit_fallback);
+                DPRINTF(2, "Best-exit fallback at unsupported (pos=%d, quality=%d)\n",
+                        ts_unsup->best_exit_buffer_pos, ts_unsup->best_exit_quality);
+                goto done;
+            }
+            // Fall back: rewind to last _SET_IP and replace with _DEOPT.
             _PyUOpInstruction *curr = uop_buffer_last(trace);
             while (curr->opcode != _SET_IP && uop_buffer_length(trace) > 2) {
                 trace->next--;
                 curr = uop_buffer_last(trace);
             }
-            assert(curr->opcode == _SET_IP || uop_buffer_length(trace) == 2);
             if (curr->opcode == _SET_IP) {
                 int32_t old_target = (int32_t)uop_get_target(curr);
                 curr->opcode = _DEOPT;
@@ -763,6 +868,39 @@ _PyJit_translate_single_bytecode_to_trace(
         return 1;
     }
 
+    // Fitness-based trace quality check (before reserving space for this instruction)
+    _PyJitTracerTranslatorState *ts = &tracer->translator_state;
+    int32_t eq = compute_exit_quality(target_instr, opcode,
+                                        &tstate->interp->opt_config);
+
+    // Record best exit candidate.
+    // Only record after minimum progress to avoid truncating to near-empty traces.
+    if (eq > ts->best_exit_quality &&
+        uop_buffer_length(trace) > CODE_SIZE_NO_PROGRESS) {
+        ts->best_exit_quality = eq;
+        ts->best_exit_buffer_pos = uop_buffer_length(trace);
+        ts->best_exit_target = target;
+    }
+
+    // Check if fitness is depleted — should we stop the trace?
+    if (ts->fitness < eq &&
+        !(progress_needed && uop_buffer_length(trace) < CODE_SIZE_NO_PROGRESS)) {
+        // Prefer stopping at the best recorded exit point
+        if (try_best_exit_fallback(trace, ts, progress_needed)) {
+            ADD_TO_TRACE(_EXIT_TRACE, 0, 0, ts->best_exit_target);
+            uop_buffer_last(trace)->operand1 = true; // is_control_flow
+        }
+        else {
+            // No valid best exit — stop at current position
+            ADD_TO_TRACE(_EXIT_TRACE, 0, 0, target);
+            uop_buffer_last(trace)->operand1 = true; // is_control_flow
+        }
+        OPT_STAT_INC(fitness_terminated_traces);
+        DPRINTF(2, "Fitness terminated: fitness=%d < exit_quality=%d\n",
+                ts->fitness, eq);
+        goto done;
+    }
+
     // One for possible _DEOPT, one because _CHECK_VALIDITY itself might _DEOPT
     trace->end -= 2;
 
@@ -793,6 +931,12 @@ _PyJit_translate_single_bytecode_to_trace(
         DPRINTF(2, "No room for expansions and guards (need %d, got %d)\n",
                 space_needed, uop_buffer_remaining_space(trace));
         OPT_STAT_INC(trace_too_long);
+        // Try best-exit fallback before giving up
+        if (try_best_exit_fallback(trace, &tracer->translator_state, progress_needed)) {
+            ADD_TO_TRACE(_EXIT_TRACE, 0, 0, tracer->translator_state.best_exit_target);
+            uop_buffer_last(trace)->operand1 = true; // is_control_flow
+            OPT_STAT_INC(best_exit_fallback);
+        }
         goto done;
     }
 
@@ -986,7 +1130,12 @@ _PyJit_translate_single_bytecode_to_trace(
         ADD_TO_TRACE(_JUMP_TO_TOP, 0, 0, 0);
         goto done;
     }
-    DPRINTF(2, "Trace continuing\n");
+    // Update fitness AFTER translation, BEFORE returning to continue tracing.
+    // This ensures the next iteration's fitness check reflects the cost of
+    // all instructions translated so far.
+    update_trace_fitness(&tracer->translator_state, opcode, target_instr,
+                         &tstate->interp->opt_config);
+    DPRINTF(2, "Trace continuing (fitness=%d)\n", tracer->translator_state.fitness);
     return 1;
 done:
     DPRINTF(2, "Trace done\n");
@@ -1069,6 +1218,18 @@ _PyJit_TryInitializeTracing(
     assert(curr_instr->op.code == JUMP_BACKWARD_JIT || curr_instr->op.code == RESUME_CHECK_JIT || (exit != NULL));
     tracer->initial_state.jump_backward_instr = curr_instr;
 
+    // Initialize fitness tracking state
+    const _PyOptimizationConfig *cfg = &tstate->interp->opt_config;
+    _PyJitTracerTranslatorState *ts = &tracer->translator_state;
+    bool is_side_trace = (exit != NULL);
+    ts->fitness = is_side_trace
+        ? (int32_t)cfg->fitness_initial_side
+        : (int32_t)cfg->fitness_initial;
+    ts->best_exit_quality = 0;
+    ts->best_exit_buffer_pos = -1;
+    ts->best_exit_target = 0;
+    ts->frame_depth = 0;
+
     tracer->is_tracing = true;
     return 1;
 }

Python/pystate.c

@@ -635,6 +635,40 @@ init_interpreter(PyInterpreterState *interp,
                 "PYTHON_JIT_SIDE_EXIT_INITIAL_BACKOFF",
                 SIDE_EXIT_INITIAL_BACKOFF, 0, MAX_BACKOFF);
 
+    // Trace fitness configuration
+    init_policy(&interp->opt_config.fitness_initial,
+                "PYTHON_JIT_FITNESS_INITIAL",
+                FITNESS_INITIAL, 100, 10000);
+    init_policy(&interp->opt_config.fitness_initial_side,
+                "PYTHON_JIT_FITNESS_INITIAL_SIDE",
+                FITNESS_INITIAL_SIDE, 50, 5000);
+    init_policy(&interp->opt_config.fitness_per_instruction,
+                "PYTHON_JIT_FITNESS_PER_INSTRUCTION",
+                FITNESS_PER_INSTRUCTION, 0, 100);
+    init_policy(&interp->opt_config.fitness_branch_biased,
+                "PYTHON_JIT_FITNESS_BRANCH_BIASED",
+                FITNESS_BRANCH_BIASED, 0, 500);
+    init_policy(&interp->opt_config.fitness_branch_unbiased,
+                "PYTHON_JIT_FITNESS_BRANCH_UNBIASED",
+                FITNESS_BRANCH_UNBIASED, 0, 500);
+    init_policy(&interp->opt_config.fitness_backward_edge,
+                "PYTHON_JIT_FITNESS_BACKWARD_EDGE",
+                FITNESS_BACKWARD_EDGE, 0, 1000);
+    init_policy(&interp->opt_config.fitness_frame_entry,
+                "PYTHON_JIT_FITNESS_FRAME_ENTRY",
+                FITNESS_FRAME_ENTRY, 0, 1000);
+
+    // Exit quality thresholds
+    init_policy(&interp->opt_config.exit_quality_enter_executor,
+                "PYTHON_JIT_EXIT_QUALITY_ENTER_EXECUTOR",
+                EXIT_QUALITY_ENTER_EXECUTOR, 0, 10000);
+    init_policy(&interp->opt_config.exit_quality_default,
+                "PYTHON_JIT_EXIT_QUALITY_DEFAULT",
+                EXIT_QUALITY_DEFAULT, 0, 10000);
+    init_policy(&interp->opt_config.exit_quality_specializable,
+                "PYTHON_JIT_EXIT_QUALITY_SPECIALIZABLE",
+                EXIT_QUALITY_SPECIALIZABLE, 0, 10000);
+
     interp->opt_config.specialization_enabled = !is_env_enabled("PYTHON_SPECIALIZATION_OFF");
     interp->opt_config.uops_optimize_enabled = !is_env_disabled("PYTHON_UOPS_OPTIMIZE");
     if (interp != &runtime->_main_interpreter) {

Python/pystats.c

@@ -274,6 +274,8 @@ print_optimization_stats(FILE *out, OptimizationStats *stats)
     fprintf(out, "Optimization low confidence: %" PRIu64 "\n", stats->low_confidence);
     fprintf(out, "Optimization unknown callee: %" PRIu64 "\n", stats->unknown_callee);
     fprintf(out, "Executors invalidated: %" PRIu64 "\n", stats->executors_invalidated);
+    fprintf(out, "Optimization fitness terminated: %" PRIu64 "\n", stats->fitness_terminated_traces);
+    fprintf(out, "Optimization best exit fallback: %" PRIu64 "\n", stats->best_exit_fallback);
 
     print_histogram(out, "Trace length", stats->trace_length_hist);
     print_histogram(out, "Trace run length", stats->trace_run_length_hist);