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

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

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

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;