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

4 files changed, 37 insertions, 5 deletions

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: