Return Array of Eigen::Matrix from C++ to Python without copying

问题

I have some C++ code that generates and manipulates arrays of Eigen matrices. In the end I want to use those matrices in python and thought this might be a job for pybind11.

Basically what I want back in python are two nested lists / numpy arrays mat_a(I, 4, 4) and mat_b(J, K, 4, 4). Because I have to do a lot of linear algebra stuff in C++ I wanted to use Eigen and the data structure I used is std::array<std::array<Eigen::Matrix4f, 2>, 3>>> mat_b // for J=3, K=2. The problem now is how to get this to python efficiently?

Additionally I want to perform those calculations for multiple inputs x = [x_0, x_1, ..., x_N] and than expect mat_a(N, I, 4, 4) and mat_b(N, J, K, 4, 4) as result. The calculations for each x_i are independent but I thought maybe it is faster to write this loop over x_i in C++ as well. If on the other hand the task gets easier if we only have fixed sized arrays in C++ this loop can also move to python.

Here is some dummy code of my problem (I=5, J=3, K=2) :

// example.cpp
#include <pybind11/pybind11.h>
#include <pybind11/eigen.h>
#include <pybind11/stl.h>
#include <pybind11/functional.h>
#include <pybind11/stl_bind.h>

#include <array>
#include <vector>
#include <Eigen/Dense>


Eigen::Matrix4f get_dummy(){
    Eigen::Matrix4f mat_a;
    mat_a << 1, 2, 3, 4,
             5, 6, 7, 8,
             9, 8, 7, 6,
             5, 4, 3, 2;
    return mat_a;
}

std::pair< std::vector<std::array<Eigen::Matrix4f, 5> >,
           std::vector<std::array<std::array<Eigen::Matrix4f, 2>, 3> > >  get_matrices(std::vector<float> & x){

    std::vector<std::array<Eigen::Matrix4f, 5> > mat_a(x.size());
    std::vector< std::array< std::array< Eigen::Matrix4f, 2>, 3> > mat_b(x.size());

    //    for (u_int i=0; i< x.size(); i++)
    //        do_stuff(x[i], mat_a[i], mat_b[i]);
    mat_a[0][0] = get_dummy();

    return std::make_pair(mat_a, mat_b);
    }


PYBIND11_MODULE(example, m) {
    m.def("get_dummy", &get_dummy, pybind11::return_value_policy::reference_internal);
    m.def("get_matrices", &get_matrices, pybind11::return_value_policy::reference_internal);
}

I compile the code via:

c++ -O3 -Wall -shared -std=c++14 -fPIC `python3 -m pybind11 --includes` example.cpp -o example`python3-config --extension-suffix`

And than use it in python:

import numpy as np
import example

x = np.zeros(1000)

mat_a, mat_b = get_matrices(x)

print(np.shape(mat_a))
print(np.shape(mat_b))
print(mat_a[0][0])

If I just want to return a single Eigen::Matrix it works fast and as far as I can tell without copying. But when I try to nest the Eigen:Matrices with std::array/std::vector pybind returns a nested list of numpy arrays instead of one multidimensional array. This is as expected and I am actually impressed how well this works but it seems rather slow to me especially as the dimensions of the arrays grow.

The question is how can I improve this to get multidimensional numpy arrays without unnecessary copying.

