Constructing ndarray by moving from a std::unique_ptr

[agtumulak]

I have the following code to dynamically allocate float array, randomly populate it (with a CUDA call which requires a raw pointer), then return a ndarray object. Is this the correct/idiomatic way to do this without any copies?

#include "sample.cuh"
#include <memory>
#include <nanobind/nanobind.h>
#include <nanobind/ndarray.h>

namespace nb = nanobind;

auto sample(size_t N) {
  auto samples_ptr = std::make_unique<float[]>(N);
  cuda_sample(samples_ptr.get(), N);
  return nb::ndarray<float, nb::numpy>(samples_ptr.get(), {N}).cast();
}

NB_MODULE(sample, m) { m.def("sample", &sample); }

As I understand it, in the return statement, the constructor for nb::ndarray returns a temporary, so calling cast() will call the T&& overload which takes ownership of the underlying pointer in samples_ptr. However, when I print out samples_ptr.get() after cast() I don't get nullptr. I would expect that samples_ptr would get std::moved at some point inside cast()?

[hpkfft]

I would follow the example in https://nanobind.readthedocs.io/en/latest/ndarray.html#data-ownership
Maybe, there's no reason to use a unique_ptr, since you're always using .get() anyway to pass the raw pointer to functions?
Though, maybe you want safety if cuda_sample throws an exception?

The nb::ndarray constructor in your example above is only given the raw pointer, so neither it nor the subsequent .cast() can use std::move to affect the unique_ptr itself.

If you need to use unique_ptr for some reason (e.g., exception considerations, other functions calls not shown in your minimal example above), then you can use samples_ptr.release() in the nb::ndarray constructor so that the unique_ptr does not free the memory when it goes out of scope at the end of the sample() function. But somebody has to own and eventually free the memory--that would be the nb::capsule owner Python object shown in the nanobind docs, which will call delete[] when it dies.

[agtumulak]

Ah! I didn't think to call release() to express that I was transferring ownership of the underlying pointer to the ndarray. I get samples_ptr.get() == nullptr true after constructing the ndarray now. There are no copies of the buffer the way this is written, right?

[hpkfft]

Correct

[wjakob]

If you are using CUDA, then I don't quite understand the reason to go through NumPy, which involves a costly (and synchronizing) GPU->CPU copy. The snippet you specified will make another copy on top of that. This is because you did not specify an owner object that will manage the lifetime of the data region. Ownership and lifetimes are described in great detail in the nanobind ndarray docs, please refer to them for more details.