pybind11 NumPy 互操作完全指南:从 Buffer Protocol 到 vectorize、unchecked 直访与 memoryview
【免费下载链接】pybind11Seamless operability between C++11 and Python项目地址: https://gitcode.com/GitHub_Trending/py/pybind11
本篇指南基于 pybind11 官方文档 NumPy 章节 整理,系统讲解在 C++11 扩展模块中操作 NumPy 数据的完整技术栈:通过 buffer protocol 让 C++ 类被零拷贝地读入/写出 NumPy 数组、用py::array_t<T>精确控制接收的数组类型与内存布局、用PYBIND11_NUMPY_DTYPE注册结构化体类型、用py::numpy_scalar区分单双精度标量、用py::vectorize把标量函数变成向量化函数、用unchecked<N>代理实现免检查的高性能元素访问,以及用py::memoryview暴露裸 C/C++ 缓冲区。读完后你将能够编写出既正确又高性能的 pybind11 × NumPy 绑定代码,并能从 numpy.h 源码层面理解每一项机制的底层行为。
1. Buffer Protocol:让任意 C++ 类型暴露为缓冲区
Python 支持一种非常通用的插件间数据交换机制:类型可以暴露一个 buffer view,提供对原始内部数据表示的快速直接访问。pybind11 的 numpy.h 与 buffer_info.h 就是对这一机制的完整封装。
1.1 把 C++ 类暴露为 buffer
假设要绑定一个简单的Matrix类:
class Matrix { public: Matrix(size_t rows, size_t cols) : m_rows(rows), m_cols(cols) { m_data = new float[rows*cols]; } float *data() { return m_data; } size_t rows() const { return m_rows; } size_t cols() const { return m_cols; } private: size_t m_rows, m_cols; float *m_data; };绑定代码将Matrix的内容暴露为 buffer 对象,从而可以把它强转为 NumPy 数组,甚至可以用np.array(matrix_instance, copy=False)这类 Python 表达式完全避免拷贝:
py::class_<Matrix>(m, "Matrix", py::buffer_protocol()) .def_buffer([](Matrix &m) -> py::buffer_info { return py::buffer_info( m.data(), /* Pointer to buffer */ sizeof(float), /* Size of one scalar */ py::format_descriptor<float>::format(), /* Python struct-style format descriptor */ 2, /* Number of dimensions */ { m.rows(), m.cols() }, /* Buffer dimensions */ { sizeof(float) * m.cols(), /* Strides (in bytes) for each index */ sizeof(float) } ); });为一个新的类型提供 buffer 协议支持需要两步:在py::class_构造函数中指定py::buffer_protocol()标签,并调用def_buffer()传入一个在需要时生成py::buffer_info描述记录的 lambda,描述给定矩阵实例。
1.2py::buffer_info的完整结构
py::buffer_info的内容镜像了 Python buffer protocol 规范。官方文档给出的结构体是:
struct buffer_info { void *ptr; py::ssize_t itemsize; std::string format; py::ssize_t ndim; std::vector<py::ssize_t> shape; std::vector<py::ssize_t> strides; };对照源码 buffer_info.h,实际实现还包含两个文档未列出的字段,值得注意:
ssize_t size:总元素数,构造函数中自动计算为所有shape之积;bool readonly:标识底层存储是否可写,默认false(可写)。
构造函数有一致性校验:若ndim与shape、strides的长度不匹配,直接pybind11_fail抛出异常(buffer_info.h#L69-L72)。此外还提供了一个模板便捷构造:buffer_info(T *ptr, shape, strides, readonly)会自动填充sizeof(T)与format_descriptor<T>::format()。
buffer_info内部还持有可选的Py_buffer *m_view;当它由Py_buffer构造且ownview为真时,析构函数会调用PyBuffer_Release释放视图(buffer_info.h#L147-L152)。这保证了从 Python 侧request()得到的缓冲区描述符生命周期正确。
一个实用成员函数是item_type_is_equivalent_to<T>(),它用于校验缓冲区元素类型是否与T等价:
template <typename T> bool item_type_is_equivalent_to() const { return detail::compare_buffer_info<T>::compare(*this); }从源码结构看(buffer_info.h#L188-L208),对于整型T的比较会额外放宽:当sizeof(T) == sizeof(long)时接受"l"/"L",当sizeof(T) == sizeof(size_t)时接受"n"/"N"——这正是int/long在不同平台上宽度不同所带来的兼容性处理。
1.3 接收 Python buffer:以 Eigen 矩阵为例
要让 C++ 函数接收 Python buffer 对象作参数,只需使用类型py::buffer。由于 buffer 可能处于各种形态,函数体内通常需要若干安全检查。下面是一个为 Eigen 双精度矩阵Eigen::MatrixXd定义自定义构造函数的例子,支持从兼容的 buffer 对象(例如 NumPy 矩阵)初始化:
/* Bind MatrixXd (or some other Eigen type) to Python */ typedef Eigen::MatrixXd Matrix; typedef Matrix::Scalar Scalar; constexpr bool rowMajor = Matrix::Flags & Eigen::RowMajorBit; py::class_<Matrix>(m, "Matrix", py::buffer_protocol()) .def(py::init([](py::buffer b) { typedef Eigen::Stride<Eigen::Dynamic, Eigen::Dynamic> Strides; /* Request a buffer descriptor from Python */ py::buffer_info info = b.request(); /* Some basic validation checks ... */ if (!info.item_type_is_equivalent_to<Scalar>()) throw std::runtime_error("Incompatible format: expected a double array!"); if (info.ndim != 2) throw std::runtime_error("Incompatible buffer dimension!"); auto strides = Strides( info.strides[rowMajor ? 0 : 1] / (py::ssize_t)sizeof(Scalar), info.strides[rowMajor ? 1 : 0] / (py::ssize_t)sizeof(Scalar)); auto map = Eigen::Map<Matrix, 0, Strides>( static_cast<Scalar *>(info.ptr), info.shape[0], info.shape[1], strides); return Matrix(map); }));注意这里直接利用info.strides构造Eigen::Stride,让Eigen::Map支持任意行列步进的视图——这是零拷贝的关键。作为参考,这个 Eigen 类型对应的def_buffer()调用应写成:
.def_buffer([](Matrix &m) -> py::buffer_info { return py::buffer_info( m.data(), /* Pointer to buffer */ sizeof(Scalar), /* Size of one scalar */ py::format_descriptor<Scalar>::format(), /* Python struct-style format descriptor */ 2, /* Number of dimensions */ { m.rows(), m.cols() }, /* Buffer dimensions */ { sizeof(Scalar) * (rowMajor ? m.cols() : 1), sizeof(Scalar) * (rowMajor ? 1 : m.rows()) } /* Strides (in bytes) for each index */ ); })若希望以更简单的方式绑定 Eigen 类型(但有一些限制),可参考文档 Eigen 章节。
完整示例参见 tests/test_buffers.cpp,它更详细地演示了如何在 pybind11 中使用 buffer protocol;对应的 Python 侧测试在 tests/test_buffers.py。
2. Arrays:py::array与py::array_t<T>类型约束
将上文片段中的py::buffer换成py::array,就可以把函数限制为只接受 NumPy 数组(而非任何满足 buffer protocol 的 Python 对象)。
2.1 限定数据类型:py::array_t<T>
很多场景下我们只想接受特定数据类型的 NumPy 数组,这通过模板py::array_t<T>实现。例如以下函数要求参数是包含双精度值的 NumPy 数组:
void f(py::array_t<double> array);当传入其它类型(如整型数组或整型列表)时,绑定代码会尝试将输入转换为所请求类型的 NumPy 数组。该特性需要包含 pybind11/numpy.h 头文件。注意:pybind11/numpy.h不依赖 NumPy 头文件,因此无需在构建期声明对 NumPy 的依赖即可使用;但NumPy ≥ 1.7.0 是运行时依赖。
从源码看,这种"必要时强制转换"的行为来自array_t的静态工厂raw_array_t(numpy.h#L1511-L1525):
static PyObject *raw_array_t(PyObject *ptr) { if (ptr == nullptr) { set_error(PyExc_ValueError, "cannot create a py::array_t from a nullptr"); return nullptr; } return detail::npy_api::get().PyArray_FromAny_(ptr, dtype::of<T>().release().ptr(), 0, 0, detail::npy_api::NPY_ARRAY_ENSUREARRAY_ | ExtraFlags, nullptr); }它最终调用 NumPy C-API 的PyArray_FromAny_,并叠加NPY_ARRAY_ENSUREARRAY_与用户指定的ExtraFlags,从而保证参数确实是所需 dtype 的数组。
2.2 限定内存布局:c_style/f_style/forcecast
NumPy 数组中的数据并不保证是密集(packed)存储的,元素之间可能存在任意的列/行步长。有时我们只希望函数接受采用 C 顺序(行主序)或 Fortran 顺序(列主序)的密集数组,这可以通过第二个模板参数实现,取值为py::array::c_style或py::array::f_style:
void f(py::array_t<double, py::array::c_style | py::array::forcecast> array);py::array::forcecast是第二个模板参数的默认值,它确保不符合要求的实参会被转换成满足指定要求的数组,而不是放弃该次转换、去尝试下一个函数重载。
这三个标志在源码中直接映射到 NumPy 的常量(numpy.h#L1029-L1033):
enum { c_style = detail::npy_api::NPY_ARRAY_C_CONTIGUOUS_, f_style = detail::npy_api::NPY_ARRAY_F_CONTIGUOUS_, forcecast = detail::npy_api::NPY_ARRAY_FORCECAST_ };另外可以推断:array_t默认的第二模板参数就是array::forcecast(numpy.h#L1376-L1377):
template <typename T, int ExtraFlags = array::forcecast> class array_t : public array { ... };一个容易忽略的细节:以f_style构造的array_t会自动使用 Fortran 步长布局。源码中的便捷构造会二选一(numpy.h#L1422-L1428):
explicit array_t(ShapeContainer shape, const T *ptr = nullptr, handle base = handle()) : array_t(private_ctor{}, std::move(shape), (ExtraFlags & f_style) != 0 ? detail::f_strides(*shape, itemsize()) : detail::c_strides(*shape, itemsize()), ptr, base) {}其中c_strides/f_strides分别按行主序/列主序规则计算字节步长(buffer_info.h#L19-L38)。
2.3py::array的实用方法
数组提供若干方法;除下方"引用获取"一节列出的方法外,以下基于 NumPy API 的函数也可直接使用:
.dtype()返回所含值的类型;.strides()返回数组步长的指针(可选传入一个整型 axis 以获得单个数值);.flags()返回标志设置,.writable()与.owndata()可直接使用;.offset_at()返回偏移量(可选传入索引);.squeeze()返回移除长度为 1 的轴后的视图;.view(dtype)返回以不同 dtype 解释的数组视图;.reshape({i, j, ...})返回不同形状的数组视图,.resize({...})也可用;.index_at(i, j, ...)计算从开头到给定索引的元素计数。
3. Structured Types:注册结构化(record)类型
为了让py::array_t支持结构化(record)类型,需要先注册该类型的内存布局。这通过在插件定义代码中调用PYBIND11_NUMPY_DTYPE宏完成,宏接收类型名后跟字段名:
struct A { int x; double y; }; struct B { int z; A a; }; // ... PYBIND11_MODULE(test, m, py::mod_gil_not_used()) { // ... PYBIND11_NUMPY_DTYPE(A, x, y); PYBIND11_NUMPY_DTYPE(B, z, a); /* now both A and B can be used as template arguments to py::array_t */ }结构体的字段只能是:基本算术类型、std::complex、此前已注册过的子结构体,以及上述任意类型的数组。C++ 原生数组与std::array都受支持。虽然存在静态断言来阻止大多数不受支持的结构体,但用户仍必须保证只使用可以安全地按裸内存操作而不违反任何不变量的"平凡"结构体。
对于类型拼写中包含逗号的类型,必须用PYBIND11_TYPE包裹,例如PYBIND11_NUMPY_DTYPE(PYBIND11_TYPE(C<int, double>), x, y)。
源码层面可以确认:该宏在未启用 NumPy 支持时会被定义为空操作(numpy.h#L1798),启用时则展开为向 pybind11 的 dtype 注册表写入字段名、格式与偏移信息的初始化代码(numpy.h#L1859 起)。
4. Scalar Types:py::numpy_scalar区分单双精度
有时我们需要接收或返回np.float32、np.float64这样的 NumPy 标量值,并希望在 C++ 侧分别处理单精度与双精度。但两者默认都绑定到 Python 内建双精度float,无法分开处理;旧的做法是用py::buffer技巧,可读性很差。
为此 pybind11 提供了辅助类型py::numpy_scalar:
m.def("add", [](py::numpy_scalar<float> a, py::numpy_scalar<float> b) { return py::make_scalar(a + b); }); m.def("add", [](py::numpy_scalar<double> a, py::numpy_scalar<double> b) { return py::make_scalar(a + b); });该类型与其包装的类型之间可平凡地互相转换。当前支持的标量类型是 NumPy 的算术类型:bool_、int8、int16、int32、int64、uint8、uint16、uint32、uint64、float32、float64、complex64、complex128,全部映射到对应的 C++ 类型。
两个重要约束(文档原样保留):
py::numpy_scalar<T>严格匹配 NumPy 标量类型。例如py::numpy_scalar<int64_t>会接受np.int64(123),但不接受普通 Pythonint(如123)。原生 C 类型到 NumPy 类型的映射是平台相关的:例如
char可能映射到np.int8或np.uint8,long依平台可能是 4 或 8 字节。除非你清楚理解差异与自己的需求,请使用<cstdint>中的定宽类型。
从源码结构看,这一"严格匹配"由 type_caster 的load实现:它先检查isinstance(src, target_type()),而target_type是通过 NumPy C-APIPyArray_TypeObjectFromType_按typenum取到的精确标量类型对象(numpy.h#L778-L805):
bool load(handle src, bool) { if (isinstance(src, target_type())) { npy_api::get().PyArray_ScalarAsCtype_(src.ptr(), &value.value); return true; } return false; }返回侧则由cast调用PyArray_Scalar_生成对应 dtype 的 NumPy 标量对象。numpy_scalar本身只是持有一个value_type value的透明包装(numpy.h#L809-L829),配合py::make_scalar工厂函数(numpy.h#L831-L834)即可在绑定函数签名中自然地表达"输入/输出必须是特定 NumPy 标量"。
5. Vectorizing Functions:py::vectorize
假设要把如下签名的函数绑定到 Python,使它能处理任意 NumPy 数组参数(向量、矩阵、一般 N 维数组):
double my_func(int x, float y, double z);包含pybind11/numpy.h后,做法极其简单:
m.def("vectorized_func", py::vectorize(my_func));以如下方式调用时,my_func会被数组的每个元素各调用一次(共 4 次)。与numpy.vectorize()等方案相比,其显著优势在于对元素的遍历完全在 C++ 侧进行,编译器可以把它压缩成一个紧凑、优化过的高性能循环。结果以numpy.dtype.float64类型的 NumPy 数组返回:
>>> x = np.array([[1, 3], [5, 7]]) >>> y = np.array([[2, 4], [6, 8]]) >>> z = 3 >>> result = vectorized_func(x, y, z)标量参数z被透明地复制 4 次。输入数组x、y会自动转换成正确类型(它们原本是numpy.dtype.int64,但分别需要numpy.dtype.int32与numpy.dtype.float32)。
注意:只有按值传递或按
const &引用的算术、复数与 POD 类型会被向量化;所有其他参数原样透传。带右值引用参数的函数不能被向量化。
这些规则在源码中有精确对应。vectorize_arg模板定义了哪些参数参与向量化(numpy.h#L2083-L2101):
template <typename T> struct vectorize_arg { static_assert(!std::is_rvalue_reference<T>::value, "Functions with rvalue reference arguments cannot be vectorized"); // Is this a vectorized argument? static constexpr bool vectorize = satisfies_any_of<call_type, std::is_arithmetic, is_complex, is_pod>::value && satisfies_none_of<call_type, std::is_pointer, std::is_array, is_std_array, std::is_enum>::value && (!std::is_reference<T>::value || (std::is_lvalue_reference<T>::value && std::is_const<call_type>::value)); // Accept this type: an array for vectorized types, otherwise the type as-is: using type = conditional_t<vectorize, array_t<remove_cv_t<call_type>, array::forcecast>, T>; };并且vectorize_helper会在编译期强制"至少一个可向量化的参数"(numpy.h#L2151-L2153):
static_assert( NVectorized >= 1, "pybind11::vectorize(...) requires a function with at least one vectorizable argument");运行时流程为:对每个向量化参数request()出buffer_info→ 用 NumPy 广播规则推断输出 shape → 若所有输入都是 0 维(单个标量)则直接返回标量结果 → 否则创建结果数组(若全部输入都是 F 序且广播平凡,则结果也按array::f_style创建,见 numpy.h#L2108-L2113)→ 以紧密循环逐元素调用原函数。
完整示例参见 tests/test_numpy_vectorize.cpp,演示了vectorize更详细的使用方式;Python 侧验证在 tests/test_numpy_vectorize.py。
5.1 手动访问缓冲区:vectorize不够用时的完整写法
如果计算太复杂、无法归约为vectorize,就需要手动创建和访问缓冲区内容。以下片段是一个完整示例(这段代码略显人为,因为它其实可以用vectorize更简单地实现):
#include <pybind11/pybind11.h> #include <pybind11/numpy.h> namespace py = pybind11; py::array_t<double> add_arrays(py::array_t<double> input1, py::array_t<double> input2) { py::buffer_info buf1 = input1.request(), buf2 = input2.request(); if (buf1.ndim != 1 || buf2.ndim != 1) throw std::runtime_error("Number of dimensions must be one"); if (buf1.size != buf2.size) throw std::runtime_error("Input shapes must match"); /* No pointer is passed, so NumPy will allocate the buffer */ auto result = py::array_t<double>(buf1.size); py::buffer_info buf3 = result.request(); double *ptr1 = static_cast<double *>(buf1.ptr); double *ptr2 = static_cast<double *>(buf2.ptr); double *ptr3 = static_cast<double *>(buf3.ptr); for (size_t idx = 0; idx < buf1.shape[0]; idx++) ptr3[idx] = ptr1[idx] + ptr2[idx]; return result; } PYBIND11_MODULE(test, m, py::mod_gil_not_used()) { m.def("add_arrays", &add_arrays, "Add two NumPy arrays"); }注意py::array_t<double>(buf1.size)这个构造:不传数据指针,由 NumPy 负责分配底层 buffer;随后对结果再次request()拿到可写的buffer_info并直接以裸指针循环写入。
6. Direct Access:unchecked<N>免检查代理
出于性能考虑——特别是处理非常大的数组时——通常希望在索引已知合法的情况下直接访问数组元素,而不在每次访问时做维度与边界检查。为此array类与array_t<T>模板类提供了 unchecked 代理对象,可通过unchecked<N>与mutable_unchecked<N>方法获得(N为数组所需的维度数):
m.def("sum_3d", [](py::array_t<double> x) { auto r = x.unchecked<3>(); // x must have ndim = 3; can be non-writeable double sum = 0; for (py::ssize_t i = 0; i < r.shape(0); i++) for (py::ssize_t j = 0; j < r.shape(1); j++) for (py::ssize_t k = 0; k < r.shape(2); k++) sum += r(i, j, k); return sum; }); m.def("increment_3d", [](py::array_t<double> x) { auto r = x.mutable_unchecked<3>(); // Will throw if ndim != 3 or flags.writeable is false for (py::ssize_t i = 0; i < r.shape(0); i++) for (py::ssize_t j = 0; j < r.shape(1); j++) for (py::ssize_t k = 0; k < r.shape(2); k++) r(i, j, k) += 1.0; }, py::arg().noconvert());使用规则与源码细节:
- 从
array对象获取代理时,必须同时把数据类型和维度数作为模板参数,例如auto r = myarray.mutable_unchecked<float, 2>(); - 如果维度数在编译期未知,可以省略维度模板参数(即调用
arr_t.unchecked()或arr.unchecked<T>())。此时得到的代理对象行为相同,但生成的代码可优化性较差,在紧密循环中会有小幅效率损失——源码中对应的代理类为unchecked_reference<T, Dims>/unchecked_mutable_reference<T, Dims>(numpy.h#L646、L733),Dims为运行时维度时偏移计算走通用路径; - 安全约束:返回的代理对象直接引用数组数据,仅在构造时读取一次 shape、strides 与 writable 标志。你必须确保在被引用数组的生命期内,其不被销毁或 reshape——通常做法是把返回实例的作用域限定在原地;
mutable_unchecked要求数组带 writable 标志,否则抛异常;unchecked只读则不要求可写。
代理对象支持py::array的部分同名方法,可在部分原有"索引检查"用法中作为直接替代:
.ndim()返回维度数;.data(1, 2, ...)与.mutable_data(1, 2, ...)分别返回给定索引处const T或T数据的指针;后者只有经a.mutable_unchecked()获得的代理才可用;.itemsize()返回元素字节数,即sizeof(T);.shape(n)返回第n维的大小;.size()返回元素总数(各维 shape 之积);.nbytes()返回被引用元素占用的字节数(即itemsize() * size())。
更多示例见 tests/test_numpy_array.cpp。
7. Ellipsis:C++ 侧的...切片
Python 提供了便捷的...省略号记法,常用于切片多维数组。例如以下代码提取张量的中间维度,首尾索引取零:
a = ... # a NumPy array b = a[0, ..., 0]函数py::ellipsis()可以在 C++ 侧完成同样的操作:
py::array a = /* A NumPy array */; py::array b = a[py::make_tuple(0, py::ellipsis(), 0)];8. Memory View:暴露裸 C/C++ 缓冲区
当只是想为 C/C++ 缓冲区提供一个直接访问器、而没有具体类对象时,可以返回memoryview对象。例如希望为一个 2x4 的uint8_t数组暴露memoryview,可以这样做:
const uint8_t buffer[] = { 0, 1, 2, 3, 4, 5, 6, 7 }; m.def("get_memoryview2d", []() { return py::memoryview::from_buffer( buffer, // buffer pointer { 2, 4 }, // shape (rows, cols) { sizeof(uint8_t) * 4, sizeof(uint8_t) } // strides in bytes ); });这种方式用于为不由 Python 管理的 C/C++ 缓冲区提供memoryview,缓冲区的生命周期由用户负责。如果先在 C++ 侧删除了缓冲区,之后仍在使用以这种方式创建的memoryview,属于未定义行为。
对于简单的一维连续缓冲区,还可以用memoryview::from_memory:
m.def("get_memoryview1d", []() { return py::memoryview::from_memory( buffer, // buffer pointer sizeof(uint8_t) * 8 // buffer size ); });memoryview::from_memory是 2.6 版本新增的。实现位于 pytypes.h:from_buffer系列构造内部先组装buffer_info再走PyMemoryView_FromBuffer,源码注释明确要求传入的ptr与format必须比所创建的 memoryview 活得更久;from_memory则直接封装PyMemoryView_FromMemory,并提供bool readonly = false的只读参数。
9. 小结与延伸阅读
本文覆盖了 pybind11 文档 NumPy 章节 的全部主题,并结合源码验证了关键机制:
| 需求 | 机制 | 核心 API | 源码位置 |
|---|---|---|---|
| C++ 类暴露给 NumPy 零拷贝读写 | buffer protocol | py::buffer_protocol()+def_buffer()+py::buffer_info | buffer_info.h |
| 接收任意 buffer | buffer 参数 | py::buffer+.request() | numpy.h |
| 只接受特定 dtype 的数组 | 类型模板 | py::array_t<T, c_style/f_style/forcecast> | numpy.h#L1376 |
| 结构化 record 类型 | dtype 注册宏 | PYBIND11_NUMPY_DTYPE(Type, fields...) | numpy.h#L1859 |
| 区分 np.float32/64 等标量 | 标量包装 | py::numpy_scalar<T>/py::make_scalar | numpy.h#L809 |
| 标量函数向量化 | vectorize | py::vectorize(f) | numpy.h#L2139 |
| 免检查高性能访问 | unchecked 代理 | .unchecked<N>()/.mutable_unchecked<N>() | numpy.h#L646 |
| 裸缓冲区暴露 | memoryview | py::memoryview::from_buffer/from_memory | pytypes.h#L2383 |
配套测试文件可作为可运行的完整参考:tests/test_buffers.cpp、tests/test_numpy_array.cpp、tests/test_numpy_vectorize.cpp、tests/test_numpy_scalars.cpp、tests/test_numpy_dtypes.cpp,以及对应的 Python 测试(如 tests/test_numpy_array.py)。
【免费下载链接】pybind11Seamless operability between C++11 and Python项目地址: https://gitcode.com/GitHub_Trending/py/pybind11
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考