Rewrite the tracing JIT to use a common opcode handler

Author: Fidget-Spinner

Include/internal/pycore_interp_structs.h

@@ -759,18 +759,23 @@ typedef _Py_CODEUNIT *(*_PyJitEntryFuncPtr)(struct _PyExecutorObject *exec, _PyI
 
 typedef struct _PyJitTracerState {
     bool dependencies_still_valid;
+    bool do_not_specialize;
+    bool dynamic_jump_taken;
     int code_max_size;
     int code_curr_size;
     int initial_stack_depth;
     int initial_chain_depth;
+    int prev_instr_oparg;
+    int prev_instr_stacklevel;
     _PyUOpInstruction *code_buffer;
     _Py_CODEUNIT *insert_exec_instr;
     _Py_CODEUNIT *close_loop_instr;
-    _Py_CODEUNIT *last_specialized_instr;
     PyCodeObject *initial_code; // Strong
     PyFunctionObject *initial_func; // Strong
-    struct _PyExitData *previous_exit;
-    _PyInterpreterFrame *current_frame;
+    _Py_CODEUNIT *prev_instr;
+    PyCodeObject *prev_instr_code; // Strong
+    struct _PyExitData *prev_exit;
+    _PyInterpreterFrame *prev_instr_frame;
     _PyBloomFilter dependencies;
 } _PyJitTracerState;
 

Include/internal/pycore_magic_number.h

@@ -286,6 +286,7 @@ Known values:
     Python 3.15a1 3653 (Fix handling of opcodes that may leave operands on the stack when optimizing LOAD_FAST)
     Python 3.15a1 3654 (Fix missing exception handlers in logical expression)
     Python 3.15a1 3655 (Fix miscompilation of some module-level annotations)
+    Python 3.15a2 3656 (Add RECORD_PREVIOUS_INST for a tracing JIT)
 
 
     Python 3.16 will start with 3700
@@ -299,7 +300,7 @@ PC/launcher.c must also be updated.
 
 */
 
-#define PYC_MAGIC_NUMBER 3655
+#define PYC_MAGIC_NUMBER 3656
 /* This is equivalent to converting PYC_MAGIC_NUMBER to 2 bytes
    (little-endian) and then appending b'\r\n'. */
 #define PYC_MAGIC_NUMBER_TOKEN \

Include/internal/pycore_opcode_metadata.h

@@ -358,22 +358,12 @@ PyAPI_FUNC(int) _PyDumpExecutors(FILE *out);
 extern void _Py_ClearExecutorDeletionList(PyInterpreterState *interp);
 #endif
 
-int
-_PyJit_translate_single_bytecode_to_trace(
-    PyThreadState *tstate,
-    _PyInterpreterFrame *frame,
-    _Py_CODEUNIT *this_instr,
-    _Py_CODEUNIT *next_instr,
-    PyCodeObject *code,
-    PyFunctionObject *func,
-    int old_stack_level,
-    int opcode,
-    int oparg,
-    int jump_taken);
+int _PyJit_translate_single_bytecode_to_trace(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *next_instr);
 
 void
-_PyJit_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame, _Py_CODEUNIT *insert_exec_instr,
-    _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit);
+_PyJit_InitializeTracing(PyThreadState *tstate, _PyInterpreterFrame *frame,
+    _Py_CODEUNIT *curr_instr, _Py_CODEUNIT *insert_exec_instr,
+    _Py_CODEUNIT *close_loop_instr, int curr_stackdepth, int chain_depth, _PyExitData *exit, int oparg);
 
 void _PyJit_FinalizeTracing(PyThreadState *tstate);
 

Include/internal/pycore_optimizer.h

@@ -2255,7 +2255,7 @@ def frame_3_jit() -> None:
                 # 1 extra iteration for tracing.
                 for i in range(_testinternalcapi.TIER2_THRESHOLD + 2):
                     # Careful, doing this in the reverse order breaks tracing:
