Modify elasticity codegen for nvcc compliance

Q-Minh · Q-Minh · commit e2edd7a4b8b9 · 2024-11-01T23:11:35.000-04:00
std::tuple is not supported by nvcc.
diff --git a/python/elasticity/potentials.py b/python/elasticity/potentials.py
@@ -55,7 +55,6 @@ def codegen(fpsi, energy_name: str):
 #include "pbat/math/linalg/mini/Matrix.h"
 
 #include <cmath>
-#include <tuple>
 
 namespace pbat {{
 namespace physics {{
@@ -82,17 +81,17 @@ def codegen(fpsi, energy_name: str):
         hesspsicode = cg.codegen(hesspsi.transpose(
         ), lhs=sp.MatrixSymbol("H", vecF.shape[0], vecF.shape[0]), scalar_type="ScalarType")
         evalgradpsi = cg.codegen([psi, gradpsi], lhs=[sp.Symbol(
-            "psi"), sp.MatrixSymbol("G", *gradpsi.shape)], scalar_type="ScalarType")
+            "psi"), sp.MatrixSymbol("gF", *gradpsi.shape)], scalar_type="ScalarType")
         evalgradhesspsi = cg.codegen([psi, gradpsi, hesspsi], lhs=[
             sp.Symbol("psi"),
-            sp.MatrixSymbol("G", *gradpsi.shape),
+            sp.MatrixSymbol("gF", *gradpsi.shape),
             sp.MatrixSymbol(
-                "H", vecF.shape[0], vecF.shape[0])
+                "HF", vecF.shape[0], vecF.shape[0])
         ], scalar_type="ScalarType")
         gradhesspsi = cg.codegen([gradpsi, hesspsi], lhs=[
-            sp.MatrixSymbol("G", *gradpsi.shape),
+            sp.MatrixSymbol("gF", *gradpsi.shape),
             sp.MatrixSymbol(
-                "H", vecF.shape[0], vecF.shape[0])
+                "HF", vecF.shape[0], vecF.shape[0])
         ], scalar_type="ScalarType")
         impl = f"""
 template <>
@@ -131,39 +130,45 @@ def codegen(fpsi, energy_name: str):
             typename TMatrix::ScalarType mu,
             typename TMatrix::ScalarType lambda) const;
 
-        template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
-        PBAT_HOST_DEVICE
-        std::tuple<
-            typename TMatrix::ScalarType,
-            SVector<typename TMatrix::ScalarType, {vecF.shape[0]}>
+        template <
+            math::linalg::mini::CReadableVectorizedMatrix TMatrix, 
+            math::linalg::mini::CWriteableVectorizedMatrix TMatrixGF
             >
+        PBAT_HOST_DEVICE
+        typename TMatrix::ScalarType
         evalWithGrad(
             TMatrix const& F,
             typename TMatrix::ScalarType mu,
-            typename TMatrix::ScalarType lambda) const;
+            typename TMatrix::ScalarType lambda,
+            TMatrixGF& gF) const;
 
-        template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
-        PBAT_HOST_DEVICE
-        std::tuple<
-            typename TMatrix::ScalarType,
-            SVector<typename TMatrix::ScalarType, {vecF.shape[0]}>,
-            SMatrix<typename TMatrix::ScalarType, {vecF.shape[0]},{vecF.shape[0]}>
+        template <
+            math::linalg::mini::CReadableVectorizedMatrix TMatrix,
+            math::linalg::mini::CWriteableVectorizedMatrix TMatrixGF, 
+            math::linalg::mini::CWriteableVectorizedMatrix TMatrixHF
             >
+        PBAT_HOST_DEVICE
+        typename TMatrix::ScalarType
         evalWithGradAndHessian(
             TMatrix const& F,
             typename TMatrix::ScalarType mu,
-            typename TMatrix::ScalarType lambda) const;
-
-        template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
-        PBAT_HOST_DEVICE
-        std::tuple<
-            SVector<typename TMatrix::ScalarType, {vecF.shape[0]}>,
-            SMatrix<typename TMatrix::ScalarType, {vecF.shape[0]},{vecF.shape[0]}>
+            typename TMatrix::ScalarType lambda,
+            TMatrixGF& gF,
+            TMatrixHF& HF) const;
+
+        template <
+            math::linalg::mini::CReadableVectorizedMatrix TMatrix,
+            math::linalg::mini::CWriteableVectorizedMatrix TMatrixGF, 
+            math::linalg::mini::CWriteableVectorizedMatrix TMatrixHF
             >
+        PBAT_HOST_DEVICE
+        void
         gradAndHessian(
             TMatrix const& F,
             typename TMatrix::ScalarType mu,
-            typename TMatrix::ScalarType lambda) const;
+            typename TMatrix::ScalarType lambda,
+            TMatrixGF& gF,
+            TMatrixHF& HF) const;
 }};
 
 template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
@@ -208,60 +213,75 @@ def codegen(fpsi, energy_name: str):
     return H;
 }}
 
-template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
-PBAT_HOST_DEVICE
-std::tuple<
-    typename TMatrix::ScalarType,
-    {energy_name}<{d}>::SVector<typename TMatrix::ScalarType, {vecF.shape[0]}>
+template <
+    math::linalg::mini::CReadableVectorizedMatrix TMatrix,
+    math::linalg::mini::CWriteableVectorizedMatrix TMatrixGF
     >
+PBAT_HOST_DEVICE
+typename TMatrix::ScalarType
 {energy_name}<{d}>::evalWithGrad(
     [[maybe_unused]] TMatrix const& F,
     [[maybe_unused]] typename TMatrix::ScalarType mu,
-    [[maybe_unused]] typename TMatrix::ScalarType lambda) const
+    [[maybe_unused]] typename TMatrix::ScalarType lambda,
+    TMatrixGF& gF) const
 {{
+    static_assert(
+        TMatrixGF::kRows == {vecF.shape[0]} and TMatrixGF::kCols == 1, 
+        "Grad w.r.t. F must have dimensions {vecF.shape[0]}x1");
     using ScalarType = typename TMatrix::ScalarType;
     ScalarType psi;
-    SVector<ScalarType, {vecF.shape[0]}> G;
 {cg.tabulate(evalgradpsi, spaces=4)}
-    return {{psi, G}};
+    return psi;
 }}
 
-template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
-PBAT_HOST_DEVICE
-std::tuple<
-    typename TMatrix::ScalarType,
-    {energy_name}<{d}>::SVector<typename TMatrix::ScalarType, {vecF.shape[0]}>,
-    {energy_name}<{d}>::SMatrix<typename TMatrix::ScalarType, {vecF.shape[0]},{vecF.shape[0]}>
+template <
+    math::linalg::mini::CReadableVectorizedMatrix TMatrix,
+    math::linalg::mini::CWriteableVectorizedMatrix TMatrixGF,
+    math::linalg::mini::CWriteableVectorizedMatrix TMatrixHF
     >
+PBAT_HOST_DEVICE
+typename TMatrix::ScalarType
 {energy_name}<{d}>::evalWithGradAndHessian(
     [[maybe_unused]] TMatrix const& F,
     [[maybe_unused]] typename TMatrix::ScalarType mu,
-    [[maybe_unused]] typename TMatrix::ScalarType lambda) const
+    [[maybe_unused]] typename TMatrix::ScalarType lambda,
+    TMatrixGF& gF,
+    TMatrixHF& HF) const
 {{
+    static_assert(
+        TMatrixGF::kRows == {vecF.shape[0]} and TMatrixGF::kCols == 1, 
+        "Grad w.r.t. F must have dimensions {vecF.shape[0]}x1");
+    static_assert(
+        TMatrixHF::kRows == {vecF.shape[0]} and TMatrixHF::kCols == {vecF.shape[0]}, 
+        "Hessian w.r.t. F must have dimensions {vecF.shape[0]}x{vecF.shape[0]}");
     using ScalarType = typename TMatrix::ScalarType;
     ScalarType psi;
-    SVector<ScalarType, {vecF.shape[0]}> G;
-    SMatrix<ScalarType, {vecF.shape[0]},{vecF.shape[0]}> H;
 {cg.tabulate(evalgradhesspsi, spaces=4)}
-    return {{psi, G, H}};
+    return psi;
 }}
 
-template <math::linalg::mini::CReadableVectorizedMatrix TMatrix>
-PBAT_HOST_DEVICE
-std::tuple<
-    {energy_name}<{d}>::SVector<typename TMatrix::ScalarType, {vecF.shape[0]}>,
-    {energy_name}<{d}>::SMatrix<typename TMatrix::ScalarType, {vecF.shape[0]},{vecF.shape[0]}>
+template <
+    math::linalg::mini::CReadableVectorizedMatrix TMatrix,
+    math::linalg::mini::CWriteableVectorizedMatrix TMatrixGF,
+    math::linalg::mini::CWriteableVectorizedMatrix TMatrixHF
     >
+PBAT_HOST_DEVICE
+void
 {energy_name}<{d}>::gradAndHessian(
     [[maybe_unused]] TMatrix const& F,
     [[maybe_unused]] typename TMatrix::ScalarType mu,
-    [[maybe_unused]] typename TMatrix::ScalarType lambda) const
+    [[maybe_unused]] typename TMatrix::ScalarType lambda,
+    TMatrixGF& gF,
+    TMatrixHF& HF) const
 {{
+    static_assert(
+        TMatrixGF::kRows == {vecF.shape[0]} and TMatrixGF::kCols == 1, 
+        "Grad w.r.t. F must have dimensions {vecF.shape[0]}x1");
+    static_assert(
+        TMatrixHF::kRows == {vecF.shape[0]} and TMatrixHF::kCols == {vecF.shape[0]}, 
+        "Hessian w.r.t. F must have dimensions {vecF.shape[0]}x{vecF.shape[0]}");
     using ScalarType = typename TMatrix::ScalarType;
-    SVector<ScalarType, {vecF.shape[0]}> G;
-    SMatrix<ScalarType, {vecF.shape[0]},{vecF.shape[0]}> H;
 {cg.tabulate(gradhesspsi, spaces=4)}
-    return {{G, H}};
 }}
 """
         source.append(impl)
