Add free-threaded specialization for lists and tuples.

Author: Yhg1s

Include/internal/pycore_list.h

@@ -13,6 +13,10 @@ 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);
+#ifdef Py_GIL_DISABLED
+// Returns -1 in case of races with other threads.
+extern int _PyList_GetItemRefNoLock(PyListObject *, Py_ssize_t i, PyObject ** result);
+#endif
 
 #define _PyList_ITEMS(op) _Py_RVALUE(_PyList_CAST(op)->ob_item)
 

Include/internal/pycore_uop_metadata.h

@@ -26,6 +26,7 @@
     'test_threadsignals',
     'test_weakref',
     'test_free_threading.test_slots',
+    'test_free_threading.test_iteration',
 ]
 
 

Lib/test/libregrtest/tsan.py

@@ -357,7 +357,7 @@ list_get_item_ref(PyListObject *op, Py_ssize_t i)
         return NULL;
     }
     Py_ssize_t cap = list_capacity(ob_item);
-    assert(cap != -1 && cap >= size);
+    assert(cap != -1);
     if (!valid_index(i, cap)) {
         return NULL;
     }
@@ -415,6 +415,32 @@ _PyList_GetItemRef(PyListObject *list, Py_ssize_t i)
     return list_get_item_ref(list, i);
 }
 
+#ifdef Py_GIL_DISABLED
+int
+_PyList_GetItemRefNoLock(PyListObject *list, Py_ssize_t i, PyObject **result)
+{
+    assert(_Py_IsOwnedByCurrentThread((PyObject *)list) ||
+           _PyObject_GC_IS_SHARED(list));
+    if (!valid_index(i, PyList_GET_SIZE(list))) {
+        return 0;
+    }
+    PyObject **ob_item = _Py_atomic_load_ptr(&list->ob_item);
+    if (ob_item == NULL) {
+        return 0;
+    }
+    Py_ssize_t cap = list_capacity(ob_item);
+    assert(cap != -1);
+    if (!valid_index(i, cap)) {
+        return 0;
+    }
+    *result = _Py_TryXGetRef(&ob_item[i]);
+    if (*result == NULL) {
+        return -1;
+    }
+    return 1;
+}
+#endif
+
 int
 PyList_SetItem(PyObject *op, Py_ssize_t i,
                PyObject *newitem)

Objects/listobject.c

@@ -1020,13 +1020,15 @@ tupleiter_next(PyObject *self)
         return NULL;
     assert(PyTuple_Check(seq));
 
-    if (it->it_index < PyTuple_GET_SIZE(seq)) {
-        item = PyTuple_GET_ITEM(seq, it->it_index);
-        ++it->it_index;
+    Py_ssize_t idx = FT_ATOMIC_LOAD_SSIZE_RELAXED(it->it_index);
+    if ((size_t)idx < (size_t)PyTuple_GET_SIZE(seq)) {
+        item = PyTuple_GET_ITEM(seq, idx++);
+        FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, idx);
         return Py_NewRef(item);
     }
 
-    it->it_seq = NULL;
+    FT_ATOMIC_STORE_SSIZE_RELAXED(it->it_index, -1);
+    FT_ATOMIC_STORE_PTR_RELAXED(it->it_seq, NULL);
     Py_DECREF(seq);
     return NULL;
 }

Objects/tupleobject.c

@@ -2929,15 +2929,15 @@ dummy_func(
         };
 
         specializing op(_SPECIALIZE_FOR_ITER, (counter/1, iter -- iter)) {
-            #if ENABLE_SPECIALIZATION
+            #if ENABLE_SPECIALIZATION_FT
             if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
                 next_instr = this_instr;
                 _Py_Specialize_ForIter(iter, next_instr, oparg);
                 DISPATCH_SAME_OPARG();
             }
             OPCODE_DEFERRED_INC(FOR_ITER);
             ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
