Skip to content

Commit

Permalink
pythongh-115999: Add free-threaded specialization for STORE_SUBSCR
Browse files Browse the repository at this point in the history 
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.
  • Loading branch information
@colesbury
colesbury committed Nov 22, 2024
1 parent a264637 commit 2b63929
Show file tree
Hide file tree
Showing 5 changed files with 109 additions and 70 deletions.
11 changes: 8 additions & 3 deletions Python/bytecodes.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -910,15 +910,15 @@ 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);
DISPATCH_SAME_OPARG();
}
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 -- )) {
Expand All @@ -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);
Expand Down
23 changes: 23 additions & 0 deletions Python/ceval_macros.h
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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)
# 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
Expand Down
12 changes: 10 additions & 2 deletions Python/executor_cases.c.h
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

11 changes: 8 additions & 3 deletions Python/generated_cases.c.h
View file

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

122 changes: 60 additions & 62 deletions Python/specialize.c
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1813,112 +1813,110 @@ _Py_Specialize_BinarySubscr(
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) {
PyFunctionObject *func = (PyFunctionObject *)descriptor;
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.
Expand Down

0 comments on commit 2b63929

Please sign in to comment.