NEP 49 — 数据分配策略#
- 作者:
Matti Picus
- 状态:
最终
- 类型:
标准轨迹
- 创建:
2021-04-18
- 决议:
摘要#
numpy.ndarray 需要额外的内存分配来保存 numpy.ndarray.strides、numpy.ndarray.shape 和 numpy.ndarray.data 属性。这些属性在 __new__ 方法中创建 Python 对象后被特别分配。
此 NEP 提出了一种机制，可以使用用户提供的替代方案来覆盖用于 ndarray->data 的内存管理策略。此分配保存数据，并且可能非常大。由于访问此数据经常成为性能瓶颈，因此自定义分配策略以保证数据对齐或将分配固定到专用内存硬件可以实现特定于硬件的优化。其他分配保持不变。
动机和范围#
用户可能希望使用他们自己的例程来覆盖内部数据内存例程。这两个用例是确保数据对齐并将某些分配固定到某些 NUMA 内核。这种对齐的需求在邮件列表中多次讨论 2005 年，以及在 问题 5312 中 2014 年，这导致了 PR 5457 和更多邮件列表讨论 此处 和此处。在关于该问题的评论中 2017 年，一位用户描述了 64 字节对齐如何将性能提高 40 倍。
相关的还有关于在 Linux 上使用 madvise 和巨页的 问题 14177。
各种跟踪和分析库，如 filprofiler 或 electric fence 覆盖了 malloc。
关于 BPO 18835 的长期 CPython 讨论始于讨论对 PyMem_Alloc32 和 PyMem_Alloc64 的需求。早期的结论是，（浪费填充的）成本与对齐内存的益处相比最好留给用户，但随后演变成关于处理内存分配的各种提案的讨论，包括 PEP 445 内存接口 到 PyTraceMalloc_Track，这显然是专门为 NumPy 添加的。
允许用户通过 NumPy C-API 实现不同的策略将能够探索这个丰富的可能优化领域。目的是创建一个足够灵活的接口，而不会给规范用户带来负担。
用法和影响#
新函数只能通过 NumPy C-API 访问。本 NEP 稍后将包含一个示例。添加的 struct 将增加 ndarray 对象的大小。这是这种方法必须付出的代价。我们可以合理地确定大小的改变对最终用户代码的影响最小，因为 NumPy 1.20 版已经改变了对象大小。
该实现保留了 PyTraceMalloc_Track 的使用，以跟踪 NumPy 中已存在的分配。
向后兼容性#
该设计不会破坏向后兼容性。将值赋给 ndarray->data 指针的项目已经破坏了当前的内存管理策略，并且应该在调用 Py_DECREF 之前恢复 ndarray->data。如上所述，大小的变化不应影响最终用户。
详细说明#
高级设计#
希望更改 NumPy 数据内存管理例程的用户将使用 PyDataMem_SetHandler()，它使用 PyDataMem_Handler 结构来保存用于管理数据内存的函数指针。为了允许 context 的生命周期管理，该结构被包装在一个 PyCapsule 中。
由于对 PyDataMem_SetHandler 的调用将更改默认函数，但该函数可能在 ndarray 对象的生命周期内被调用，因此每个 ndarray 将在其实例化时携带它所使用的 PyDataMem_Handler 包装的 PyCapsule，并且这些将用于重新分配或释放实例的数据内存。在内部，NumPy 可能会对数据内存的指针使用 memcpy 或 memset。
处理程序的名称将在 Python 层面通过 numpy.core.multiarray.get_handler_name(arr) 函数公开。如果调用为 numpy.core.multiarray.get_handler_name()，它将返回将用于为下一个新的 ndarrray 分配数据的处理程序的名称。
处理程序的版本将在 Python 层面通过 numpy.core.multiarray.get_handler_version(arr) 函数公开。如果调用为 numpy.core.multiarray.get_handler_version()，它将返回将用于为下一个新的 ndarrray 分配数据的处理程序的版本。
当前版本 1 允许将来增强 PyDataMemAllocator。如果添加字段，则必须将其添加到末尾。
NumPy C-API 函数#
-
type PyDataMem_Handler#
一个结构体，用于保存用于操作内存的函数指针
typedef struct { char name[127]; /* multiple of 64 to keep the struct aligned */ uint8_t version; /* currently 1 */ PyDataMemAllocator allocator; } PyDataMem_Handler;
其中分配器结构是
/* The declaration of free differs from PyMemAllocatorEx */ typedef struct { void *ctx; void* (*malloc) (void *ctx, size_t size); void* (*calloc) (void *ctx, size_t nelem, size_t elsize); void* (*realloc) (void *ctx, void *ptr, size_t new_size); void (*free) (void *ctx, void *ptr, size_t size); } PyDataMemAllocator;
在
free函数中使用size参数将此结构体与 Python 中的PyMemAllocatorEx结构体区分开来。此调用签名目前在 NumPy 内部使用，也在其他地方使用，例如 C++98 <https://cppreference.cn/w/cpp/memory/allocator/deallocate>，C++11 <https://cppreference.cn/w/cpp/memory/allocator_traits/deallocate> 和 Rust (allocator_api) <https://doc.rust-lang.net.cn/std/alloc/trait.Allocator.html#tymethod.deallocate>。PyDataMemAllocator 接口的使用者必须跟踪
size并确保它与传递给(m|c|re)alloc函数的参数一致。当请求的数组的形状包含
0时，NumPy 本身可能会违反此要求，因此 PyDataMemAllocators 的编写者应该将size参数视为最佳猜测。修复此问题的相关工作正在 PR 15780 和 15788 中进行，但尚未解决。解决后，应重新审视此 NEP。
PyDataMem_Handler 线程安全性和生命周期#
活动处理程序通过 ContextVar 存储在当前的 Context 中。这确保它可以针对每个线程和每个异步协程进行配置。
目前没有 PyDataMem_Handler 的生命周期管理。 PyDataMem_SetHandler 的使用者必须确保参数在其分配的任何对象的生命周期内以及在其作为活动处理程序期间保持活动状态。实际上，这意味着处理程序必须是永久存在的。
作为实现细节，当前此 ContextVar 包含一个 PyCapsule 对象，该对象存储指向 PyDataMem_Handler 的指针，且没有析构函数，但不应依赖于此。
示例代码#
此代码为每个 data 指针添加 64 字节的头部，并在头部中存储有关分配的信息。在调用 free 之前，会进行检查以确保 sz 参数正确。
#define NPY_NO_DEPRECATED_API NPY_1_7_API_VERSION
#include <numpy/arrayobject.h>
NPY_NO_EXPORT void *
typedef struct {
void *(*malloc)(size_t);
void *(*calloc)(size_t, size_t);
void *(*realloc)(void *, size_t);
void (*free)(void *);
} Allocator;
NPY_NO_EXPORT void *
shift_alloc(Allocator *ctx, size_t sz) {
char *real = (char *)ctx->malloc(sz + 64);
if (real == NULL) {
return NULL;
}
snprintf(real, 64, "originally allocated %ld", (unsigned long)sz);
return (void *)(real + 64);
}
NPY_NO_EXPORT void *
shift_zero(Allocator *ctx, size_t sz, size_t cnt) {
char *real = (char *)ctx->calloc(sz + 64, cnt);
if (real == NULL) {
return NULL;
}
snprintf(real, 64, "originally allocated %ld via zero",
(unsigned long)sz);
return (void *)(real + 64);
}
NPY_NO_EXPORT void
shift_free(Allocator *ctx, void * p, npy_uintp sz) {
if (p == NULL) {
return ;
}
char *real = (char *)p - 64;
if (strncmp(real, "originally allocated", 20) != 0) {
fprintf(stdout, "uh-oh, unmatched shift_free, "
"no appropriate prefix\\n");
/* Make C runtime crash by calling free on the wrong address */
ctx->free((char *)p + 10);
/* ctx->free(real); */
}
else {
npy_uintp i = (npy_uintp)atoi(real +20);
if (i != sz) {
fprintf(stderr, "uh-oh, unmatched shift_free"
"(ptr, %ld) but allocated %ld\\n", sz, i);
/* This happens when the shape has a 0, only print */
ctx->free(real);
}
else {
ctx->free(real);
}
}
}
NPY_NO_EXPORT void *
shift_realloc(Allocator *ctx, void * p, npy_uintp sz) {
if (p != NULL) {
char *real = (char *)p - 64;
if (strncmp(real, "originally allocated", 20) != 0) {
fprintf(stdout, "uh-oh, unmatched shift_realloc\\n");
return realloc(p, sz);
}
return (void *)((char *)ctx->realloc(real, sz + 64) + 64);
}
else {
char *real = (char *)ctx->realloc(p, sz + 64);
if (real == NULL) {
return NULL;
}
snprintf(real, 64, "originally allocated "
"%ld via realloc", (unsigned long)sz);
return (void *)(real + 64);
}
}
static Allocator new_handler_ctx = {
malloc,
calloc,
realloc,
free
};
static PyDataMem_Handler new_handler = {
"secret_data_allocator",
1,
{
&new_handler_ctx,
shift_alloc, /* malloc */
shift_zero, /* calloc */
shift_realloc, /* realloc */
shift_free /* free */
}
};
实现#
此 NEP 已在 PR 17582 中实现。
替代方案#
这些在 issue 17467 中进行了讨论。 PR 5457 和 PR 5470 提出了一种用于指定对齐分配的全局接口。
PyArray_malloc_aligned 及其相关函数已随 numpy.random 模块 API 重构添加到 NumPy 中，并在其中用于性能提升。
PR 390 包含两部分：通过 NumPy C-API 公开 PyDataMem_*，以及一种钩子机制。该 PR 已合并，但没有提供使用这些功能的示例代码。
讨论#
邮件列表上的讨论导致了具有类似于 PyMemAllocatorEx 的 context 字段的 PyDataMemAllocator 结构体，但 free 函数的签名不同。
参考文献和脚注#
版权#
本文档已置于公共领域。 [1]