mirror/cpython
1
0
Fork 0
mirror of https://github.com/python/cpython.git synced 2026-01-26 12:55:08 +00:00
Code Issues Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Ken Jin 2026-01-07 00:39:57 +08:00 committed by GitHub
parent faa3dc7c64
commit 90c44bc803
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 102 additions and 27 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
Include/internal/pycore_tstate.h
View File

@ -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

57
Lib/test/test_capi/test_opt.py
View File

@ -1168,22 +1168,6 @@ class TestUopsOptimization(unittest.TestCase):
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 @@ class TestUopsOptimization(unittest.TestCase):
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

1
Misc/NEWS.d/next/Core_and_Builtins/2026-01-01-23-41-50.gh-issue-131798.QUqDdK.rst Normal file
View File

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

34
Python/optimizer_bytecodes.c
View File

@ -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])) {

32
Python/optimizer_cases.c.h generated
View File

@ -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;