如何将多维点积矩阵改写为矩阵乘法（Fortran/NumPy优化）

问题背景 #

我维护着一个数值模拟Fortran代码库，当前核心计算逻辑是通过嵌套concurrent循环结合dot_product完成多维点积计算，代码如下：

n_dimensions = 2 ! 二维场景
n_faces = n_dimensions + 1 ! 二维三角形对应3个面

do concurrent(l=1:n_different_orders, k=1:n_different_orders)
    do concurrent(j=1:size(W(k,l)%array, 3), i=1:size(W(k,l)%array, 2), &
                  face=1:n_faces, jdim=1:n_dimensions, kdim=1:n_dimensions)
    ! 对所有点维度执行内积计算
    A(k,l)%array(kdim,jdim,face,i,j) &
        = dot_product(C(kdim,jdim,face,:), W(k,l)%array(face,i,j,:))
    end do
end do

该操作在OpenMP循环中重复数十万次，且C随迭代动态变化，是程序的核心耗时模块。我尝试通过BLAS/LAPACK的zgemm进行张量收缩改写，但结果不正确，改写代码如下：

m = 2 * 2 * 3
i = num_points
! 提取W数组的最大尺寸
max_n = maxval(ndof_per_order)**2 * 3

allocate(Amat(m, max_n), Wmat(max_n, i))

Cmat = reshape(C, [m, i])

do l=1, n_different_orders
    do k=1, n_different_orders
        n = 3 * size(W(k,l)%array, 3) * size(W(k,l)%array, 2)

        Wmat(1:n, 1:i) = reshape(W(k,l)%array, [n, i])

        call zgemm('N','T',m,n,i,1.0,Cmat,m,Wmat,n,0.0,Amat,m)

        A(k,l)%array = reshape(Amat(1:m, 1:n), [2, 2, 3, size(W(k,l)%array, 2), size(W(k,l)%array, 3)])
    end do
end do

参数规模 #

• num_points：60~100
• ndof_per_order：最大60
• n_different_orders：1~9

Python验证参考 #

我在Python中用numpy.einsum和矩阵乘法实现了正确逻辑，但存在矩阵冗余，效率不如直接点积；按面拆分调用矩阵乘法的方案，效率仍未达预期。

完整验证代码 #

import numpy as np

np.random.seed(1)

ndofi = ndofj = 55
jdim = kdim = 2
faces = jdim + 1
points = 120

C = np.random.rand(kdim, jdim, faces, points)
W = np.random.rand(faces, ndofi, ndofj, points)
AA = np.einsum("abcd,cefd->abcef", C, W)

path_info = np.einsum_path("abcd,cefd->abcef", C, W, optimize="optimal")
print(path_info[0])
print(path_info[1])

AB = np.zeros((kdim, jdim, faces, ndofi, ndofj))
for k in range(kdim):
    for d in range(jdim):
        for f in range(faces):
            for i in range(ndofi):
                for j in range(ndofj):
                    AB[k, d, f, i, j] = np.dot(C[k, d, f, :], W[f, i, j, :])

C_mat = C.reshape(kdim * jdim * faces, points)
W_mat = W.reshape(faces * ndofi * ndofj, points)
A_mat = C_mat @ W_mat.T
A_tmp = A_mat.reshape(kdim, jdim, faces, faces, ndofi, ndofj)
AC = np.zeros((kdim, jdim, faces, ndofi, ndofj))
for f in range(faces):
    AC[:, :, f, :, :] = A_tmp[:, :, f, f, :, :]

assert np.allclose(AA, AB)
assert np.allclose(AC, AB)

按面拆分的实现代码 #

A = np.zeros((kdim, jdim, faces, ndofi, ndofj))

for f in range(faces):
    C_f = C[:, :, f, :].reshape(kdim * jdim, points)
    W_f = W[f, :, :, :].reshape(ndofi * ndofj, points)

    A_f = C_f @ W_f.T

    A[:, :, f, :, :] = A_f.reshape(kdim, jdim, ndofi, ndofj)

需求 #

寻求更高效的Fortran实现方案，降低该操作的运行时占比。

内容的提问来源于stack exchange，提问作者glowl