Add serialization for sparse matrix and sparse vector.

This commit is contained in:
Antonio Sánchez 2022-11-21 19:43:07 +00:00
parent 044f3f6234
commit e7b1ad0315
4 changed files with 267 additions and 5 deletions

View File

@ -1496,6 +1496,125 @@ struct evaluator<SparseMatrix<Scalar_,Options_,StorageIndex_> >
}
// Specialization for SparseMatrix.
// Serializes [rows, cols, isCompressed, outerSize, numNonZeros, innerNonZeros,
// outerIndices, innerIndices, values].
template <typename Scalar, int Options, typename StorageIndex>
class Serializer<SparseMatrix<Scalar, Options, StorageIndex>, void> {
public:
typedef SparseMatrix<Scalar, Options, StorageIndex> SparseMat;
struct Header {
typename SparseMat::Index rows;
typename SparseMat::Index cols;
bool compressed;
Index outer_size;
Index num_non_zeros;
};
EIGEN_DEVICE_FUNC size_t size(const SparseMat& value) const {
// innerNonZeros.
std::size_t num_storage_indices =
value.isCompressed() ? 0 : value.outerSize();
// Outer indices.
num_storage_indices += value.outerSize() + 1;
// Inner indices.
num_storage_indices += value.nonZeros();
// Values.
std::size_t num_values = value.nonZeros();
return sizeof(Header) + sizeof(Scalar) * num_values +
sizeof(StorageIndex) * num_storage_indices;
}
EIGEN_DEVICE_FUNC uint8_t* serialize(uint8_t* dest, uint8_t* end,
const SparseMat& value) {
if (EIGEN_PREDICT_FALSE(dest == nullptr)) return nullptr;
if (EIGEN_PREDICT_FALSE(dest + size(value) > end)) return nullptr;
const size_t header_bytes = sizeof(Header);
Header header = {value.rows(), value.cols(), value.isCompressed(),
value.outerSize(), value.nonZeros()};
EIGEN_USING_STD(memcpy)
memcpy(dest, &header, header_bytes);
dest += header_bytes;
// innerNonZeros.
size_t data_bytes = sizeof(StorageIndex) * header.outer_size;
if (!header.compressed) {
memcpy(dest, value.innerNonZeroPtr(), data_bytes);
dest += data_bytes;
}
// Outer indices.
data_bytes = sizeof(StorageIndex) * (header.outer_size + 1);
memcpy(dest, value.outerIndexPtr(), data_bytes);
dest += data_bytes;
// Inner indices.
data_bytes = sizeof(StorageIndex) * header.num_non_zeros;
memcpy(dest, value.innerIndexPtr(), data_bytes);
dest += data_bytes;
// Values.
data_bytes = sizeof(Scalar) * header.num_non_zeros;
memcpy(dest, value.valuePtr(), data_bytes);
dest += data_bytes;
return dest;
}
EIGEN_DEVICE_FUNC const uint8_t* deserialize(const uint8_t* src,
const uint8_t* end,
SparseMat& value) const {
if (EIGEN_PREDICT_FALSE(src == nullptr)) return nullptr;
if (EIGEN_PREDICT_FALSE(src + sizeof(Header) > end)) return nullptr;
const size_t header_bytes = sizeof(Header);
Header header;
EIGEN_USING_STD(memcpy)
memcpy(&header, src, header_bytes);
src += header_bytes;
value.setZero();
value.resize(header.rows, header.cols);
// Initialize compressed state and inner non-zeros.
size_t data_bytes = sizeof(StorageIndex) * header.outer_size;
if (EIGEN_PREDICT_FALSE(src + data_bytes > end)) return nullptr;
if (header.compressed) {
value.makeCompressed();
value.resizeNonZeros(header.num_non_zeros);
} else {
// Temporarily load inner sizes, then reserve.
std::vector<StorageIndex> inner_sizes(header.outer_size);
memcpy(inner_sizes.data(), src, data_bytes);
src += data_bytes;
value.uncompress();
value.reserve(inner_sizes);
memcpy(value.innerNonZeroPtr(), inner_sizes.data(), data_bytes);
}
// Outer indices.
data_bytes = sizeof(StorageIndex) * (header.outer_size + 1);
memcpy(value.outerIndexPtr(), src, data_bytes);
src += data_bytes;
// Inner indices.
data_bytes = sizeof(StorageIndex) * header.num_non_zeros;
if (EIGEN_PREDICT_FALSE(src + data_bytes > end)) return nullptr;
memcpy(value.innerIndexPtr(), src, data_bytes);
src += data_bytes;
// Values.
data_bytes = sizeof(Scalar) * header.num_non_zeros;
if (EIGEN_PREDICT_FALSE(src + data_bytes > end)) return nullptr;
memcpy(value.valuePtr(), src, data_bytes);
src += data_bytes;
return src;
}
};
} // end namespace Eigen
#endif // EIGEN_SPARSEMATRIX_H