-                    expected = {enabled} and i >= _testinternalcapi.TIER2_THRESHOLD
+                    expected = {enabled} and i >= _testinternalcapi.TIER2_THRESHOLD + 1
                     assert sys._jit.is_active() is expected
                     frame_2_jit(expected)
                     assert sys._jit.is_active() is expected

Include/opcode_ids.h

@@ -2141,7 +2141,7 @@ Objects/mimalloc/page.o: $(srcdir)/Objects/mimalloc/page-queue.c
 regen-cases: \
         regen-opcode-ids regen-opcode-targets regen-uop-ids regen-opcode-metadata-py \
 		regen-generated-cases regen-executor-cases regen-optimizer-cases \
-		regen-opcode-metadata regen-uop-metadata regen-tracer-cases
+		regen-opcode-metadata regen-uop-metadata
 
 .PHONY: regen-opcode-ids
 regen-opcode-ids:
@@ -2187,13 +2187,6 @@ regen-optimizer-cases:
 	    $(srcdir)/Python/bytecodes.c
 	$(UPDATE_FILE) $(srcdir)/Python/optimizer_cases.c.h $(srcdir)/Python/optimizer_cases.c.h.new
 
-.PHONY: regen-tracer-cases
-regen-tracer-cases:
-	$(PYTHON_FOR_REGEN) $(srcdir)/Tools/cases_generator/tracer_generator.py \
-	    -o $(srcdir)/Python/generated_tracer_cases.c.h.new $(srcdir)/Python/bytecodes.c
-	$(UPDATE_FILE) $(srcdir)/Python/generated_tracer_cases.c.h $(srcdir)/Python/generated_tracer_cases.c.h.new
-
-
 .PHONY: regen-opcode-metadata
 regen-opcode-metadata:
 	$(PYTHON_FOR_REGEN) $(srcdir)/Tools/cases_generator/opcode_metadata_generator.py \

Lib/_opcode_metadata.py

@@ -45,6 +45,8 @@
 
 #define USE_COMPUTED_GOTOS 0
 #include "ceval_macros.h"
+#include "ceval_macros.h"
+#include "ceval_macros.h"
 #include "../Include/internal/pycore_code.h"
 #include "../Include/internal/pycore_stackref.h"
 
@@ -2985,14 +2987,9 @@ dummy_func(
                         oparg >>= 8;
                         insert_exec_at--;
                     }
-                    _PyJit_InitializeTracing(tstate, frame, insert_exec_at, next_instr, STACK_LEVEL(), 0, NULL);
+                    _PyJit_InitializeTracing(tstate, frame, this_instr, insert_exec_at, next_instr, STACK_LEVEL(), 0, NULL, oparg);
                     ENTER_TRACING();
                 }
