如何在Fortran OpenMP代码中实现指定层级的线程分配

原代码 #

program main

USE OMP_LIB
IMPLICIT NONE

 INTEGER :: i,j,mm,m1,a,k,a1,b,bvv
 logical :: nn

call omp_set_nested(.true.)
call omp_set_max_active_levels(2)
CALL omp_set_num_threads(2)

!$omp parallel do private(j,mm,i,bvv) &
!$omp reduction (+:b)

   DO j=1, 2
   mm=omp_get_thread_num()

!$omp parallel do shared(mm,j) private(i,bvv) 

   DO i = 1, 2
    b=b+1
    bvv=omp_get_thread_num()
   WRITE(*,'(A,I2,2x,A,I2,I2,I2)')'DO1_THREAD',mm, 'DO2_Thread ',bvv  ,i,j
   END DO

!$omp end parallel do

  enddo

!$omp end parallel do 

print*,b
end program main

需求说明 #

当前代码已设置CALL omp_set_num_threads(2)，希望实现以下线程分配结构：

0------------------- 1
            |                     |
            |                     |
       ----------            ------------
     |            |          |           |
     |            |          |           |                                                        
     2            3          4           5 

即外层并行区域使用线程0、1，每个外层线程分别衍生出2个子线程（线程2、3和线程4、5）。

修改方案 #

关键调整点 #

1. 初始化未定义变量：原代码中b未初始化，需在声明时添加初始值b=0，避免计算结果异常。
2. 强制内层并行区域线程数：在内层!$omp parallel do指令中添加num_threads(2)子句，确保每个外层线程衍生的内层并行区域固定启动2个线程。外层线程数已通过omp_set_num_threads(2)设置为2，配合已开启的嵌套并行（omp_set_nested(.true.)），即可实现目标线程层级。

修改后的完整代码 #

program main

USE OMP_LIB
IMPLICIT NONE

 INTEGER :: i,j,mm,m1,a,k,a1,b=0,bvv  ! 初始化b为0
 logical :: nn

call omp_set_nested(.true.)
call omp_set_max_active_levels(2)
CALL omp_set_num_threads(2)

!$omp parallel do private(j,mm,i,bvv) &
!$omp reduction (+:b)

   DO j=1, 2
   mm=omp_get_thread_num()

!$omp parallel do shared(mm,j) private(i,bvv) num_threads(2)  ! 添加num_threads(2)子句

   DO i = 1, 2
    b=b+1
    bvv=omp_get_thread_num()
   WRITE(*,'(A,I2,2x,A,I2,I2,I2)')'DO1_THREAD',mm, 'DO2_Thread ',bvv  ,i,j
   END DO

!$omp end parallel do

  enddo

!$omp end parallel do 

print*,b
end program main

效果说明 #

修改后，外层并行do会启动线程0和1；每个外层线程进入内层并行do时，会各自启动2个新线程（线程2、3由线程0衍生，线程4、5由线程1衍生），完全匹配需求的线程结构。同时初始化b后，最终输出的b值会稳定为4（2个外层线程 × 2次内层循环）。

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