Initial implementation of tagged stackrefs for GIL build. Requires fi…

…xes to immortal objects.
python · markshannon · Mar 10, 2025 · Oct 9, 2024 · Oct 11, 2024 · Oct 11, 2024
commit 2ec65dc331054cf68f7a3b740c19db7f64533ff5
@@ -54,12 +54,13 @@ typedef union _PyStackRef {
 } _PyStackRef;
 
 
+#ifdef Py_GIL_DISABLED
+
 #define Py_TAG_DEFERRED (1)
 
 #define Py_TAG_PTR      ((uintptr_t)0)
 #define Py_TAG_BITS     ((uintptr_t)1)
 
-#ifdef Py_GIL_DISABLED
 
 static const _PyStackRef PyStackRef_NULL = { .bits = Py_TAG_DEFERRED};
 #define PyStackRef_IsNull(stackref) ((stackref).bits == PyStackRef_NULL.bits)
@@ -153,36 +154,104 @@ PyStackRef_AsStrongReference(_PyStackRef stackref)
     return PyStackRef_FromPyObjectSteal(PyStackRef_AsPyObjectSteal(stackref));
 }
 
-
 #else // Py_GIL_DISABLED
 
 // With GIL
+
+#define Py_TAG_BITS 1
+#define Py_TAG_REFCNT 1
+#define BITS_TO_PTR(REF) ((PyObject *)((REF).bits))
+#define BITS_TO_PTR_MASKED(REF) ((PyObject *)(((REF).bits) & (~Py_TAG_BITS)))
+
 static const _PyStackRef PyStackRef_NULL = { .bits = 0 };
-#define PyStackRef_IsNull(stackref) ((stackref).bits == 0)
-#define PyStackRef_True ((_PyStackRef){.bits = (uintptr_t)&_Py_TrueStruct })
-#define PyStackRef_False ((_PyStackRef){.bits = ((uintptr_t)&_Py_FalseStruct) })
-#define PyStackRef_None ((_PyStackRef){.bits = ((uintptr_t)&_Py_NoneStruct) })
 
-#define PyStackRef_AsPyObjectBorrow(stackref) ((PyObject *)(stackref).bits)
+#define PyStackRef_IsNull(ref) ((ref).bits == 0)
+#define PyStackRef_True ((_PyStackRef){.bits = ((uintptr_t)&_Py_TrueStruct) | Py_TAG_REFCNT })
+#define PyStackRef_False ((_PyStackRef){.bits = ((uintptr_t)&_Py_FalseStruct) | Py_TAG_REFCNT })
+#define PyStackRef_None ((_PyStackRef){.bits = ((uintptr_t)&_Py_NoneStruct) | Py_TAG_REFCNT })
 
-#define PyStackRef_AsPyObjectSteal(stackref) PyStackRef_AsPyObjectBorrow(stackref)
+static inline int
+PyStackRef_HasCount(_PyStackRef ref)
+{
+    return ref.bits & Py_TAG_REFCNT;
+}
 
-#define PyStackRef_FromPyObjectSteal(obj) ((_PyStackRef){.bits = ((uintptr_t)(obj))})
+static inline PyObject *
+PyStackRef_AsPyObjectBorrow(_PyStackRef ref)
+{
+    return BITS_TO_PTR_MASKED(ref);
+}
 
-#define PyStackRef_FromPyObjectNew(obj) ((_PyStackRef){ .bits = (uintptr_t)(Py_NewRef(obj)) })
+static inline PyObject *
+PyStackRef_AsPyObjectSteal(_PyStackRef ref)
+{
+    if (PyStackRef_HasCount(ref)) {
+        return Py_NewRef(BITS_TO_PTR_MASKED(ref));
+    }
+    else {
+        return BITS_TO_PTR(ref);
+    }
+}
 
-#define PyStackRef_FromPyObjectImmortal(obj) ((_PyStackRef){ .bits = (uintptr_t)(obj) })
+/* We will want to extend this to a larger set of objects in the future */
+#define _Py_IsDeferrable _Py_IsImmortal
 
