gh-112529: Use GC heaps for GC allocations in free-threaded builds (gh-114157)

* gh-112529: Use GC heaps for GC allocations in free-threaded builds

The free-threaded build's garbage collector implementation will need to
find GC objects by traversing mimalloc heaps. This hooks up the
allocation calls with the correct heaps by using a thread-local
"current_obj_heap" variable.

* Refactor out setting heap based on type

Author: colesbury

Include/internal/pycore_object_alloc.h

@@ -0,0 +1,71 @@
+#ifndef Py_INTERNAL_OBJECT_ALLOC_H
+#define Py_INTERNAL_OBJECT_ALLOC_H
+
+#include "pycore_object.h"      // _PyType_HasFeature()
+#include "pycore_pystate.h"     // _PyThreadState_GET()
+#include "pycore_tstate.h"      // _PyThreadStateImpl
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#ifndef Py_BUILD_CORE
+#  error "this header requires Py_BUILD_CORE define"
+#endif
+
+#ifdef Py_GIL_DISABLED
+static inline mi_heap_t *
+_PyObject_GetAllocationHeap(_PyThreadStateImpl *tstate, PyTypeObject *tp)
+{
+    struct _mimalloc_thread_state *m = &tstate->mimalloc;
+    if (_PyType_HasFeature(tp, Py_TPFLAGS_PREHEADER)) {
+        return &m->heaps[_Py_MIMALLOC_HEAP_GC_PRE];
+    }
+    else if (_PyType_IS_GC(tp)) {
+        return &m->heaps[_Py_MIMALLOC_HEAP_GC];
+    }
+    else {
+        return &m->heaps[_Py_MIMALLOC_HEAP_OBJECT];
+    }
+}
+#endif
+
+// Sets the heap used for PyObject_Malloc(), PyObject_Realloc(), etc. calls in
+// Py_GIL_DISABLED builds. We use different heaps depending on if the object
+// supports GC and if it has a pre-header. We smuggle the choice of heap
+// through the _mimalloc_thread_state. In the default build, this simply
+// calls PyObject_Malloc().
+static inline void *
+_PyObject_MallocWithType(PyTypeObject *tp, size_t size)
+{
+#ifdef Py_GIL_DISABLED
+    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
+    struct _mimalloc_thread_state *m = &tstate->mimalloc;
+    m->current_object_heap = _PyObject_GetAllocationHeap(tstate, tp);
+#endif
+    void *mem = PyObject_Malloc(size);
+#ifdef Py_GIL_DISABLED
+    m->current_object_heap = &m->heaps[_Py_MIMALLOC_HEAP_OBJECT];
+#endif
+    return mem;
+}
+
+static inline void *
+_PyObject_ReallocWithType(PyTypeObject *tp, void *ptr, size_t size)
+{
+#ifdef Py_GIL_DISABLED
+    _PyThreadStateImpl *tstate = (_PyThreadStateImpl *)_PyThreadState_GET();
+    struct _mimalloc_thread_state *m = &tstate->mimalloc;
+    m->current_object_heap = _PyObject_GetAllocationHeap(tstate, tp);
+#endif
+    void *mem = PyObject_Realloc(ptr, size);
+#ifdef Py_GIL_DISABLED
+    m->current_object_heap = &m->heaps[_Py_MIMALLOC_HEAP_OBJECT];
+#endif
+    return mem;
+}
+
+#ifdef __cplusplus
+}
+#endif
+#endif  // !Py_INTERNAL_OBJECT_ALLOC_H

Makefile.pre.in

@@ -1852,6 +1852,7 @@ PYTHON_HEADERS= \
 		$(srcdir)/Include/internal/pycore_moduleobject.h \
 		$(srcdir)/Include/internal/pycore_namespace.h \
 		$(srcdir)/Include/internal/pycore_object.h \
+		$(srcdir)/Include/internal/pycore_object_alloc.h \
 		$(srcdir)/Include/internal/pycore_object_state.h \
 		$(srcdir)/Include/internal/pycore_obmalloc.h \
 		$(srcdir)/Include/internal/pycore_obmalloc_init.h \

Objects/typeobject.c

@@ -11,6 +11,7 @@
 #include "pycore_modsupport.h"    // _PyArg_NoKwnames()
 #include "pycore_moduleobject.h"  // _PyModule_GetDef()
 #include "pycore_object.h"        // _PyType_HasFeature()
+#include "pycore_object_alloc.h"  // _PyObject_MallocWithType()
 #include "pycore_pyerrors.h"      // _PyErr_Occurred()
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_symtable.h"      // _Py_Mangle()
@@ -1729,7 +1730,7 @@ _PyType_AllocNoTrack(PyTypeObject *type, Py_ssize_t nitems)
     const size_t size = _PyObject_VAR_SIZE(type, nitems+1);
 
     const size_t presize = _PyType_PreHeaderSize(type);
