Keeping underlying Fortran-contiguous ndarray buffer valid on return from C++ code

[Chronum94]

Hello,

Great library, pleasure to use it.

I'm currently creating a Python API to a C++ library which has a pattern that like so:

.def_static(
          "from_ndarray",
          [](nb::ndarray<double, nb::device::cpu> &a)
          {
            int local_m = a.shape(0);
            int local_n = a.shape(1);

            double *local_tmp = new double[local_m * local_n];
            std::copy(a.data(), a.data() + local_m * local_n, local_tmp);

            return MatrixObject::from_ndarray(
                local_tmp);
          })

I am currently copying the ndarray's contents into a raw pointer but if possible, I'd like to avoid that copy operation (it's not too bad, but I'd like to avoid it).

The Numpy arrays that are fed into this C++ code will definitely remain valid (non-GC'd) for the lifetime of the returned MatrixObject. The MatrixObject requires that the memory at a.data() remains valid for the lifetime of the MatrixObject. Is there a way to do this without the copy op?

I don't know if there's any more information I can provide but I'd be more than happy to provide as much detail as I can.

[wjakob]

Two ways to do this: modify your MatrixObject to stash the nb::ndarray somewhere. Or use a keep_alive call annotation to ensure that the argument (a) is kept alive at least as long as the return value of the function.

[Chronum94]

I'm accepting this as the answer since a nb::keep_alive<0, 1>() did do the needful. The np.asfortranarray situation later on in the thread is mostly a different issue.

[Chronum94]

The above example is with a .def_static but I've been trying the keep_alive annotation by making a standalone function. It looks like (slightly simplified but the core logic is here) so:

m.def(
    "matrix_from_ndarray",
    [](nb::ndarray<double, nb::device::cpu> a)
    {
        return MatrixObject::from_ndarray((double *)a.data());
    }, nb::keep_alive<0, 1>());

My understanding from the keep_alive documentation is that this means that as long as the returned MatrixObject (the 0) is alive, a (the 1) is kept alive. But this isn't working as far as I can tell. I still get segfaults, nondeterministic corruption of the numerical data.

I think this is from the return value policy (https://nanobind.readthedocs.io/en/latest/api_core.html#_CPPv4N8nanobind9rv_policy9automaticE) being automatically set to move since I am returning an rvalue, although I could be wrong. But if the rv policy is move, then I don't know what the behaviour of keep_alive would be. I've tried changing the rv_policy to reference_internal but that also gives me a segfault.

Still very new to the library and a bit rusty with C++ memory management so I really appreciate the help through this.

[wjakob]

And here is an even better version using nb::pointer_and_handle that simultaneously accepts the C++ and Python version of an object to avoid the need for a subsequent nb::find call.

        .def(
            "numpy",
            [](nb::pointer_and_handle<A> pah)
            {
                A *a = pah.p;
                nb::handle h = pah.h;

                size_t shape[2] = {a->n_, a->n_};
                int64_t strides[2] = {1, (Py_ssize_t) a->n_};

                // When returning an nd-array, all of the annotations below
                // (except *nb::numpy*) are just decorations for the
                // docstring. In contrast to accepting an nd-array as input,
                // no conversions are performed. In other words, it's your job
                // to provide a numpy array of the promised properties. You
                // must also specify a strides vector if you use non-C strides.
                return nb::ndarray<nb::numpy, double, nb::shape<nb::any>,
                                   nb::f_contig>(a->a_, 2, shape, h, strides);
            });

[Chronum94]

After fairly aggressive testing of the exact code I have in production (and fixing some import errors I'm not proud of), I can confirm that master fixes the problem.

Also really appreciate the code review of the mwe. Learning good practices/new things as I go.

I deeply appreciate all of the time you've spent on this, and I think master fixes the subtlety with keep_alive.

[wjakob]

Great, thanks for the kind words and for double-checking.

[Chronum94]

I've changed the title slightly to make the discussion somewhat easier to find in the future should anyone need it. Unless you think otherwise, we can close this discussion since all issues have been addressed.

[wjakob]

Sounds good!