Skip to content

Commit 0f9f236

Browse files
DinoVDinoV
committed
mimalloc: enable as allocator option
1 parent 96487df commit 0f9f236

File tree

6 files changed

+193
-10
lines changed

6 files changed

+193
-10
lines changed

‎Include/cpython/pymem.h

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,10 @@ typedef enum {
2929
PYMEM_ALLOCATOR_PYMALLOC = 5,
3030
PYMEM_ALLOCATOR_PYMALLOC_DEBUG = 6,
3131
#endif
32+
#ifdef WITH_MIMALLOC
33+
PYMEM_ALLOCATOR_MIMALLOC = 7,
34+
PYMEM_ALLOCATOR_MIMALLOC_DEBUG = 8,
35+
#endif
3236
} PyMemAllocatorName;
3337

3438

‎Include/internal/pycore_pymem_init.h

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -18,17 +18,31 @@ extern void * _PyMem_RawRealloc(void *, void *, size_t);
1818
extern void _PyMem_RawFree(void *, void *);
1919
#define PYRAW_ALLOC {NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree}
2020

21-
#ifdef WITH_PYMALLOC
21+
#ifdef WITH_MIMALLOC
22+
extern void *_PyMem_MiMalloc(void *ctx, size_t size);
23+
extern void * _PyMem_MiCalloc(void *ctx, size_t nelem, size_t elsize);
24+
extern void *_PyMem_MiRealloc(void *ctx, void *ptr, size_t size);
25+
extern void _PyMem_MiFree(void *ctx, void *ptr);
26+
extern void * _PyObject_MiMalloc(void *ctx, size_t nbytes);
27+
extern void* _PyObject_MiCalloc(void *ctx, size_t nelem, size_t elsize);
28+
extern void* _PyObject_MiRealloc(void *ctx, void *ptr, size_t nbytes);
29+
extern void _PyObject_MiFree(void *ctx, void *ptr);
30+
31+
# define PYMEM_ALLOC {NULL, _PyMem_MiMalloc, _PyMem_MiCalloc, _PyMem_MiRealloc, _PyMem_MiFree}
32+
# define PYOBJ_ALLOC {NULL, _PyObject_MiMalloc, _PyObject_MiCalloc, _PyObject_MiRealloc, _PyObject_MiFree}
33+
34+
#elif defined(WITH_PYMALLOC)
2235
extern void* _PyObject_Malloc(void *, size_t);
2336
extern void* _PyObject_Calloc(void *, size_t, size_t);
2437
extern void _PyObject_Free(void *, void *);
2538
extern void* _PyObject_Realloc(void *, void *, size_t);
2639
# define PYOBJ_ALLOC {NULL, _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free}
40+
# define PYMEM_ALLOC PYOBJ_ALLOC
2741
#else
2842
# define PYOBJ_ALLOC PYRAW_ALLOC
43+
# define PYMEM_ALLOC PYOBJ_ALLOC
2944
#endif // WITH_PYMALLOC
3045

31-
#define PYMEM_ALLOC PYOBJ_ALLOC
3246

3347
extern void* _PyMem_DebugRawMalloc(void *, size_t);
3448
extern void* _PyMem_DebugRawCalloc(void *, size_t, size_t);

‎Lib/test/support/__init__.py

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2025,6 +2025,11 @@ def with_pymalloc():
20252025
return _testcapi.WITH_PYMALLOC
20262026

20272027

