Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit 0bfbca2

Browse files
kumaraditya303kumaraditya303
authored
ENH: use PyMem_RawMalloc instead of system allocator for better performance on free-threading (#31503)
1 parent dfb0fc6 commit 0bfbca2

22 files changed

Lines changed: 155 additions & 150 deletions

‎doc/release/upcoming_changes/31503.change.rst‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,4 @@
1+
NumPy's internal memory allocations now use ``PyMem_RawMalloc``
2+
---------------------------------------------------------------
3+
NumPy's internal memory allocations now use ``PyMem_RawMalloc`` instead of
4+
``malloc`` and can be tracked by ``tracemalloc``.

‎numpy/_core/src/_simd/_simd_convert.inc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ simd_sequence_new(Py_ssize_t len, simd_data_type dtype)
6060
const simd_data_info *info = simd_data_getinfo(dtype);
6161
assert(len > 0 && info->is_sequence && info->lane_size > 0);
6262
size_t size = sizeof(simd__alloc_data) + len * info->lane_size + NPY_SIMD_WIDTH;
63-
void *ptr = malloc(size);
63+
void *ptr = PyMem_RawMalloc(size);
6464
if (ptr == NULL) {
6565
return PyErr_NoMemory();
6666
}
@@ -82,7 +82,7 @@ simd_sequence_len(void const *ptr)
8282
static void
8383
simd_sequence_free(void *ptr)
8484
{
85-
free(((simd__alloc_data *)ptr)[-1].ptr);
85+
PyMem_RawFree(((simd__alloc_data *)ptr)[-1].ptr);
8686
}
8787

8888
static void *

‎numpy/_core/src/common/mem_overlap.c‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -544,9 +544,9 @@ solve_diophantine(unsigned int n, diophantine_term_t *E, npy_int64 b,
544544
diophantine_term_t *Ep = NULL;
545545
npy_int64 *Epsilon = NULL, *Gamma = NULL;
546546

547-
Ep = malloc(n * sizeof(diophantine_term_t));
548-
Epsilon = malloc(n * sizeof(npy_int64));
549-
Gamma = malloc(n * sizeof(npy_int64));
547+
Ep = PyMem_RawMalloc(n * sizeof(diophantine_term_t));
548+
Epsilon = PyMem_RawMalloc(n * sizeof(npy_int64));
549+
Gamma = PyMem_RawMalloc(n * sizeof(npy_int64));
550550
if (Ep == NULL || Epsilon == NULL || Gamma == NULL) {
551551
res = MEM_OVERLAP_ERROR;
552552
}
@@ -557,9 +557,9 @@ solve_diophantine(unsigned int n, diophantine_term_t *E, npy_int64 b,
557557
res = diophantine_dfs(n, n-1, E, Ep, Gamma, Epsilon, b, max_work,
558558
require_ub_nontrivial, x, &count);
559559
}
560-
free(Ep);
561-
free(Gamma);
562-
free(Epsilon);
560+
PyMem_RawFree(Ep);
561+
PyMem_RawFree(Gamma);
562+
PyMem_RawFree(Epsilon);
563563
return res;
564564
}
565565
}

‎numpy/_core/src/common/npy_cpuinfo_parser.h‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -188,7 +188,7 @@ extract_cpuinfo_field(const char* buffer, int buflen, const char* field)
188188