-    char *alloc = PyObject_Malloc(size + presize);
+    char *alloc = _PyObject_MallocWithType(type, size + presize);
     if (alloc  == NULL) {
         return PyErr_NoMemory();
     }

PCbuild/pythoncore.vcxproj

@@ -255,6 +255,7 @@
     <ClInclude Include="..\Include\internal\pycore_moduleobject.h" />
     <ClInclude Include="..\Include\internal\pycore_namespace.h" />
     <ClInclude Include="..\Include\internal\pycore_object.h" />
+    <ClInclude Include="..\Include\internal\pycore_object_alloc.h" />
     <ClInclude Include="..\Include\internal\pycore_object_state.h" />
     <ClInclude Include="..\Include\internal\pycore_obmalloc.h" />
     <ClInclude Include="..\Include\internal\pycore_obmalloc_init.h" />

PCbuild/pythoncore.vcxproj.filters

@@ -693,6 +693,9 @@
     <ClInclude Include="..\Include\internal\pycore_object.h">
       <Filter>Include\internal</Filter>
     </ClInclude>
+    <ClInclude Include="..\Include\internal\pycore_object_alloc.h">
+      <Filter>Include\internal</Filter>
+    </ClInclude>
     <ClInclude Include="..\Include\internal\pycore_object_state.h">
       <Filter>Include\internal</Filter>
     </ClInclude>

Python/gc.c

@@ -9,6 +9,7 @@
 #include "pycore_initconfig.h"
 #include "pycore_interp.h"        // PyInterpreterState.gc
 #include "pycore_object.h"
+#include "pycore_object_alloc.h"  // _PyObject_MallocWithType()
 #include "pycore_pyerrors.h"
 #include "pycore_pystate.h"       // _PyThreadState_GET()
 #include "pycore_weakref.h"       // _PyWeakref_ClearRef()
@@ -1795,14 +1796,14 @@ _Py_RunGC(PyThreadState *tstate)
 }
 
 static PyObject *
-gc_alloc(size_t basicsize, size_t presize)
+gc_alloc(PyTypeObject *tp, size_t basicsize, size_t presize)
 {
     PyThreadState *tstate = _PyThreadState_GET();
     if (basicsize > PY_SSIZE_T_MAX - presize) {
         return _PyErr_NoMemory(tstate);
     }
     size_t size = presize + basicsize;
-    char *mem = PyObject_Malloc(size);
+    char *mem = _PyObject_MallocWithType(tp, size);
     if (mem == NULL) {
         return _PyErr_NoMemory(tstate);
     }
@@ -1817,7 +1818,7 @@ PyObject *
 _PyObject_GC_New(PyTypeObject *tp)
 {
     size_t presize = _PyType_PreHeaderSize(tp);
-    PyObject *op = gc_alloc(_PyObject_SIZE(tp), presize);
+    PyObject *op = gc_alloc(tp, _PyObject_SIZE(tp), presize);
     if (op == NULL) {
         return NULL;
     }
@@ -1836,7 +1837,7 @@ _PyObject_GC_NewVar(PyTypeObject *tp, Py_ssize_t nitems)
     }
     size_t presize = _PyType_PreHeaderSize(tp);
     size_t size = _PyObject_VAR_SIZE(tp, nitems);
-    op = (PyVarObject *)gc_alloc(size, presize);
+    op = (PyVarObject *)gc_alloc(tp, size, presize);
     if (op == NULL) {
         return NULL;
     }
@@ -1848,7 +1849,7 @@ PyObject *
 PyUnstable_Object_GC_NewWithExtraData(PyTypeObject *tp, size_t extra_size)
 {
     size_t presize = _PyType_PreHeaderSize(tp);
-    PyObject *op = gc_alloc(_PyObject_SIZE(tp) + extra_size, presize);
+    PyObject *op = gc_alloc(tp, _PyObject_SIZE(tp) + extra_size, presize);
     if (op == NULL) {
         return NULL;
     }
@@ -1867,7 +1868,7 @@ _PyObject_GC_Resize(PyVarObject *op, Py_ssize_t nitems)
         return (PyVarObject *)PyErr_NoMemory();
     }
     char *mem = (char *)op - presize;
-    mem = (char *)PyObject_Realloc(mem,  presize + basicsize);
+    mem = (char *)_PyObject_ReallocWithType(Py_TYPE(op), mem, presize + basicsize);
     if (mem == NULL) {
         return (PyVarObject *)PyErr_NoMemory();
     }