gh-136183: Deal with escapes in JIT optimizer's constant evaluator

.github/workflows/jit.yml

@@ -5,6 +5,7 @@ on:
       - '**jit**'
       - 'Python/bytecodes.c'
       - 'Python/optimizer*.c'
+      - 'Python/executor_cases.c.h'
       - '!Python/perf_jit_trampoline.c'
       - '!**/*.md'
       - '!**/*.ini'
@@ -13,6 +14,7 @@ on:
       - '**jit**'
       - 'Python/bytecodes.c'
       - 'Python/optimizer*.c'
+      - 'Python/executor_cases.c.h'
       - '!Python/perf_jit_trampoline.c'
       - '!**/*.md'
       - '!**/*.ini'

Lib/test/test_capi/test_opt.py

@@ -2451,6 +2451,21 @@ def testfunc(n):
         self.assertNotIn("_GUARD_TOS_FLOAT", uops)
         self.assertNotIn("_GUARD_NOS_FLOAT", uops)
 
+    def test_binary_op_constant_evaluate(self):
+        def testfunc(n):
+            for _ in range(n):
+                2 ** 65
+
+        testfunc(TIER2_THRESHOLD)
+
+        ex = get_first_executor(testfunc)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        # For now.. until we constant propagte it away.
+        self.assertIn("_BINARY_OP", uops)
+
+
 def global_identity(x):
     return x
 

Lib/test/test_generated_cases.py

@@ -2398,6 +2398,53 @@ def test_replace_opcode_uop_body_copied_in_complex(self):
         """
         self.run_cases_test(input, input2, output)
 
+    def test_replace_opcode_escaping_uop_body_copied_in_complex(self):
+        input = """
+        pure op(OP, (foo -- res)) {
+            if (foo) {
+                res = ESCAPING_CODE(foo);
+            }
+            else {
+                res = 1;
+            }
+        }
+        """
+        input2 = """
+        op(OP, (foo -- res)) {
+            REPLACE_OPCODE_IF_EVALUATES_PURE(foo);
+            res = sym_new_known(ctx, foo);
+        }
+        """
+        output = """
+        case OP: {
+            JitOptRef foo;
+            JitOptRef res;
+            foo = stack_pointer[-1];
+            if (
+                sym_is_safe_const(ctx, foo)
+            ) {
+                JitOptRef foo_sym = foo;
+                _PyStackRef foo = sym_get_const_as_stackref(ctx, foo_sym);
+                _PyStackRef res_stackref;
+                /* Start of uop copied from bytecodes for constant evaluation */
+                if (foo) {
+                    res_stackref = ESCAPING_CODE(foo);
+                }
+                else {
+                    res_stackref = 1;
+                }
+                /* End of uop copied from bytecodes for constant evaluation */
+                res = sym_new_const_steal(ctx, PyStackRef_AsPyObjectSteal(res_stackref));
+                stack_pointer[-1] = res;
+                break;
+            }
+            res = sym_new_known(ctx, foo);
+            stack_pointer[-1] = res;
+            break;
+        }
+        """
+        self.run_cases_test(input, input2, output)
+
     def test_replace_opocode_uop_reject_array_effects(self):
         input = """
         pure op(OP, (foo[2] -- res)) {

Tools/cases_generator/generators_common.py

@@ -106,8 +106,9 @@ class Emitter:
     out: CWriter
     labels: dict[str, Label]
     _replacers: dict[str, ReplacementFunctionType]
+    cannot_escape: bool
 
-    def __init__(self, out: CWriter, labels: dict[str, Label]):
+    def __init__(self, out: CWriter, labels: dict[str, Label], cannot_escape: bool = False):
         self._replacers = {
             "EXIT_IF": self.exit_if,
             "DEOPT_IF": self.deopt_if,
@@ -127,6 +128,7 @@ def __init__(self, out: CWriter, labels: dict[str, Label]):
         }
         self.out = out
         self.labels = labels
+        self.cannot_escape = cannot_escape
 
     def dispatch(
         self,
@@ -238,7 +240,8 @@ def decref_inputs(
         next(tkn_iter)
         self._print_storage("DECREF_INPUTS", storage)
         try:
-            storage.close_inputs(self.out)
+            if not self.cannot_escape:
+                storage.close_inputs(self.out)
         except StackError as ex:
             raise analysis_error(ex.args[0], tkn)
         except Exception as ex:
@@ -476,7 +479,7 @@ def emit_SimpleStmt(
         reachable = True
         tkn = stmt.contents[-1]
         try:
-            if stmt in uop.properties.escaping_calls:
+            if stmt in uop.properties.escaping_calls and not self.cannot_escape:
                 escape = uop.properties.escaping_calls[stmt]
                 if escape.kills is not None:
                     self.stackref_kill(escape.kills, storage, True)
@@ -513,7 +516,7 @@ def emit_SimpleStmt(
                         self.out.emit(tkn)
                 else:
                     self.out.emit(tkn)
-            if stmt in uop.properties.escaping_calls:
+            if stmt in uop.properties.escaping_calls and not self.cannot_escape:
                 self.emit_reload(storage)
             return reachable, None, storage
         except StackError as ex:

Tools/cases_generator/optimizer_generator.py

@@ -245,6 +245,7 @@ def __init__(self, out: CWriter, labels: dict[str, Label], original_uop: Uop, st
             outp.name: self.emit_stackref_override for outp in self.original_uop.stack.outputs
         }
         self._replacers = {**self._replacers, **overrides}
+        self.cannot_escape = True
 
     def emit_to_with_replacement(
         self,