diff --git a/python/examples/elasticity.py b/python/examples/elasticity.py
@@ -53,7 +53,7 @@
     # Create hyper elastic potential
     Y, nu, energy = args.Y, args.nu, pbat.fem.HyperElasticEnergy.StableNeoHookean
     hep, detJeU, GNeU = pbat.fem.hyper_elastic_potential(
-        mesh, Y=Y, nu=nu, energy=energy, detJ=detJeF)
+        mesh, Y=Y, nu=nu, energy=energy, detJe=detJeF)
 
     # Set Dirichlet boundary conditions
     Xmin = mesh.X.min(axis=1)
diff --git a/python/examples/ipc.py b/python/examples/ipc.py
@@ -379,7 +379,7 @@ def __call__(self, xk: np.ndarray):
     # Create hyper elastic potential
     Y, nu, psi = args.Y, args.nu, pbat.fem.HyperElasticEnergy.StableNeoHookean
     hep, detJeU, GNeU = pbat.fem.hyper_elastic_potential(
-        mesh, Y=Y, nu=nu, energy=psi, detJ=detJeF)
+        mesh, Y=Y, nu=nu, energy=psi, detJe=detJeF)
 
     # Setup IPC contact handling
     F = igl.boundary_facets(C)
diff --git a/python/examples/vbd.py b/python/examples/vbd.py
@@ -118,7 +118,7 @@ def partition_vertices(GVT, dbcs):
     GNeU = pbat.fem.shape_function_gradients(mesh, quadrature_order=1)
     g = np.zeros(mesh.dims)
     g[-1] = -9.81
