Fortran等维四维数组快速交换第二、第三维度的高效方法

场景说明 #

针对每维尺寸均为n的四维数组A(:,:,:,:)，需要完成第二维度与第三维度的转置交换。最初的实现采用内置reshape函数完成，对应代码为：
A2 = reshape(A, (/ n, n, n, n/), order = (/1,3,2,4/) )
由于操作仅涉及部分维度交换，目标是找到运行效率更高的等价实现。

性能测试代码 #

Program reshape_test

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

  real :: A(32,32,32,32), A2(32,32,32,32)
  integer :: i,j,k,l,n,m, m_iter
  Integer( li ) :: start, finish, rate
  
  
  m_iter = 3200
  n = 32
  do l = 1, n
    do k = 1, n
       do j = 1, n
         do i = 1, n
           A(i,j,k,l) = i + j * k**2 - l
         end do 
       end do   
    end do      
  end do        
 
  Call System_clock( start, rate )
  do m=1,m_iter
    A2 = reshape(A, (/n, n, n, n/), order = (/1,3,2,4/) ) 
  end do
  
 ! write (*,*) A(1,2,3,4), A(1,3,2,4), A2(1,2,3,4), A2(1,3,2,4)
  Call System_clock( finish, rate )  
  
  Write(*,*) 'Time for reshape-1', Real( finish - start, wp ) / rate  
  
  
  Call System_clock( start, rate )
  do m=1,m_iter

  do j = 1, n
    do i = 1, n  
      A2(:,i,j,:) =  A(:,j,i,:)
    end do      
  end do      
  end do
  
  !write (*,*) A(1,2,3,4), A(1,3,2,4), A2(1,2,3,4), A2(1,3,2,4)
  Call System_clock( finish, rate )  
  
  Write(*,*) 'Time for reshape-2', Real( finish - start, wp ) / rate  
    
    
  Call System_clock( start, rate )
  do m=1,m_iter 
  do l = 1, n
    do k = 1, n
       do j = 1, n
         do i = 1, n
           A2(i,k,j,l) = A(i,j,k,l)
         end do 
       end do   
    end do      
  end do   
  end do
  
 ! write (*,*) A(1,2,3,4), A(1,3,2,4), A2(1,2,3,4), A2(1,3,2,4)
  Call System_clock( finish, rate )      
  
  Write(*,*) 'Time for reshape-3', Real( finish - start, wp ) / rate  
  

end program reshape_test

测试结果 #

仅使用gfortran -O3参数编译 #

三轮独立运行的耗时结果（单位：秒）：

Time for reshape-1   13.500307800000000
Time for reshape-2   10.146418400000000
Time for reshape-3   20.489294800000000

Time for reshape-1   11.421597100000000
Time for reshape-2   9.3823936999999997
Time for reshape-3   19.856221900000001

Time for reshape-1   14.376207500000000
Time for reshape-2   10.756465400000000
Time for reshape-3   21.301044500000000

追加-march=native参数编译 #

三轮独立运行的耗时结果（单位：秒）：

Time for reshape-1   12.517529200000000
Time for reshape-2   10.240939200000000
Time for reshape-3   26.637998799999998

Time for reshape-1   12.436479700000000
Time for reshape-2   9.9803838000000002
Time for reshape-3   21.760061700000001

Time for reshape-1   11.419214300000000
Time for reshape-2   10.908747200000001
Time for reshape-3   21.445951399999998

测试结论 #

• 三类实现中，两层循环驱动连续内存切片赋值的方案（对应测试中的reshape-2）运行效率最高，相比内置reshape实现（reshape-1）提速约20%~30%
• 逐元素遍历的全嵌套循环实现（对应测试中的reshape-3）性能最差，耗时接近内置reshape实现的2倍
• 追加-march=native编译参数对整体性能提升不显著，三种实现的性能排序未发生变化

本次测试的计时代码规范参考了Stack Overflow平台关于Fortran高阶张量转置性能优化的相关讨论内容。

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