Make specialization thread-safe

Include/internal/pycore_code.h

@@ -326,7 +326,7 @@ extern void _PyCode_Clear_Executors(PyCodeObject *code);
 #define ENABLE_SPECIALIZED_CONTAINS_OP 0 && ENABLE_SPECIALIZATION
 #define ENABLE_SPECIALIZED_FOR_ITER 0 && ENABLE_SPECIALIZATION
 #define ENABLE_SPECIALIZED_LOAD_ATTR 0 && ENABLE_SPECIALIZATION
-#define ENABLE_SPECIALIZED_LOAD_GLOBAL 0 && ENABLE_SPECIALIZATION
+#define ENABLE_SPECIALIZED_LOAD_GLOBAL ENABLE_SPECIALIZATION
 #define ENABLE_SPECIALIZED_LOAD_SUPER_ATTR 0 && ENABLE_SPECIALIZATION
 #define ENABLE_SPECIALIZED_SEND 0 && ENABLE_SPECIALIZATION
 #define ENABLE_SPECIALIZED_STORE_ATTR 0 && ENABLE_SPECIALIZATION

Objects/dictobject.c

@@ -1682,7 +1682,7 @@ insert_combined_dict(PyInterpreterState *interp, PyDictObject *mp,
 
     uint64_t new_version = _PyDict_NotifyEvent(
         interp, PyDict_EVENT_ADDED, mp, key, value);
-    mp->ma_keys->dk_version = 0;
+    FT_ATOMIC_STORE_UINT32_RELAXED(mp->ma_keys->dk_version, 0);
 
     Py_ssize_t hashpos = find_empty_slot(mp->ma_keys, hash);
     dictkeys_set_index(mp->ma_keys, hashpos, mp->ma_keys->dk_nentries);
@@ -2634,7 +2634,7 @@ delitem_common(PyDictObject *mp, Py_hash_t hash, Py_ssize_t ix,
         ASSERT_CONSISTENT(mp);
     }
     else {
-        mp->ma_keys->dk_version = 0;
+        FT_ATOMIC_STORE_UINT32_RELAXED(mp->ma_keys->dk_version, 0);
         dictkeys_set_index(mp->ma_keys, hashpos, DKIX_DUMMY);
         if (DK_IS_UNICODE(mp->ma_keys)) {
             PyDictUnicodeEntry *ep = &DK_UNICODE_ENTRIES(mp->ma_keys)[ix];
@@ -4461,7 +4461,7 @@ dict_popitem_impl(PyDictObject *self)
             return NULL;
         }
     }
-    self->ma_keys->dk_version = 0;
+    FT_ATOMIC_STORE_UINT32_RELAXED(self->ma_keys->dk_version, 0);
 
     /* Pop last item */
     PyObject *key, *value;
@@ -7340,17 +7340,30 @@ _PyDictKeys_DecRef(PyDictKeysObject *keys)
     dictkeys_decref(interp, keys, false);
 }
 
