From a9c244ae9199c41471af6767e194173632a883e0 Mon Sep 17 00:00:00 2001 From: Ken Jin Date: Mon, 12 Jan 2026 16:40:24 +0000 Subject: [PATCH 1/2] JIT: Support custom binary op and property frames --- Lib/test/test_capi/test_opt.py | 38 ++++++++++++++++++++++++++++++++++ Python/optimizer_bytecodes.c | 26 +++++++++++++++++------ Python/optimizer_cases.c.h | 30 ++++++++++++++++++++++----- 3 files changed, 83 insertions(+), 11 deletions(-) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index 79c5c22c13f557..faee49d3787eda 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -3487,6 +3487,44 @@ def testfunc(n): # _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_binary_op_subscr_init_frame(self): + class B: + def __getitem__(self, other): + return other + 1 + def testfunc(*args): + n, b = args[0] + for _ in range(n): + y = b[2] + + res, ex = self._run_with_optimizer(testfunc, (TIER2_THRESHOLD, B())) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + self.assertIn("_BINARY_OP_SUBSCR_INIT_CALL", 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_load_attr_property_frame(self): + class B: + @property + def prop(self): + return 3 + def testfunc(*args): + n, b = args[0] + for _ in range(n): + y = b.prop + b.prop + + testfunc((3, B())) + import dis + dis.dis(testfunc, adaptive=True) + res, ex = self._run_with_optimizer(testfunc, (TIER2_THRESHOLD, B())) + self.assertIsNotNone(ex) + uops = get_opnames(ex) + + self.assertIn("_LOAD_ATTR_PROPERTY_FRAME", uops) + # This is a sign the optimizer ran and didn't hit bottom. + self.assertIn("_INSERT_2_LOAD_CONST_INLINE_BORROW", uops) + def test_unary_negative(self): def testfunc(n): a = 3 diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index e85536bfc3a493..a70527f7725d1a 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -327,9 +327,17 @@ dummy_func(void) { GETLOCAL(this_instr->operand0) = sym_new_null(ctx); } - op(_BINARY_OP_SUBSCR_INIT_CALL, (container, sub, getitem -- new_frame)) { - new_frame = PyJitRef_NULL; - ctx->done = true; + op(_BINARY_OP_SUBSCR_INIT_CALL, (container, sub, getitem -- new_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 *f = frame_new(ctx, co, 0, NULL, 0); + f->locals[0] = container; + f->locals[1] = sub; + new_frame = PyJitRef_Wrap((JitOptSymbol *)f); } op(_BINARY_OP_SUBSCR_STR_INT, (str_st, sub_st -- res, s, i)) { @@ -753,9 +761,15 @@ dummy_func(void) { } op(_LOAD_ATTR_PROPERTY_FRAME, (fget/4, owner -- new_frame)) { - (void)fget; - new_frame = PyJitRef_NULL; - ctx->done = true; + assert((this_instr + 2)->opcode == _PUSH_FRAME); + PyCodeObject *co = get_code_with_logging(this_instr + 2); + if (co == NULL) { + ctx->done = true; + break; + } + _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0); + f->locals[0] = owner; + new_frame = PyJitRef_Wrap((JitOptSymbol *)f); } op(_INIT_CALL_BOUND_METHOD_EXACT_ARGS, (callable, self_or_null, unused[oparg] -- callable, self_or_null, unused[oparg])) { diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index 898930a01b16af..1aa70aa5cf5eff 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -1102,9 +1102,21 @@ } case _BINARY_OP_SUBSCR_INIT_CALL: { + JitOptRef sub; + JitOptRef container; JitOptRef new_frame; - new_frame = PyJitRef_NULL; - ctx->done = true; + sub = stack_pointer[-2]; + container = stack_pointer[-3]; + 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 *f = frame_new(ctx, co, 0, NULL, 0); + f->locals[0] = container; + f->locals[1] = sub; + new_frame = PyJitRef_Wrap((JitOptSymbol *)f); CHECK_STACK_BOUNDS(-2); stack_pointer[-3] = new_frame; stack_pointer += -2; @@ -1935,11 +1947,19 @@ } case _LOAD_ATTR_PROPERTY_FRAME: { + JitOptRef owner; JitOptRef new_frame; + owner = stack_pointer[-1]; PyObject *fget = (PyObject *)this_instr->operand0; - (void)fget; - new_frame = PyJitRef_NULL; - ctx->done = true; + assert((this_instr + 2)->opcode == _PUSH_FRAME); + PyCodeObject *co = get_code_with_logging(this_instr + 2); + if (co == NULL) { + ctx->done = true; + break; + } + _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0); + f->locals[0] = owner; + new_frame = PyJitRef_Wrap((JitOptSymbol *)f); stack_pointer[-1] = new_frame; break; } From 0f3e071e06a0b7f7f244e5785ce84119baf88c0d Mon Sep 17 00:00:00 2001 From: Ken Jin Date: Tue, 13 Jan 2026 08:48:30 +0000 Subject: [PATCH 2/2] Address code review --- Lib/test/test_capi/test_opt.py | 6 ++---- Python/optimizer_bytecodes.c | 7 +++++++ Python/optimizer_cases.c.h | 6 ++++++ 3 files changed, 15 insertions(+), 4 deletions(-) diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py index 9fe27af8eca017..71a525a9534d57 100644 --- a/Lib/test/test_capi/test_opt.py +++ b/Lib/test/test_capi/test_opt.py @@ -3552,7 +3552,7 @@ def testfunc(*args): uops = get_opnames(ex) self.assertIn("_BINARY_OP_SUBSCR_INIT_CALL", uops) - # _POP_TOP_NOP is a sign the optimizer ran and didn't hit bottom. + # _POP_TOP_NOP is a sign the optimizer ran and didn't hit contradiction. self.assertGreaterEqual(count_ops(ex, "_POP_TOP_NOP"), 1) def test_load_attr_property_frame(self): @@ -3566,14 +3566,12 @@ def testfunc(*args): y = b.prop + b.prop testfunc((3, B())) - import dis - dis.dis(testfunc, adaptive=True) res, ex = self._run_with_optimizer(testfunc, (TIER2_THRESHOLD, B())) self.assertIsNotNone(ex) uops = get_opnames(ex) self.assertIn("_LOAD_ATTR_PROPERTY_FRAME", uops) - # This is a sign the optimizer ran and didn't hit bottom. + # This is a sign the optimizer ran and didn't hit contradiction. self.assertIn("_INSERT_2_LOAD_CONST_INLINE_BORROW", uops) def test_unary_negative(self): diff --git a/Python/optimizer_bytecodes.c b/Python/optimizer_bytecodes.c index 5e3031c16b0b9e..f76d0d5fe51c99 100644 --- a/Python/optimizer_bytecodes.c +++ b/Python/optimizer_bytecodes.c @@ -335,6 +335,9 @@ dummy_func(void) { break; } _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0); + if (f == NULL) { + break; + } f->locals[0] = container; f->locals[1] = sub; new_frame = PyJitRef_Wrap((JitOptSymbol *)f); @@ -767,6 +770,7 @@ dummy_func(void) { } op(_LOAD_ATTR_PROPERTY_FRAME, (fget/4, owner -- new_frame)) { + // + 1 for _SAVE_RETURN_OFFSET assert((this_instr + 2)->opcode == _PUSH_FRAME); PyCodeObject *co = get_code_with_logging(this_instr + 2); if (co == NULL) { @@ -774,6 +778,9 @@ dummy_func(void) { break; } _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0); + if (f == NULL) { + break; + } f->locals[0] = owner; new_frame = PyJitRef_Wrap((JitOptSymbol *)f); } diff --git a/Python/optimizer_cases.c.h b/Python/optimizer_cases.c.h index 4b1e89422ca038..3db380e3cc2e32 100644 --- a/Python/optimizer_cases.c.h +++ b/Python/optimizer_cases.c.h @@ -1114,6 +1114,9 @@ break; } _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0); + if (f == NULL) { + break; + } f->locals[0] = container; f->locals[1] = sub; new_frame = PyJitRef_Wrap((JitOptSymbol *)f); @@ -1958,6 +1961,9 @@ break; } _Py_UOpsAbstractFrame *f = frame_new(ctx, co, 0, NULL, 0); + if (f == NULL) { + break; + } f->locals[0] = owner; new_frame = PyJitRef_Wrap((JitOptSymbol *)f); stack_pointer[-1] = new_frame;