基于BLAS实现高效张量收缩：优化C[a,d,e]=A[a,b,c]*B[d,b,c,e]计算

当前实现通过循环nd次调用DGEMM完成张量收缩，存在调用开销大、非连续内存访问影响性能的问题。以下是通过矩阵重塑实现单次DGEMM调用的优化方案：

核心思路 #

将张量收缩转化为二维矩阵乘法：

1. 合并收缩索引：把需要求和的双索引(b,c)合并为单一维度k = nb * nc；把输出的双索引(d,e)合并为单一维度m = nd * ne。
2. 张量重塑：将三维张量A、四维张量B分别重塑为二维矩阵A_2D(na, k)、B_2D(k, m)，确保矩阵内存连续。
3. 单次矩阵乘法：调用DGEMM完成C_2D(na, m) = A_2D(na, k) * B_2D(k, m)，最后将C_2D重塑回三维张量C(na, nd, ne)。

优化后的Fortran代码 #

Program double_contractions_single_gemm

  Use, Intrinsic :: iso_fortran_env, Only :  wp => real64, li => int64
  External :: dgemm

  Implicit None

  Real( wp ), Dimension( :, :, : ), Allocatable :: a, c
  Real( wp ), Dimension( :, :, :, : ), Allocatable :: b
  Real( wp ), Dimension( :, : ), Allocatable :: a_2d, b_2d, c_2d

  Integer :: na, nb, nc, nd, ne, k, m
  Integer( li ) :: start, finish, rate

  Write( *, * ) 'na, nb, nc, nd, ne ?'
  Read( *, * ) na, nb, nc, nd, ne

  k = nb * nc
  m = nd * ne

  allocate( a( na, nb, nc ) )
  allocate( b( nd, nb, nc, ne ) )
  allocate( c( na, nd, ne ) )

  allocate( a_2d( na, k ) )
  allocate( b_2d( k, m ) )
  allocate( c_2d( na, m ) )

  Call Random_number( a )
  Call Random_number( b )

  ! 将A重塑为二维矩阵：A_2D[a, k]，k对应合并后的b、c索引
  a_2d = reshape( a, [na, k] )

  ! 将B重塑为二维矩阵：B_2D[k, m]，k对应b、c，m对应d、e
  ! order参数指定按b→c→d→e的顺序读取原B的元素，保证重塑后矩阵连续
  b_2d = reshape( b, [k, m], order = [2, 3, 1, 4] )

  Call System_clock( start, rate )

  ! 单次DGEMM完成所有张量收缩
  Call dgemm( 'N', 'N', na, m, k, 1.0_wp, &
              a_2d, size(a_2d, 1), &
              b_2d, size(b_2d, 1), &
              0.0_wp, c_2d, size(c_2d, 1) )

  ! 将结果重塑回三维张量C[a,d,e]
  c = reshape( c_2d, [na, nd, ne] )

  Call System_clock( finish, rate )

  write( *, * ) c

  deallocate( a_2d, b_2d, c_2d )

end program double_contractions_single_gemm

关键细节说明 #

• 内存连续性保证：通过reshape的order参数调整B的元素读取顺序，避免原张量非连续切片导致的BLAS性能损耗。
• 维度映射关系：
  • A_2D[a, k]对应原张量A[a,b,c]，其中k = (b-1)*nc + c（Fortran列优先索引）。
  • B_2D[k, m]对应原张量B[d,b,c,e]，其中m = (d-1)*ne + e。
  • 矩阵乘法后，C_2D[a, m]直接映射到C[a,d,e]。
• 性能优势：消除循环调用DGEMM的开销，充分利用BLAS库的高度优化实现，大幅提升大张量场景下的计算效率。

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