+// In free-threaded builds the caller must ensure that the keys object is not
+// being mutated concurrently by another thread.
 uint32_t _PyDictKeys_GetVersionForCurrentState(PyInterpreterState *interp,
                                                PyDictKeysObject *dictkeys)
 {
-    if (dictkeys->dk_version != 0) {
-        return dictkeys->dk_version;
+    uint32_t dk_version = FT_ATOMIC_LOAD_UINT32_RELAXED(dictkeys->dk_version);
+    if (dk_version != 0) {
+        return dk_version;
     }
+#ifdef Py_GIL_DISABLED
+    uint32_t v;
+    do {
+        v = _Py_atomic_load_uint32_relaxed(&interp->dict_state.next_keys_version);
+        if (v == 0) {
+            return 0;
+        }
+    } while(!_Py_atomic_compare_exchange_uint32(&interp->dict_state.next_keys_version, &v, v + 1));
+#else
     if (interp->dict_state.next_keys_version == 0) {
         return 0;
     }
     uint32_t v = interp->dict_state.next_keys_version++;
-    dictkeys->dk_version = v;
+#endif
+    FT_ATOMIC_STORE_UINT32_RELAXED(dictkeys->dk_version, v);
     return v;
 }
 

Python/bytecodes.c

@@ -1575,34 +1575,42 @@ dummy_func(
         op(_GUARD_GLOBALS_VERSION, (version/1 -- globals_keys: PyDictKeysObject*)) {
             PyDictObject *dict = (PyDictObject *)GLOBALS();
             DEOPT_IF(!PyDict_CheckExact(dict));
-            globals_keys = dict->ma_keys;
-            DEOPT_IF(globals_keys->dk_version != version);
+            globals_keys = FT_ATOMIC_LOAD_PTR_ACQUIRE(dict->ma_keys);
+            DEOPT_IF(FT_ATOMIC_LOAD_UINT32_RELAXED(globals_keys->dk_version) != version);
             assert(DK_IS_UNICODE(globals_keys));
         }
 
         op(_GUARD_BUILTINS_VERSION, (version/1 -- builtins_keys: PyDictKeysObject*)) {
             PyDictObject *dict = (PyDictObject *)BUILTINS();
             DEOPT_IF(!PyDict_CheckExact(dict));
-            builtins_keys = dict->ma_keys;
-            DEOPT_IF(builtins_keys->dk_version != version);
+            builtins_keys = FT_ATOMIC_LOAD_PTR_ACQUIRE(dict->ma_keys);
+            DEOPT_IF(FT_ATOMIC_LOAD_UINT32_RELAXED(builtins_keys->dk_version) != version);
             assert(DK_IS_UNICODE(builtins_keys));
         }
 
         op(_LOAD_GLOBAL_MODULE, (index/1, globals_keys: PyDictKeysObject* -- res, null if (oparg & 1))) {
             PyDictUnicodeEntry *entries = DK_UNICODE_ENTRIES(globals_keys);
-            PyObject *res_o = entries[index].me_value;
+            PyObject *res_o = FT_ATOMIC_LOAD_PTR_RELAXED(entries[index].me_value);
             DEOPT_IF(res_o == NULL);
+            #if Py_GIL_DISABLED
+            DEOPT_IF(!_Py_TryIncrefCompare(&entries[index].me_value, res_o));
+            #else
             Py_INCREF(res_o);
+            #endif
             STAT_INC(LOAD_GLOBAL, hit);
             null = PyStackRef_NULL;
             res = PyStackRef_FromPyObjectSteal(res_o);
         }
 
         op(_LOAD_GLOBAL_BUILTINS, (index/1, globals_keys: PyDictKeysObject*, builtins_keys: PyDictKeysObject* -- res, null if (oparg & 1))) {
             PyDictUnicodeEntry *entries = DK_UNICODE_ENTRIES(builtins_keys);
-            PyObject *res_o = entries[index].me_value;
+            PyObject *res_o = FT_ATOMIC_LOAD_PTR_RELAXED(entries[index].me_value);
             DEOPT_IF(res_o == NULL);
+            #if Py_GIL_DISABLED
+            DEOPT_IF(!_Py_TryIncrefCompare(&entries[index].me_value, res_o));
+            #else
             Py_INCREF(res_o);
+            #endif
             STAT_INC(LOAD_GLOBAL, hit);
             null = PyStackRef_NULL;
             res = PyStackRef_FromPyObjectSteal(res_o);

Python/specialize.c

@@ -1464,7 +1464,7 @@ PyObject *descr, DescriptorClassification kind, bool is_method)
 }
 
 void
-_Py_Specialize_LoadGlobal(
+specialize_load_global_lock_held(
     PyObject *globals, PyObject *builtins,
     _Py_CODEUNIT *instr, PyObject *name)
 {
@@ -1487,6 +1487,8 @@ _Py_Specialize_LoadGlobal(
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_EXPECTED_ERROR);
         goto fail;
     }
+    uint8_t specialized_opcode;
+    _PyLoadGlobalCache new_cache;
     PyInterpreterState *interp = _PyInterpreterState_GET();
     if (index != DKIX_EMPTY) {
         if (index != (uint16_t)index) {
@@ -1503,9 +1505,9 @@ _Py_Specialize_LoadGlobal(
             SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_OUT_OF_RANGE);
             goto fail;
         }
-        cache->index = (uint16_t)index;
-        cache->module_keys_version = (uint16_t)keys_version;
-        instr->op.code = LOAD_GLOBAL_MODULE;
+        new_cache.index = (uint16_t)index;
+        new_cache.module_keys_version = (uint16_t)keys_version;
+        specialized_opcode = LOAD_GLOBAL_MODULE;
         goto success;
     }
     if (!PyDict_CheckExact(builtins)) {
@@ -1546,23 +1548,35 @@ _Py_Specialize_LoadGlobal(
         SPECIALIZATION_FAIL(LOAD_GLOBAL, SPEC_FAIL_OUT_OF_RANGE);
         goto fail;
     }
-    cache->index = (uint16_t)index;
-    cache->module_keys_version = (uint16_t)globals_version;
-    cache->builtin_keys_version = (uint16_t)builtins_version;
-    instr->op.code = LOAD_GLOBAL_BUILTIN;
+    new_cache.index = (uint16_t)index;
+    new_cache.module_keys_version = (uint16_t)globals_version;
+    new_cache.builtin_keys_version = (uint16_t)builtins_version;
+    specialized_opcode = LOAD_GLOBAL_BUILTIN;
     goto success;
 fail:
+    SET_OPCODE_OR_RETURN(instr, LOAD_GLOBAL);
     STAT_INC(LOAD_GLOBAL, failure);
     assert(!PyErr_Occurred());
-    instr->op.code = LOAD_GLOBAL;
     cache->counter = adaptive_counter_backoff(cache->counter);
     return;
 success:
+    SET_OPCODE_OR_RETURN(instr, specialized_opcode);
     STAT_INC(LOAD_GLOBAL, success);
+    *cache = new_cache;
     assert(!PyErr_Occurred());
     cache->counter = adaptive_counter_cooldown();
 }
 
+void
+_Py_Specialize_LoadGlobal(
+    PyObject *globals, PyObject *builtins,
+    _Py_CODEUNIT *instr, PyObject *name)
+{
+    Py_BEGIN_CRITICAL_SECTION2(globals, builtins);
+    specialize_load_global_lock_held(globals, builtins, instr, name);
+    Py_END_CRITICAL_SECTION2();
+}
+
 #ifdef Py_STATS
 static int
 binary_subscr_fail_kind(PyTypeObject *container_type, PyObject *sub)