gh-135832: Implement _Py_DECREF_SPECIALIZED and PyStackRef_CLOSE_SPECIALIZED for free-threading

This should help single-threaded performance.

Author: brijkapadia

Include/internal/pycore_object.h

@@ -253,11 +253,25 @@ _Py_DECREF_SPECIALIZED(PyObject *op, const destructor destruct)
 }
 
 #else
-// TODO: implement Py_DECREF specializations for Py_GIL_DISABLED build
 static inline void
 _Py_DECREF_SPECIALIZED(PyObject *op, const destructor destruct)
 {
-    Py_DECREF(op);
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
+    if (local == _Py_IMMORTAL_REFCNT_LOCAL) {
+        _Py_DECREF_IMMORTAL_STAT_INC();
+        return;
+    }
+    _Py_DECREF_STAT_INC();
+    if (_Py_IsOwnedByCurrentThread(op)) {
+        local--;
+        _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local);
+        if (local == 0) {
+            _Py_MergeZeroLocalRefcountSpecialized(op, destruct);
+        }
+    }
+    else {
+        _Py_DecRefShared(op);
+    }
 }
 
 static inline int
@@ -465,7 +479,7 @@ static inline void Py_DECREF_MORTAL_SPECIALIZED(PyObject *op, destructor destruc
 #endif
 #else  // Py_GIL_DISABLED
 # define Py_DECREF_MORTAL(op) Py_DECREF(op)
-# define Py_DECREF_MORTAL_SPECIALIZED(op, destruct) Py_DECREF(op)
+# define Py_DECREF_MORTAL_SPECIALIZED(op, destruct) _Py_DECREF_SPECIALIZED(op, destruct)
 #endif
 
 /* Inline functions trading binary compatibility for speed:

Include/refcount.h

@@ -321,6 +321,7 @@ PyAPI_FUNC(void) _Py_DecRefSharedDebug(PyObject *, const char *, int);
 // zero. Otherwise, the thread gives up ownership and merges the reference
 // count fields.
 PyAPI_FUNC(void) _Py_MergeZeroLocalRefcount(PyObject *);
+PyAPI_FUNC(void) _Py_MergeZeroLocalRefcountSpecialized(PyObject *, const destructor);
 #endif  // Py_GIL_DISABLED
 #endif  // Py_LIMITED_API
 

Objects/object.c

@@ -464,6 +464,30 @@ _Py_MergeZeroLocalRefcount(PyObject *op)
     }
 }
 
+void
+_Py_MergeZeroLocalRefcountSpecialized(PyObject *op, const destructor destruct)
+{
+    assert(op->ob_ref_local == 0);
+
+    Py_ssize_t shared = _Py_atomic_load_ssize_acquire(&op->ob_ref_shared);
+    if (shared == 0) {
+        destruct(op);
+        return;
+    }
+
+    _Py_atomic_store_uintptr_relaxed(&op->ob_tid, 0);
+
+    Py_ssize_t new_shared;
+    do {
+        new_shared = (shared & ~_Py_REF_SHARED_FLAG_MASK) | _Py_REF_MERGED;
+    } while (!_Py_atomic_compare_exchange_ssize(&op->ob_ref_shared,
+                                                &shared, new_shared));
+
+    if (new_shared == _Py_REF_MERGED) {
+        destruct(op);
+    }
+}
+
 Py_ssize_t
 _Py_ExplicitMergeRefcount(PyObject *op, Py_ssize_t extra)
 {

Python/ceval.h

@@ -132,7 +132,26 @@
     } while (0)
 
 #undef _Py_DECREF_SPECIALIZED
-#define _Py_DECREF_SPECIALIZED(arg, dealloc) Py_DECREF(arg)
+#define _Py_DECREF_SPECIALIZED(arg, dealloc) \
+    do { \
+        PyObject *op = _PyObject_CAST(arg); \
+        uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local); \
+        if (local == _Py_IMMORTAL_REFCNT_LOCAL) { \
+            _Py_DECREF_IMMORTAL_STAT_INC(); \
+            break; \
+        } \
+        _Py_DECREF_STAT_INC(); \
+        if (_Py_IsOwnedByCurrentThread(op)) { \
+            local--; \
+            _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local); \
+            if (local == 0) { \
+                _Py_MergeZeroLocalRefcountSpecialized(op, dealloc); \
+            } \
+        } \
+        else { \
+            _Py_DecRefShared(op); \
+        } \
+    } while (0)
 
 #endif
 #endif