Passing tensor buffer to methods taking const Eigen::ArrayXd& without copying

[ianhbell]

This library is incredible. The overhead on very often-called, but fast functions (runtime of sub-microsecond) on a derived type is something like an order of magnitude faster. I'm moving from pybind11 to nanobind as we speak.

One thing I was unclear of is how STL containers, namely std::vector, work in the interface. If I have a function like so that I want to wrap:

auto indexer (int i, const std::vector<double> &e){
    return e[i];
}

does nanobind still make a temporary? What if I make a numpy array and pass that to the function? It seems like the copies are avoided, but I am not 100% sure.

Related: it would be nice to add std::valarray, as that allows for concise math stuff without the weight of Eigen.

In pybind11 we had the ability to specify .noconvert() on an argument to avoid copies, and that compiles with nanobind, but I wanted to confirm that the same thing is working in nanobind without copying.

Alternatively, how would one access a numpy array buffer (1d of double) in nanobind? The tensor machinery seems overkill.

[wjakob]

Hi @ianhbell.

issues are reserved for things that are serious problems with the library. I have converted this into a discussion.

Now with regards to the questions you raised:

• Just like in pybind11, anything that involves a normal std::* caster involves a full copy of the data. You must create bindings for the array types in Python-land to enable zero-copy sharing of data structures.

• std::valarray is not currently supported, but a PR to add it would be welcomed. If support is added through the form of a caster, it will necessarily involve a conversion from STL data structure into a Python list or similar.

• .noconvert() continues to be supported, but it does something different from what I think you had in mind. It does not prevent the use of casters. Instead, it disables implicit conversions involving multiple bound types. For example, if you call a function f(a: A) with a B() instance, but A has an implicit constructor that knows how to create an A from a B. Adding .noconvert() will prevent this implicit conversion path from being used.

• I don't understand the issue with buffer/DLPack access via nb::tensor. This is the way to get zero-copy access to NumPy arrays (except that it is now much more general). Under the hood, what happens will be quite similar to passing around py::array* types in pybind11.

Wenzel

[ianhbell]

Sorry about that, I couldn't figure out where to ask the question other than via issues and this is clearly not a bug. Could you add that to the issue template to make it more clear somehow? Or add the discussion link to the main README?

You are right, I thought .noconvert() was like in pybind11 where copies are not allowed.

Is it possible to put together an example of zero-cost access to a 1D array from Python? I am sure that would be widely useful for people. I couldn't find such an example going through all the examples and tests.

[wjakob]

Did you look at https://github.com/wjakob/nanobind/blob/master/docs/tensor.md? It gives an example of a function named process that demonstrates how it works.

[ianhbell]

Here was my best attempt to call a function that indexes a const Eigen::ArrayXd& with a zero-copy tensor passed in. But I don't think I succeeded, as the access time is rather huge (order of milliseconds), and should be more like nanoseconds, and scales strongly with the length of the array, further evidence that some sort of copy is happening. Is there a way, for debugging, to disable copying of any kind in nanobind?

#include <nanobind/nanobind.h>
#include <nanobind/tensor.h>

#include <Eigen/Dense>

namespace nb = nanobind;
using namespace nb::literals;

using tensor1d = nb::tensor<double, nb::shape<nb::any>, nb::c_contig, nb::device::cpu>;

auto construct_Eigen_tensor_view(const tensor1d& tensor){
    return Eigen::Map<const Eigen::ArrayXd>(static_cast<const double*>(tensor.data()), tensor.shape(0));
}

auto indexer(const Eigen::ArrayXd &x, int i){
    return x(i);
}

NB_MODULE(n2flann, m) {
    m.def("indexer", [](const tensor1d& ten, int i){ return indexer(construct_Eigen_tensor_view(ten), i);}, "x"_a.noconvert(), "i"_a);
    m.attr("__version__") = "0.0.0";
}

with the calling code:

In [1]: import n2flann as n2, numpy as np

In [2]: z = np.linspace(0,1,1000000)

In [3]: %timeit n2.indexer(z, 3)
420 µs ± 20.5 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

[wjakob]

