Changes from review:

* Add some code comments about _Py_INCREF_TYPE/CODE/DICT
* Fix typo
* Move size check to `_PyDict_EnablePerThreadRefcounting`

Include/cpython/dictobject.h

@@ -17,7 +17,9 @@ typedef struct {
     /* This is a private field for CPython's internal use.
      * Bits 0-7 are for dict watchers.
      * Bits 8-11 are for the watched mutation counter (used by tier2 optimization)
-     * The remaining bits are not currently used. */
+     * Bits 12-31 are currently unused
+     * Bits 32-63 are a unique id in the free threading build (used for per-thread refcounting)
+     */
     uint64_t _ma_watcher_tag;
 
     PyDictKeysObject *ma_keys;

Include/internal/pycore_dict.h

@@ -229,8 +229,8 @@ static inline PyDictUnicodeEntry* DK_UNICODE_ENTRIES(PyDictKeysObject *dk) {
 #define DICT_VERSION_INCREMENT (1 << (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS))
 #define DICT_WATCHER_MASK ((1 << DICT_MAX_WATCHERS) - 1)
 #define DICT_WATCHER_AND_MODIFICATION_MASK ((1 << (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS)) - 1)
-#define DICT_UNIQUE_ID_SHIFT (DICT_MAX_WATCHERS + DICT_WATCHED_MUTATION_BITS)
-#define DICT_UNIQUE_ID_MAX (UINT64_MAX >> DICT_UNIQUE_ID_SHIFT)
+#define DICT_UNIQUE_ID_SHIFT (32)
+#define DICT_UNIQUE_ID_MAX ((UINT64_C(1) << (64 - DICT_UNIQUE_ID_SHIFT)) - 1)
 
 
 PyAPI_FUNC(void)
@@ -314,14 +314,15 @@ extern void _PyDict_EnablePerThreadRefcounting(PyObject *op);
 
 PyDictObject *_PyObject_MaterializeManagedDict_LockHeld(PyObject *);
 
+// See `_Py_INCREF_TYPE()` in pycore_object.h
 #ifndef Py_GIL_DISABLED
 #  define _Py_INCREF_DICT Py_INCREF
 #  define _Py_DECREF_DICT Py_DECREF
 #else
 static inline Py_ssize_t
 _PyDict_UniqueId(PyDictObject *mp)
 {
-    // Offset by on so that _ma_watcher_tag=0 represents an unassigned id
+    // Offset by one so that _ma_watcher_tag=0 represents an unassigned id
     return (Py_ssize_t)(mp->_ma_watcher_tag >> DICT_UNIQUE_ID_SHIFT) - 1;
 }
 

Include/internal/pycore_object.h

@@ -288,6 +288,20 @@ extern PyStatus _PyObject_InitState(PyInterpreterState *interp);
 extern void _PyObject_FiniState(PyInterpreterState *interp);
 extern bool _PyRefchain_IsTraced(PyInterpreterState *interp, PyObject *obj);
 
+// Macros used for per-thread reference counting in the free threading build.
+// They resolve to normal Py_INCREF/DECREF calls in the default build.
+//
+// The macros are used for only a few references that would otherwise cause
+// scaling bottlenecks in the free threading build:
+// - The reference from an object to `ob_type`.
+// - The reference from a function to `func_code`.
+// - The reference from a function to `func_globals` and `func_builtins`.
+//
+// It's safe, but not performant or necessary, to use these macros for other
+// references to code, type, or dict objects. It's also safe to mix their
+// usage with normal Py_INCREF/DECREF calls.
+//
+// See also Include/internal/pycore_dict.h for _Py_INCREF_DICT/_Py_DECREF_DICT.
 #ifndef Py_GIL_DISABLED
 #  define _Py_INCREF_TYPE Py_INCREF
 #  define _Py_DECREF_TYPE Py_DECREF

Include/internal/pycore_uniqueid.h

@@ -48,6 +48,9 @@ struct _Py_unique_id_pool {
 // Assigns the next id from the pool of ids.
 extern Py_ssize_t _PyObject_AssignUniqueId(PyObject *obj);
 
+// Releases the allocated id back to the pool.
+extern void _PyObject_ReleaseUniqueId(Py_ssize_t unique_id);
+
 // Releases the allocated id back to the pool.
 extern void _PyObject_DisablePerThreadRefcounting(PyObject *obj);
 

Objects/dictobject.c

@@ -1642,6 +1642,10 @@ _PyDict_EnablePerThreadRefcounting(PyObject *op)
     assert(PyDict_Check(op));
 #ifdef Py_GIL_DISABLED
     Py_ssize_t id = _PyObject_AssignUniqueId(op);
+    if ((uint64_t)id >= (uint64_t)DICT_UNIQUE_ID_MAX) {
+        _PyObject_ReleaseUniqueId(id);
+        return;
+    }
 
     PyDictObject *mp = (PyDictObject *)op;
     assert((mp->_ma_watcher_tag >> DICT_UNIQUE_ID_SHIFT) == 0);

Python/uniqueid.c

@@ -1,6 +1,6 @@
 #include "Python.h"
 
-#include "pycore_dict.h"        // DICT_UNIQUE_ID_MAX
+#include "pycore_dict.h"        // _PyDict_UniqueId()
 #include "pycore_lock.h"        // PyMutex_LockFlags()
 #include "pycore_pystate.h"     // _PyThreadState_GET()
 #include "pycore_object.h"      // _Py_IncRefTotal
@@ -91,21 +91,16 @@ _PyObject_AssignUniqueId(PyObject *obj)
     }
 
     _Py_unique_id_entry *entry = pool->freelist;
-    Py_ssize_t unique_id = (entry - pool->table);
-    if ((uint64_t)unique_id >= DICT_UNIQUE_ID_MAX) {
-        UNLOCK_POOL(pool);
-        return -1;
-    }
-
     pool->freelist = entry->next;
     entry->obj = obj;
     _PyObject_SetDeferredRefcount(obj);
+    Py_ssize_t unique_id = (entry - pool->table);
     UNLOCK_POOL(pool);
     return unique_id;
 }
 
-static void
-release_unique_id(Py_ssize_t unique_id)
+void
+_PyObject_ReleaseUniqueId(Py_ssize_t unique_id)
 {
     PyInterpreterState *interp = _PyInterpreterState_GET();
     struct _Py_unique_id_pool *pool = &interp->unique_ids;
@@ -147,7 +142,7 @@ _PyObject_DisablePerThreadRefcounting(PyObject *obj)
 {
     Py_ssize_t id = clear_unique_id(obj);
     if (id >= 0) {
-        release_unique_id(id);
+        _PyObject_ReleaseUniqueId(id);
     }
 }