Centralize eigen matrix to/from buffer conversions

Author: Q-Minh

source/pbat/gpu/common/Buffer.cuh

@@ -40,6 +40,7 @@ class Buffer
     Data() const;
 
     void Resize(std::size_t count);
+    void SetConstant(T value);
 
     std::conditional_t<(D > 1), std::array<T*, D>, T*> Raw();
     std::conditional_t<(D > 1), std::array<T const*, D>, T const*> Raw() const;
@@ -159,6 +160,15 @@ void Buffer<T, D>::Resize(std::size_t count)
     }
 }
 
+template <class T, int D>
+void Buffer<T, D>::SetConstant(T value)
+{
+    for (auto d = 0; d < Dimensions(); ++d)
+    {
+        thrust::fill(mBuffers[d].begin(), mBuffers[d].end(), value);
+    }
+}
+
 template <class T, int D>
 std::conditional_t<(D > 1), std::array<T*, D>, T*> Buffer<T, D>::Raw()
 {

source/pbat/gpu/common/CMakeLists.txt

@@ -4,6 +4,7 @@ target_sources(PhysicsBasedAnimationToolkit_PhysicsBasedAnimationToolkit
     FILES
     "Buffer.cuh"
     "Cuda.cuh"
+    "Eigen.cuh"
     "SynchronizedList.cuh"
     "Var.cuh"
 )
@@ -12,6 +13,7 @@ target_sources(PhysicsBasedAnimationToolkit_PhysicsBasedAnimationToolkit
     PRIVATE
     "Buffer.cu"
     "Cuda.cu"
+    "Eigen.cu"
     "SynchronizedList.cu"
     "Var.cu"
 )

source/pbat/gpu/common/Eigen.cu

@@ -0,0 +1 @@
+#include "Eigen.cuh"

source/pbat/gpu/common/Eigen.cuh

@@ -0,0 +1,61 @@
+#ifndef PBAT_GPU_COMMON_EIGEN_CUH
+#define PBAT_GPU_COMMON_EIGEN_CUH
+
+#include "Buffer.cuh"
+#include "pbat/common/Eigen.h"
+
+#include <Eigen/Core>
+
+namespace pbat {
+namespace gpu {
+namespace common {
+
+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)
+{
+    if constexpr (D > 1)
+    {
+        if (A.rows() != buf.Dimensions() and A.cols() != buf.Size())
+        {
+            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());
+        }
+        for (auto d = 0; d < buf.Dimensions(); ++d)
+        {
+            thrust::copy(A.row(d).begin(), A.row(d).end(), buf[d].begin());
+        }
+    }
+    else
+    {
+        if (A.size() != buf.Size())
+        {
+            std::ostringstream ss{};
+            ss << "Expected input dimensions " << buf.Dimensions() << "x" << buf.Size()
+               << " or its transpose, but got " << A.rows() << "x" << A.cols() << "\n";
+            throw std::invalid_argument(ss.str());
+        }
+        thrust::copy(A.data(), A.data() + A.size(), buf.Data());
+    }
+}
+
+template <class ScalarType, auto D, auto StorageOrder = Eigen::RowMajor>
+Eigen::Matrix<ScalarType, Eigen::Dynamic, Eigen::Dynamic, StorageOrder>
+ToEigen(Buffer<ScalarType, D> const& buf)
+{
+    return pbat::common::ToEigen(buf.Get()).reshaped(buf.Size(), buf.Dimensions()).transpose();
+}
+
+} // namespace common
+} // namespace gpu
+} // namespace pbat
+
+#endif // PBAT_GPU_COMMON_EIGEN_CUH

source/pbat/gpu/geometry/BvhImpl.cu

@@ -4,7 +4,7 @@
 
 #include "BvhImpl.cuh"
 #include "BvhImplKernels.cuh"
-#include "pbat/common/Eigen.h"
+#include "pbat/gpu/common/Eigen.cuh"
 
 #include <exception>
 #include <string>
@@ -30,7 +30,7 @@ BvhImpl::BvhImpl(std::size_t nPrimitives, std::size_t nOverlaps)
       visits(nPrimitives - 1),
       overlaps(nOverlaps)
 {
-    thrust::fill(thrust::device, parent.Data(), parent.Data() + parent.Size(), GpuIndex{-1});
+    parent.SetConstant(GpuIndex(-1));
 }
 
 void BvhImpl::Build(
@@ -50,7 +50,7 @@ void BvhImpl::Build(
     }
 
     // 0. Reset intermediate data
-    thrust::fill(thrust::device, visits.Raw(), visits.Raw() + visits.Size(), GpuIndex{0});
+    visits.SetConstant(GpuIndex(0));
 
     // 1. Construct leaf node (i.e. simplex) bounding boxes
     auto const leafBegin        = n - 1;
@@ -155,50 +155,42 @@ std::size_t BvhImpl::NumberOfAllocatedBoxes() const
 
 Eigen::Matrix<GpuScalar, Eigen::Dynamic, Eigen::Dynamic> BvhImpl::Min() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(b.Get()).reshaped(b.Size(), b.Dimensions());
+    return common::ToEigen(b);
 }
 
 Eigen::Matrix<GpuScalar, Eigen::Dynamic, Eigen::Dynamic> BvhImpl::Max() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(e.Get()).reshaped(e.Size(), e.Dimensions());
+    return common::ToEigen(e);
 }
 
 Eigen::Vector<GpuIndex, Eigen::Dynamic> BvhImpl::SimplexOrdering() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(simplex.Get());
