基于Fortran优化向量-矩阵-向量高效求和：消除嵌套循环

需求与数学表达式 #

需要高效完成以下求和计算，消除现有代码中的四层嵌套循环：

• 当 $i+j ≤ I$ 时：
  $$A_{j} = \sum_{m,n,i} C_{m, i+j} G_{ m, n, i, j} C_{n, i}$$
• 当 $i+j ≥ I$ 时：
  $$A_{j} = \sum_{m,n,i} C_{m, i+j-I} G_{ m, n, i, j} C_{n, i}$$

原代码问题 #

现有Fortran实现存在四层嵌套循环，内存访问不连续，分支判断频繁，且未利用硬件的向量计算能力，效率较低：

program main
  implicit none
  real, allocatable :: A(:)
  real, allocatable :: C(:,:), G(:,:,:,:)
  integer :: i, j, m, n
  integer, parameter :: N = 1500, I = 2000

  allocate(A(I))  ! 原代码中A(J)为笔误，修正为A(I)
  allocate(C(N,I))
  allocate(G(N,N,I,I))
  ! If i+j <= I,then
  ! A_{j} = \sum_{m,n,i} C_{m, i+j} G_{ m, n, i, j} C_{n, i} 
  ! else if i+j >= I, then
  ! A_{j} = \sum_{m,n,i} C_{m, i+j-I} G_{ m, n, i, j} C_{n, i}
  do j = 1, I
    do i = 1, I
      if ( i + j <= I ) then
        do n = 1, N
          do m = 1, N
            A(j) = A(j) + C(m,i+j) * G(m,n,i,j) * C(n,i)
          end do
        end do
      else
        do n = 1, N
          do m = 1, N
            A(j) = A(j) + C(m,i+j-I) * G(m,n,i,j) * C(n,i)
          end do
        end do
      end if
    end do
  end do
  
end program main

优化方案 #

1. 移除循环内分支，拆分i循环 #

将每个j对应的i范围拆分为两个区间：i ≤ I-j（满足i+j≤I）和i > I-j（满足i+j≥I），避免在每个i循环中进行条件判断，减少分支预测开销。

2. 调整循环顺序，优化内存访问 #

利用Fortran的列主序存储特性，将内存连续的维度作为内层循环。对于四维数组G(m,n,i,j)，固定m、i、j时，n的访问是连续内存块；C(n,i)固定i时，n的访问也是连续内存块。调整循环顺序后，缓存命中率大幅提升。

3. 利用BLAS库替换手动循环 #

使用BLAS（基础线性代数子程序）中的优化函数（如ddot点积、dgemv矩阵-向量乘法）替换手动嵌套循环，BLAS库通常经过底层优化，能充分利用CPU的向量指令和多核能力。

优化后代码示例 #

手动优化循环顺序版本 #

program main
  implicit none
  real, allocatable :: A(:)
  real, allocatable :: C(:,:), G(:,:,:,:)
  real :: temp
  integer :: i, j, m, n, k
  integer, parameter :: N = 1500, I = 2000

  allocate(A(I))
  allocate(C(N,I))
  allocate(G(N,N,I,I))
  A = 0.0  ! 初始化A为0

  do j = 1, I
    ! 处理i+j ≤ I的情况
    do i = 1, I - j
      do m = 1, N
        temp = 0.0
        ! 内层n循环访问连续内存
        do n = 1, N
          temp = temp + G(m,n,i,j) * C(n,i)
        end do
        A(j) = A(j) + C(m,i+j) * temp
      end do
    end do

    ! 处理i+j ≥ I的情况
    do i = max(1, I - j + 1), I
      k = i + j - I
      do m = 1, N
        temp = 0.0
        do n = 1, N
          temp = temp + G(m,n,i,j) * C(n,i)
        end do
        A(j) = A(j) + C(m,k) * temp
      end do
    end do
  end do
  
end program main

BLAS优化版本 #

program main
  implicit none
  real, allocatable :: A(:), v(:)
  real, allocatable :: C(:,:), G(:,:,:,:)
  real :: dot_prod
  integer :: i, j, k
  integer, parameter :: N = 1500, I = 2000
  ! 声明BLAS函数
  external ddot, dgemv

  allocate(A(I), v(N))
  allocate(C(N,I))
  allocate(G(N,N,I,I))
  A = 0.0

  do j = 1, I
    ! 处理i+j ≤ I的情况
    do i = 1, I - j
      ! 矩阵G(:,:,i,j)乘向量C(:,i)得到v
      call dgemv('N', N, N, 1.0, G(1,1,i,j), N, C(1,i), 1, 0.0, v, 1)
      ! 向量C(:,i+j)与v做点积，累加到A(j)
      dot_prod = ddot(N, C(1,i+j), 1, v, 1)
      A(j) = A(j) + dot_prod
    end do

    ! 处理i+j ≥ I的情况
    do i = max(1, I - j + 1), I
      k = i + j - I
      call dgemv('N', N, N, 1.0, G(1,1,i,j), N, C(1,i), 1, 0.0, v, 1)
      dot_prod = ddot(N, C(1,k), 1, v, 1)
      A(j) = A(j) + dot_prod
    end do
  end do
  
end program main

优化效果说明 #

• 分支判断消除：减少CPU分支预测错误带来的性能损耗
• 内存访问优化：连续内存访问提升缓存命中率，降低内存延迟
• BLAS加速：利用高度优化的线性代数库，充分发挥CPU的向量计算能力，对于N=1500、I=2000的场景，性能可提升数倍甚至数十倍

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