-                int _jump_taken = false;
-                PyCodeObject *old_code = _PyFrame_GetCode(frame);
-                PyFunctionObject *old_func = (PyFunctionObject *)PyStackRef_AsPyObjectBorrow(frame->f_funcobj);
-                int _old_stack_level = 0;
-                TRACING_DISPATCH();
             }
             else {
                 ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
@@ -3057,9 +3054,6 @@ dummy_func(
             }
             assert(executor != tstate->interp->cold_executor);
             tstate->jit_exit = NULL;
-            #if TRACING_JIT
-            RECORD_TRACE_NO_DISPATCH();
-            #endif
             TIER1_TO_TIER2(executor);
             #else
             Py_FatalError("ENTER_EXECUTOR is not supported in this build");
@@ -5224,6 +5218,25 @@ dummy_func(
             Py_FatalError("Executing RESERVED instruction.");
         }
 
+        tier1 no_save_ip inst(RECORD_PREVIOUS_INST, (--)) {
+            assert(IS_JIT_TRACING());
+            int opcode = next_instr->op.code;
+            int full = !_PyJit_translate_single_bytecode_to_trace(tstate, frame, next_instr);
+            if (full) {
+                LEAVE_TRACING();
+                int err = bail_tracing_and_jit(tstate, frame);
+                ERROR_IF(err < 0);
+                DISPATCH_GOTO_NON_TRACING();
+            }
+            tstate->interp->jit_state.do_not_specialize = false;
+            PyCodeObject *prev_code = (PyCodeObject *)Py_NewRef(_PyFrame_GetCode(frame));
+            Py_SETREF(tstate->interp->jit_state.prev_instr_code, prev_code);
+            tstate->interp->jit_state.prev_instr = next_instr;
+            tstate->interp->jit_state.prev_instr_frame = frame;
+            tstate->interp->jit_state.prev_instr_oparg = oparg;
+            tstate->interp->jit_state.prev_instr_stacklevel = STACK_LEVEL();
+            DISPATCH_GOTO_NON_TRACING();
+        }
         ///////// Tier-2 only opcodes /////////
 
         op (_GUARD_IS_TRUE_POP, (flag -- )) {
@@ -5456,7 +5469,10 @@ dummy_func(
                 _PyExecutorObject *previous_executor = _PyExecutor_FromExit(exit);
                 assert(tstate->current_executor == (PyObject *)previous_executor);
                 int chain_depth = previous_executor->vm_data.chain_depth + 1;
-                _PyJit_InitializeTracing(tstate, frame, target, target, STACK_LEVEL(), chain_depth, exit);
+                // Note: it's safe to use target->op.arg here instead of the oparg given by EXTENDED_ARG.
+                // The invariant in the optimizer is the deopt target always points back to the first EXTENDED_ARG.
+                // So setting it to anything else is wrong.
+                _PyJit_InitializeTracing(tstate, frame, target, target, target, STACK_LEVEL(), chain_depth, exit, target->op.arg);
                 exit->temperature = initial_temperature_backoff_counter();
                 GOTO_TIER_ONE(target, 1);
             }

Lib/test/test_sys.py

@@ -984,17 +984,6 @@ _PyObjectArray_Free(PyObject **array, PyObject **scratch)
     }
 }
 
-#if _Py_TIER2
-// 1 for trace full, 0 for successful write.
-static inline int
-add_to_code_trace(PyThreadState *tstate, _PyInterpreterFrame *frame, PyCodeObject *old_code, PyFunctionObject *old_func, int old_stack_level, _Py_CODEUNIT *this_instr, _Py_CODEUNIT *next_instr, int opcode, int oparg, int jump_taken)
-{
-    assert(frame != NULL);
-    assert(tstate->interp->jit_state.code_curr_size < UOP_MAX_TRACE_LENGTH);
-    return !_PyJit_translate_single_bytecode_to_trace(tstate, frame, this_instr, next_instr, old_code, old_func, old_stack_level, opcode, oparg, jump_taken);
-}
-
-
 // 0 for success, -1  for error.
 static int
 bail_tracing_and_jit(PyThreadState *tstate, _PyInterpreterFrame *frame)
@@ -1007,7 +996,6 @@ bail_tracing_and_jit(PyThreadState *tstate, _PyInterpreterFrame *frame)
     _PyJit_FinalizeTracing(tstate);
     return err;
 }
-#endif
 
 /* _PyEval_EvalFrameDefault is too large to optimize for speed with PGO on MSVC.
  */
@@ -1025,7 +1013,6 @@ bail_tracing_and_jit(PyThreadState *tstate, _PyInterpreterFrame *frame)
 #if _Py_TAIL_CALL_INTERP
 #include "opcode_targets.h"
 #include "generated_cases.c.h"
-#include "generated_tracer_cases.c.h"
 #endif
 
 #if (defined(__GNUC__) && __GNUC__ >= 10 && !defined(__clang__)) && defined(__x86_64__)
@@ -1154,7 +1141,6 @@ _PyEval_EvalFrameDefault(PyThreadState *tstate, _PyInterpreterFrame *frame, int
 #else
     goto start_frame;
 #   include "generated_cases.c.h"
-    #include "generated_tracer_cases.c.h"
 #endif