-            #endif  /* ENABLE_SPECIALIZATION */
+            #endif  /* ENABLE_SPECIALIZATION_FT */
         }
 
         replaced op(_FOR_ITER, (iter -- iter, next)) {
@@ -3015,10 +3015,20 @@ dummy_func(
 
 
         op(_ITER_CHECK_LIST, (iter -- iter)) {
-            EXIT_IF(Py_TYPE(PyStackRef_AsPyObjectBorrow(iter)) != &PyListIter_Type);
+            PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
+            EXIT_IF(Py_TYPE(iter_o) != &PyListIter_Type);
+#ifdef Py_GIL_DISABLED
+            _PyListIterObject *it = (_PyListIterObject *)iter_o;
+            EXIT_IF(it->it_seq == NULL ||
+                    !_Py_IsOwnedByCurrentThread((PyObject *)it->it_seq) ||
+                    !_PyObject_GC_IS_SHARED(it->it_seq));
+#endif
         }
 
         replaced op(_ITER_JUMP_LIST, (iter -- iter)) {
+// For free-threaded Python, the loop exit can happen at any point during item
+// retrieval, so separate ops don't make much sense.
+#ifndef Py_GIL_DISABLED
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyListIterObject *it = (_PyListIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyListIter_Type);
@@ -3036,10 +3046,12 @@ dummy_func(
                 JUMPBY(oparg + 1);
                 DISPATCH();
             }
+#endif
         }
 
         // Only used by Tier 2
         op(_GUARD_NOT_EXHAUSTED_LIST, (iter -- iter)) {
+#ifndef Py_GIL_DISABLED
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyListIterObject *it = (_PyListIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyListIter_Type);
@@ -3049,6 +3061,7 @@ dummy_func(
                 it->it_index = -1;
                 EXIT_IF(1);
             }
+#endif
         }
 
         op(_ITER_NEXT_LIST, (iter -- iter, next)) {
@@ -3057,8 +3070,30 @@ dummy_func(
             assert(Py_TYPE(iter_o) == &PyListIter_Type);
             PyListObject *seq = it->it_seq;
             assert(seq);
+#ifdef Py_GIL_DISABLED
+            assert(_Py_IsOwnedByCurrentThread((PyObject *)seq) ||
+                   _PyObject_GC_IS_SHARED(seq));
+            STAT_INC(FOR_ITER, hit);
+            Py_ssize_t idx = _Py_atomic_load_ssize_relaxed(&it->it_index);
+            PyObject *item;
+            int result = _PyList_GetItemRefNoLock(it->it_seq, idx, &item);
+            // A negative result means we lost a race with another thread
+            // and we need to take the slow path.
+            EXIT_IF(result < 0);
+            if (result == 0) {
+                _Py_atomic_store_ssize_relaxed(&it->it_index, -1);
+                PyStackRef_CLOSE(iter);
+                STACK_SHRINK(1);
+                /* Jump forward oparg, then skip following END_FOR and POP_TOP instructions */
+                JUMPBY(oparg + 2);
+                DISPATCH();
+            }
+            _Py_atomic_store_ssize_relaxed(&it->it_index, idx + 1);
+            next = PyStackRef_FromPyObjectSteal(item);
+#else
             assert(it->it_index < PyList_GET_SIZE(seq));
             next = PyStackRef_FromPyObjectNew(PyList_GET_ITEM(seq, it->it_index++));
+#endif
         }
 
         macro(FOR_ITER_LIST) =
@@ -3073,8 +3108,11 @@ dummy_func(
 
         replaced op(_ITER_JUMP_TUPLE, (iter -- iter)) {
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
-            _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
+// For free-threaded Python, the loop exit can happen at any point during item
+// retrieval, so separate ops don't make much sense.
+#ifndef Py_GIL_DISABLED
+            _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             STAT_INC(FOR_ITER, hit);
             PyTupleObject *seq = it->it_seq;
             if (seq == NULL || it->it_index >= PyTuple_GET_SIZE(seq)) {
@@ -3086,26 +3124,44 @@ dummy_func(
                 JUMPBY(oparg + 1);
                 DISPATCH();
             }
+#endif
         }
 
         // Only used by Tier 2
         op(_GUARD_NOT_EXHAUSTED_TUPLE, (iter -- iter)) {
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
+#ifndef Py_GIL_DISABLED
             PyTupleObject *seq = it->it_seq;
             EXIT_IF(seq == NULL);
             EXIT_IF(it->it_index >= PyTuple_GET_SIZE(seq));
+#ifndef Py_GIL_DISABLED
         }
 
         op(_ITER_NEXT_TUPLE, (iter -- iter, next)) {
             PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
             _PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
             assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
             PyTupleObject *seq = it->it_seq;
+#ifdef Py_GIL_DISABLED
+            STAT_INC(FOR_ITER, hit);
+            Py_ssize_t idx = _Py_atomic_load_ssize_relaxed(&it->it_index);
+            if (seq == NULL || (size_t)idx >= (size_t)PyTuple_GET_SIZE(seq)) {
+                _Py_atomic_store_ssize_relaxed(&it->it_index, -1);
+                PyStackRef_CLOSE(iter);
+                STACK_SHRINK(1);
+                /* Jump forward oparg, then skip following END_FOR and POP_TOP instructions */
+                JUMPBY(oparg + 2);
+                DISPATCH();
+            }
+            _Py_atomic_store_ssize_relaxed(&it->it_index, idx + 1);
+            next = PyStackRef_FromPyObjectNew(PyTuple_GET_ITEM(seq, idx));
+#else
             assert(seq);
             assert(it->it_index < PyTuple_GET_SIZE(seq));
             next = PyStackRef_FromPyObjectNew(PyTuple_GET_ITEM(seq, it->it_index++));
+#endif
         }
 
         macro(FOR_ITER_TUPLE) =