[QUESTION] Segfault when exposing class member variable with buffer protocol

[rcffc]

I want to expose the following class, specifically I am interested in the *voxelBlocks:

template<class ITMVoxel>
    class ITMLocalVBA
    {
    private:
        ORUtils::MemoryBlock<TVoxel> *voxelBlocks;
        MemoryDeviceType memoryType;

    public:
        inline TVoxel *GetVoxelBlocks(void) { return voxelBlocks->GetData(memoryType); }
        inline const TVoxel *GetVoxelBlocks(void) const { return voxelBlocks->GetData(memoryType); }
        int lastFreeBlockId;
        int allocatedSize;
  }

GetData(MemoryDeviceType memoryType) looks like this:

/** Get the data pointer on CPU or GPU. */
inline DEVICEPTR(T)* GetData(MemoryDeviceType memoryType)
{
	switch (memoryType)
	{
	case MEMORYDEVICE_CPU: return data_cpu;
	case MEMORYDEVICE_CUDA: return data_cuda;
	}

	return 0;
}

Reading the documentation, it seems that the following buffer protocol should suffice:

  PYBIND11_NUMPY_DTYPE(ITMVoxel, sdf, w_depth);  // sdf, w_depth are uchar and short

  py::class_<ITMLocalVBA<ITMVoxel>>(m, "ITMLocalVBA_ITMVoxel", pybind11::buffer_protocol())
    .def_readonly("allocated_size", &ITMLocalVBA<ITMVoxel>::allocatedSize)
    .def_buffer([](ITMLocalVBA<ITMVoxel>& v) -> pybind11::buffer_info {
      return pybind11::buffer_info(
        v.GetVoxelBlocks(), 
        sizeof(ITMVoxel), 
        py::format_descriptor<ITMVoxel>::format(), 
        1, 
        { v.allocatedSize }, 
        { sizeof(ITMVoxel) }
      );
    })
  ;

However, when I try to access it, the program crashes (with a segmentation fault). Does anyone have a clue? Help would be much appreciated, my Master thesis depends on it.

[ljluestc]

#include <pybind11/pybind11.h>
#include <pybind11/numpy.h>
#include <pybind11/stl.h>

namespace py = pybind11;

template <class ITMVoxel>
class ITMLocalVBA {
private:
    ORUtils::MemoryBlock<ITMVoxel> *voxelBlocks;
    MemoryDeviceType memoryType;

public:
    int lastFreeBlockId;
    int allocatedSize;

    ITMLocalVBA(ORUtils::MemoryBlock<ITMVoxel> *voxelBlocks, MemoryDeviceType memoryType)
        : voxelBlocks(voxelBlocks), memoryType(memoryType), lastFreeBlockId(0), allocatedSize(0) {}

    inline ITMVoxel *GetVoxelBlocks(void) {
        return voxelBlocks->GetData(memoryType);
    }

    inline const ITMVoxel *GetVoxelBlocks(void) const {
        return voxelBlocks->GetData(memoryType);
    }
};

template <typename ITMVoxel>
py::array_t<ITMVoxel> voxel_blocks_to_numpy(ITMLocalVBA<ITMVoxel>& v) {
    // Get data pointer
    ITMVoxel *data_ptr = v.GetVoxelBlocks();
    
    if (v.memoryType == MEMORYDEVICE_CUDA) {
        // Handle GPU memory - Copy the data to CPU
        // This can depend on how you interface with CUDA (this is just an example)
        // You can copy data from GPU to CPU here if required
        // For example, use cudaMemcpy to copy from GPU to a temporary CPU buffer.
        std::vector<ITMVoxel> cpu_data(v.allocatedSize);
        cudaMemcpy(cpu_data.data(), data_ptr, v.allocatedSize * sizeof(ITMVoxel), cudaMemcpyDeviceToHost);
        data_ptr = cpu_data.data();  // Now we have the CPU pointer.
    }

    // Expose the data to numpy array via buffer protocol
    return py::array_t<ITMVoxel>(
        { v.allocatedSize },                        // shape
        { sizeof(ITMVoxel) },                      // strides
        data_ptr                                  // pointer to data
    );
}

PYBIND11_MODULE(ur_mpc_bindings, m) {
    py::class_<ITMLocalVBA<ITMVoxel>>(m, "ITMLocalVBA_ITMVoxel")
        .def(py::init<ORUtils::MemoryBlock<ITMVoxel>*, MemoryDeviceType>(),
            py::arg("voxelBlocks"), py::arg("memoryType"))
        .def("get_voxel_blocks", &ITMLocalVBA<ITMVoxel>::GetVoxelBlocks)
        .def_property("allocated_size", &ITMLocalVBA<ITMVoxel>::allocatedSize, nullptr)
        .def_buffer([](ITMLocalVBA<ITMVoxel>& v) -> pybind11::buffer_info {
            // Expose voxel blocks via buffer protocol
            return voxel_blocks_to_numpy(v);
        });
}