gh-115999: Add partial free-thread specialization for BINARY_SUBSCR (gh-127227)

Author: corona10

Include/internal/pycore_list.h

@@ -10,6 +10,9 @@ extern "C" {
 
 PyAPI_FUNC(PyObject*) _PyList_Extend(PyListObject *, PyObject *);
 extern void _PyList_DebugMallocStats(FILE *out);
+// _PyList_GetItemRef should be used only when the object is known as a list
+// because it doesn't raise TypeError when the object is not a list, whereas PyList_GetItemRef does.
+extern PyObject* _PyList_GetItemRef(PyListObject *, Py_ssize_t i);
 
 #define _PyList_ITEMS(op) _Py_RVALUE(_PyList_CAST(op)->ob_item)
 

Include/internal/pycore_opcode_metadata.h

@@ -1260,27 +1260,6 @@ def test_super_instructions(self):
         got = self.get_disassembly(load_test, adaptive=True)
         self.do_disassembly_compare(got, dis_load_test_quickened_code)
 
-    @cpython_only
-    @requires_specialization
-    def test_binary_subscr_specialize(self):
-        binary_subscr_quicken = """\
-  0           RESUME_CHECK             0
-
-  1           LOAD_NAME                0 (a)
-              LOAD_SMALL_INT           0
-              %s
-              RETURN_VALUE
-"""
-        co_list = compile('a[0]', "<list>", "eval")
-        self.code_quicken(lambda: exec(co_list, {}, {'a': [0]}))
-        got = self.get_disassembly(co_list, adaptive=True)
-        self.do_disassembly_compare(got, binary_subscr_quicken % "BINARY_SUBSCR_LIST_INT")
-
-        co_dict = compile('a[0]', "<dict>", "eval")
-        self.code_quicken(lambda: exec(co_dict, {}, {'a': {0: '1'}}))
-        got = self.get_disassembly(co_dict, adaptive=True)
-        self.do_disassembly_compare(got, binary_subscr_quicken % "BINARY_SUBSCR_DICT")
-
     @cpython_only
     @requires_specialization
     def test_load_attr_specialize(self):

Include/internal/pycore_uop_metadata.h

@@ -617,7 +617,7 @@ def write(items):
         opname = "BINARY_SUBSCR_GETITEM"
         self.assert_races_do_not_crash(opname, get_items, read, write)
 
-    @requires_specialization
+    @requires_specialization_ft
     def test_binary_subscr_list_int(self):
         def get_items():
             items = []
@@ -1023,7 +1023,7 @@ def write(items):
         opname = "STORE_ATTR_WITH_HINT"
         self.assert_races_do_not_crash(opname, get_items, read, write)
 
-    @requires_specialization
+    @requires_specialization_ft
     def test_store_subscr_list_int(self):
         def get_items():
             items = []
@@ -1229,48 +1229,48 @@ class TestSpecializer(TestBase):
     @cpython_only
     @requires_specialization_ft
     def test_binary_op(self):
-        def f():
+        def binary_op_add_int():
             for _ in range(100):
                 a, b = 1, 2
                 c = a + b
                 self.assertEqual(c, 3)
 
-        f()
-        self.assert_specialized(f, "BINARY_OP_ADD_INT")
-        self.assert_no_opcode(f, "BINARY_OP")
+        binary_op_add_int()
+        self.assert_specialized(binary_op_add_int, "BINARY_OP_ADD_INT")
+        self.assert_no_opcode(binary_op_add_int, "BINARY_OP")
 
-        def g():
+        def binary_op_add_unicode():
             for _ in range(100):
                 a, b = "foo", "bar"
                 c = a + b
                 self.assertEqual(c, "foobar")
 
-        g()
-        self.assert_specialized(g, "BINARY_OP_ADD_UNICODE")
-        self.assert_no_opcode(g, "BINARY_OP")
+        binary_op_add_unicode()
+        self.assert_specialized(binary_op_add_unicode, "BINARY_OP_ADD_UNICODE")
+        self.assert_no_opcode(binary_op_add_unicode, "BINARY_OP")
 
     @cpython_only
     @requires_specialization_ft
     def test_contain_op(self):
-        def f():
+        def contains_op_dict():
             for _ in range(100):
                 a, b = 1, {1: 2, 2: 5}
                 self.assertTrue(a in b)
                 self.assertFalse(3 in b)
 
-        f()
-        self.assert_specialized(f, "CONTAINS_OP_DICT")
-        self.assert_no_opcode(f, "CONTAINS_OP")
+        contains_op_dict()
+        self.assert_specialized(contains_op_dict, "CONTAINS_OP_DICT")
+        self.assert_no_opcode(contains_op_dict, "CONTAINS_OP")
 
-        def g():
+        def contains_op_set():
             for _ in range(100):
                 a, b = 1, {1, 2}
                 self.assertTrue(a in b)
                 self.assertFalse(3 in b)
 
-        g()
-        self.assert_specialized(g, "CONTAINS_OP_SET")
-        self.assert_no_opcode(g, "CONTAINS_OP")
+        contains_op_set()
+        self.assert_specialized(contains_op_set, "CONTAINS_OP_SET")
+        self.assert_no_opcode(contains_op_set, "CONTAINS_OP")
 
     @cpython_only
     @requires_specialization_ft
@@ -1342,34 +1342,81 @@ def to_bool_str():
     @cpython_only
     @requires_specialization_ft
     def test_unpack_sequence(self):
-        def f():
+        def unpack_sequence_two_tuple():
             for _ in range(100):
                 a, b = 1, 2
                 self.assertEqual(a, 1)
                 self.assertEqual(b, 2)
 
-        f()
-        self.assert_specialized(f, "UNPACK_SEQUENCE_TWO_TUPLE")
-        self.assert_no_opcode(f, "UNPACK_SEQUENCE")
+        unpack_sequence_two_tuple()
+        self.assert_specialized(unpack_sequence_two_tuple,
+                                "UNPACK_SEQUENCE_TWO_TUPLE")
+        self.assert_no_opcode(unpack_sequence_two_tuple, "UNPACK_SEQUENCE")
 
-        def g():
+        def unpack_sequence_tuple():
             for _ in range(100):
                 a, = 1,
                 self.assertEqual(a, 1)
 
-        g()
-        self.assert_specialized(g, "UNPACK_SEQUENCE_TUPLE")
-        self.assert_no_opcode(g, "UNPACK_SEQUENCE")
+        unpack_sequence_tuple()
+        self.assert_specialized(unpack_sequence_tuple, "UNPACK_SEQUENCE_TUPLE")
+        self.assert_no_opcode(unpack_sequence_tuple, "UNPACK_SEQUENCE")
 
-        def x():
+        def unpack_sequence_list():
             for _ in range(100):
                 a, b = [1, 2]
                 self.assertEqual(a, 1)
                 self.assertEqual(b, 2)
 
-        x()
-        self.assert_specialized(x, "UNPACK_SEQUENCE_LIST")
-        self.assert_no_opcode(x, "UNPACK_SEQUENCE")
+        unpack_sequence_list()
+        self.assert_specialized(unpack_sequence_list, "UNPACK_SEQUENCE_LIST")
+        self.assert_no_opcode(unpack_sequence_list, "UNPACK_SEQUENCE")
+
+    @cpython_only
+    @requires_specialization_ft
+    def test_binary_subscr(self):
+        def binary_subscr_list_int():
+            for _ in range(100):
+                a = [1, 2, 3]
+                for idx, expected in enumerate(a):
+                    self.assertEqual(a[idx], expected)
+
+        binary_subscr_list_int()
+        self.assert_specialized(binary_subscr_list_int,
+                                "BINARY_SUBSCR_LIST_INT")
+        self.assert_no_opcode(binary_subscr_list_int, "BINARY_SUBSCR")
+
+        def binary_subscr_tuple_int():
+            for _ in range(100):
+                a = (1, 2, 3)
+                for idx, expected in enumerate(a):
+                    self.assertEqual(a[idx], expected)
+
+        binary_subscr_tuple_int()
+        self.assert_specialized(binary_subscr_tuple_int,
+                                "BINARY_SUBSCR_TUPLE_INT")
+        self.assert_no_opcode(binary_subscr_tuple_int, "BINARY_SUBSCR")
+
+        def binary_subscr_dict():
+            for _ in range(100):
+                a = {1: 2, 2: 3}
+                self.assertEqual(a[1], 2)
+                self.assertEqual(a[2], 3)
+
+        binary_subscr_dict()
+        self.assert_specialized(binary_subscr_dict, "BINARY_SUBSCR_DICT")
+        self.assert_no_opcode(binary_subscr_dict, "BINARY_SUBSCR")
+
+        def binary_subscr_str_int():
+            for _ in range(100):
+                a = "foobar"
+                for idx, expected in enumerate(a):
+                    self.assertEqual(a[idx], expected)
+
+        binary_subscr_str_int()
+        self.assert_specialized(binary_subscr_str_int, "BINARY_SUBSCR_STR_INT")
+        self.assert_no_opcode(binary_subscr_str_int, "BINARY_SUBSCR")
+
 
 if __name__ == "__main__":
     unittest.main()

Lib/test/test_dis.py

@@ -391,6 +391,12 @@ PyList_GetItemRef(PyObject *op, Py_ssize_t i)
     return item;
 }
 
