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 访问。本文档后面包含一个示例。添加的 struct 将增加 ndarray 对象的大小。这是此方法必须付出的必要代价。我们可以 reasonably 确定对象大小的变化对最终用户代码的影响很小，因为 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函数的参数一致。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 中实现。
替代方案#
这些在 issue 17467 中进行了讨论。 PR 5457 和 PR 5470 提出了一个用于指定对齐分配的全局接口。
PyArray_malloc_aligned 及其相关函数已随 numpy.random 模块 API 重构添加到 NumPy 中，并在其中用于提高性能。
PR 390 分为两部分：通过 NumPy C-API 暴露 PyDataMem_*，以及一个钩子机制。PR 被合并，但没有使用这些功能的示例代码。
讨论#
邮件列表上的讨论导致了 PyDataMemAllocator 结构，其中有一个 context 字段，类似于 PyMemAllocatorEx，但 free 的签名不同。
参考文献和脚注#
版权#
本文档已放入公共领域。 [1]