189189
/* Copy the line into a heap-allocated buffer */
190190
len = q - p;
191-
result = malloc(len + 1);
191+
result = PyMem_RawMalloc(len + 1);
192192
if (result == NULL) {
193193
goto EXIT;
194194
}
@@ -247,15 +247,15 @@ get_feature_from_proc_cpuinfo(unsigned long *hwcap, unsigned long *hwcap2) {
247247
if (cpuinfo_len < 0) {
248248
return 0;
249249
}
250-
char *cpuinfo = malloc(cpuinfo_len);
250+
char *cpuinfo = PyMem_RawMalloc(cpuinfo_len);
251251
if (cpuinfo == NULL) {
252252
return 0;
253253
}
254254

255255
cpuinfo_len = read_file("/proc/cpuinfo", cpuinfo, cpuinfo_len);
256256
char *cpuFeatures = extract_cpuinfo_field(cpuinfo, cpuinfo_len, "Features");
257257
if (cpuFeatures == NULL) {
258-
free(cpuinfo);
258+
PyMem_RawFree(cpuinfo);
259259
return 0;
260260
}
261261
*hwcap |= has_list_item(cpuFeatures, "fphp") ? NPY__HWCAP_FPHP : 0;
@@ -281,8 +281,8 @@ get_feature_from_proc_cpuinfo(unsigned long *hwcap, unsigned long *hwcap2) {
281281
*hwcap |= has_list_item(cpuFeatures, "sha2") ? NPY__HWCAP_SHA2 : 0;
282282
*hwcap |= has_list_item(cpuFeatures, "crc32") ? NPY__HWCAP_CRC32 : 0;
283283
#endif
284-
free(cpuinfo);
285-
free(cpuFeatures);
284+
PyMem_RawFree(cpuinfo);
285+
PyMem_RawFree(cpuFeatures);
286286
return 1;
287287
}
288288
#endif /* NUMPY_CORE_SRC_COMMON_NPY_CPUINFO_PARSER_H_ */

‎numpy/_core/src/multiarray/arraytypes.c.src‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -652,7 +652,7 @@ UNICODE_getitem(void *ip, void *vap)
652652

653653
/* swap and align if needed */
654654
if (swap || align) {
655-
buf = (npy_ucs4 *)malloc(size);
655+
buf = (npy_ucs4 *)PyMem_RawMalloc(size);
656656
if (buf == NULL) {
657657
PyErr_NoMemory();
658658
return NULL;
@@ -669,7 +669,7 @@ UNICODE_getitem(void *ip, void *vap)
669669
ucs4len--;
670670
}
671671
PyObject *ret = PyUnicode_FromKindAndData(PyUnicode_4BYTE_KIND, src, ucs4len);
672-
free(buf);
672+
PyMem_RawFree(buf);
673673
return ret;
674674
}
675675

‎numpy/_core/src/multiarray/buffer.c‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -934,7 +934,7 @@ _descriptor_from_pep3118_format(char const *s)
934934
}
935935

936936
/* Strip whitespace, except from field names */
937-
buf = malloc(strlen(s) + 1);
937+
buf = PyMem_RawMalloc(strlen(s) + 1);
938938
if (buf == NULL) {
939939
PyErr_NoMemory();
940940
return NULL;
@@ -956,15 +956,15 @@ _descriptor_from_pep3118_format(char const *s)
956956

957957
str = PyUnicode_FromStringAndSize(buf, strlen(buf));
958958
if (str == NULL) {
959-
free(buf);
959+
PyMem_RawFree(buf);
960960
return NULL;
961961
}
962962

963963
/* Convert */
964964
_numpy_internal = PyImport_ImportModule("numpy._core._internal");
965965
if (_numpy_internal == NULL) {
966966
Py_DECREF(str);
967-
free(buf);
967+
PyMem_RawFree(buf);
968968
return NULL;
969969
}
970970
descr = PyObject_CallMethod(
@@ -977,18 +977,18 @@ _descriptor_from_pep3118_format(char const *s)
977977
PyErr_Format(PyExc_ValueError,
978978
"'%s' is not a valid PEP 3118 buffer format string", buf);
979979
npy_PyErr_ChainExceptionsCause(exc, val, tb);
980-
free(buf);
980+
PyMem_RawFree(buf);
981981
return NULL;
982982
}
983983
if (!PyArray_DescrCheck(descr)) {
984984
PyErr_Format(PyExc_RuntimeError,
985985
"internal error: numpy._core._internal._dtype_from_pep3118 "
986986
"did not return a valid dtype, got %s", buf);
987987
Py_DECREF(descr);
988-
free(buf);
988+
PyMem_RawFree(buf);
989989
return NULL;
990990
}
991-
free(buf);
991+
PyMem_RawFree(buf);
992992
return (PyArray_Descr*)descr;
993993
}
994994

‎numpy/_core/src/multiarray/ctors.c‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -135,7 +135,7 @@ swab_separator(const char *sep)
135135
int skip_space = 0;
136136
char *s, *start;
137137

