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 年的 问题 5312，这导致了 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。
处理程序的名称将通过 numpy.core.multiarray.get_handler_name(arr) 函数在 Python 层暴露。如果以 numpy.core.multiarray.get_handler_name() 方式调用，它将返回用于为下一个新 ndarray 分配数据的处理程序名称。
处理程序的版本将通过 numpy.core.multiarray.get_handler_version(arr) 函数在 Python 层暴露。如果以 numpy.core.multiarray.get_handler_version() 方式调用，它将返回用于为下一个新 ndarray 分配数据的处理程序版本。
版本号（当前为 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函数的参数保持一致。NumPy 本身在请求的数组形状包含
0时可能会违反此要求，因此 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 中实现。
替代方案#
这些内容已在 问题 17467 中讨论过。PR 5457 和 PR 5470 提出了一个用于指定对齐分配的全局接口。
PyArray_malloc_aligned 及其相关函数是随着 numpy.random 模块 API 重构被添加到 NumPy 中的，并在该模块中用于提升性能。
PR 390 包含两部分：通过 NumPy C-API 暴露 PyDataMem_*，以及一种钩子机制。该 PR 合并时未包含使用这些功能的示例代码。
讨论#
邮件列表上的讨论最终确定了 PyDataMemAllocator 结构体，该结构体带有一个 context 字段，类似于 PyMemAllocatorEx，但 free 的签名不同。
参考文献和脚注#
版权#
本文档已进入公共领域。[1]