GH-112354: Initial implementation of warm up on exits and trace-stitching (GH-114142)

Author: markshannon

.gitattributes

@@ -77,6 +77,7 @@ Include/internal/pycore_opcode.h                    generated
 Include/internal/pycore_opcode_metadata.h           generated
 Include/internal/pycore_*_generated.h               generated
 Include/internal/pycore_uop_ids.h                   generated
+Include/internal/pycore_uop_metadata.h              generated
 Include/opcode.h                                    generated
 Include/opcode_ids.h                                generated
 Include/token.h                                     generated

Include/cpython/optimizer.h

@@ -33,16 +33,28 @@ typedef struct {
 typedef struct {
     uint16_t opcode;
     uint16_t oparg;
-    uint32_t target;
+    union {
+        uint32_t target;
+        uint32_t exit_index;
+    };
     uint64_t operand;  // A cache entry
 } _PyUOpInstruction;
 
+typedef struct _exit_data {
+    uint32_t target;
+    int16_t temperature;
+    const struct _PyExecutorObject *executor;
+} _PyExitData;
+
 typedef struct _PyExecutorObject {
     PyObject_VAR_HEAD
+    const _PyUOpInstruction *trace;
     _PyVMData vm_data; /* Used by the VM, but opaque to the optimizer */
-    void *jit_code;
+    uint32_t exit_count;
+    uint32_t code_size;
     size_t jit_size;
-    _PyUOpInstruction trace[1];
+    void *jit_code;
+    _PyExitData exits[1];
 } _PyExecutorObject;
 
 typedef struct _PyOptimizerObject _PyOptimizerObject;
@@ -59,6 +71,7 @@ typedef struct _PyOptimizerObject {
     /* These thresholds are treated as signed so do not exceed INT16_MAX
      * Use INT16_MAX to indicate that the optimizer should never be called */
     uint16_t resume_threshold;
+    uint16_t side_threshold;
     uint16_t backedge_threshold;
     /* Data needed by the optimizer goes here, but is opaque to the VM */
 } _PyOptimizerObject;
@@ -73,16 +86,16 @@ PyAPI_FUNC(int) PyUnstable_Replace_Executor(PyCodeObject *code, _Py_CODEUNIT *in
 
 _PyOptimizerObject *_Py_SetOptimizer(PyInterpreterState *interp, _PyOptimizerObject* optimizer);
 
-PyAPI_FUNC(void) PyUnstable_SetOptimizer(_PyOptimizerObject* optimizer);
+PyAPI_FUNC(int) PyUnstable_SetOptimizer(_PyOptimizerObject* optimizer);
 
 PyAPI_FUNC(_PyOptimizerObject *) PyUnstable_GetOptimizer(void);
 
 PyAPI_FUNC(_PyExecutorObject *) PyUnstable_GetExecutor(PyCodeObject *code, int offset);
 
 int
-_PyOptimizer_Optimize(struct _PyInterpreterFrame *frame, _Py_CODEUNIT *start, PyObject **stack_pointer);
+_PyOptimizer_Optimize(struct _PyInterpreterFrame *frame, _Py_CODEUNIT *start, PyObject **stack_pointer, _PyExecutorObject **exec_ptr);
 
-void _Py_ExecutorInit(_PyExecutorObject *, _PyBloomFilter *);
+void _Py_ExecutorInit(_PyExecutorObject *, const _PyBloomFilter *);
 void _Py_ExecutorClear(_PyExecutorObject *);
 void _Py_BloomFilter_Init(_PyBloomFilter *);
 void _Py_BloomFilter_Add(_PyBloomFilter *bloom, void *obj);

Include/cpython/pystate.h

@@ -212,6 +212,8 @@ struct _ts {
     /* The thread's exception stack entry.  (Always the last entry.) */
     _PyErr_StackItem exc_state;
 
+    PyObject *previous_executor;
+
 };
 
 #ifdef Py_DEBUG

Include/internal/pycore_interp.h

@@ -237,10 +237,13 @@ struct _is {
     struct callable_cache callable_cache;
     _PyOptimizerObject *optimizer;
     _PyExecutorObject *executor_list_head;
-    /* These values are shifted and offset to speed up check in JUMP_BACKWARD */
+
+    /* These two values are shifted and offset to speed up check in JUMP_BACKWARD */
     uint32_t optimizer_resume_threshold;
     uint32_t optimizer_backedge_threshold;
 
+    uint16_t optimizer_side_threshold;
+
     uint32_t next_func_version;
     _rare_events rare_events;
     PyDict_WatchCallback builtins_dict_watcher;

Include/internal/pycore_jit.h

@@ -13,7 +13,7 @@ extern "C" {
 
 typedef _Py_CODEUNIT *(*jit_func)(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *tstate);
 
-int _PyJIT_Compile(_PyExecutorObject *executor, _PyUOpInstruction *trace, size_t length);
+int _PyJIT_Compile(_PyExecutorObject *executor, const _PyUOpInstruction *trace, size_t length);
 void _PyJIT_Free(_PyExecutorObject *executor);
 
 #endif  // _Py_JIT

Include/internal/pycore_opcode_metadata.h

@@ -331,6 +331,7 @@ def f():
             # on the *very next* allocation:
             gc.collect()
             gc.set_threshold(1, 0, 0)
+            sys._clear_internal_caches()
             # Okay, so here's the nightmare scenario:
             # - We're tracing the resumption of a generator, which creates a new
             #   frame object.

Include/internal/pycore_uop_ids.h

@@ -794,6 +794,17 @@ def test_annotated_op(self):
         self.run_cases_test(input, output)
 
 
+    def test_deopt_and_exit(self):
+        input = """
+        pure op(OP, (arg1 -- out)) {
+            DEOPT_IF(1);
+            EXIT_IF(1);
+        }
+        """
+        output = ""
+        with self.assertRaises(Exception):
+            self.run_cases_test(input, output)
+
 class TestGeneratedAbstractCases(unittest.TestCase):
     def setUp(self) -> None:
         super().setUp()