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 倍。
与之相关的还有 问题 14177，围绕在 Linux 上使用 madvise 和巨型页面。
各种跟踪和分析库，如 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 本身可能会违反此要求，因此 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 中实现。
替代方案#
这些在 问题 17467 中进行了讨论。 PR 5457 和 PR 5470 提出了一个用于指定对齐分配的全局接口。
PyArray_malloc_aligned 及其相关函数是在 NumPy 的 numpy.random 模块 API 重构中添加的，并在其中用于性能优化。
PR 390 包含两个部分：通过 NumPy C-API 公开 PyDataMem_* 和一个钩子机制。此 PR 已合并，但没有提供使用这些功能的示例代码。
讨论#
邮件列表上的讨论导致了具有类似于 PyMemAllocatorEx 的 context 字段的 PyDataMemAllocator 结构体，但 free 的签名不同。
参考文献和脚注#
版权#
本文档已置于公共领域。 [1]