138-
s = start = malloc(strlen(sep)+3);
138+
s = start = PyMem_RawMalloc(strlen(sep)+3);
139139
if (s == NULL) {
140140
PyErr_NoMemory();
141141
return NULL;
@@ -3602,7 +3602,7 @@ array_from_text(PyArray_Descr *dtype, npy_intp num, char const *sep, size_t *nre
36023602
}
36033603
NPY_END_ALLOW_THREADS;
36043604

3605-
free(clean_sep);
3605+
PyMem_RawFree(clean_sep);
36063606

36073607
if (stop_reading_flag == -2) {
36083608
if (PyErr_Occurred()) {

‎numpy/_core/src/multiarray/dtypemeta.c‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -417,7 +417,7 @@ dtypemeta_initialize_struct_from_spec(
417417
* Like PyArray_RegisterDataType, this mutates global userdescrs state
418418
* and is expected to run during module import while single-threaded.
419419
*/
420-
_PyArray_LegacyDescr **tmp = realloc(userdescrs,
420+
_PyArray_LegacyDescr **tmp = PyMem_RawRealloc(userdescrs,
421421
(NPY_NUMUSERTYPES + 1) * sizeof(void *));
422422
if (tmp == NULL) {
423423
PyErr_NoMemory();

‎numpy/_core/src/multiarray/item_selection.c‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1960,7 +1960,7 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
19601960
PyDataMem_FREE(valbuffer);
19611961
goto fail;
19621962
}
1963-
swaps = malloc(NPY_LIKELY(n > 0) ? n * sizeof(int) : 1);
1963+
swaps = PyMem_RawMalloc(NPY_LIKELY(n > 0) ? n * sizeof(int) : 1);
19641964
if (swaps == NULL) {
19651965
PyDataMem_FREE(valbuffer);
19661966
PyDataMem_FREE(indbuffer);
@@ -1993,7 +1993,7 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
19931993
&& PyErr_Occurred())) {
19941994
PyDataMem_FREE(valbuffer);
19951995
PyDataMem_FREE(indbuffer);
1996-
free(swaps);
1996+
PyMem_RawFree(swaps);
19971997
goto fail;
19981998
}
19991999
PyArray_ITER_NEXT(its[j]);
@@ -2004,7 +2004,7 @@ PyArray_LexSort(PyObject *sort_keys, int axis)
20042004
}
20052005
PyDataMem_FREE(valbuffer);
20062006
PyDataMem_FREE(indbuffer);
2007-
free(swaps);
2007+
PyMem_RawFree(swaps);
20082008
}
20092009
else {
20102010
while (size--) {

‎numpy/_core/src/multiarray/usertypes.c‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -57,7 +57,7 @@ _append_new(int **p_types, int insert)
5757
while (types[n] != NPY_NOTYPE) {
5858
n++;
5959
}
60-
newtypes = (int *)realloc(types, (n + 2)*sizeof(int));
60+
newtypes = (int *)PyMem_RawRealloc(types, (n + 2)*sizeof(int));
6161
if (newtypes == NULL) {
6262
PyErr_NoMemory();
6363
return -1;
@@ -281,7 +281,7 @@ PyArray_RegisterDataType(PyArray_DescrProto *descr_proto)
281281
}
282282
}
283283

284-
userdescrs = realloc(userdescrs,
284+
userdescrs = PyMem_RawRealloc(userdescrs,
285285
(NPY_NUMUSERTYPES+1)*sizeof(void *));
286286
if (userdescrs == NULL) {
287287
PyErr_SetString(PyExc_MemoryError, "RegisterDataType");
@@ -484,7 +484,7 @@ PyArray_RegisterCanCast(PyArray_Descr *descr, int totype,
484484
* -- they become part of the data-type
485485
*/
486486
if (PyDataType_GetArrFuncs(descr)->cancastto == NULL) {
487-
PyDataType_GetArrFuncs(descr)->cancastto = (int *)malloc(1*sizeof(int));
487+
PyDataType_GetArrFuncs(descr)->cancastto = (int *)PyMem_RawMalloc(1*sizeof(int));
488488
if (PyDataType_GetArrFuncs(descr)->cancastto == NULL) {
489489
PyErr_NoMemory();
490490
return -1;
@@ -498,7 +498,7 @@ PyArray_RegisterCanCast(PyArray_Descr *descr, int totype,
498498
if (PyDataType_GetArrFuncs(descr)->cancastscalarkindto == NULL) {
499499
int i;
500500
PyDataType_GetArrFuncs(descr)->cancastscalarkindto =
501-
(int **)malloc(NPY_NSCALARKINDS* sizeof(int*));
501+
(int **)PyMem_RawMalloc(NPY_NSCALARKINDS* sizeof(int*));
502502
if (PyDataType_GetArrFuncs(descr)->cancastscalarkindto == NULL) {
503503
PyErr_NoMemory();
504504
return -1;
@@ -509,7 +509,7 @@ PyArray_RegisterCanCast(PyArray_Descr *descr, int totype,
509509
}
510510
if (PyDataType_GetArrFuncs(descr)->cancastscalarkindto[scalar] == NULL) {
511511
PyDataType_GetArrFuncs(descr)->cancastscalarkindto[scalar] =
512-
(int *)malloc(1*sizeof(int));
512+
(int *)PyMem_RawMalloc(1*sizeof(int));
513513
if (PyDataType_GetArrFuncs(descr)->cancastscalarkindto[scalar] == NULL) {
514514
PyErr_NoMemory();
515515
return -1;

0 commit comments

Comments
 (0)