gh-131798: Support generator frames in the JIT optimizer (GH-143340)

Author: Fidget-Spinner

Include/internal/pycore_tstate.h

@@ -47,10 +47,15 @@ typedef struct _PyJitTracerPreviousState {
     _PyBloomFilter dependencies;
 } _PyJitTracerPreviousState;
 
+typedef struct _PyJitTracerTranslatorState {
+    int jump_backward_seen;
+} _PyJitTracerTranslatorState;
+
 typedef struct _PyJitTracerState {
     _PyUOpInstruction *code_buffer;
     _PyJitTracerInitialState initial_state;
     _PyJitTracerPreviousState prev_state;
+    _PyJitTracerTranslatorState translator_state;
 } _PyJitTracerState;
 
 #endif

Lib/test/test_capi/test_opt.py

@@ -1168,22 +1168,6 @@ def testfunc(n):
         self.assertIsNotNone(ex)
         self.assertIn("_FOR_ITER_TIER_TWO", get_opnames(ex))
 
-    @unittest.skip("Tracing into generators currently isn't supported.")
-    def test_for_iter_gen(self):
-        def gen(n):
-            for i in range(n):
-                yield i
-        def testfunc(n):
-            g = gen(n)
-            s = 0
-            for i in g:
-                s += i
-            return s
-        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD)
-        self.assertEqual(res, sum(range(TIER2_THRESHOLD)))
-        self.assertIsNotNone(ex)
-        self.assertIn("_FOR_ITER_GEN_FRAME", get_opnames(ex))
-
     def test_modified_local_is_seen_by_optimized_code(self):
         l = sys._getframe().f_locals
         a = 1
@@ -3404,6 +3388,47 @@ def test_is_none(n):
         self.assertIn("_POP_TOP_NOP", uops)
         self.assertLessEqual(count_ops(ex, "_POP_TOP"), 2)
 
+    def test_for_iter_gen_frame(self):
+        def f(n):
+            for i in range(n):
+                # Should be optimized to POP_TOP_NOP
+                yield i + i
+        def testfunc(n):
+            for _ in f(n):
+                pass
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD*2)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_FOR_ITER_GEN_FRAME", uops)
+        # _POP_TOP_NOP is a sign the optimizer ran and didn't hit bottom.
+        self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
+
+    def test_send_gen_frame(self):
+
+        def gen(n):
+            for i in range(n):
+                yield i + i
+        def send_gen(n):
+            yield from gen(n)
+        def testfunc(n):
+            for _ in send_gen(n):
+                pass
+
+        for _ in range(_testinternalcapi.SPECIALIZATION_THRESHOLD):
+            # Ensure SEND is specialized to SEND_GEN
+            send_gen(10)
+
+        res, ex = self._run_with_optimizer(testfunc, TIER2_THRESHOLD*2)
+        self.assertIsNotNone(ex)
+        uops = get_opnames(ex)
+
+        self.assertIn("_FOR_ITER_GEN_FRAME", uops)
+        self.assertIn("_SEND_GEN_FRAME", uops)
+        # _POP_TOP_NOP is a sign the optimizer ran and didn't hit bottom.
+        self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1)
+
     def test_143026(self):
         # https://github.com/python/cpython/issues/143026
 

Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst

@@ -0,0 +1 @@
+The JIT optimizer now understands more generator instructions.

Python/optimizer_bytecodes.c

@@ -939,15 +939,35 @@ dummy_func(void) {
     }
 
     op(_FOR_ITER_GEN_FRAME, (unused, unused -- unused, unused, gen_frame)) {
-        gen_frame = PyJitRef_NULL;
-        /* We are about to hit the end of the trace */
-        ctx->done = true;
+        assert((this_instr + 1)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging((this_instr + 1));
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+        _Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
+        if (new_frame == NULL) {
+            ctx->done = true;
+            break;
+        }
+        new_frame->stack[0] = sym_new_const(ctx, Py_None);
+        gen_frame = PyJitRef_Wrap((JitOptSymbol *)new_frame);
     }
 
-    op(_SEND_GEN_FRAME, (unused, unused -- unused, gen_frame)) {
-        gen_frame = PyJitRef_NULL;
-        // We are about to hit the end of the trace:
-        ctx->done = true;
+    op(_SEND_GEN_FRAME, (unused, v -- unused, gen_frame)) {
+        assert((this_instr + 1)->opcode == _PUSH_FRAME);
+        PyCodeObject *co = get_code_with_logging((this_instr + 1));
+        if (co == NULL) {
+            ctx->done = true;
+            break;
+        }
+        _Py_UOpsAbstractFrame *new_frame = frame_new(ctx, co, 1, NULL, 0);
+        if (new_frame == NULL) {
+            ctx->done = true;
+            break;
+        }
+        new_frame->stack[0] = PyJitRef_StripReferenceInfo(v);
+        gen_frame = PyJitRef_Wrap((JitOptSymbol *)new_frame);
     }
 
     op(_CHECK_STACK_SPACE, (unused, unused, unused[oparg] -- unused, unused, unused[oparg])) {