Templated viscous stress models

Problem

The Newtonian and potential viscosity models currently run very slow, accounting for nearly half of the assembly runtime for structural mechanics simulations. It looks like this slowdown is mostly due to heap allocation of several runtime-sized arrays for every gauss point that should be moved to the stack. Furthermore, the viscosity models can be templated on nsd, similarly for what is already done for compute_pk2cc. Fixing this cuts assembly runtime by ~30% (20% of total runtime)

Solution

I implemented nsd templated viscosity models in this branch with fixed-size Eigen arrays used in place of the previous heap-allocated arrays.

Other minor improvements include:

• The column count of dynamic arrays can be capped at 27, as this is the largest possible value for eNoN:

/// @brief Largest element node count the fixed-size views below allow (HEX27).
constexpr int MAX_ELEMENT_NODES = 27;

/// @brief An nsd x eNoN matrix whose column count is bounded at compile time,
/// so that it is held on the stack rather than heap allocated.
template <int nsd>
using MatNodes = Eigen::Matrix<double, nsd, Eigen::Dynamic, 0, nsd, MAX_ELEMENT_NODES>;

• nsd templated overloads of mat_dev and mat_symm

template <int nsd>
Matrix<nsd> mat_symm(const Matrix<nsd>& A) {
    return 0.5 * (A + A.transpose());
}

template <int nsd>
Matrix<nsd> mat_dev(const Matrix<nsd>& A) {
    return A - (A.trace() / nsd) * Matrix<nsd>::Identity();
}

Below is the proposed potential viscosity implementation (Newtonian viscosity follows a similar pattern):

template <int nsd>
void compute_visc_stress_potential_impl(const double mu, const int eNoN, const Array<double>& Nx,
                        const Array<double>& vx, const Array<double>& F,
                        Array<double>& Svis, Array3<double>& Kvis_u, Array3<double>& Kvis_v) {
    using MatNsd = mat_fun::Matrix<nsd>;

    // Initialize Svis, Kvis_u, Kvis_v to zero
    Svis = 0.0;
    Kvis_u = 0.0;
    Kvis_v = 0.0;

    Eigen::Map<const MatNsd> F_map(F.data());
    Eigen::Map<const MatNsd> vx_map(vx.data());
    Eigen::Map<const MatNodes<nsd>> Nx_map(Nx.data(), nsd, eNoN);

    const MatNsd F_Ft  = F_map * F_map.transpose();
    const MatNsd Ft_vx = F_map.transpose() * vx_map;
    const MatNsd F_vxt = F_map * vx_map.transpose();

    // F_Nx(i,a) = sum_j F(i,j) * Nx(j,a), and likewise for vx.
    const MatNodes<nsd> F_Nx  = F_map  * Nx_map;
    const MatNodes<nsd> vx_Nx = vx_map * Nx_map;

    // 2nd Piola-Kirchhoff stress due to viscosity,
    // Svis = mu * 1/2 * ( (F^T * dv/dX) + (F^T * dv/dX)^T )
    Eigen::Map<MatNsd> Svis_map(Svis.data());
    Svis_map.noalias() = mu * mat_fun::mat_symm<nsd>(Ft_vx);

    // Tangent matrix contributions due to viscosity
    for (int b = 0; b < eNoN; ++b) {
        for (int a = 0; a < eNoN; ++a) {
            double Nx_Nx = 0.0;
            for (int i = 0; i < nsd; ++i) {
                Nx_Nx += Nx(i,a) * Nx(i,b);
            }

            for (int i = 0; i < nsd; ++i) {
                for (int j = 0; j < nsd; ++j) {
                    int ii = i * nsd + j;
                    Kvis_u(ii,a,b) = 0.5 * mu * (F_Nx(i,b) * vx_Nx(j,a) + Nx_Nx * F_vxt(i,j));
                    Kvis_v(ii,a,b) = 0.5 * mu * (Nx_Nx * F_Ft(i,j) + F_Nx(i,b) * F_Nx(j,a));
                }
            }
        }
    }
}

This is dispatched by nsd in:

/// @brief Dispatches on the spatial dimension.
void compute_visc_stress_potential(const double mu, const int eNoN, const Array<double>& Nx, const Array<double>& vx, const Array<double>& F,
                        Array<double>& Svis, Array3<double>& Kvis_u, Array3<double>& Kvis_v) {
    if (F.nrows() == 3) {
        compute_visc_stress_potential_impl<3>(mu, eNoN, Nx, vx, F, Svis, Kvis_u, Kvis_v);
    } else if (F.nrows() == 2) {
        compute_visc_stress_potential_impl<2>(mu, eNoN, Nx, vx, F, Svis, Kvis_u, Kvis_v);
    }
}

Additional context

No response

Code of Conduct

• I agree to follow this project's Code of Conduct and Contributing Guidelines