gh-115999: Add free-threaded specialization for `STORE_SUBSCR` (#127169)

The specialization only depends on the type, so no special thread-safety
considerations there.

STORE_SUBSCR_LIST_INT needs to lock the list before modifying it.

`_PyDict_SetItem_Take2` already internally locks the dictionary using a
critical section.
This commit is contained in:
Sam Gross 2024-11-26 21:46:06 +00:00 committed by GitHub
parent f0d3f10c43
commit 71ede1142d
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 109 additions and 70 deletions

View File

@ -910,7 +910,7 @@ dummy_func(
};
specializing op(_SPECIALIZE_STORE_SUBSCR, (counter/1, container, sub -- container, sub)) {
#if ENABLE_SPECIALIZATION
#if ENABLE_SPECIALIZATION_FT
if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
next_instr = this_instr;
_Py_Specialize_StoreSubscr(container, sub, next_instr);
@ -918,7 +918,7 @@ dummy_func(
}
OPCODE_DEFERRED_INC(STORE_SUBSCR);
ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
#endif /* ENABLE_SPECIALIZATION */
#endif /* ENABLE_SPECIALIZATION_FT */
}
op(_STORE_SUBSCR, (v, container, sub -- )) {
@ -940,13 +940,18 @@ dummy_func(
// Ensure nonnegative, zero-or-one-digit ints.
DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub));
Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
DEOPT_IF(!LOCK_OBJECT(list));
// Ensure index < len(list)
DEOPT_IF(index >= PyList_GET_SIZE(list));
if (index >= PyList_GET_SIZE(list)) {
UNLOCK_OBJECT(list);
DEOPT_IF(true);
}
STAT_INC(STORE_SUBSCR, hit);
PyObject *old_value = PyList_GET_ITEM(list, index);
PyList_SET_ITEM(list, index, PyStackRef_AsPyObjectSteal(value));
assert(old_value != NULL);
UNLOCK_OBJECT(list); // unlock before decrefs!
Py_DECREF(old_value);
PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
DEAD(sub_st);

View File

@ -284,6 +284,29 @@ GETITEM(PyObject *v, Py_ssize_t i) {
}
// Try to lock an object in the free threading build, if it's not already
// locked. Use with a DEOPT_IF() to deopt if the object is already locked.
// These are no-ops in the default GIL build. The general pattern is:
//
// DEOPT_IF(!LOCK_OBJECT(op));
// if (/* condition fails */) {
// UNLOCK_OBJECT(op);
// DEOPT_IF(true);
// }
// ...
// UNLOCK_OBJECT(op);
//
// NOTE: The object must be unlocked on every exit code path and you should
// avoid any potentially escaping calls (like PyStackRef_CLOSE) while the
// object is locked.
#ifdef Py_GIL_DISABLED
# define LOCK_OBJECT(op) PyMutex_LockFast(&(_PyObject_CAST(op))->ob_mutex._bits)
# define UNLOCK_OBJECT(op) PyMutex_Unlock(&(_PyObject_CAST(op))->ob_mutex)
#else
# define LOCK_OBJECT(op) (1)
# define UNLOCK_OBJECT(op) ((void)0)
#endif
#define GLOBALS() frame->f_globals
#define BUILTINS() frame->f_builtins
#define LOCALS() frame->f_locals

View File

@ -1235,15 +1235,23 @@
JUMP_TO_JUMP_TARGET();
}
Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
// Ensure index < len(list)
if (index >= PyList_GET_SIZE(list)) {
if (!LOCK_OBJECT(list)) {
UOP_STAT_INC(uopcode, miss);
JUMP_TO_JUMP_TARGET();
}
// Ensure index < len(list)
if (index >= PyList_GET_SIZE(list)) {
UNLOCK_OBJECT(list);
if (true) {
UOP_STAT_INC(uopcode, miss);
JUMP_TO_JUMP_TARGET();
}
}
STAT_INC(STORE_SUBSCR, hit);
PyObject *old_value = PyList_GET_ITEM(list, index);
PyList_SET_ITEM(list, index, PyStackRef_AsPyObjectSteal(value));
assert(old_value != NULL);
UNLOCK_OBJECT(list); // unlock before decrefs!
Py_DECREF(old_value);
PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
PyStackRef_CLOSE(list_st);

View File

@ -7629,7 +7629,7 @@
container = stack_pointer[-2];
uint16_t counter = read_u16(&this_instr[1].cache);
(void)counter;
#if ENABLE_SPECIALIZATION
#if ENABLE_SPECIALIZATION_FT
if (ADAPTIVE_COUNTER_TRIGGERS(counter)) {
next_instr = this_instr;
_PyFrame_SetStackPointer(frame, stack_pointer);
@ -7639,7 +7639,7 @@
}
OPCODE_DEFERRED_INC(STORE_SUBSCR);
ADVANCE_ADAPTIVE_COUNTER(this_instr[1].counter);
#endif /* ENABLE_SPECIALIZATION */
#endif /* ENABLE_SPECIALIZATION_FT */
}
// _STORE_SUBSCR
{
@ -7704,12 +7704,17 @@
// Ensure nonnegative, zero-or-one-digit ints.
DEOPT_IF(!_PyLong_IsNonNegativeCompact((PyLongObject *)sub), STORE_SUBSCR);
Py_ssize_t index = ((PyLongObject*)sub)->long_value.ob_digit[0];
DEOPT_IF(!LOCK_OBJECT(list), STORE_SUBSCR);
// Ensure index < len(list)
DEOPT_IF(index >= PyList_GET_SIZE(list), STORE_SUBSCR);
if (index >= PyList_GET_SIZE(list)) {
UNLOCK_OBJECT(list);
DEOPT_IF(true, STORE_SUBSCR);
}
STAT_INC(STORE_SUBSCR, hit);
PyObject *old_value = PyList_GET_ITEM(list, index);
PyList_SET_ITEM(list, index, PyStackRef_AsPyObjectSteal(value));
assert(old_value != NULL);
UNLOCK_OBJECT(list); // unlock before decrefs!
Py_DECREF(old_value);
PyStackRef_CLOSE_SPECIALIZED(sub_st, (destructor)PyObject_Free);
PyStackRef_CLOSE(list_st);

View File

@ -1814,86 +1814,54 @@ success:
cache->counter = adaptive_counter_cooldown();
}
void
_Py_Specialize_StoreSubscr(_PyStackRef container_st, _PyStackRef sub_st, _Py_CODEUNIT *instr)
{
PyObject *container = PyStackRef_AsPyObjectBorrow(container_st);
PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
assert(ENABLE_SPECIALIZATION);
_PyStoreSubscrCache *cache = (_PyStoreSubscrCache *)(instr + 1);
PyTypeObject *container_type = Py_TYPE(container);
if (container_type == &PyList_Type) {
if (PyLong_CheckExact(sub)) {
if (_PyLong_IsNonNegativeCompact((PyLongObject *)sub)
&& ((PyLongObject *)sub)->long_value.ob_digit[0] < (size_t)PyList_GET_SIZE(container))
{
instr->op.code = STORE_SUBSCR_LIST_INT;
goto success;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
goto fail;
}
}
else if (PySlice_Check(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_LIST_SLICE);
goto fail;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
goto fail;
}
}
if (container_type == &PyDict_Type) {
instr->op.code = STORE_SUBSCR_DICT;
goto success;
}
#ifdef Py_STATS
static int
store_subscr_fail_kind(PyObject *container_type)
{
PyMappingMethods *as_mapping = container_type->tp_as_mapping;
if (as_mapping && (as_mapping->mp_ass_subscript
== PyDict_Type.tp_as_mapping->mp_ass_subscript)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_DICT_SUBCLASS_NO_OVERRIDE);
goto fail;
return SPEC_FAIL_SUBSCR_DICT_SUBCLASS_NO_OVERRIDE;
}
if (PyObject_CheckBuffer(container)) {
if (PyLong_CheckExact(sub) && (!_PyLong_IsNonNegativeCompact((PyLongObject *)sub))) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
return SPEC_FAIL_OUT_OF_RANGE;
}
else if (strcmp(container_type->tp_name, "array.array") == 0) {
if (PyLong_CheckExact(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_ARRAY_INT);
return SPEC_FAIL_SUBSCR_ARRAY_INT;
}
else if (PySlice_Check(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_ARRAY_SLICE);
return SPEC_FAIL_SUBSCR_ARRAY_SLICE;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
return SPEC_FAIL_OTHER;
}
}
else if (PyByteArray_CheckExact(container)) {
if (PyLong_CheckExact(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BYTEARRAY_INT);
return SPEC_FAIL_SUBSCR_BYTEARRAY_INT;
}
else if (PySlice_Check(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BYTEARRAY_SLICE);
return SPEC_FAIL_SUBSCR_BYTEARRAY_SLICE;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
return SPEC_FAIL_OTHER;
}
}
else {
if (PyLong_CheckExact(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BUFFER_INT);
return SPEC_FAIL_SUBSCR_BUFFER_INT;
}
else if (PySlice_Check(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_BUFFER_SLICE);
return SPEC_FAIL_SUBSCR_BUFFER_SLICE;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
return SPEC_FAIL_OTHER;
}
}
goto fail;
return SPEC_FAIL_OTHER;
}
PyObject *descriptor = _PyType_Lookup(container_type, &_Py_ID(__setitem__));
if (descriptor && Py_TYPE(descriptor) == &PyFunction_Type) {
@ -1901,25 +1869,55 @@ _Py_Specialize_StoreSubscr(_PyStackRef container_st, _PyStackRef sub_st, _Py_COD
PyCodeObject *code = (PyCodeObject *)func->func_code;
int kind = function_kind(code);
if (kind == SIMPLE_FUNCTION) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_PY_SIMPLE);
return SPEC_FAIL_SUBSCR_PY_SIMPLE;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_PY_OTHER);
return SPEC_FAIL_SUBSCR_PY_OTHER;
}
goto fail;
}
#endif // Py_STATS
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
fail:
STAT_INC(STORE_SUBSCR, failure);
assert(!PyErr_Occurred());
instr->op.code = STORE_SUBSCR;
cache->counter = adaptive_counter_backoff(cache->counter);
return;
success:
STAT_INC(STORE_SUBSCR, success);
assert(!PyErr_Occurred());
cache->counter = adaptive_counter_cooldown();
return SPEC_FAIL_OTHER;
}
#endif
void
_Py_Specialize_StoreSubscr(_PyStackRef container_st, _PyStackRef sub_st, _Py_CODEUNIT *instr)
{
PyObject *container = PyStackRef_AsPyObjectBorrow(container_st);
PyObject *sub = PyStackRef_AsPyObjectBorrow(sub_st);
assert(ENABLE_SPECIALIZATION_FT);
PyTypeObject *container_type = Py_TYPE(container);
if (container_type == &PyList_Type) {
if (PyLong_CheckExact(sub)) {
if (_PyLong_IsNonNegativeCompact((PyLongObject *)sub)
&& ((PyLongObject *)sub)->long_value.ob_digit[0] < (size_t)PyList_GET_SIZE(container))
{
specialize(instr, STORE_SUBSCR_LIST_INT);
return;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OUT_OF_RANGE);
unspecialize(instr);
return;
}
}
else if (PySlice_Check(sub)) {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_SUBSCR_LIST_SLICE);
unspecialize(instr);
return;
}
else {
SPECIALIZATION_FAIL(STORE_SUBSCR, SPEC_FAIL_OTHER);
unspecialize(instr);
return;
}
}
if (container_type == &PyDict_Type) {
specialize(instr, STORE_SUBSCR_DICT);
return;
}
SPECIALIZATION_FAIL(STORE_SUBSCR, store_subscr_fail_kind(container_type));
unspecialize(instr);
}
/* Returns a borrowed reference.