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

7 files changed, 68 insertions, 10 deletions

diff --git a/Tools/c-analyzer/cpython/_parser.py b/Tools/c-analyzer/cpython/_parser.py
index 61cd41e..8482bf8 100644
--- a/Tools/c-analyzer/cpython/_parser.py
+++ b/Tools/c-analyzer/cpython/_parser.py
@@ -321,6 +321,7 @@ MAX_SIZES = {
     _abs('Objects/stringlib/unicode_format.h'): (10_000, 400),
     _abs('Objects/typeobject.c'): (35_000, 200),
     _abs('Python/compile.c'): (20_000, 500),
+    _abs('Python/optimizer.c'): (100_000, 5_000),
     _abs('Python/parking_lot.c'): (40_000, 1000),
     _abs('Python/pylifecycle.c'): (500_000, 5000),
     _abs('Python/pystate.c'): (500_000, 5000),
diff --git a/Tools/c-analyzer/cpython/ignored.tsv b/Tools/c-analyzer/cpython/ignored.tsv
index 14bcd85..94be337 100644
--- a/Tools/c-analyzer/cpython/ignored.tsv
+++ b/Tools/c-analyzer/cpython/ignored.tsv
@@ -382,6 +382,11 @@ Python/optimizer.c	-	_PyCounterOptimizer_Type	-
 Python/optimizer.c	-	_PyUOpExecutor_Type	-
 Python/optimizer.c	-	_PyUOpOptimizer_Type	-
 Python/optimizer.c	-	_PyOptimizer_Default	-
+Python/optimizer.c	-	_ColdExit_Type	-
+Python/optimizer.c	-	COLD_EXITS	-
+Python/optimizer.c	-	Py_FatalErrorExecutor	-
+Python/optimizer.c	-	EMPTY_FILTER	-
+Python/optimizer.c	-	cold_exits_initialized	-
 
 ##-----------------------
 ## test code
diff --git a/Tools/cases_generator/analyzer.py b/Tools/cases_generator/analyzer.py
index 3497b7f..bcffd75 100644
--- a/Tools/cases_generator/analyzer.py
+++ b/Tools/cases_generator/analyzer.py
@@ -21,6 +21,7 @@ class Properties:
     uses_co_names: bool
     uses_locals: bool
     has_free: bool
+    side_exit: bool
 
     pure: bool
     passthrough: bool
@@ -48,6 +49,7 @@ class Properties:
             uses_co_names=any(p.uses_co_names for p in properties),
             uses_locals=any(p.uses_locals for p in properties),
             has_free=any(p.has_free for p in properties),
+            side_exit=any(p.side_exit for p in properties),
             pure=all(p.pure for p in properties),
             passthrough=all(p.passthrough for p in properties),
         )
@@ -69,6 +71,7 @@ SKIP_PROPERTIES = Properties(
     uses_co_names=False,
     uses_locals=False,
     has_free=False,
+    side_exit=False,
     pure=False,
     passthrough=False,
 )
@@ -269,9 +272,7 @@ def override_error(
 
 
 def convert_stack_item(item: parser.StackEffect) -> StackItem:
-    return StackItem(
-        item.name, item.type, item.cond, (item.size or "1")
-    )
+    return StackItem(item.name, item.type, item.cond, (item.size or "1"))
 
 
 def analyze_stack(op: parser.InstDef) -> StackEffect:
@@ -448,13 +449,24 @@ def compute_properties(op: parser.InstDef) -> Properties:
         or variable_used(op, "PyCell_GET")
         or variable_used(op, "PyCell_SET")
     )
+    deopts_if = variable_used(op, "DEOPT_IF")
+    exits_if = variable_used(op, "EXIT_IF")
+    if deopts_if and exits_if:
+        tkn = op.tokens[0]
+        raise lexer.make_syntax_error(
+            "Op cannot contain both EXIT_IF and DEOPT_IF",
+            tkn.filename,
+            tkn.line,
+            tkn.column,
+            op.name,
+        )
     infallible = is_infallible(op)
