Exposing Eigen::SparseMatrix class causes segfault in setFromTriplets

[MarcelFerrari]

Hi all,

I am trying to expose some bindings for the Eigen::SparseMatrix class so that I can use Eigen instead of Scipy from Python.

I defined the following templated function to initialise a SparseMatrix class with given datatype and ordering.

#include <nanobind/nanobind.h>
#include <nanobind/operators.h>
#include <nanobind/eigen/dense.h>
#include <nanobind/stl/string.h>
#include <Eigen/Sparse>
#include <string>
#include "include/utils.hpp"
#include<iostream>

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

template <typename Scalar, typename Index, int Options>
void bind_sparse_matrix(nb::module_ &m, const std::string &class_name) {
    using SparseMatrix = Eigen::SparseMatrix<Scalar, Options, Index>;

    nb::class_<SparseMatrix>(m, class_name.c_str())
        // Constructors
        .def(nb::init<>(), "Default constructor")
        .def(nb::init<Index, Index>(), "rows"_a, "cols"_a)

        // Properties
        .def_prop_ro("rows", &SparseMatrix::rows, "Number of rows")
        .def_prop_ro("cols", &SparseMatrix::cols, "Number of columns")
        .def_prop_ro("non_zeros", &SparseMatrix::nonZeros, "Number of non-zero elements")
        .def_prop_ro("is_compressed", &SparseMatrix::isCompressed, "Check if the matrix is compressed")

        // Matrix initialization using triplets
        .def("set_from_triplets", [](SparseMatrix &self,
                                     nb::DRef<Eigen::Array<Index, Eigen::Dynamic, 1>> idx_i,
                                     nb::DRef<Eigen::Array<Index, Eigen::Dynamic, 1>> idx_j,
                                     nb::DRef<Eigen::Array<Scalar, Eigen::Dynamic, 1>> values) {
            // Determine matrix size
            Index rows = idx_i.maxCoeff() + 1;
            Index cols = idx_j.maxCoeff() + 1;

            if (rows != cols) {
                throw std::runtime_error("Matrix must be square.");
            }

            // Create ArrayTripletIterator object
            ArrayTripletIterator<Scalar, Index> begin(idx_i, idx_j, values);
            ArrayTripletIterator<Scalar, Index> end(idx_i, idx_j, values, values.rows());
            
            // Set the matrix from triplets
            self.resize(rows, cols);
            self.setFromTriplets(begin, end);
        }, "idx_i"_a, "idx_j"_a, "values"_a,
           "Set the matrix from arrays of triplets idx_i, idx_j, values in COO format.")
}

I then use this to create bindings for my sparse matrix objects:

#include <nanobind/nanobind.h>
#include <Eigen/Core>
#include <Eigen/Sparse>

#include "include/sparse_matrix.hpp"

namespace nb = nanobind;

NB_MODULE(phasma, m) {
    // Bind for various scalar types and storage options
    bind_sparse_matrix<double, int, Eigen::ColMajor>(m, "SparseMatrixDCol");
    bind_sparse_matrix<double, int, Eigen::RowMajor>(m, "SparseMatrixDRow");
    bind_sparse_matrix<float,  int, Eigen::ColMajor>(m, "SparseMatrixFCol");
    bind_sparse_matrix<float,  int, Eigen::RowMajor>(m, "SparseMatrixFRow");
}

The problem I am facing is that whenever I try to run the set_from_triplets method, I hit a segfault or bus error depending on what datatypes I use. Specifically, the problem seems to be the setFromTriplets method of the sparse matrix.

The ArrayTripletIterator class is just an iterator implementation in order to construct the sparse matrix without having to rely on a std::vector<Eigen::Triplet<Scalar>> object.

template <typename Scalar, typename Index>
class ArrayTripletIterator {
    const Eigen::Array<Index, Eigen::Dynamic, 1>& rows_;
    const Eigen::Array<Index, Eigen::Dynamic, 1>& cols_;
    const Eigen::Array<Scalar, Eigen::Dynamic, 1>& values_;
    Index index_;

public:
    ArrayTripletIterator(const Eigen::Array<Index, Eigen::Dynamic, 1>& rows,
                         const Eigen::Array<Index, Eigen::Dynamic, 1>& cols,
                         const Eigen::Array<Scalar, Eigen::Dynamic, 1>& values, Index index = 0)
        : rows_(rows), cols_(cols), values_(values), index_(index) {}

    Index row() const { return rows_[index_]; }
    Index col() const { return cols_[index_]; }
    Scalar value() const { return values_[index_]; }

    ArrayTripletIterator& operator++() {
        ++index_;
        return *this;
    }

    bool operator!=(const ArrayTripletIterator& other) const { return index_ != other.index_; }

    // Pointer-like interface for Eigen compatibility
    const ArrayTripletIterator* operator->() const { return this; }
};

When I test this exact same class in a normal isolated C++ snippet it works absolutely fine.

Example:

#include <Eigen/Sparse>
#include <iostream>
#include <vector>
#include <chrono>

