以便携/快速方式获取指向Numpy/Numpypy数据的指针

Question

我最近尝试过,PyPy并对这种方法很感兴趣.我有很多Python的C扩展,它们都PyArray_DATA()用来获取指向numpy数组数据部分的指针.不幸的是,PyPy似乎没有numpypy在他们的cpyext模块中导出他们的数组的等价物,所以我尝试按照他们网站上的推荐使用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来说,两者都很慢,我的应用程序完全成为瓶颈!是否有快速便携的方法来获取此指针？或者是否有相当于PyArray_DATA()PyPy(可能没有文档)？

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 用于我的科学计算应用程序之后，只要他们（相当顽固地）拒绝支持，我就看不到它的未来完整的 CPython API。

更新

引入后我发现ctypes.cast()现在成为了瓶颈accel._get_ptr()。可以通过将接口中的所有指针声明为 来摆脱强制转换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()s，ctypes 的速度实际上都可以与本机 CPython 扩展竞争。所以我只需要等到有人重写h5py，matplotlib并scipy使用 ctypes 就可以尝试 PyPy 来做任何严肃的事情......