Just a quick thought: Eigen IIRC has column-contiguous storage by default (https://eigen.tuxfamily.org/dox/group__TopicStorageOrders.html), which would correspond to nb::f_contig in your tensor type. Otherwise the system will make a copy of the data to transpose the contents.

Can you try making that change to see if it improves performance?

[ianhbell]

Eigen is column-major by default: https://eigen.tuxfamily.org/dox/group__TopicStorageOrders.html "If the storage order is not specified, then Eigen defaults to storing the entry in column-major."

[ianhbell]

Sorry I read your response too quickly. The f_contig didn't seem to help (in my case, which is 1D, that makes sense, since column-major and row-major are the same for a 1D memory chunk).

I tried to break down the problem into parts and it seems the problem is the passing of the Eigen::Map to the indexing function. Making the map is fast, both locally and in another function that returns it. It is only when passing the map to another function that the copy happens. This isn't is a nanobind problem, but it would be good to have some of these examples documented somewhere so that others can avoid re-inventing the wheel if Eigen casters are no longer available.

In [1]: import n2flann as n2, numpy as np; z = np.linspace(0,1,1000000)

In [2]: %timeit n2.indexer(z, 3)
474 µs ± 6.01 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

In [3]: %timeit n2.indexer2(z, 3)
204 ns ± 4.55 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

In [4]: %timeit n2.indexer3(z, 3)
213 ns ± 0.753 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

from

#include <nanobind/nanobind.h>
#include <nanobind/tensor.h>

#include <Eigen/Dense>

namespace nb = nanobind;
using namespace nb::literals;

using tensor1d = nb::tensor<double, nb::shape<nb::any>, nb::f_contig, nb::device::cpu>;

Eigen::Map<const Eigen::ArrayXd> construct_Eigen_tensor_view(const tensor1d& tensor){
    return Eigen::Map<const Eigen::ArrayXd>(static_cast<const double*>(tensor.data()), tensor.shape(0));
}

auto indexer(const Eigen::ArrayXd &x, int i){
    return x(i);
}

NB_MODULE(n2flann, m) {
    m.def("indexer", [](const tensor1d& ten, int i){ 
        return indexer(construct_Eigen_tensor_view(ten), i);}, 
        "x"_a.noconvert(), "i"_a);
    m.def("indexer2", [](const tensor1d& ten, int i){ 
        Eigen::Map<const Eigen::ArrayXd> view(static_cast<const double*>(ten.data()), ten.shape(0));
        return view(i);}, 
        "x"_a.noconvert(), "i"_a);
    m.def("indexer3", [](const tensor1d& ten, int i){ 
        auto view = construct_Eigen_tensor_view(ten);
        return view(i);}, 
        "x"_a.noconvert(), "i"_a);
    m.attr("__version__") = "0.0.0";
}

[ianhbell]

A bit more digging: this doesn't seem to work properly in pybind11 either. It seems that any code that takes an Eigen::Ref<const Eigen::ArrayXd>, when passing that argument to another function taking a const Eigen::ArrayXd&, invokes a copy. I don't really understand why, as if you have a contiguous allocated chunk of memory, a copy should not be required. And if it is 1D, the row- and column-major problem doesn't exist.

[wjakob]

Aha. So it sounds like the cost occurs in the part of the code involving Eigen and not the nb::tensor abstraction. Perhaps the solution is to rewrite your code to take Eigen::Map everywhere?

[wjakob]

By the way: you can respond to the individual messages to make a large thread instead of creating separate "answers" .That way, the discussion is easier to follow for others visiting this topic later.

[ianhbell]

Yes, sorry, just now realized that..

[ianhbell]

Aha. So it sounds like the cost occurs in the part of the code involving Eigen and not the nb::tensor abstraction. Perhaps the solution is to rewrite your code to take Eigen::Map everywhere?

I REALLY don't want to do that. I already have aggressively templated the code to allow for various numerical types, so I suppose the Eigen::Ref template is not MUCH worse, but it is a pity because then one loses auto-complete.

The fundamental problem is that it seems one cannot mix and match Eigen::Map/Ref and Eigen::ArrayXd in the way that I had thought.

[ianhbell]

My previous point stands, Eigen::Ref<...> -> const Eigen::ArrayXd& always invokes a copy. So, how should one make a function that calls to a function in an API that takes const Eigen::ArrayXd& as an argument?

[wjakob]

I think this more of a question for the Eigen folks than something I can answer. My guess is that you have to rewrite all API so that it takes an Eigen::Ref (those functions should still be callable using Eigen::ArrayXd instances).

[ianhbell]

But I think you are right.

[ianhbell]

I think this more of a question for the Eigen folks than something I can answer. My guess is that you have to rewrite all API so that it takes an Eigen::Ref (those functions should still be callable using Eigen::ArrayXd instances).

I hadn't appreciated that the inverse was still possible, but I wonder if that still invokes a copy.

[wjakob]

Again, something to discuss with the eigen folks :). I will mark this discussion as answered, since it does not seem to be a nanobind issue.

[ianhbell]

Yes, thank you very much for your time.

[nachovizzo]

This would be the missing feature in my case to migrate all my pybind projects to nanobind!

I'm still getting to know nanobind, and trying to adapt my existing codebase to have np.array on the python side and std::vector<Eigen::Vector3d> on the C++ api,.

In pybind11 I have an ugly adapter between both, similar to Open3D but I'm still tring to figure out how to do the same for nb::tensor.

With nanobind I have the following draft:

// tensor view of pointcloud
using vector3dvector = nb::tensor<nb::numpy, double, nb::shape<nb::any, 3>, nb::device::cpu>;

std::vector<Eigen::Vector3d> _to_eigen(vector3dvector tensor) {
    std::vector<Eigen::Vector3d> eigen_vectors(tensor.shape(0));
    for (auto i = 0; i < tensor.shape(0); ++i) {
        eigen_vectors[i] = Eigen::Map<Eigen::Vector3d>(&tensor(i, 0));
    }
    return eigen_vectors;
}

auto _to_numpy(std::vector<Eigen::Vector3d>&& points) {
    size_t shape[2] = {points.size(), 3};
    return vector3dvector(points.data(), 2, shape);
}

Any ideas?