template <typename Scalar, typename Index>
class ArrayTripletIterator {
    const Eigen::Array<Index, Eigen::Dynamic, 1>& rows_;
    const Eigen::Array<Index, Eigen::Dynamic, 1>& cols_;
    const Eigen::Array<Scalar, Eigen::Dynamic, 1>& values_;
    Index index_;

public:
    ArrayTripletIterator(const Eigen::Array<Index, Eigen::Dynamic, 1>& rows,
                         const Eigen::Array<Index, Eigen::Dynamic, 1>& cols,
                         const Eigen::Array<Scalar, Eigen::Dynamic, 1>& values, Index index = 0)
        : rows_(rows), cols_(cols), values_(values), index_(index) {}

    Index row() const { return rows_[index_]; }
    Index col() const { return cols_[index_]; }
    Scalar value() const { return values_[index_]; }

    ArrayTripletIterator& operator++() {
        ++index_;
        return *this;
    }

    bool operator!=(const ArrayTripletIterator& other) const { return index_ != other.index_; }

    // Pointer-like interface for Eigen compatibility
    const ArrayTripletIterator* operator->() const { return this; }
};

void generate_random_coo(int rows, int cols, int nnz, 
                         Eigen::ArrayXi& row_indices,
                         Eigen::ArrayXi& col_indices,
                         Eigen::ArrayXd& values) {
    row_indices = Eigen::ArrayXi::Random(nnz).cwiseAbs().eval().unaryExpr([&rows](int x) { return x % rows; });
    col_indices = Eigen::ArrayXi::Random(nnz).cwiseAbs().eval().unaryExpr([&cols](int x) { return x % cols; });
    values = Eigen::ArrayXd::Random(nnz);
}


// Benchmark function
void benchmark_setFromTriplets(int rows, int cols, int nnz) {
    // Generate random COO data
    Eigen::ArrayXi row_indices, col_indices;
    Eigen::ArrayXd values;
    generate_random_coo(rows, cols, nnz, row_indices, col_indices, values);

    // Benchmark with custom iterator
    auto start = std::chrono::high_resolution_clock::now();
    Eigen::SparseMatrix<double> matrix1(rows, cols);
    matrix1.setFromTriplets(ArrayTripletIterator<double, int>(row_indices, col_indices, values),
                            ArrayTripletIterator<double, int>(row_indices, col_indices, values, nnz));
    auto end = std::chrono::high_resolution_clock::now();
    auto custom_iterator_time = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
    std::cout << "Custom iterator setFromTriplets: " << custom_iterator_time << " microseconds\n";

    // Benchmark with std::vector<Eigen::Triplet>
    start = std::chrono::high_resolution_clock::now();
    std::vector<Eigen::Triplet<double>> triplets;
    triplets.reserve(nnz);
    for (int i = 0; i < nnz; ++i) {
        triplets.emplace_back(row_indices[i], col_indices[i], values[i]);
    }
    Eigen::SparseMatrix<double> matrix2(rows, cols);
    matrix2.setFromTriplets(triplets.begin(), triplets.end());
    end = std::chrono::high_resolution_clock::now();
    auto triplet_vector_time = std::chrono::duration_cast<std::chrono::microseconds>(end - start).count();
    std::cout << "std::vector<Eigen::Triplet> setFromTriplets: " << triplet_vector_time << " microseconds\n";
}


int main() {
    // Matrix dimensions and number of non-zeros
    int rows = 20000;
    int cols = 20000;
    int nnz = int(1e-3 * rows * cols);

    // Run benchmark
    benchmark_setFromTriplets(rows, cols, nnz);

    return 0;
}

Even worse, if I create a temporary vector of triplets and copy the data from the input arrays, then the nanobind bindings work absolutely great.

I have no clue what is going wrong, but I hit a brick wall and I am stuck. I am also not sure whether this is a problem with my code or with nanobind.

Thanks in advance!

Update 1

If I compile the code with -march=native then the segfault disappears, but I get the wrong results:

idx_i dtype:  int32
idx_j dtype:  int32
vals dtype:  float64
idx_i: 0 1 2 3
idx_j: 0 1 2 3
values: 1 2 3 4

Output:
0 1 0 0 
0 2 0 0 
0 3 0 0 
0 4 0 0 

Given this, I immediately thought about checking memory alignment, but it seems like all is good.

I also got in touch with the developers of Eigen and it looks like the way I am initialising the matrix is correct. This seems to be the case as the code works correctly when run from C++.

This seems to be a nanobind specific issue or some other incompatibility between Numpy and Eigen that I cannot figure out.

[MarcelFerrari]

It seem like the core issue lies with how the nb::DRef class interacts with regular Eigen objects.

My iterator class originally stored const references to Eigen objects:

const Eigen::Array<Index, Eigen::Dynamic, 1>& rows_;
const Eigen::Array<Index, Eigen::Dynamic, 1>& cols_;
const Eigen::Array<Scalar, Eigen::Dynamic, 1>& values_;

When passing nb::Dref objects to the iterator, no errors are raised and everything looks fine.
However, accessing the values results in completely corrupted data.

Replacing the Eigen types with nb::Dref results in the correct and expected behaviour.

using IndexArray = nb::DRef<Eigen::Array<Index, Eigen::Dynamic, 1>>;
using ScalarArray = nb::DRef<Eigen::Array<Scalar, Eigen::Dynamic, 1>>;

IndexArray rows_;
IndexArray cols_;
ScalarArray values_;

This is not an issue with row-major vs column-major ordering as we are dealing with a simple linear arrays.
If this is not an intended use of nb::Dref, then an error should be raised at compile time.

I will open an issue related to this.