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

8 投票

2 回答

2018 浏览

提问于 2025-04-17 17:55

我最近尝试了 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() 的东西（可能没有文档说明）？

pypy numpy compatibility performance optimization array manipulation memory management c extensions data pointer

2 个回答

这可能不是一个完整的答案，但希望能给你一些好的提示。我在代码的某些部分使用了scipy.weave.inline()。我对这个接口的速度了解不多，因为我执行的函数比较复杂，只依赖几个指针和数组，但对我来说似乎运行得很快。也许你可以从scipy.weave的代码中获得一些灵感，特别是attempt_function_call这个部分。

https://github.com/scipy/scipy/blob/master/scipy/weave/inline_tools.py#L390

如果你想看看scipy.weave生成的C++代码，

可以从这里生成一个简单的例子：http://docs.scipy.org/doc/scipy/reference/tutorial/weave.html，
运行这个Python脚本，

找到scipy.weave的缓存文件夹：

import scipy.weave.catalog as ctl
ctl.default_dir()
Out[5]: '/home/user/.python27_compiled'

查看文件夹里生成的C++代码。

回答于 2025-04-17 由 Python大师

分享举报

我还没有找到一个完全满意的解决方案，但有一些方法可以在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。

更新

我发现引入accel._get_ptr()后，ctypes.cast()成了瓶颈。通过将接口中的所有指针声明为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了……

回答于 2025-04-17 由 Python大师

分享举报

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

2 个回答

撰写回答