+PyObject *
+_PyList_GetItemRef(PyListObject *list, Py_ssize_t i)
+{
+    return list_get_item_ref(list, i);
+}
+
 int
 PyList_SetItem(PyObject *op, Py_ssize_t i,
                PyObject *newitem)

Lib/test/test_opcache.py

@@ -704,7 +704,7 @@ dummy_func(
         };
 
         specializing op(_SPECIALIZE_BINARY_SUBSCR, (counter/1, container, sub -- container, sub)) {
-            #if ENABLE_SPECIALIZATION
+            #if ENABLE_SPECIALIZATION_FT
             assert(frame->stackpointer == NULL);
             if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                 next_instr = this_instr;
@@ -713,7 +713,7 @@ dummy_func(
             }
             OPCODE_DEFERRED_INC(BINARY_SUBSCR);
             ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-            #endif  /* ENABLE_SPECIALIZATION */
+            #endif  /* ENABLE_SPECIALIZATION_FT */
         }
 
         op(_BINARY_SUBSCR, (container, sub -- res)) {
@@ -790,11 +790,17 @@ dummy_func(
             // Deopt unless 0 <= sub < PyList_Size(list)
             DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub));
             Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
+#ifdef Py_GIL_DISABLED
+            PyObject *res_o = _PyList_GetItemRef((PyListObject*)list, index);
+            DEOPT_IF(res_o == NULL);
+            STAT_INC(BINARY_SUBSCR, hit);
+#else
             DEOPT_IF(index >= PyList_GET_SIZE(list));
             STAT_INC(BINARY_SUBSCR, hit);
             PyObject *res_o = PyList_GET_ITEM(list, index);
             assert(res_o != NULL);
             Py_INCREF(res_o);
+#endif
             PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
             DEAD(sub_st);
             PyStackRef_CLOSE(list_st);