gh-131798: JIT optimizer: Support custom binary op and property frames (GH-143735)

diff --git a/Lib/test/test_capi/test_opt.py b/Lib/test/test_capi/test_opt.py
index a2c041eb4a67d0d5271a517096ba5edf416f43ec..757e5e6ef53574a6486a37e0ca03f1c8a823e603 100644 (file)
--- a/Lib/test/test_capi/test_opt.py
+++ b/Lib/test/test_capi/test_opt.py
@@ -3575,6 +3575,42 @@ class TestUopsOptimization(unittest.TestCase):
         # _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 contradiction.
+        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()))
+        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 contradiction.
+        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 7b386454b76c097480cc23f8864982d5e18ff096..1a64810b50a3a4a689a6de00eacd8ca57ec7e1eb 100644 (file)
--- a/Python/optimizer_bytecodes.c
+++ b/Python/optimizer_bytecodes.c
@@ -327,9 +327,20 @@ 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);
+        if (f == NULL) {
+            break;
+        }
+        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)) {
@@ -761,9 +772,19 @@ dummy_func(void) {
     }
 
     op(_LOAD_ATTR_PROPERTY_FRAME, (fget/4, owner -- new_frame)) {
-        (void)fget;
-        new_frame = PyJitRef_NULL;
-        ctx->done = true;
+        // + 1 for _SAVE_RETURN_OFFSET
+        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);
+        if (f == NULL) {
+            break;
+        }
+        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 accb3e620ebc4e2c4360526df3858f567db75ee2..f3bc7213fcce3f71aa138b87bc52d4739ff96b28 100644 (file)
--- a/Python/optimizer_cases.c.h
+++ b/Python/optimizer_cases.c.h
@@ -1114,9 +1114,24 @@
         }
 
         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);
+            if (f == NULL) {
+                break;
+            }
+            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;
@@ -1947,11 +1962,22 @@
         }
 
         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);
+            if (f == NULL) {
+                break;
+            }
+            f->locals[0] = owner;
+            new_frame = PyJitRef_Wrap((JitOptSymbol *)f);
             stack_pointer[-1] = new_frame;
             break;
         }