Eigen.cuh

#ifndef PBAT_GPU_IMPL_COMMON_EIGEN_H
#define PBAT_GPU_IMPL_COMMON_EIGEN_H
 
#include "Buffer.cuh"
#include "pbat/common/Eigen.h"
 
#include <Eigen/Core>
 
namespace pbat {
namespace gpu {
namespace impl {
namespace common {
 
template <class ScalarType, auto D, class TDerived>
void ToBuffer(Eigen::PlainObjectBase<TDerived> const& A, Buffer<ScalarType, D>& buf)
{
    if constexpr (not std::is_same_v<typename TDerived::Scalar, ScalarType>)
    {
        auto A2 = A.template cast<ScalarType>().eval();
        ToBuffer(A2, buf);
    }
 
    using SizeType = decltype(buf.Size());
    auto ncols     = static_cast<SizeType>(A.cols());
    auto nrows     = static_cast<SizeType>(A.rows());
    if constexpr (D > 1)
    {
        if (nrows != buf.Dimensions())
        {
            std::ostringstream ss{};
            ss << "Expected input dimensions " << buf.Dimensions() << "x" << buf.Size()
               << ", but got " << A.rows() << "x" << A.cols() << "\n";
            throw std::invalid_argument(ss.str());
        }
        if (buf.Size() != ncols)
            buf.Resize(ncols);
        for (auto d = 0; d < buf.Dimensions(); ++d)
        {
            thrust::copy(A.row(d).begin(), A.row(d).end(), buf[d].begin());
        }
    }
    else
    {
        auto nsize = static_cast<SizeType>(A.size());
        if (buf.Size() != nsize)
            buf.Resize(nsize);
        thrust::copy(A.data(), A.data() + A.size(), buf.Data());
    }
}
 
template <
    class ScalarType,
    auto D,
    auto StorageOrder = Eigen::ColMajor,
    auto EigenRows    = Eigen::Dynamic,
    auto EigenCols    = Eigen::Dynamic>
void ToBuffer(
    Eigen::Ref<Eigen::Matrix<ScalarType, EigenRows, EigenCols, StorageOrder> const> const& A,
    Buffer<ScalarType, D>& buf)
{
    using SizeType = decltype(buf.Size());
    if constexpr (D > 1)
    {
        auto nrows = static_cast<SizeType>(A.rows());
        auto ncols = static_cast<SizeType>(A.cols());
        if (nrows != buf.Dimensions())
        {
            std::ostringstream ss{};
            ss << "Expected input dimensions " << buf.Dimensions() << "x" << buf.Size()
               << ", but got " << A.rows() << "x" << A.cols() << "\n";
            throw std::invalid_argument(ss.str());
        }
        if (buf.Size() != ncols)
            buf.Resize(ncols);
        for (auto d = 0; d < buf.Dimensions(); ++d)
        {
            thrust::copy(A.row(d).begin(), A.row(d).end(), buf[d].begin());
        }
    }
    else
    {
        auto nsize = static_cast<SizeType>(A.size());
        if (buf.Size() != nsize)
            buf.Resize(nsize);
        thrust::copy(A.data(), A.data() + A.size(), buf.Data());
    }
}
 
template <class ScalarType, auto D, auto StorageOrder = (D > 1) ? Eigen::RowMajor : Eigen::ColMajor>
Eigen::Matrix<ScalarType, Eigen::Dynamic, Eigen::Dynamic, StorageOrder>
ToEigen(Buffer<ScalarType, D> const& buf)
{
    if constexpr (D > 1)
    {
        return pbat::common::ToEigen(buf.Get()).reshaped(buf.Size(), buf.Dimensions()).transpose();
    }
    else
    {
        return pbat::common::ToEigen(buf.Get());
    }
}
 
template <class ScalarType, auto D, class TDerived>
Buffer<ScalarType, D> EigenToBuffer(Eigen::PlainObjectBase<TDerived> const& A)
{
    Buffer<ScalarType, D> buf;
    ToBuffer(A, buf);
    return buf;
}
 
template <
    class ScalarType,
    auto D,
    auto StorageOrder = Eigen::ColMajor,
    auto EigenRows    = Eigen::Dynamic,
    auto EigenCols    = Eigen::Dynamic>
Buffer<ScalarType, D> EigenToBuffer(
    Eigen::Ref<Eigen::Matrix<ScalarType, EigenRows, EigenCols, StorageOrder> const> const& A)
{
    Buffer<ScalarType, D> buf;
    ToBuffer(A, buf);
    return buf;
}
 
} // namespace common
} // namespace impl
} // namespace gpu
} // namespace pbat
 
#endif // PBAT_GPU_IMPL_COMMON_EIGEN_H