View File

@ -495,6 +495,78 @@ struct sparse_vector_assign_selector<Dest,Src,SVA_RuntimeSwitch> {
}
// Specialization for SparseVector.
// Serializes [size, numNonZeros, innerIndices, values].
template <typename Scalar, int Options, typename StorageIndex>
class Serializer<SparseVector<Scalar, Options, StorageIndex>, void> {
public:
typedef SparseVector<Scalar, Options, StorageIndex> SparseMat;
struct Header {
typename SparseMat::Index size;
Index num_non_zeros;
};
EIGEN_DEVICE_FUNC size_t size(const SparseMat& value) const {
return sizeof(Header) +
(sizeof(Scalar) + sizeof(StorageIndex)) * value.nonZeros();
}
EIGEN_DEVICE_FUNC uint8_t* serialize(uint8_t* dest, uint8_t* end,
const SparseMat& value) {
if (EIGEN_PREDICT_FALSE(dest == nullptr)) return nullptr;
if (EIGEN_PREDICT_FALSE(dest + size(value) > end)) return nullptr;
const size_t header_bytes = sizeof(Header);
Header header = {value.innerSize(), value.nonZeros()};
EIGEN_USING_STD(memcpy)
memcpy(dest, &header, header_bytes);
dest += header_bytes;
// Inner indices.
std::size_t data_bytes = sizeof(StorageIndex) * header.num_non_zeros;
memcpy(dest, value.innerIndexPtr(), data_bytes);
dest += data_bytes;
// Values.
data_bytes = sizeof(Scalar) * header.num_non_zeros;
memcpy(dest, value.valuePtr(), data_bytes);
dest += data_bytes;
return dest;
}
EIGEN_DEVICE_FUNC const uint8_t* deserialize(const uint8_t* src,
const uint8_t* end,
SparseMat& value) const {
if (EIGEN_PREDICT_FALSE(src == nullptr)) return nullptr;
if (EIGEN_PREDICT_FALSE(src + sizeof(Header) > end)) return nullptr;
const size_t header_bytes = sizeof(Header);
Header header;
EIGEN_USING_STD(memcpy)
memcpy(&header, src, header_bytes);
src += header_bytes;
value.setZero();
value.resize(header.size);
value.resizeNonZeros(header.num_non_zeros);
// Inner indices.
std::size_t data_bytes = sizeof(StorageIndex) * header.num_non_zeros;
if (EIGEN_PREDICT_FALSE(src + data_bytes > end)) return nullptr;
memcpy(value.innerIndexPtr(), src, data_bytes);
src += data_bytes;
// Values.
data_bytes = sizeof(Scalar) * header.num_non_zeros;
if (EIGEN_PREDICT_FALSE(src + data_bytes > end)) return nullptr;
memcpy(value.valuePtr(), src, data_bytes);
src += data_bytes;
return src;
}
};
} // end namespace Eigen
#endif // EIGEN_SPARSEVECTOR_H

