NEP 49 — 数据分配策略#

作者:

Matti Picus

状态:

最终

类型:

标准跟踪

创建:

2021-04-18

解决:

https://mail.python.org/archives/list/[email protected]/thread/YZ3PNTXZUT27B6ITFAD3WRSM3T3SRVK4/#PKYXCTG4R5Q6LIRZC4SEWLNBM6GLRF26

摘要#

 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 倍。

还与 问题 14177 相关的使用 madvise 和 Linux 上的巨大页面。

各种跟踪和分析库,例如 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://en.cppreference.com/w/cpp/memory/allocator/deallocate>C++11 <https://en.cppreference.com/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。

PyObject *PyDataMem_SetHandler(PyObject *handler)#

设置新的分配策略。如果输入值为 NULL,将策略重置为默认策略。返回先前的策略,或者如果发生错误则返回 NULL。我们将用户提供的策略封装起来,以便它们仍会调用 Python 和 NumPy 内存管理回调挂钩。所有函数指针都必须填写,不接受 NULL

const PyObject *PyDataMem_GetHandler()#

返回将用于为下一个 PyArrayObject 分配数据的当前策略。如果失败,则返回 NULL

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 合并后没有使用这些功能的示例代码。

讨论#

邮件列表中的讨论导致了 PyDataMemAllocator 结构,它带有类似于 PyMemAllocatorEx 的 context 字段,但 free 的签名不同。

参考文献和脚注#