https://github.com/python/cpython/commit/90c44bc803f132b4ea4a09b228f9d5f97ee213f6
commit: 90c44bc803f132b4ea4a09b228f9d5f97ee213f6
branch: main
author: Ken Jin <[email protected]>
committer: Fidget-Spinner <[email protected]>
date: 2026-01-06T16:39:57Z
summary:
gh-131798: Support generator frames in the JIT optimizer (GH-143340)
files:
A
Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst
M Include/internal/pycore_tstate.h
M Lib/test/test_capi/test_opt.py
M Python/optimizer_bytecodes.c
M Python/optimizer_cases.c.h
diff --git a/Include/internal/pycore_tstate.h b/Include/internal/pycore_tstate.h
index 27299d1ebcb57a..d8f4bfef98af7e 100644
--- a/Include/internal/pycore_tstate.h
+++ b/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
diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index 92b3b3454e4bf7..4f092beeea179b 100644
--- a/Lib/test/test_capi/test_opt.py
+++ b/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
diff --git
a/Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst
b/Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst
new file mode 100644
index 00000000000000..528eb70433576e
--- /dev/null
+++
b/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.
diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c
index 7f8d2dbd8e83b8..cbcae5075a279d 100644
--- a/Python/optimizer_bytecodes.c
+++ b/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])) {
diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h
index 37e3ba29fa20f0..96aee7830b9e24 100644
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -1081,9 +1081,22 @@
/* _SEND is not a viable micro-op for tier 2 */
case _SEND_GEN_FRAME: {
+ JitOptRef v;
JitOptRef gen_frame;
- gen_frame = PyJitRef_NULL;
- ctx->done = true;
+ v = stack_pointer[-1];
+ 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);
stack_pointer[-1] = gen_frame;
break;
}
@@ -2267,8 +2280,19 @@
case _FOR_ITER_GEN_FRAME: {
JitOptRef gen_frame;
- gen_frame = PyJitRef_NULL;
- 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);
CHECK_STACK_BOUNDS(1);
stack_pointer[0] = gen_frame;
stack_pointer += 1;
_______________________________________________
Python-checkins mailing list -- [email protected]
To unsubscribe send an email to [email protected]
https://mail.python.org/mailman3//lists/python-checkins.python.org
Member address: [email protected]