View File

@ -670,6 +670,12 @@ bool test_isCwiseApprox(const DenseBase<Derived1>& m1,
return true;
}
template <typename Derived1, typename Derived2>
bool test_isCwiseApprox(const SparseMatrixBase<Derived1>& m1,
const SparseMatrixBase<Derived2>& m2, bool exact) {
return test_isCwiseApprox(m1.toDense(), m2.toDense(), exact);
}
template<typename T, typename U>
bool test_is_equal(const T& actual, const U& expected, bool expect_equal)
{

View File

@ -9,8 +9,71 @@
#include "main.h"
#include <vector>
#include <Eigen/Core>
#include <Eigen/SparseCore>
#include <vector>
template <typename T>
struct RandomImpl {
static auto Create(Eigen::Index rows, Eigen::Index cols) {
return T::Random(rows, cols);
}
};
template <typename Scalar, int Options, typename DenseIndex>
struct RandomImpl<Eigen::SparseMatrix<Scalar, Options, DenseIndex>> {
using T = Eigen::SparseMatrix<Scalar, Options, DenseIndex>;
static auto Create(Eigen::Index rows, Eigen::Index cols) {
Eigen::SparseMatrix<Scalar, Options, DenseIndex> M(rows, cols);
M.setZero();
double density = 0.1;
// Reserve some space along each inner dim.
int nnz = static_cast<int>(density * 1.5 * M.innerSize());
M.reserve(Eigen::VectorXi::Constant(M.outerSize(), nnz));
for (int j = 0; j < M.outerSize(); j++) {
for (int i = 0; i < M.innerSize(); i++) {
bool zero = (Eigen::internal::random<double>(0, 1) > density);
if (!zero) {
M.insertByOuterInner(j, i) = internal::random<Scalar>();
}
}
}
// 50-50 whether to compress or not.
if (Eigen::internal::random<double>(0, 1) >= 0.5) {
M.makeCompressed();
}
return M;
}
};
template <typename Scalar, int Options, typename DenseIndex>
struct RandomImpl<Eigen::SparseVector<Scalar, Options, DenseIndex>> {
using T = Eigen::SparseVector<Scalar, Options, DenseIndex>;
static auto Create(Eigen::Index rows, Eigen::Index cols) {
Eigen::SparseVector<Scalar, Options, DenseIndex> M(rows, cols);
M.setZero();
double density = 0.1;
// Reserve some space along each inner dim.
int nnz = static_cast<int>(density * 1.5 * M.innerSize());
M.reserve(nnz);
for (int i = 0; i < M.innerSize(); i++) {
bool zero = (Eigen::internal::random<double>(0, 1) > density);
if (!zero) {
M.insert(i) = internal::random<Scalar>();
}
}
return M;
}
};
struct MyPodType {
double x;
@ -67,9 +130,9 @@ template<typename T>
void test_eigen_type(const T& type) {
const Index rows = type.rows();
const Index cols = type.cols();
const T initial = T::Random(rows, cols);
const T initial = RandomImpl<T>::Create(rows, cols);
// Serialize.
Eigen::Serializer<T> serializer;
size_t buffer_size = serializer.size(initial);
@ -160,7 +223,9 @@ EIGEN_DECLARE_TEST(serializer)
CALL_SUBTEST( test_eigen_type(Eigen::Vector3f()) );
CALL_SUBTEST( test_eigen_type(Eigen::Matrix4d()) );
CALL_SUBTEST( test_eigen_type(Eigen::MatrixXd(15, 17)) );
CALL_SUBTEST(test_eigen_type(Eigen::SparseMatrix<float>(13, 12)));
CALL_SUBTEST(test_eigen_type(Eigen::SparseVector<float>(17)));
CALL_SUBTEST( test_dense_types( Eigen::Array33f(),
Eigen::ArrayXd(10),
Eigen::MatrixXd(15, 17)) );