Templated viscous stress models

[dseyler]

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

[ktbolt]

@dseyler Excellent work improving performance ! Please have a look at the fluids code next. 😄

It would be interesting to see the profiling data here.

[dseyler]

Haha no promises on the fluids, but I'll add it to the list!

Profiling data from a slightly older and ~5% slower branch can be seen in this plot labeled "templated viscosity"

[ktbolt]

@dseyler Very interesting to see this data.

How do the results change ?

[dseyler]

@ktbolt All test cases pass. For my profiled case, I checked the values of Kvis_u and Kvis_v pre- and post-, and the relative difference was less than 1e-15.

[michelebucelli]

Thanks @dseyler! I think the code looks fine, and the performance gains are very nice.

My only question is whether you think it still makes sense to have the dispatching function compute_visc_stress_newtonian, or we could just replace it with the templated one that uses Eigen. If the caller also uses Eigen, at least for the local terms (e.g. the local stress tensor), the code could probably be a bit simpler and with less indirection layers.

[dseyler]

@michelebucelli Sounds good! I was also thinking we could skip recomputing the viscous stress and tangent at each Gauss point for linear elements since it's only dependent on the shape function gradient. This just requires passing an extra bool into struct_3d() and then returning early (Kvis_u, Kvis_v, and Svis aren't zeroed) when a given element has that flag:

switch (lDmn.solid_visc.viscType) {
      case consts::SolidViscosityModelType::viscType_Newtonian:
        // Viscosity is constant at all Gauss points for linear elements
        if (!recompute_visc) {
          return;
        }
        if (F.nrows() == 3) {
          compute_visc_stress_newtonian<3>(lDmn.solid_visc.mu, eNoN, Nx, vx, F, Svis, Kvis_u, Kvis_v);
        } else if (F.nrows() == 2) {
          compute_visc_stress_newtonian<2>(lDmn.solid_visc.mu, eNoN, Nx, vx, F, Svis, Kvis_u, Kvis_v);
        }
      break;

[michelebucelli]

@dseyler That might be a good idea, although it has a couple of design downsides in my opinion:

1. it introduces coupling between the choice of finite element space and the assembly code (whereas ideally we'd want the assembly code to be FE-space-agnostic);
2. it relies on preserving the value of some variables between calls, although those variables appear on first read to be local and "temporary" (as in, they store scratch information).

Because of this, I'd suggest only keeping this optimization if you measured that it has a noticeable and significant performance impact (that is, if it's less than say a 5% improvement I would recommend against it).

If it does have a major impact, I'd suggest making these choices prominent in the code, by e.g. using explicit variable names and detailed comments (for example, recompute_visc might be replaced by something like using_linear_tet_elements, which would make the connection more explicit). If we go down this path, I'll try giving more precise suggestions in the code review.

[dseyler]

@dseyler That might be a good idea, although it has a couple of design downsides in my opinion:

1. it introduces coupling between the choice of finite element space and the assembly code (whereas ideally we'd want the assembly code to be FE-space-agnostic);
2. it relies on preserving the value of some variables between calls, although those variables appear on first read to be local and "temporary" (as in, they store scratch information).

Because of this, I'd suggest only keeping this optimization if you measured that it has a noticeable and significant performance impact (that is, if it's less than say a 5% improvement I would recommend against it).

If it does have a major impact, I'd suggest making these choices prominent in the code, by e.g. using explicit variable names and detailed comments (for example, recompute_visc might be replaced by something like using_linear_tet_elements, which would make the connection more explicit). If we go down this path, I'll try giving more precise suggestions in the code review.

Makes sense! I measured it, and it comes out to about a 5% improvement. It looks like this pattern of only computing shape function derivatives for the first Gauss point for linear elements is already used about 30 times throughout the code, all triggered by the lShpF flag