-    f, detJeF = pbat.fem.load_vector(mesh, rho*g, detJ=detJeU, flatten=False)
+    f, detJeF = pbat.fem.load_vector(mesh, rho*g, detJe=detJeU, flatten=False)
     a = f / m
 
     # Compute material (Lame) constants
diff --git a/source/pbat/fem/HyperElasticPotential.h b/source/pbat/fem/HyperElasticPotential.h
@@ -241,8 +241,9 @@ HyperElasticPotential<TMesh, THyperElasticEnergy, QuadratureOrder>::ComputeEleme
     auto constexpr kNodesPerElement = ElementType::kNodes;
     auto constexpr kDofsPerElement  = kNodesPerElement * kDims;
     auto const wg                   = common::ToEigen(QuadratureRuleType::weights);
-    using math::linalg::mini::FromEigen;
-    using math::linalg::mini::ToEigen;
+    namespace mini                  = math::linalg::mini;
+    using mini::FromEigen;
+    using mini::ToEigen;
     if (not bWithGradient and not bWithHessian)
     {
         tbb::parallel_for(Index{0}, Index{numberOfElements}, [&](Index e) {
@@ -275,7 +276,8 @@ HyperElasticPotential<TMesh, THyperElasticEnergy, QuadratureOrder>::ComputeEleme
                     e * kStride + g * MeshType::kDims);
                 Matrix<kDims, kDims> const F = xe * gradPhi;
                 auto vecF                    = FromEigen(F);
-                auto [psiF, gradPsiF]        = Psi.evalWithGrad(vecF, mue(g, e), lambdae(g, e));
+                mini::SVector<Scalar, kDims * kDims> gradPsiF;
+                auto psiF = Psi.evalWithGrad(vecF, mue(g, e), lambdae(g, e), gradPsiF);
                 Ue(e) += (wg(g) * detJe(g, e)) * psiF;
                 auto const GP = FromEigen(gradPhi);
                 auto GPsix    = GradientWrtDofs<ElementType, kDims>(gradPsiF, GP);
@@ -321,8 +323,10 @@ HyperElasticPotential<TMesh, THyperElasticEnergy, QuadratureOrder>::ComputeEleme
                     e * kStride + g * MeshType::kDims);
                 Matrix<kDims, kDims> const F = xe * gradPhi;
                 auto vecF                    = FromEigen(F);
-                auto [psiF, gradPsiF, hessPsiF] =
-                    Psi.evalWithGradAndHessian(vecF, mue(g, e), lambdae(g, e));
+                mini::SVector<Scalar, kDims * kDims> gradPsiF;
+                mini::SMatrix<Scalar, kDims * kDims, kDims * kDims> hessPsiF;
+                auto psiF =
+                    Psi.evalWithGradAndHessian(vecF, mue(g, e), lambdae(g, e), gradPsiF, hessPsiF);
                 Ue(e) += (wg(g) * detJe(g, e)) * psiF;
                 auto const GP = FromEigen(gradPhi);
                 auto GPsix    = GradientWrtDofs<ElementType, kDims>(gradPsiF, GP);
diff --git a/source/pbat/gpu/vbd/VbdImplKernels.cu b/source/pbat/gpu/vbd/VbdImplKernels.cu
@@ -44,12 +44,9 @@ __global__ void MinimizeBackwardEuler(BackwardEulerMinimization BDF)
     // 2. Accumulate results into vertex hessian and gradient
     SVector<GpuScalar, 3> xti     = FromBuffers<3, 1>(BDF.xt, i);
     SVector<GpuScalar, 3> xitilde = FromBuffers<3, 1>(BDF.xtilde, i);
-    SVector<GpuScalar, 3> xi{
-        BDF.x[0][i],
-        BDF.x[1][i],
-        BDF.x[2][i]} /*= FromBuffers<3, 1>(BDF.x, i)*/;
-    SMatrix<GpuScalar, 3, 3> Hi = Zeros<GpuScalar, 3, 3>{};
-    SVector<GpuScalar, 3> gi    = Zeros<GpuScalar, 3, 1>{};
+    SVector<GpuScalar, 3> xi      = FromBuffers<3, 1>(BDF.x, i);
+    SMatrix<GpuScalar, 3, 3> Hi   = Zeros<GpuScalar, 3, 3>{};
+    SVector<GpuScalar, 3> gi      = Zeros<GpuScalar, 3, 1>{};
     // Add elastic energy derivatives
     auto const nActiveThreads = min(nAdjacentElements, nThreadsPerBlock);
     for (auto j = 0; j < nActiveThreads; ++j)
@@ -93,14 +90,10 @@ PBAT_DEVICE void BackwardEulerMinimization::ComputeStableNeoHookeanDerivatives(
     SMatrix<GpuScalar, 3, 4> xe  = FromBuffers(x, v.Transpose());
     SMatrix<GpuScalar, 4, 3> GPe = FromFlatBuffer<4, 3>(GP, e);
     SMatrix<GpuScalar, 3, 3> Fe  = xe * GPe;
+    physics::StableNeoHookeanEnergy<3> Psi{};
     SVector<GpuScalar, 9> gF;
     SMatrix<GpuScalar, 9, 9> HF;
-    physics::StableNeoHookeanEnergy<3> Psi{};
-    // NOTE:
-    // For some reason, nvcc doesn't like structured bindings, so I can't write
-    // auto [gF, HF] = Psi.gradAndHessian(Fe, lamee(0), lamee(1));
-    // or else I get nan values?? Weird.
-    std::tie(gF, HF) = Psi.gradAndHessian(Fe, lamee(0), lamee(1));
+    Psi.gradAndHessian(Fe, lamee(0), lamee(1), gF, HF);
     // Write vertex-specific derivatives into output memory HGe
     SMatrixView<GpuScalar, 3, 4> HGei(Hge);
     auto Hi = HGei.Slice<3, 3>(0, 0);
diff --git a/source/pbat/physics/HyperElasticity.h b/source/pbat/physics/HyperElasticity.h
@@ -9,7 +9,6 @@
 #include <fmt/core.h>
 #include <pbat/Aliases.h>
 #include <string>
-#include <tuple>
 
 namespace pbat {
 namespace physics {
@@ -40,26 +39,25 @@ concept CHyperElasticEnergy = requires(T t)
         t.evalWithGrad(
             math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims>{},
             Scalar{},
-            Scalar{})
-    } -> std::convertible_to<
-        std::tuple<Scalar, math::linalg::mini::SVector<Scalar, T::kDims * T::kDims>>>;
+            Scalar{},
+            std::declval<math::linalg::mini::SVector<Scalar, T::kDims * T::kDims>&>())
+    } -> std::convertible_to<Scalar>;
     {
         t.evalWithGradAndHessian(
             math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims>{},
             Scalar{},
-            Scalar{})
-    } -> std::convertible_to<std::tuple<
-        Scalar,
-        math::linalg::mini::SVector<Scalar, T::kDims * T::kDims>,
-        math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims, T::kDims * T::kDims>>>;
-    {
-        t.gradAndHessian(
-            math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims>{},
             Scalar{},
-            Scalar{})
-    } -> std::convertible_to<std::tuple<
-        math::linalg::mini::SVector<Scalar, T::kDims * T::kDims>,
-        math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims, T::kDims * T::kDims>>>;
+            std::declval<math::linalg::mini::SVector<Scalar, T::kDims * T::kDims>&>(),
+            std::declval<
+                math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims, T::kDims * T::kDims>&>())
+    } -> std::convertible_to<Scalar>;
+    {t.gradAndHessian(
+        math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims>{},
+        Scalar{},
+        Scalar{},
+        std::declval<math::linalg::mini::SVector<Scalar, T::kDims * T::kDims>&>(),
+        std::declval<
+            math::linalg::mini::SMatrix<Scalar, T::kDims * T::kDims, T::kDims * T::kDims>&>())};
 };
 
 template <class TDerivedY, class TDerivednu>