-    deopts = variable_used(op, "DEOPT_IF")
     passthrough = stack_effect_only_peeks(op) and infallible
     return Properties(
         escapes=makes_escaping_api_call(op),
         infallible=infallible,
-        deopts=deopts,
+        deopts=deopts_if or exits_if,
+        side_exit=exits_if,
         oparg=variable_used(op, "oparg"),
         jumps=variable_used(op, "JUMPBY"),
         eval_breaker=variable_used(op, "CHECK_EVAL_BREAKER"),
diff --git a/Tools/cases_generator/generators_common.py b/Tools/cases_generator/generators_common.py
index 2fc2ab1..54ffea7 100644
--- a/Tools/cases_generator/generators_common.py
+++ b/Tools/cases_generator/generators_common.py
@@ -154,6 +154,7 @@ def replace_check_eval_breaker(
 
 
 REPLACEMENT_FUNCTIONS = {
+    "EXIT_IF": replace_deopt,
     "DEOPT_IF": replace_deopt,
     "ERROR_IF": replace_error,
     "DECREF_INPUTS": replace_decrefs,
@@ -205,6 +206,8 @@ def cflags(p: Properties) -> str:
         flags.append("HAS_EVAL_BREAK_FLAG")
     if p.deopts:
         flags.append("HAS_DEOPT_FLAG")
+    if p.side_exit:
+        flags.append("HAS_EXIT_FLAG")
     if not p.infallible:
         flags.append("HAS_ERROR_FLAG")
     if p.escapes:
diff --git a/Tools/cases_generator/opcode_metadata_generator.py b/Tools/cases_generator/opcode_metadata_generator.py
index 3e9fa3e..52dc09e 100644
--- a/Tools/cases_generator/opcode_metadata_generator.py
+++ b/Tools/cases_generator/opcode_metadata_generator.py
@@ -50,6 +50,7 @@ FLAGS = [
     "DEOPT",
     "ERROR",
     "ESCAPES",
+    "EXIT",
     "PURE",
     "PASSTHROUGH",
 ]
diff --git a/Tools/cases_generator/tier2_generator.py b/Tools/cases_generator/tier2_generator.py
index 7897b89..8b4d164 100644
--- a/Tools/cases_generator/tier2_generator.py
+++ b/Tools/cases_generator/tier2_generator.py
@@ -98,9 +98,25 @@ def tier2_replace_deopt(
     out.emit(") goto deoptimize;\n")
 
 
+def tier2_replace_exit_if(
+    out: CWriter,
+    tkn: Token,
+    tkn_iter: Iterator[Token],
+    uop: Uop,
+    unused: Stack,
+    inst: Instruction | None,
+) -> None:
+    out.emit_at("if ", tkn)
+    out.emit(next(tkn_iter))
+    emit_to(out, tkn_iter, "RPAREN")
+    next(tkn_iter)  # Semi colon
+    out.emit(") goto side_exit;\n")
+
+
 TIER2_REPLACEMENT_FUNCTIONS = REPLACEMENT_FUNCTIONS.copy()
 TIER2_REPLACEMENT_FUNCTIONS["ERROR_IF"] = tier2_replace_error
 TIER2_REPLACEMENT_FUNCTIONS["DEOPT_IF"] = tier2_replace_deopt
+TIER2_REPLACEMENT_FUNCTIONS["EXIT_IF"] = tier2_replace_exit_if
 
 
 def write_uop(uop: Uop, out: CWriter, stack: Stack) -> None:
diff --git a/Tools/jit/template.c b/Tools/jit/template.c
index 12303a5..d79c6ef 100644
--- a/Tools/jit/template.c
+++ b/Tools/jit/template.c
@@ -38,6 +38,20 @@
         goto LABEL ## _tier_two; \
     } while (0)
 
+#undef GOTO_TIER_TWO
+#define GOTO_TIER_TWO(EXECUTOR) \
+do {  \
+    __attribute__((musttail))                     \
+    return ((jit_func)((EXECUTOR)->jit_code))(frame, stack_pointer, tstate); \
+} while (0)
+
+#undef GOTO_TIER_ONE
+#define GOTO_TIER_ONE(TARGET) \
+do {  \
+    _PyFrame_SetStackPointer(frame, stack_pointer); \
+    return TARGET; \
+} while (0)
+
 #undef LOAD_IP
 #define LOAD_IP(UNUSED) \
     do {                \
@@ -59,7 +73,6 @@ _JIT_ENTRY(_PyInterpreterFrame *frame, PyObject **stack_pointer, PyThreadState *
     PATCH_VALUE(_PyExecutorObject *, current_executor, _JIT_EXECUTOR)
     int oparg;
     int opcode = _JIT_OPCODE;
-    _PyUOpInstruction *next_uop;
     // Other stuff we need handy:
     PATCH_VALUE(uint16_t, _oparg, _JIT_OPARG)
     PATCH_VALUE(uint64_t, _operand, _JIT_OPERAND)
@@ -90,9 +103,16 @@ pop_2_error_tier_two:
 pop_1_error_tier_two:
     STACK_SHRINK(1);
 error_tier_two:
-    _PyFrame_SetStackPointer(frame, stack_pointer);
-    return NULL;
+    tstate->previous_executor = (PyObject *)current_executor;
+    GOTO_TIER_ONE(NULL);
 deoptimize:
-    _PyFrame_SetStackPointer(frame, stack_pointer);
-    return _PyCode_CODE(_PyFrame_GetCode(frame)) + _target;
+    tstate->previous_executor = (PyObject *)current_executor;
+    GOTO_TIER_ONE(_PyCode_CODE(_PyFrame_GetCode(frame)) + _target);
+side_exit:
+    {
+        _PyExitData *exit = &current_executor->exits[_target];
+        Py_INCREF(exit->executor);
+        tstate->previous_executor = (PyObject *)current_executor;
+        GOTO_TIER_TWO(exit->executor);
+    }
 }