Some roads I tried but did not work (for me, what doesn't mean that they do not work in general; I just could not figure it out):

use Eigen::Tensor instead of the arrays of Eigen:Matrix
create the matrices in python and pass it to C++ by reference
building a custom wrapper for array<array<Matrix4f, K>, J>

回答1:

Your best option may be to create the data on python side so it gets refcounted and garbage collected.

test.py

import example
import numpy as np

array = np.zeros((3, 2, 4, 4), 'f4')

example.do_math(array, 3, 2)
print(array[0, 0])

example.cpp

#define PY_SSIZE_T_CLEAN
#include <Python.h>

#include <Eigen/Dense>

Eigen::Matrix4f get_dummy() {
    Eigen::Matrix4f mat_a;
    mat_a << 1, 2, 3, 4,
             5, 6, 7, 8,
             9, 8, 7, 6,
             5, 4, 3, 2;
    return mat_a;
}

PyObject * example_meth_do_math(PyObject * self, PyObject * args, PyObject * kwargs) {
    static char * keywords[] = {"array", "rows", "cols", NULL};

    PyObject * array;
    int rows, cols;

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "Oii", keywords, &array, &rows, &cols)) {
        return NULL;
    }

    Py_buffer view = {};
    if (PyObject_GetBuffer(array, &view, PyBUF_SIMPLE)) {
        return NULL;
    }

    Eigen::Matrix4f * ptr = (Eigen::Matrix4f *)view.buf;

    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            ptr[i * cols + j] = get_dummy();
        }
    }

    PyBuffer_Release(&view);
    Py_RETURN_NONE;
}

PyMethodDef module_methods[] = {
    {"do_math", (PyCFunction)example_meth_do_math, METH_VARARGS | METH_KEYWORDS, NULL},
    {},
};

PyModuleDef module_def = {PyModuleDef_HEAD_INIT, "example", NULL, -1, module_methods};

extern "C" PyObject * PyInit_example() {
    PyObject * module = PyModule_Create(&module_def);
    return module;
}

setup.py

from setuptools import Extension, setup

ext = Extension(
    name='example',
    sources=['./example.cpp'],
    extra_compile_args=['-fpermissive'],
    include_dirs=['.'], # add the path of Eigen
    library_dirs=[],
    libraries=[],
)

setup(
    name='example',
    version='0.1.0',
    ext_modules=[ext],
)

It should be trivial from here to add a second parameter and use the two arrays for the calculation.

You can build this with python setup.py develop.

if you want to distribute it you can create a wheel file with python setup.py bdist_wheel.

I used numpy to create the data, this ensures the underlying memory of the data is C contiguous.

This example was kept simple and it uses a Matrix4f pointer to iterate a 3x2 array of matrices. Feel free to cast the ptr to an Eigen::Array<Eigen::Matrix4f>, 3, 2>. You cannot cast it to an std::vector since the internal data of an std::vector contains a pointer.

Please note that std::vector<std::array<...>> does not have a single contiguous array in the memory. Use Eigen::Array instead.

edit:

Here is a function that uses an Eigen Array Map:

PyObject * example_meth_do_math(PyObject * self, PyObject * args, PyObject * kwargs) {
    static char * keywords[] = {"array", NULL};

    PyObject * array;

    if (!PyArg_ParseTupleAndKeywords(args, kwargs, "O", keywords, &array)) {
        return NULL;
    }

    Py_buffer view = {};
    if (PyObject_GetBuffer(array, &view, PyBUF_SIMPLE)) {
        return NULL;
    }

    Eigen::Map<Eigen::Array<Eigen::Matrix4f, 2, 3>> array_map((Eigen::Matrix4f *)view.buf, 2, 3);

    for (int i = 0; i < 2; ++i) {
        for (int j = 0; j < 3; ++j) {
            array_map(i, j) = get_dummy();
        }
    }

    PyBuffer_Release(&view);
    Py_RETURN_NONE;
}

回答2:

One other possibility if you're not too tied to Eigen is xtensor ( found here). I've used their python bindings before which give an example of communicating directly with python(found here). This will have the advantage of being able to handle larger, multi-dimensional arrays. The linear algebra won't be as slick (hard to beat Eigen there), but will be similar to what you would do in numpy (np.dot(A,B) for example.

If you want to stick with Eigen, note there is some technical specifics with using the STL. As your std::array is no longer able to contain a fixed number of Matrices, when you move to std::vector you will hit alignment issues (admittedly, that I don't totally understand). Will get a working implementation of xtensor for you shortly.

来源：https://stackoverflow.com/questions/63159807/return-array-of-eigenmatrix-from-c-to-python-without-copying

标签

python

c++

numpy

eigen

pybind11