diff --git a/source/pbat/physics/SaintVenantKirchhoffEnergy.cpp b/source/pbat/physics/SaintVenantKirchhoffEnergy.cpp
@@ -9,19 +9,21 @@
 TEST_CASE("[physics] SaintVenantKirchhoffEnergy")
 {
     using namespace pbat;
+    namespace mini = pbat::math::linalg::mini;
     common::ForValues<1, 2, 3>([]<auto Dims>() {
-        using math::linalg::mini::FromEigen;
+        using mini::FromEigen;
         physics::SaintVenantKirchhoffEnergy<Dims> psi{};
         Matrix<Dims, Dims> const F = Matrix<Dims, Dims>::Identity();
         Scalar constexpr Y         = 1e6;
         Scalar constexpr nu        = 0.45;
         auto const [mu, lambda]    = physics::LameCoefficients(Y, nu);
         auto vecF                  = FromEigen(F.reshaped());
         Scalar const ePsi          = psi.eval(vecF, mu, lambda);
-        auto const eGradPsi        = psi.evalWithGrad(vecF, mu, lambda);
-        Scalar const ePsiFromGrad  = std::get<0>(eGradPsi);
-        auto const eGradHessPsi    = psi.evalWithGradAndHessian(vecF, mu, lambda);
-        Scalar const ePsiFromHess  = std::get<0>(eGradHessPsi);
+        mini::SVector<Scalar, Dims * Dims> gF;
+        Scalar const ePsiFromGrad = psi.evalWithGrad(vecF, mu, lambda, gF);
+        gF.SetZero();
+        mini::SMatrix<Scalar, Dims * Dims, Dims * Dims> HF;
+        Scalar const ePsiFromHess = psi.evalWithGradAndHessian(vecF, mu, lambda, gF, HF);
         bool const bIsEnergyNonNegative =
             (ePsi >= 0.) && (ePsiFromGrad >= 0.) && (ePsiFromHess >= 0.);
         CHECK(bIsEnergyNonNegative);
@@ -30,8 +32,8 @@ TEST_CASE("[physics] SaintVenantKirchhoffEnergy")
         Scalar const trE           = E.trace();
         Scalar const ePsiExpected  = mu * (E.array() * E.array()).sum() + 0.5 * lambda * trE * trE;
         Scalar const ePsiError     = std::abs(ePsi - ePsiExpected) +
-                                 std::abs(std::get<0>(eGradPsi) - ePsiExpected) +
-                                 std::abs(std::get<0>(eGradHessPsi) - ePsiExpected);
+                                 std::abs(ePsiFromGrad - ePsiExpected) +
+                                 std::abs(ePsiFromHess - ePsiExpected);
         Scalar constexpr zero = 1e-15;
         CHECK_LE(ePsiError, zero);
     });
diff --git a/source/pbat/physics/SaintVenantKirchhoffEnergy.h b/source/pbat/physics/SaintVenantKirchhoffEnergy.h
diff --git a/source/pbat/physics/StableNeoHookeanEnergy.cpp b/source/pbat/physics/StableNeoHookeanEnergy.cpp
diff --git a/source/pbat/physics/StableNeoHookeanEnergy.h b/source/pbat/physics/StableNeoHookeanEnergy.h
