python - 获取指向 Numpy/Numpypy 数据的指针的便携/快速方法

Question

我最近尝试了 PyPy 并且对这种方法很感兴趣。我有很多 Python 的 C 扩展，它们都使用 PyArray_DATA() 来获取指向 numpy 数组数据部分的指针。不幸的是，PyPy 似乎没有在他们的 cpyext 模块中为他们的 numpypy 数组导出等效项，所以我尝试按照他们网站上的建议使用 ctypes。这将获取指针的任务推向了Python层面。

似乎有两种方式:

import ctypes as C
p_t = C.POINTER(C.c_double)

def get_ptr_ctypes(x):
    return x.ctypes.data_as(p_t)

def get_ptr_array(x):
    return C.cast(x.__array_interface__['data'][0], p_t)

只有第二个适用于 PyPy，因此为了兼容性，选择很明确。对于 CPython，两者都非常慢，并且是我应用程序的完全瓶颈!有没有一种快速便携的方法来获取这个指针？或者 PyPy 是否有等效的 PyArray_DATA()(可能未记录)？

Answer 1

我仍然没有找到一个完全令人满意的解决方案，但是在 CPython 中可以做一些事情来以更少的开销获取指针。首先，上面提到的两种方式之所以这么慢，是因为 .ctypes 和 .__array_interface__ 都是按需属性，由 array_ctypes_get( ) 和 array_interface_get() 在 numpy/numpy/core/src/multiarray/getset.c 中。第一个导入 ctypes 并创建一个 numpy.core._internal._ctypes 实例，而第二个创建一个新字典并在其中填充除数据指针之外的许多不必要的内容。

对于这种开销，在 Python 级别上我们无能为力，但是可以在 C 级别上编写一个微模块来绕过大部分开销:

#include <Python.h>
#include <numpy/arrayobject.h>

PyObject *_get_ptr(PyObject *self, PyObject *obj) {
    return PyLong_FromVoidPtr(PyArray_DATA(obj));
}

static PyMethodDef methods[] = {
    {"_get_ptr", _get_ptr, METH_O, "Wrapper to PyArray_DATA()"},
    {NULL, NULL, 0, NULL}
};

PyMODINIT_FUNC initaccel(void) {
    Py_InitModule("accel", methods);
}

像往常一样在 setup.py 中作为扩展编译，并导入为

try:
    from accel import _get_ptr
    def get_ptr(x):
        return C.cast(_get_ptr(x), p_t)
except ImportError:
    get_ptr = get_ptr_array

在 PyPy 上，from accel import _get_ptr 将失败，get_ptr 将回退到 get_ptr_array，它与 Numpypy 一起工作。

就性能而言，对于轻量级 C 函数调用，ctypes + accel._get_ptr() 仍然比原生 CPython 扩展慢很多，后者基本上没有开销。它当然比上面的 get_ptr_ctypes() 和 get_ptr_array() 快得多，因此对于中等重量的 C 函数调用来说，开销可能变得微不足道。

一个人已经获得了与 PyPy 的兼容性，尽管我不得不说，在花了相当多的时间尝试为我的科学计算应用程序评估 PyPy 之后，只要他们(非常顽固)，我就看不到它的 future 拒绝支持完整的 CPython API。

更新

我发现 ctypes.cast() 在引入 accel._get_ptr() 后成为了瓶颈。可以通过将接口(interface)中的所有指针声明为 ctypes.c_void_p 来摆脱强制转换。这就是我最终得到的:

def get_ptr_ctypes2(x):
    return x.ctypes._data

def get_ptr_array(x):
    return x.__array_interface__['data'][0]

try:
    from accel import _get_ptr as get_ptr
except ImportError:
    get_ptr = get_ptr_array

此处，get_ptr_ctypes2() 通过直接访问隐藏的 ndarray.ctypes._data 属性来避免转换。以下是从 Python 调用重量级和轻量级 C 函数的一些计时结果:

                             heavy C (few calls)      light C (many calls)
ctypes + get_ptr_ctypes():         0.71 s                   15.40 s
ctypes + get_ptr_ctypes2():        0.68 s                   13.30 s
ctypes + get_ptr_array():          0.65 s                   11.50 s
ctypes + accel._get_ptr():         0.63 s                    9.47 s

native CPython:                    0.62 s                    8.54 s
Cython (no decorators):            0.64 s                    9.96 s

因此，使用 accel._get_ptr() 并且没有 ctypes.cast()，ctypes 的速度实际上可以与原生 CPython 扩展相媲美。所以我只需要等到有人用 ctypes 重写 h5py、matplotlib 和 scipy 就可以尝试 PyPy 做任何严肃的事情......