NEP 49 — 資料配置策略#

作者:

Matti Picus

狀態:

最終

類型:

標準追蹤

建立於:

2021-04-18

決議:

https://mail.python.org/archives/list/numpy-discussion@python.org/thread/YZ3PNTXZUT27B6ITFAD3WRSM3T3SRVK4/#PKYXCTG4R5Q6LIRZC4SEWLNBM6GLRF26

摘要#

numpy.ndarray 需要額外的記憶體配置來保存 numpy.ndarray.stridesnumpy.ndarray.shapenumpy.ndarray.data 屬性。這些屬性是在 __new__ 方法中建立 python 物件後特別配置的。

此 NEP 提案一種機制,以使用者提供的替代方案來覆寫用於 ndarray->data 的記憶體管理策略。此配置保存資料,而且可能非常大。由於存取此資料經常成為效能瓶頸,因此自訂配置策略以保證資料對齊或將配置釘選到專用記憶體硬體可以實現硬體特定的最佳化。其他配置保持不變。

動機與範疇#

使用者可能希望使用自己的常式來覆寫內部資料記憶體常式。其中兩個用例是確保資料對齊,以及將某些配置釘選到某些 NUMA 核心。在郵件清單 2005 年以及 issue 5312(2014 年)中多次討論了對齊的需求,這促成了 PR 5457 和更多的郵件清單討論 此處 和此處。在 2017 年的 issue 上的評論中,一位使用者描述了 64 位元組對齊如何將效能提升 40 倍。

與此相關的還有關於在 Linux 上使用 madvise 和巨頁的 issue 14177

各種追蹤和效能分析程式庫(例如 filprofilerelectric fence)會覆寫 malloc

BPO 18835 的長期 CPython 討論始於討論對 PyMem_Alloc32PyMem_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 可以在資料記憶體指標上使用 memcpymemset

處理常式的名稱將透過 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.org/std/alloc/trait.Allocator.html#tymethod.deallocate>

PyDataMemAllocator 介面的消費者必須追蹤 size,並確保它與傳遞至 (m|c|re)alloc 函式的參數一致。

當請求陣列的形狀包含 0 時,NumPy 本身可能會違反此要求,因此 PyDataMemAllocator 的作者應將 size 參數視為最佳猜測。修復此問題的工作正在 PR 1578015788 中進行,但尚未解決。解決後應重新檢視此 NEP。

PyObject *PyDataMem_SetHandler(PyObject *handler)#

設定新的配置策略。如果輸入值為 NULL,則會將策略重設為預設值。傳回先前的策略,或在發生錯誤時傳回 NULL。我們包裝使用者提供的策略,讓它們仍然會呼叫 Python 和 NumPy 記憶體管理回呼掛鉤。所有函式指標都必須填寫,不接受 NULL

const PyObject *PyDataMem_GetHandler()#

傳回將用於為下一個 PyArrayObject 配置資料的目前策略。如果失敗,則傳回 NULL

PyDataMem_Handler 執行緒安全性和生命週期#

作用中的處理常式透過 ContextVar 儲存在目前的 Context 中。這可確保可以針對每個執行緒和每個非同步協程進行組態。

目前沒有 PyDataMem_Handler 的生命週期管理。PyDataMem_SetHandler 的使用者必須確保引數在以它配置的任何物件的生命週期內以及在其為作用中處理常式時保持存活。實際上,這表示處理常式必須是不朽的。

作為實作細節,目前此 ContextVar 包含一個 PyCapsule 物件,該物件儲存指向沒有解構函式的 PyDataMem_Handler 的指標,但不應依賴此物件。

範例程式碼#

此程式碼將 64 位元組標頭新增至每個 data 指標,並將關於配置的資訊儲存在標頭中。在呼叫 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 5457PR 5470 提案一個全域介面,用於指定對齊的配置。

PyArray_malloc_aligned 和相關功能已隨著 numpy.random 模組 API 重構新增至 NumPy。並在那裡用於提升效能。

PR 390 有兩個部分:透過 NumPy C-API 公開 PyDataMem_*,以及掛鉤機制。PR 已合併,但沒有使用這些功能的範例程式碼。

討論#

郵件清單上的討論促成了具有 context 欄位的 PyDataMemAllocator 結構,類似於 PyMemAllocatorEx,但在 free 的簽名方面有所不同。

參考資料與註腳#