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 年，以及 2014 年的 issue 5312，这导致了 PR 5457 以及更多的邮件列表讨论，此处和此处。在 2017 年的该 issue 评论中，一位用户描述了 64 字节对齐如何将性能提高 40 倍。
相关内容还有 issue 14177，涉及在 Linux 上使用 madvise 和大页内存。
各种跟踪和分析库，如 filprofiler 或 electric fence，会覆盖 malloc。
CPython 关于 BPO 18835 的长期讨论始于对 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。
处理器的名称将通过 numpy.core.multiarray.get_handler_name(arr) 函数在 Python 级别暴露。如果以 numpy.core.multiarray.get_handler_name() 形式调用，它将返回用于为下一个新 ndarrray 分配数据的处理器的名称。
处理器的版本将通过 numpy.core.multiarray.get_handler_version(arr) 函数在 Python 级别暴露。如果以 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 本身可能会违反此要求，因此 PyDataMemAllocator 的作者应将size参数视为一个最佳猜测。修复此问题的工作正在 PR 15780 和 15788 中进行，但尚未解决。一旦解决，应重新审视本 NEP。
PyDataMem_Handler 线程安全和生命周期#
活动处理器通过 Context 中的 ContextVar 存储。这确保了它可以按线程和按异步协程进行配置。
目前没有对 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 已合并，但没有提供使用这些功能的示例代码。
讨论#
邮件列表上的讨论导致了 PyDataMemAllocator 结构体，其中包含一个与 PyMemAllocatorEx 类似的 context 字段，但 free 函数签名不同。
参考文献和脚注#
版权#
本文档已置于公共领域。 [1]