2028+
def with_mimalloc():
2029+
import _testcapi
2030+
return _testcapi.WITH_MIMALLOC
2031+
2032+
20282033
class _ALWAYS_EQ:
20292034
"""
20302035
Object that is equal to anything.

‎Lib/test/test_cmd_line.py

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -724,7 +724,9 @@ def test_xdev(self):
724724
code = "import _testinternalcapi; print(_testinternalcapi.pymem_getallocatorsname())"
725725
with support.SuppressCrashReport():
726726
out = self.run_xdev("-c", code, check_exitcode=False)
727-
if support.with_pymalloc():
727+
if support.with_mimalloc():
728+
alloc_name = "mimalloc_debug"
729+
elif support.with_pymalloc():
728730
alloc_name = "pymalloc_debug"
729731
else:
730732
alloc_name = "malloc_debug"
@@ -803,7 +805,11 @@ def check_pythonmalloc(self, env_var, name):
803805
def test_pythonmalloc(self):
804806
# Test the PYTHONMALLOC environment variable
805807
pymalloc = support.with_pymalloc()
806-
if pymalloc:
808+
mimalloc = support.with_mimalloc()
809+
if mimalloc:
810+
default_name = 'mimalloc_debug' if support.Py_DEBUG else 'mimalloc'
811+
default_name_debug = 'mimalloc_debug'
812+
elif pymalloc:
807813
default_name = 'pymalloc_debug' if support.Py_DEBUG else 'pymalloc'
808814
default_name_debug = 'pymalloc_debug'
809815
else:
@@ -821,6 +827,11 @@ def test_pythonmalloc(self):
821827
('pymalloc', 'pymalloc'),
822828
('pymalloc_debug', 'pymalloc_debug'),
823829
))
830+
if mimalloc:
831+
tests.extend((
832+
('mimalloc', 'mimalloc'),
833+
('mimalloc_debug', 'mimalloc_debug'),
834+
))
824835

825836
for env_var, name in tests:
826837
with self.subTest(env_var=env_var, name=name):

‎Modules/_testcapi/mem.c

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -613,5 +613,14 @@ _PyTestCapi_Init_Mem(PyObject *mod)
613613
return -1;
614614
}
615615

616+
#ifdef WITH_MIMALLOC
617+
v = Py_True;
618+
#else
619+
v = Py_False;
620+
#endif
621+
if (PyModule_AddObjectRef(mod, "WITH_MIMALLOC", v) < 0) {
622+
return -1;
623+
}
624+
616625
return 0;
617626
}

‎Objects/obmalloc.c

Lines changed: 146 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,9 @@
1010

1111
#include <stdlib.h> // malloc()
1212
#include <stdbool.h>
13+
#ifdef WITH_MIMALLOC
14+
#include "mimalloc.h"
15+
#endif
1316

1417
#undef uint
1518
#define uint pymem_uint
@@ -74,25 +77,100 @@ _PyMem_RawFree(void *Py_UNUSED(ctx), void *ptr)
7477
free(ptr);
7578
}
7679

80+
#ifdef WITH_MIMALLOC
81+
82+
void *
83+
_PyMem_MiMalloc(void *ctx, size_t size)
84+
{
85+
PyThreadState *tstate = _PyThreadState_GET();
86+
return mi_heap_malloc(tstate->heaps[mi_heap_tag_default], size);
87+
}
88+
89+
void *
90+
_PyMem_MiCalloc(void *ctx, size_t nelem, size_t elsize)
91+
{
92+
PyThreadState *tstate = _PyThreadState_GET();
93+
return mi_heap_calloc(tstate->heaps[mi_heap_tag_default], nelem, elsize);
94+
}
95+
96+
void *
97+
_PyMem_MiRealloc(void *ctx, void *ptr, size_t size)
98+
{
99+
PyThreadState *tstate = _PyThreadState_GET();
100+
return mi_heap_realloc(tstate->heaps[mi_heap_tag_default], ptr, size);
101+
}
102+
103+
void
104+
_PyMem_MiFree(void *ctx, void *ptr)
105+
{
106+
mi_free(ptr);
107+
}
108+
109+
void *
110+
_PyObject_MiMalloc(void *ctx, size_t nbytes)
111+
{
112+
PyThreadState *tstate = _PyThreadState_GET();
113+
return mi_heap_malloc(tstate->heaps[mi_heap_tag_obj], nbytes);
114+
}
115+
116+
void *
117+
_PyObject_MiCalloc(void *ctx, size_t nelem, size_t elsize)
118+
{
119+
PyThreadState *tstate = _PyThreadState_GET();
120+
return mi_heap_calloc(tstate->heaps[mi_heap_tag_obj], nelem, elsize);
121+
}
122+
123+
124+
void *
125+
_PyObject_MiRealloc(void *ctx, void *ptr, size_t nbytes)
126+
{
127+
PyThreadState *tstate = _PyThreadState_GET();
128+
return mi_heap_realloc(tstate->heaps[mi_heap_tag_obj], ptr, nbytes);
129+
}
130+
131+
void
132+
_PyObject_MiFree(void *ctx, void *ptr)
133+
{
134+
mi_free(ptr);
135+
}
136+
137+
#endif // WITH_MIMALLOC
138+
139+
77140
#define MALLOC_ALLOC {NULL, _PyMem_RawMalloc, _PyMem_RawCalloc, _PyMem_RawRealloc, _PyMem_RawFree}
78-
#define PYRAW_ALLOC MALLOC_ALLOC
79141

80-
/* the default object allocator */
142+
143+
#ifdef WITH_MIMALLOC
144+
# define MIMALLOC_ALLOC {NULL, _PyMem_MiMalloc, _PyMem_MiCalloc, _PyMem_MiRealloc, _PyMem_MiFree}
145+
# define MIMALLOC_OBJALLOC {NULL, _PyObject_MiMalloc, _PyObject_MiCalloc, _PyObject_MiRealloc, _PyObject_MiFree}
146+
#endif
147+
148+
/* the pymalloc allocator */
81149

82150
// The actual implementation is further down.
83151

84-
#ifdef WITH_PYMALLOC
152+
#if defined(WITH_PYMALLOC)
85153
void* _PyObject_Malloc(void *ctx, size_t size);
86154
void* _PyObject_Calloc(void *ctx, size_t nelem, size_t elsize);
87155
void _PyObject_Free(void *ctx, void *p);
88156
void* _PyObject_Realloc(void *ctx, void *ptr, size_t size);
89157
# define PYMALLOC_ALLOC {NULL, _PyObject_Malloc, _PyObject_Calloc, _PyObject_Realloc, _PyObject_Free}
158+
#endif // WITH_PYMALLOC
159+
160+
#ifdef WITH_MIMALLOC
161+
# define PYRAW_ALLOC MALLOC_ALLOC
162+
# define PYMEM_ALLOC MIMALLOC_ALLOC
163+
# define PYOBJ_ALLOC MIMALLOC_OBJALLOC
164+
#elif defined(WITH_PYMALLOC)
165+
# define PYRAW_ALLOC MALLOC_ALLOC
166+
# define PYMEM_ALLOC PYMALLOC_ALLOC
90167
# define PYOBJ_ALLOC PYMALLOC_ALLOC
91168
#else
169+
# define PYRAW_ALLOC MALLOC_ALLOC
170+
# define PYMEM_ALLOC MALLOC_ALLOC
92171
# define PYOBJ_ALLOC MALLOC_ALLOC
93-
#endif // WITH_PYMALLOC
172+
#endif
94173

95-
#define PYMEM_ALLOC PYOBJ_ALLOC
96174

97175
/* the default debug allocators */
98176

@@ -291,6 +369,14 @@ _PyMem_GetAllocatorName(const char *name, PyMemAllocatorName *allocator)
291369
else if (strcmp(name, "pymalloc_debug") == 0) {
292370
*allocator = PYMEM_ALLOCATOR_PYMALLOC_DEBUG;
293371
}
372+
#endif
373+
#ifdef WITH_MIMALLOC
374+
else if (strcmp(name, "mimalloc") == 0) {
375+
*allocator = PYMEM_ALLOCATOR_MIMALLOC;
376+
}
377+
else if (strcmp(name, "mimalloc_debug") == 0) {
378+
*allocator = PYMEM_ALLOCATOR_MIMALLOC_DEBUG;
379+
}
294380
#endif
295381
else if (strcmp(name, "malloc") == 0) {
296382
*allocator = PYMEM_ALLOCATOR_MALLOC;
@@ -343,6 +429,26 @@ set_up_allocators_unlocked(PyMemAllocatorName allocator)
343429
break;
344430
}
345431
#endif
432+
#ifdef WITH_MIMALLOC
433+
case PYMEM_ALLOCATOR_MIMALLOC:
434+
case PYMEM_ALLOCATOR_MIMALLOC_DEBUG:
435+
{
436+
PyMemAllocatorEx malloc_alloc = MALLOC_ALLOC;
437+
set_allocator_unlocked(PYMEM_DOMAIN_RAW, &malloc_alloc);
438+
439+
PyMemAllocatorEx pymalloc = MIMALLOC_ALLOC;
440+
set_allocator_unlocked(PYMEM_DOMAIN_MEM, &pymalloc);
441+
442+
PyMemAllocatorEx objmalloc = MIMALLOC_OBJALLOC;
443+
set_allocator_unlocked(PYMEM_DOMAIN_OBJ, &objmalloc);
444+
445+
if (allocator == PYMEM_ALLOCATOR_MIMALLOC_DEBUG) {
446+
set_up_debug_hooks_unlocked();
447+
}
448+
449+
break;
450+
}
451+
#endif
346452

347453
case PYMEM_ALLOCATOR_MALLOC:
348454
case PYMEM_ALLOCATOR_MALLOC_DEBUG:
@@ -390,6 +496,10 @@ get_current_allocator_name_unlocked(void)
390496
#ifdef WITH_PYMALLOC
391497
PyMemAllocatorEx pymalloc = PYMALLOC_ALLOC;
392498
#endif
499+
#ifdef WITH_MIMALLOC
500+
PyMemAllocatorEx mimalloc = MIMALLOC_ALLOC;
501+
PyMemAllocatorEx mimalloc_obj = MIMALLOC_OBJALLOC;
502+
#endif
393503

394504
if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&
395505
pymemallocator_eq(&_PyMem, &malloc_alloc) &&
@@ -405,6 +515,14 @@ get_current_allocator_name_unlocked(void)
405515
return "pymalloc";
406516
}
407517
#endif
518+
#ifdef WITH_MIMALLOC
519+
if (pymemallocator_eq(&_PyMem_Raw, &malloc_alloc) &&
520+
pymemallocator_eq(&_PyMem, &mimalloc) &&
521+
pymemallocator_eq(&_PyObject, &mimalloc_obj))
522+
{
523+
return "mimalloc";
524+
}
525+
#endif
408526

409527
PyMemAllocatorEx dbg_raw = PYDBGRAW_ALLOC;
410528
PyMemAllocatorEx dbg_mem = PYDBGMEM_ALLOC;
@@ -428,6 +546,14 @@ get_current_allocator_name_unlocked(void)
428546
{
429547
return "pymalloc_debug";
430548
}
549+
#endif
550+
#ifdef WITH_MIMALLOC
551+
if (pymemallocator_eq(&_PyMem_Debug.raw.alloc, &malloc_alloc) &&
552+
pymemallocator_eq(&_PyMem_Debug.mem.alloc, &mimalloc) &&
553+
pymemallocator_eq(&_PyMem_Debug.obj.alloc, &mimalloc_obj))
554+
{
555+
return "mimalloc_debug";
556+
}
431557
#endif
432558
}
433559
return NULL;
@@ -443,13 +569,14 @@ _PyMem_GetCurrentAllocatorName(void)
443569
}
444570

445571

446-
#ifdef WITH_PYMALLOC
572+
#if defined(WITH_PYMALLOC) || defined(WITH_MIMALLOC)
447573
static int
448574
_PyMem_DebugEnabled(void)
449575
{
450576
return (_PyObject.malloc == _PyMem_DebugMalloc);
451577
}
452578

579+
#ifdef WITH_PYMALLOC
453580
static int
454581
_PyMem_PymallocEnabled(void)
455582
{
@@ -461,6 +588,19 @@ _PyMem_PymallocEnabled(void)
461588
}
462589
}
463590
#endif
591+
#ifdef WITH_MIMALLOC
592+
static int
593+
_PyMem_MimallocEnabled(void)
594+
{
595+
if (_PyMem_DebugEnabled()) {
596+
return (_PyMem_Debug.obj.alloc.malloc == _PyObject_MiMalloc);
597+
}
598+
else {
599+
return (_PyObject.malloc == _PyObject_MiMalloc);
600+
}
601+
}
602+
#endif
603+
#endif // defined(WITH_PYMALLOC) || defined(WITH_MIMALLOC)
464604

465605

466606
static void

0 commit comments

Comments
 (0)