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 19d9ea9eab289f..8771faf7f47fde 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 @@ -3385,6 +3369,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 e92eb2b75a291a..e56a013a2b0d0c 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 52360f8119df4c..ababa13aca6bbc 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -1073,9 +1073,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; } @@ -2259,8 +2272,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;