-#define PyStackRef_CLOSE(stackref) Py_DECREF(PyStackRef_AsPyObjectBorrow(stackref))
+static inline _PyStackRef
+PyStackRef_FromPyObjectSteal(PyObject *obj)
+{
+    assert(obj != NULL);
+    unsigned int tag = _Py_IsDeferrable(obj) ? Py_TAG_REFCNT : 0;
+    _PyStackRef ref = ((_PyStackRef){.bits = ((uintptr_t)(obj)) | tag});
+    return ref;
+}
+
+static inline _PyStackRef
+_PyStackRef_FromPyObjectNew(PyObject *obj)
+{
+    if (_Py_IsDeferrable(obj)) {
+        return (_PyStackRef){ .bits = ((uintptr_t)obj) | Py_TAG_REFCNT};
+    }
+    Py_INCREF(obj);
+    _PyStackRef ref = (_PyStackRef){ .bits = (uintptr_t)obj };
+    return ref;
+}
+#define PyStackRef_FromPyObjectNew(obj) _PyStackRef_FromPyObjectNew(_PyObject_CAST(obj))
+
+/* Create a new reference from an object with an embedded reference count */
+static inline _PyStackRef
+_PyStackRef_FromPyObjectWithCount(PyObject *obj)
+{
+    return (_PyStackRef){ .bits = (uintptr_t)obj | Py_TAG_REFCNT};
+}
+#define PyStackRef_FromPyObjectWithCount(obj) _PyStackRef_FromPyObjectWithCount(_PyObject_CAST(obj))
 
-#define PyStackRef_DUP(stackref) PyStackRef_FromPyObjectSteal(Py_NewRef(PyStackRef_AsPyObjectBorrow(stackref)))
+#define PyStackRef_FromPyObjectImmortal PyStackRef_FromPyObjectWithCount
 
 
+static inline _PyStackRef
+PyStackRef_DUP(_PyStackRef ref)
+{
+    assert(!PyStackRef_IsNull(ref));
+    if (!PyStackRef_HasCount(ref)) {
+        Py_INCREF_MORTAL(BITS_TO_PTR(ref));
+    }
+    return ref;
+}
+
+static inline void
+PyStackRef_CLOSE(_PyStackRef ref)
+{
+    assert(!PyStackRef_IsNull(ref));
+    if (!PyStackRef_HasCount(ref)) {
+        Py_DECREF_MORTAL(BITS_TO_PTR(ref));
+    }
+}
+
 #endif // Py_GIL_DISABLED
 
 // Note: this is a macro because MSVC (Windows) has trouble inlining it.
 
-#define PyStackRef_Is(a, b) ((a).bits == (b).bits)
+#define PyStackRef_Is(a, b) (((a).bits & (~Py_TAG_BITS)) == ((b).bits & (~Py_TAG_BITS)))
 
 // Converts a PyStackRef back to a PyObject *, converting the
 // stackref to a new reference.

diff --git a/Include/refcount.h b/Include/refcount.h
@@ -34,7 +34,7 @@ having all the lower 32 bits set, which will avoid the reference count to go
 beyond the refcount limit. Immortality checks for reference count decreases will
 be done by checking the bit sign flag in the lower 32 bits.
 */
