Optimize double dot product and dyadic product

Problem

Many of the operations in mat_fun.h could be sped up with a better use of Eigen and by reordering nested loops such that the innermost loop steps along the index in which memory is stored consecutively.

Specifically, double_dot_product(), dyadic_product, and symmetric_dyadic_product are good candidates for these improvements as they are called several times for each gauss point in struct assembly.

Solution

I propose the following changes, resulting in a ~12% decrease in struct/ustruct assembly runtime, which I've implemented in this branch.

double_dot_product()

Add a fast path for double dot products that are contracted along the last two dimensions, which occurs frequently when computing the isochoric elasticity tensor in bar_to_iso. The mechanism is as follows: nsd x nsd x nsd x nsd tensors are stored as a (nsd x nsd) x (nsd x nsd) matrix. As a result, if contracting along the last two dimensions, this can be expressed as a * b.transpose rather than using the general Eigen .contract(). Similarly, this could also be applied for contraction along the first two dimensions as a.transpose * b, but this pattern currently doesn't appear anywhere in the code, so I only implemented it for contraction along the last two dimensions.

// Fast path for dimsA = dimsB = {2,3}: C_ijmn = A_ijkl * B_mnkl.
        if (dimsA[0] == 2 && dimsA[1] == 3 && dimsB[0] == 2 && dimsB[1] == 3) {
            constexpr int N = nsd * nsd;
            Tensor<nsd> C;
            Eigen::Map<const Eigen::Matrix<double, N, N>> a(A.data());
            Eigen::Map<const Eigen::Matrix<double, N, N>> b(B.data());
            Eigen::Map<Eigen::Matrix<double, N, N>> c(C.data());
            c.noalias() = a * b.transpose();
            return C;
        }

dyadic_product()

Similarly, the nested for loops in the dyadic product can be written as an Eigen outer product:

dyadic_product(const Matrix<nsd>& A, const Matrix<nsd>& B) {
        // Initialize the result tensor
        Tensor<nsd> C;
        constexpr int N = nsd * nsd;

        Eigen::Map<const Eigen::Matrix<double, N, 1>> a(A.data());
        Eigen::Map<const Eigen::Matrix<double, N, 1>> b(B.data());
        Eigen::Map<Eigen::Matrix<double, N, N>> c(C.data());
        c.noalias() = a * b.transpose();

        return C;

symmetric_dyadic_product()

I couldn't figure out how to remove the nested for-loops entirely for this function, since the two matrices A and B each supply one of the first two indices and one of the last two, so the result doesn't reduce to a single outer product. Two of the four for loops can be eliminated, though, leaving only the last two indices, such that the indexed block is stored consecutively in memory. For context, most of the speedup in the other functions can be achieved just by correctly ordering the loops without switching to Eigen, so I think eliminating the other two loops wouldn't drastically improve performance.

template <int nsd>
    Tensor<nsd>
    symmetric_dyadic_product(const Matrix<nsd>& A, const Matrix<nsd>& B) {
        Tensor<nsd> C;

        // Compute the symmetric product: C_ijkl = 0.5 * (A_ik * B_jl + A_il * B_jk)
        for (int l = 0; l < nsd; ++l) {
            for (int k = 0; k < nsd; ++k) {
                Eigen::Map<Eigen::Matrix<double, nsd, nsd>> blk(C.data() + nsd * nsd * (k + nsd * l));
                blk.noalias() = 0.5 * (A.col(k) * B.col(l).transpose()
                                     + A.col(l) * B.col(k).transpose());
            }
        }

        // Return the symmetric product
        return C;

Additional context

No response

Code of Conduct

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