GPU卸载三角循环性能优化：移除Atomic指令解决性能瓶颈

我有一个用于计算力的三角循环，已通过OpenMP卸载移植到GPU，代码如下：

!$omp target teams distribute parallel do simd &
!$omp        private(i,j,ii,jj,rinvdenom,tp,index,fix0) &
!$omp        map(to:ap1,ap2,ap2_2,ap1_2,ap1p2,ap2p2,ri_2,gi3,mass) &
!$omp        map(tofrom:ugrav,ugrav_aux)

                do kk=0,nn1*(nn1-1)/2-1
                  ii=kk/nn1
                  jj=mod(kk,nn1)
                  if(jj.le.ii)then
                    ii=nn1-ii-2
                    jj=nn1-jj-1
                  endif
                  i=ii+1
                  j=jj+1
                  rinvdenom=1.d0/(ap1_2(j)+ri_2(i)+ap2_2(j)-ap2p2(i)*ap2(j))
                  tp=min(0.9995d0,ap1p2(i)*ap1(j)*rinvdenom)*10000.d0
                  index=1+int(tp)
                  fix0=gi3(index)*sqrt(rinvdenom)
                  ugrav(i)=ugrav(i)+mass(j)*fix0
                  !$omp atomic
                  ugrav_aux(j)=ugrav_aux(j)+mass(i)*fix0
                enddo

!$omp end target teams distribute parallel do simd

                do i=1,n1
                  ugrav(i)=-g/2./pi*(ugrav(i)+ugrav_aux(i))
                enddo

该代码能得到ugrav数组的正确结果，但依赖atomic指令，导致循环序列化，性能极低。我尝试对ugrav和ugrav_aux使用reduction指令，代码如下：

!$omp target teams distribute parallel do simd &
!$omp        private(i,j,ii,jj,rinvdenom,tp,index,fix0) reduction(+:ugrav_aux,ugrav)&
!$omp        map(to:ap1,ap2,ap2_2,ap1_2,ap1p2,ap2p2,ri_2,gi3,mass) &
!$omp        map(tofrom:ugrav,ugrav_aux)

但执行时出现错误：

libgomp: cuCtxSynchronize error: an illegal memory access was encountered

编译命令为：gfortran -fopenmp -foffload=nvptx-none <source>，请问如何移除atomic指令以提升计算速度？

1. 修正数组Reduction与内存映射的冲突 #

直接使用reduction(+:ugrav, ugrav_aux)时出错，核心原因是map(tofrom)与reduction的自动内存管理冲突——reduction会在GPU端创建私有副本并自动合并结果，无需手动指定tofrom。修改后的并行区域指令如下：

!$omp target teams distribute parallel do simd &
!$omp        private(i,j,ii,jj,rinvdenom,tp,index,fix0) &
!$omp        reduction(+:ugrav, ugrav_aux) &
!$omp        map(to:ap1,ap2,ap2_2,ap1_2,ap1p2,ap2p2,ri_2,gi3,mass, ugrav, ugrav_aux)

同时要确保进入并行区域前，ugrav和ugrav_aux已被初始化为0，因为reduction是在初始值基础上累加。

2. 手动分块归约（编译器Reduction支持不完善时的替代方案） #

如果gfortran版本对GPU数组reduction支持有限，可以通过减少原子操作次数来优化：

• 为每个线程/线程组创建私有累加数组，先将结果累加到私有数组
• 最后仅对每个数组元素执行一次原子合并操作，将私有结果合并到全局数组

示例代码框架：

!$omp target teams distribute parallel do simd &
!$omp        private(i,j,ii,jj,rinvdenom,tp,index,fix0) &
!$omp        private(ugrav_priv, ugrav_aux_priv) &
!$omp        map(to:ap1,ap2,ap2_2,ap1_2,ap1p2,ap2p2,ri_2,gi3,mass) &
!$omp        map(tofrom:ugrav, ugrav_aux)
    ! 初始化私有累加数组
    ugrav_priv = 0.d0
    ugrav_aux_priv = 0.d0

    do kk=0,nn1*(nn1-1)/2-1
      ! 原循环计算逻辑不变
      ii=kk/nn1
      jj=mod(kk,nn1)
      if(jj.le.ii)then
        ii=nn1-ii-2
        jj=nn1-jj-1
      endif
      i=ii+1
      j=jj+1
      rinvdenom=1.d0/(ap1_2(j)+ri_2(i)+ap2_2(j)-ap2p2(i)*ap2(j))
      tp=min(0.9995d0,ap1p2(i)*ap1(j)*rinvdenom)*10000.d0
      index=1+int(tp)
      fix0=gi3(index)*sqrt(rinvdenom)
      
      ! 累加到私有数组
      ugrav_priv(i) = ugrav_priv(i) + mass(j)*fix0
      ugrav_aux_priv(j) = ugrav_aux_priv(j) + mass(i)*fix0
    enddo

    ! 合并私有数组到全局数组，仅O(N)次原子操作
    do i=1,n1
      !$omp atomic
      ugrav(i) = ugrav(i) + ugrav_priv(i)
      !$omp atomic
      ugrav_aux(i) = ugrav_aux(i) + ugrav_aux_priv(i)
    enddo
!$omp end target teams distribute parallel do simd

这种方式将原子操作次数从O(N²)降到O(N)，性能会有显著提升。

3. 编译器优化建议 #

• 添加-O3优化级别，确保编译器生成高效的GPU代码
• 升级gfortran到12及以上版本，新版本对OpenMP GPU卸载和数组reduction的支持更完善

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