diff --git a/Include/internal/pycore_object_alloc.h b/Include/internal/pycore_object_alloc.h
new file mode 100644
index 00000000000..8cc7a444bc9
--- /dev/null
+++ b/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
diff --git a/Makefile.pre.in b/Makefile.pre.in
index d251e7c481b..1107259b5ae 100644
--- a/Makefile.pre.in
+++ b/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 \
diff --git a/Objects/typeobject.c b/Objects/typeobject.c
index ea29a38d74a..3a35a5b5975 100644
--- a/Objects/typeobject.c
+++ b/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();
}
diff --git a/PCbuild/pythoncore.vcxproj b/PCbuild/pythoncore.vcxproj
index 64738b1bbf2..57275fb2039 100644
--- a/PCbuild/pythoncore.vcxproj
+++ b/PCbuild/pythoncore.vcxproj
@@ -255,6 +255,7 @@
+
diff --git a/PCbuild/pythoncore.vcxproj.filters b/PCbuild/pythoncore.vcxproj.filters
index b37ca2dfed5..51cbb079b5b 100644
--- a/PCbuild/pythoncore.vcxproj.filters
+++ b/PCbuild/pythoncore.vcxproj.filters
@@ -693,6 +693,9 @@
Include\internal
+
+ Include\internal
+
Include\internal
diff --git a/Python/gc.c b/Python/gc.c
index 9f9a755f6ac..14870505ef1 100644
--- a/Python/gc.c
+++ b/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();
}