gh-115999: Add free-threaded specialization for FOR_ITER (#128798)

Add free-threaded versions of existing specialization for FOR_ITER (list, tuples, fast range iterators and generators), without significantly affecting their thread-safety. (Iterating over shared lists/tuples/ranges should be fine like before. Reusing iterators between threads is not fine, like before. Sharing generators between threads is a recipe for significant crashes, like before.)
This commit is contained in:
T. Wouters 2025-03-12 16:21:46 +01:00 committed by GitHub
parent db27aee2fe
commit de2f7da77d
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
13 changed files with 469 additions and 125 deletions

View file

@ -3005,7 +3005,7 @@ 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);
@ -3013,7 +3013,7 @@ dummy_func(
}
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)) {
@ -3091,31 +3091,46 @@ 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
EXIT_IF(!_PyObject_IsUniquelyReferenced(iter_o));
_PyListIterObject *it = (_PyListIterObject *)iter_o;
EXIT_IF(!_Py_IsOwnedByCurrentThread((PyObject *)it->it_seq) ||
!_PyObject_GC_IS_SHARED(it->it_seq));
#endif
}
replaced op(_ITER_JUMP_LIST, (iter -- iter)) {
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
_PyListIterObject *it = (_PyListIterObject *)iter_o;
assert(Py_TYPE(iter_o) == &PyListIter_Type);
// For free-threaded Python, the loop exit can happen at any point during
// item retrieval, so it doesn't make much sense to check and jump
// separately before item retrieval. Any length check we do here can be
// invalid by the time we actually try to fetch the item.
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced(iter_o));
(void)iter_o;
#else
_PyListIterObject *it = (_PyListIterObject *)iter_o;
STAT_INC(FOR_ITER, hit);
PyListObject *seq = it->it_seq;
if (seq == NULL || (size_t)it->it_index >= (size_t)PyList_GET_SIZE(seq)) {
it->it_index = -1;
#ifndef Py_GIL_DISABLED
if (seq != NULL) {
it->it_seq = NULL;
Py_DECREF(seq);
}
#endif
/* Jump forward oparg, then skip following END_FOR instruction */
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);
@ -3125,16 +3140,62 @@ dummy_func(
it->it_index = -1;
EXIT_IF(1);
}
#endif
}
op(_ITER_NEXT_LIST, (iter -- iter, next)) {
replaced op(_ITER_NEXT_LIST, (iter -- iter, next)) {
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
_PyListIterObject *it = (_PyListIterObject *)iter_o;
assert(Py_TYPE(iter_o) == &PyListIter_Type);
PyListObject *seq = it->it_seq;
assert(seq);
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced(iter_o));
assert(_Py_IsOwnedByCurrentThread((PyObject *)seq) ||
_PyObject_GC_IS_SHARED(seq));
STAT_INC(FOR_ITER, hit);
int result = _PyList_GetItemRefNoLock(seq, it->it_index, &next);
// A negative result means we lost a race with another thread
// and we need to take the slow path.
DEOPT_IF(result < 0);
if (result == 0) {
it->it_index = -1;
/* Jump forward oparg, then skip following END_FOR instruction */
JUMPBY(oparg + 1);
DISPATCH();
}
it->it_index++;
#else
assert(it->it_index < PyList_GET_SIZE(seq));
next = PyStackRef_FromPyObjectNew(PyList_GET_ITEM(seq, it->it_index++));
#endif
}
// Only used by Tier 2
op(_ITER_NEXT_LIST_TIER_TWO, (iter -- iter, next)) {
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
_PyListIterObject *it = (_PyListIterObject *)iter_o;
assert(Py_TYPE(iter_o) == &PyListIter_Type);
PyListObject *seq = it->it_seq;
assert(seq);
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced(iter_o));
assert(_Py_IsOwnedByCurrentThread((PyObject *)seq) ||
_PyObject_GC_IS_SHARED(seq));
STAT_INC(FOR_ITER, hit);
int result = _PyList_GetItemRefNoLock(seq, it->it_index, &next);
// 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) {
it->it_index = -1;
EXIT_IF(1);
}
it->it_index++;
#else
assert(it->it_index < PyList_GET_SIZE(seq));
next = PyStackRef_FromPyObjectNew(PyList_GET_ITEM(seq, it->it_index++));
#endif
}
macro(FOR_ITER_LIST) =
@ -3144,20 +3205,30 @@ dummy_func(
_ITER_NEXT_LIST;
op(_ITER_CHECK_TUPLE, (iter -- iter)) {
EXIT_IF(Py_TYPE(PyStackRef_AsPyObjectBorrow(iter)) != &PyTupleIter_Type);
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
EXIT_IF(Py_TYPE(iter_o) != &PyTupleIter_Type);
#ifdef Py_GIL_DISABLED
EXIT_IF(!_PyObject_IsUniquelyReferenced(iter_o));
#endif
}
replaced op(_ITER_JUMP_TUPLE, (iter -- iter)) {
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
_PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
(void)iter_o;
assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced(iter_o));
#endif
_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)) {
if (seq == NULL || (size_t)it->it_index >= (size_t)PyTuple_GET_SIZE(seq)) {
#ifndef Py_GIL_DISABLED
if (seq != NULL) {
it->it_seq = NULL;
Py_DECREF(seq);
}
#endif
/* Jump forward oparg, then skip following END_FOR instruction */
JUMPBY(oparg + 1);
DISPATCH();
@ -3169,6 +3240,9 @@ dummy_func(
PyObject *iter_o = PyStackRef_AsPyObjectBorrow(iter);
_PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced(iter_o));
#endif
PyTupleObject *seq = it->it_seq;
EXIT_IF(seq == NULL);
EXIT_IF(it->it_index >= PyTuple_GET_SIZE(seq));
@ -3179,6 +3253,9 @@ dummy_func(
_PyTupleIterObject *it = (_PyTupleIterObject *)iter_o;
assert(Py_TYPE(iter_o) == &PyTupleIter_Type);
PyTupleObject *seq = it->it_seq;
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced(iter_o));
#endif
assert(seq);
assert(it->it_index < PyTuple_GET_SIZE(seq));
next = PyStackRef_FromPyObjectNew(PyTuple_GET_ITEM(seq, it->it_index++));
@ -3193,11 +3270,17 @@ dummy_func(
op(_ITER_CHECK_RANGE, (iter -- iter)) {
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
EXIT_IF(Py_TYPE(r) != &PyRangeIter_Type);
#ifdef Py_GIL_DISABLED
EXIT_IF(!_PyObject_IsUniquelyReferenced((PyObject *)r));
#endif
}
replaced op(_ITER_JUMP_RANGE, (iter -- iter)) {
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
assert(Py_TYPE(r) == &PyRangeIter_Type);
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced((PyObject *)r));
#endif
STAT_INC(FOR_ITER, hit);
if (r->len <= 0) {
// Jump over END_FOR instruction.
@ -3216,6 +3299,9 @@ dummy_func(
op(_ITER_NEXT_RANGE, (iter -- iter, next)) {
_PyRangeIterObject *r = (_PyRangeIterObject *)PyStackRef_AsPyObjectBorrow(iter);
assert(Py_TYPE(r) == &PyRangeIter_Type);
#ifdef Py_GIL_DISABLED
assert(_PyObject_IsUniquelyReferenced((PyObject *)r));
#endif
assert(r->len > 0);
long value = r->start;
r->start = value + r->step;
@ -3234,6 +3320,13 @@ dummy_func(
op(_FOR_ITER_GEN_FRAME, (iter -- iter, gen_frame: _PyInterpreterFrame*)) {
PyGenObject *gen = (PyGenObject *)PyStackRef_AsPyObjectBorrow(iter);
DEOPT_IF(Py_TYPE(gen) != &PyGen_Type);
#ifdef Py_GIL_DISABLED
// Since generators can't be used by multiple threads anyway we
// don't need to deopt here, but this lets us work on making
// generators thread-safe without necessarily having to
// specialize them thread-safely as well.
DEOPT_IF(!_PyObject_IsUniquelyReferenced((PyObject *)gen));
#endif
DEOPT_IF(gen->gi_frame_state >= FRAME_EXECUTING);
STAT_INC(FOR_ITER, hit);
gen_frame = &gen->gi_iframe;