-#define _Py_IMMORTAL_REFCNT _Py_CAST(Py_ssize_t, UINT_MAX)
+#define _Py_IMMORTAL_REFCNT ((Py_ssize_t)0xb0000000)
 
 #else
 /*
@@ -274,6 +274,32 @@ static inline Py_ALWAYS_INLINE void Py_INCREF(PyObject *op)
 #  define Py_INCREF(op) Py_INCREF(_PyObject_CAST(op))
 #endif
 
+static inline Py_ALWAYS_INLINE void Py_INCREF_MORTAL(PyObject *op)
+{
+#if defined(Py_GIL_DISABLED)
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
+    uint32_t new_local = local + 1;
+    assert (new_local != 0);
+    if (_Py_IsOwnedByCurrentThread(op)) {
+        _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, new_local);
+    }
+    else {
+        _Py_atomic_add_ssize(&op->ob_ref_shared, (1 << _Py_REF_SHARED_SHIFT));
+    }
+#elif SIZEOF_VOID_P > 4
+    PY_UINT32_T cur_refcnt = op->ob_refcnt_split[PY_BIG_ENDIAN];
+    PY_UINT32_T new_refcnt = cur_refcnt + 1;
+    op->ob_refcnt_split[PY_BIG_ENDIAN] = new_refcnt;
+#else
+    op->ob_refcnt++;
+#endif
+    _Py_INCREF_STAT_INC();
+#if defined(Py_REF_DEBUG) && !defined(Py_LIMITED_API)
+    if (!_Py_IsImmortal(op)) {
+        _Py_INCREF_IncRefTotal();
+    }
+#endif
+}
 
 #if !defined(Py_LIMITED_API) && defined(Py_GIL_DISABLED)
 // Implements Py_DECREF on objects not owned by the current thread.
@@ -302,6 +328,27 @@ static inline void Py_DECREF(PyObject *op) {
 #define Py_DECREF(op) Py_DECREF(_PyObject_CAST(op))
 
 #elif defined(Py_GIL_DISABLED) && defined(Py_REF_DEBUG)
+static inline void Py_DECREF_MORTAL(const char *filename, int lineno, PyObject *op)
+{
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
+    _Py_DECREF_STAT_INC();
+    _Py_DECREF_DecRefTotal();
+    if (_Py_IsOwnedByCurrentThread(op)) {
+        if (local == 0) {
+            _Py_NegativeRefcount(filename, lineno, op);
+        }
+        local--;
+        _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local);
+        if (local == 0) {
+            _Py_MergeZeroLocalRefcount(op);
+        }
+    }
+    else {
+        _Py_DecRefSharedDebug(op, filename, lineno);
+    }
+}
+#define Py_DECREF_MORTAL(op) Py_DECREF_MORTAL(__FILE__, __LINE__, _PyObject_CAST(op))
+
 static inline void Py_DECREF(const char *filename, int lineno, PyObject *op)
 {
     uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
@@ -328,6 +375,22 @@ static inline void Py_DECREF(const char *filename, int lineno, PyObject *op)
 #define Py_DECREF(op) Py_DECREF(__FILE__, __LINE__, _PyObject_CAST(op))
 
 #elif defined(Py_GIL_DISABLED)
+static inline void Py_DECREF_MORTAL(PyObject *op)
+{
+    uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
+    _Py_DECREF_STAT_INC();
+    if (_Py_IsOwnedByCurrentThread(op)) {
+        local--;
+        _Py_atomic_store_uint32_relaxed(&op->ob_ref_local, local);
+        if (local == 0) {
+            _Py_MergeZeroLocalRefcount(op);
+        }
+    }
+    else {
+        _Py_DecRefShared(op);
+    }
+}
+
 static inline void Py_DECREF(PyObject *op)
 {
     uint32_t local = _Py_atomic_load_uint32_relaxed(&op->ob_ref_local);
@@ -350,6 +413,15 @@ static inline void Py_DECREF(PyObject *op)
 #define Py_DECREF(op) Py_DECREF(_PyObject_CAST(op))
 
 #elif defined(Py_REF_DEBUG)
+static inline Py_ALWAYS_INLINE void Py_DECREF_MORTAL(PyObject *op)
+{
+    _Py_DECREF_STAT_INC();
+    _Py_DECREF_DecRefTotal();
+    if (--op->ob_refcnt == 0) {
+        _Py_Dealloc(op);
+    }
+}
+
 static inline void Py_DECREF(const char *filename, int lineno, PyObject *op)
 {
     if (op->ob_refcnt <= 0) {
@@ -359,15 +431,21 @@ static inline void Py_DECREF(const char *filename, int lineno, PyObject *op)
         _Py_DECREF_IMMORTAL_STAT_INC();
         return;
     }
-    _Py_DECREF_STAT_INC();
-    _Py_DECREF_DecRefTotal();
+    Py_DECREF_MORTAL(op);
+}
+#define Py_DECREF(op) Py_DECREF(__FILE__, __LINE__, _PyObject_CAST(op))
+
+#else
+static inline Py_ALWAYS_INLINE void Py_DECREF_MORTAL(PyObject *op)
+{
+    if (!_Py_IsImmortal(op)) {
+        _Py_DECREF_STAT_INC();
+    }
     if (--op->ob_refcnt == 0) {
         _Py_Dealloc(op);
     }
 }
-#define Py_DECREF(op) Py_DECREF(__FILE__, __LINE__, _PyObject_CAST(op))
 
-#else
 static inline Py_ALWAYS_INLINE void Py_DECREF(PyObject *op)
 {
     // Non-limited C API and limited C API for Python 3.9 and older access
@@ -376,10 +454,7 @@ static inline Py_ALWAYS_INLINE void Py_DECREF(PyObject *op)
         _Py_DECREF_IMMORTAL_STAT_INC();
         return;
     }
-    _Py_DECREF_STAT_INC();
-    if (--op->ob_refcnt == 0) {
-        _Py_Dealloc(op);
-    }
+    Py_DECREF_MORTAL(op);
 }
 #define Py_DECREF(op) Py_DECREF(_PyObject_CAST(op))
 #endif

@@ -3627,11 +3627,12 @@ dummy_func(
             DEOPT_IF(!PyStackRef_IsNull(null));
             DEOPT_IF(callable_o != (PyObject *)&PyUnicode_Type);
             STAT_INC(CALL, hit);
-            res = PyStackRef_FromPyObjectSteal(PyObject_Str(arg_o));
+            PyObject *str = PyObject_Str(arg_o);
             DEAD(null);
             DEAD(callable);
             PyStackRef_CLOSE(arg);
-            ERROR_IF(PyStackRef_IsNull(res), error);
+            ERROR_IF(str == NULL, error);
+            res = PyStackRef_FromPyObjectSteal(str);
         }
 
         macro(CALL_STR_1) =
@@ -3648,11 +3649,12 @@ dummy_func(
             DEOPT_IF(!PyStackRef_IsNull(null));
             DEOPT_IF(callable_o != (PyObject *)&PyTuple_Type);
             STAT_INC(CALL, hit);
-            res = PyStackRef_FromPyObjectSteal(PySequence_Tuple(arg_o));
+            PyObject *tuple = PySequence_Tuple(arg_o);
             DEAD(null);
             DEAD(callable);
             PyStackRef_CLOSE(arg);
-            ERROR_IF(PyStackRef_IsNull(res), error);
+            ERROR_IF(tuple == NULL, error);
+            res = PyStackRef_FromPyObjectSteal(tuple);
         }
 
         macro(CALL_TUPLE_1) =

@@ -1765,7 +1765,6 @@ _PyEvalFramePushAndInit_UnTagged(PyThreadState *tstate, _PyStackRef func,
                         PyObject *locals, PyObject *const* args,
                         size_t argcount, PyObject *kwnames, _PyInterpreterFrame *previous)
 {
-#if defined(Py_GIL_DISABLED)
     size_t kw_count = kwnames == NULL ? 0 : PyTuple_GET_SIZE(kwnames);
     size_t total_argcount = argcount + kw_count;
     _PyStackRef *tagged_args_buffer = PyMem_Malloc(sizeof(_PyStackRef) * total_argcount);
@@ -1782,9 +1781,6 @@ _PyEvalFramePushAndInit_UnTagged(PyThreadState *tstate, _PyStackRef func,
     _PyInterpreterFrame *res = _PyEvalFramePushAndInit(tstate, func, locals, (_PyStackRef const *)tagged_args_buffer, argcount, kwnames, previous);
     PyMem_Free(tagged_args_buffer);
     return res;
-#else
-    return _PyEvalFramePushAndInit(tstate, func, locals, (_PyStackRef const *)args, argcount, kwnames, previous);
-#endif
 }
 
 /* Same as _PyEvalFramePushAndInit but takes an args tuple and kwargs dict.

@@ -442,25 +442,14 @@ do { \
 /* How much scratch space to give stackref to PyObject* conversion. */
 #define MAX_STACKREF_SCRATCH 10
 
-#ifdef Py_GIL_DISABLED
 #define STACKREFS_TO_PYOBJECTS(ARGS, ARG_COUNT, NAME) \
     /* +1 because vectorcall might use -1 to write self */ \
     PyObject *NAME##_temp[MAX_STACKREF_SCRATCH+1]; \
     PyObject **NAME = _PyObjectArray_FromStackRefArray(ARGS, ARG_COUNT, NAME##_temp + 1);
-#else
-#define STACKREFS_TO_PYOBJECTS(ARGS, ARG_COUNT, NAME) \
-    PyObject **NAME = (PyObject **)ARGS; \
-    assert(NAME != NULL);
-#endif
 
-#ifdef Py_GIL_DISABLED
 #define STACKREFS_TO_PYOBJECTS_CLEANUP(NAME) \
     /* +1 because we +1 previously */ \
     _PyObjectArray_Free(NAME - 1, NAME##_temp);
-#else
-#define STACKREFS_TO_PYOBJECTS_CLEANUP(NAME) \
-    (void)(NAME);
-#endif
 
 #ifdef Py_GIL_DISABLED
 #define CONVERSION_FAILED(NAME) ((NAME) == NULL)

diff --git a/Python/executor_cases.c.h b/Python/executor_cases.c.h