+    return common::ToEigen(simplex);
 }
 
 Eigen::Vector<typename BvhImpl::MortonCodeType, Eigen::Dynamic> BvhImpl::MortonCodes() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(morton.Get());
+    return common::ToEigen(morton);
 }
 
 Eigen::Matrix<GpuIndex, Eigen::Dynamic, 2> BvhImpl::Child() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(child.Get()).reshaped(child.Size(), child.Dimensions());
+    return common::ToEigen(child);
 }
 
 Eigen::Vector<GpuIndex, Eigen::Dynamic> BvhImpl::Parent() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(parent.Get());
+    return common::ToEigen(parent);
 }
 
 Eigen::Matrix<GpuIndex, Eigen::Dynamic, 2> BvhImpl::Rightmost() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(rightmost.Get()).reshaped(rightmost.Size(), rightmost.Dimensions());
+    return common::ToEigen(rightmost);
 }
 
 Eigen::Vector<GpuIndex, Eigen::Dynamic> BvhImpl::Visits() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(visits.Get());
+    return common::ToEigen(visits);
 }
 
 } // namespace geometry

source/pbat/gpu/geometry/Primitives.cu

@@ -4,7 +4,7 @@
 
 #include "Primitives.h"
 #include "PrimitivesImpl.cuh"
-#include "pbat/common/Eigen.h"
+#include "pbat/gpu/common/Eigen.cuh"
 
 #include <array>
 #include <thrust/copy.h>
@@ -46,12 +46,7 @@ PointsImpl const* Points::Impl() const
 
 GpuMatrixX Points::Get() const
 {
-    GpuMatrixX V(mImpl->x.Dimensions(), mImpl->x.Size());
-    for (auto d = 0; d < V.rows(); ++d)
-    {
-        thrust::copy(mImpl->x[d].begin(), mImpl->x[d].end(), V.row(d).begin());
-    }
-    return V;
+    return common::ToEigen(mImpl->x);
 }
 
 Points::~Points()
@@ -126,8 +121,7 @@ Bodies& geometry::Bodies::operator=(Bodies&& other) noexcept
 
 GpuIndexMatrixX geometry::Bodies::Get() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(mImpl->body.Get());
+    return common::ToEigen(mImpl->body);
 }
 
 std::size_t Bodies::NumberOfBodies() const

source/pbat/gpu/geometry/PrimitivesImpl.cu

@@ -3,6 +3,7 @@
 // clang-format on
 
 #include "PrimitivesImpl.cuh"
+#include "pbat/gpu/common/Eigen.cuh"
 
 #include <array>
 #include <exception>
@@ -31,12 +32,8 @@ std::size_t PointsImpl::Dimensions() const
 
 void PointsImpl::Update(Eigen::Ref<GpuMatrixX const> const& V)
 {
-    for (auto d = 0; d < 3; ++d)
-        x[d].resize(V.cols());
-    for (auto d = 0; d < 3; ++d)
-    {
-        thrust::copy(V.row(d).begin(), V.row(d).end(), x[d].begin());
-    }
+    x.Resize(V.cols());
+    common::ToBuffer(V, x);
 }
 
 SimplicesImpl::SimplicesImpl(Eigen::Ref<GpuIndexMatrixX const> const& C) : eSimplexType(), inds()
@@ -74,7 +71,7 @@ GpuIndex SimplicesImpl::NumberOfSimplices() const
 BodiesImpl::BodiesImpl(Eigen::Ref<GpuIndexVectorX const> const& B)
     : body(B.size()), nBodies(static_cast<GpuIndex>(B.maxCoeff() + 1))
 {
-    thrust::copy(B.data(), B.data() + B.size(), body.Data());
+    common::ToBuffer(B, body);
 }
 
 GpuIndex BodiesImpl::NumberOfBodies() const

source/pbat/gpu/geometry/SweepAndPruneImpl.cu

@@ -45,8 +45,8 @@ void SweepAndPruneImpl::SortAndSweep(
     }
 
     // 0. Preprocess internal data
-    thrust::fill(mu.Data(), mu.Data() + mu.Size(), 0.f);
-    thrust::fill(sigma.Data(), sigma.Data() + sigma.Size(), 0.f);
+    mu.SetConstant(GpuScalar(0));
+    sigma.SetConstant(GpuScalar(0));
     overlaps.Clear();
     thrust::sequence(thrust::device, binds.Data(), binds.Data() + nBoxes);
 

source/pbat/gpu/vbd/Vbd.cu

@@ -4,7 +4,7 @@
 
 #include "Vbd.h"
 #include "VbdImpl.cuh"
-#include "pbat/common/Eigen.h"
+#include "pbat/gpu/common/Eigen.cuh"
 
 namespace pbat {
 namespace gpu {
@@ -114,17 +114,12 @@ void Vbd::SetBlockSize(GpuIndex blockSize)
 
 GpuMatrixX Vbd::GetPositions() const
 {
-    using pbat::common::ToEigen;
-    return ToEigen(mImpl->X.x.Get())
-        .reshaped(mImpl->X.x.Size(), mImpl->X.x.Dimensions())
-        .transpose();
+    return common::ToEigen(mImpl->X.x);
 }
 
 GpuMatrixX Vbd::GetVelocities() const
 {
-    using pbat::common::ToEigen;
-    auto const& velocity = mImpl->GetVelocity();
-    return ToEigen(velocity.Get()).reshaped(velocity.Size(), velocity.Dimensions()).transpose();
+    return common::ToEigen(mImpl->GetVelocity());
 }
 
 } // namespace vbd