Add type-local type cache

Author: DinoV

Include/cpython/object.h

@@ -259,6 +259,9 @@ struct _specialization_cache {
     PyObject *getitem;
     uint32_t getitem_version;
     PyObject *init;
+#ifdef Py_GIL_DISABLED
+    struct local_type_cache *local_type_cache;
+#endif
 };
 
 /* The *real* layout of a type object when allocated on the heap */

Include/internal/pycore_typeobject.h

@@ -71,6 +71,25 @@ struct type_cache {
     struct type_cache_entry hashtable[1 << MCACHE_SIZE_EXP];
 };
 
+#ifdef Py_GIL_DISABLED
+
+// Type attribute lookup cache which is type-specific. Only used
+// for heap types where we store a small additional cache in free-threaded
+// builds which can be accessed without any locking.
+#define LOCAL_CACHE_TYPE_SIZE 67
+
+struct local_type_cache_entry {
+    PyObject *name;     // reference to exactly a str or NULL
+    PyObject *value;    // borrowed reference or NULL
+};
+
+struct local_type_cache {
+    unsigned int tp_version_tag;
+    struct local_type_cache_entry entries[LOCAL_CACHE_TYPE_SIZE];
+};
+
+#endif
+
 typedef struct {
     PyTypeObject *type;
     int isbuiltin;

Objects/typeobject.c

@@ -1021,6 +1021,22 @@ set_version_unlocked(PyTypeObject *tp, unsigned int version)
 #endif
 }
 
+static void
+clear_spec_cache(PyTypeObject *type)
+{
+    if (PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
+        // This field *must* be invalidated if the type is modified (see the
+        // comment on struct _specialization_cache):
+        PyHeapTypeObject *heap_type = (PyHeapTypeObject *)type;
+        heap_type->_spec_cache.getitem = NULL;
+        struct local_type_cache *cache = heap_type->_spec_cache.local_type_cache;
+        if (cache != NULL) {
+            _PyMem_FreeDelayed(cache);
+            heap_type->_spec_cache.local_type_cache = NULL;
+        }
+    }
+}
+
 static void
 type_modified_unlocked(PyTypeObject *type)
 {
@@ -1082,11 +1098,7 @@ type_modified_unlocked(PyTypeObject *type)
     }
 
     set_version_unlocked(type, 0); /* 0 is not a valid version tag */
-    if (PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
-        // This field *must* be invalidated if the type is modified (see the
-        // comment on struct _specialization_cache):
-        ((PyHeapTypeObject *)type)->_spec_cache.getitem = NULL;
-    }
+    clear_spec_cache(type);
 }
 
 void
@@ -1163,11 +1175,7 @@ type_mro_modified(PyTypeObject *type, PyObject *bases) {
     assert(!(type->tp_flags & _Py_TPFLAGS_STATIC_BUILTIN));
     set_version_unlocked(type, 0);  /* 0 is not a valid version tag */
     type->tp_versions_used = _Py_ATTR_CACHE_UNUSED;
-    if (PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
-        // This field *must* be invalidated if the type is modified (see the
-        // comment on struct _specialization_cache):
-        ((PyHeapTypeObject *)type)->_spec_cache.getitem = NULL;
-    }
+    clear_spec_cache(type);
 }
 
 /*
@@ -5553,6 +5561,19 @@ _PyType_LookupRefAndVersion(PyTypeObject *type, PyObject *name, unsigned int *ve
     struct type_cache *cache = get_type_cache();
     struct type_cache_entry *entry = &cache->hashtable[h];
 #ifdef Py_GIL_DISABLED
+    PyHeapTypeObject *heap_type = (PyHeapTypeObject *)type;
+    struct local_type_cache *local_cache = heap_type->_spec_cache.local_type_cache;
+
+    if (cache != NULL) {
+        Py_ssize_t index = (((Py_ssize_t)(name)) >> 3) % LOCAL_CACHE_TYPE_SIZE;
+        struct local_type_cache_entry *entry = &local_cache->entries[index];
+        PyObject *value;
+        if (entry->name == name && (value = _Py_TryXGetRef(&entry->value)) != NULL) {
+            *version = local_cache->tp_version_tag;
+            return value;
+        }
+    }
+
     // synchronize-with other writing threads by doing an acquire load on the sequence
     while (1) {
         uint32_t sequence = _PySeqLock_BeginRead(&entry->sequence);
@@ -5638,6 +5659,28 @@ _PyType_LookupRefAndVersion(PyTypeObject *type, PyObject *name, unsigned int *ve
     if (has_version) {
 #if Py_GIL_DISABLED
         update_cache_gil_disabled(entry, name, assigned_version, res);
+        if (res != NULL) {
+            if (PyType_HasFeature(type, Py_TPFLAGS_HEAPTYPE)) {
+                PyHeapTypeObject *heap_type = (PyHeapTypeObject *)type;
+                struct local_type_cache *local_cache = heap_type->_spec_cache.local_type_cache;
+                if (local_cache == NULL) {
+                    BEGIN_TYPE_LOCK();
+                    local_cache = PyMem_Calloc(1, sizeof(struct local_type_cache));
+                    if (local_cache != NULL) {
+                        local_cache->tp_version_tag = assigned_version;
+                        heap_type->_spec_cache.local_type_cache = local_cache;
+                    }
+                    END_TYPE_LOCK();
+                }
+                if (local_cache != NULL) {
+                    struct local_type_cache_entry *entry =
+                        &local_cache->entries[(((Py_ssize_t)(name)) >> 3) % LOCAL_CACHE_TYPE_SIZE];
+                    entry->name = name;
+                    entry->value = res;
+                }
+            }
+        }
+
 #else
         PyObject *old_value = update_cache(entry, name, assigned_version, res);
         Py_DECREF(old_value);