如何在 OpenMP Fortran 代码中使用多线程 FFTW

Question

我有一个工作正常的代码，它是用 gcc/gfortran-14 编译的（通过brew 安装）。绝大多数时间都花在通过 FFTW 进行 FFT 上。在代码的关键部分我有这个：

!$omp parallel
!$omp sections
!$omp section
       call fft(xi3d,xo3d,.false.,xip,layered_p)
!$omp section
       call fft(yi3d,yo3d,.false.,yip,layered_p)
!$omp section
       call fft(zi3d,zo3d,.false.,zip,layered_p)
!$omp end sections
!$omp sections
!$omp section
       x3d=n3d(:,:,:,1,1)*xo3d+n3d(:,:,:,2,1)*yo3d+n3d(:,:,:,3,1)*zo3d
!$omp section
       y3d=n3d(:,:,:,4,1)*xo3d+n3d(:,:,:,5,1)*yo3d+n3d(:,:,:,6,1)*zo3d
!$omp section
       z3d=n3d(:,:,:,7,1)*xo3d+n3d(:,:,:,8,1)*yo3d+n3d(:,:,:,9,1)*zo3d
!$omp end sections
!$omp sections
!$omp section
       call fft(rx3d,x3d,.true.,xp,layered_p)
!$omp section
       call fft(ry3d,y3d,.true.,yp,layered_p)
!$omp section
       call fft(rz3d,z3d,.true.,zp,layered_p)
!$omp end sections
!$omp end parallel

其中

x/y/zi3d

和

rx/y/z3d

是真正的 3d 数组，所有其他都是复杂的 3d 数组（全部使用

fftw_alloc_real/complex

创建）。

fft

子例程仅根据第三个参数（即

fftw_plan_dft_c2r_3d

或

fftw_plan_dft_r2c_3d

）调用

.true.

或

.false.

。

我通常使用三个线程运行。所以 FFT 是并行运行的。这非常有效，代码的性能确实提高了近 3 倍。

现在我尝试将 OpenMP 与 FFTW 本身结合使用。也就是说，告诉 FFTW 在每个 fft 中使用 2 或 3 个线程，而不是一个。因此，在代码开始时，我进行了所需的初始化：

     if (fftw_init_threads().eq.0) then
       print *,'fftw has problem: fftw_init_threads'
     else
       call fftw_plan_with_nthreads(max(1,omp_get_max_threads()/3))
       print *,'each fftw will use ',max(1,omp_get_max_threads()/3),' threads'
   !   call omp_set_nested(.true.)
     endif

但是，当我运行代码时（即使设置

export OMP_NUM_THREADS=6

或

export OMP_NUM_THREADS=3,2

后，代码运行速度也会变慢。设置为 6，

top

告诉我只有 3 个处理器正在运行。其他设置显示为 4.5 个处理器正在运行，但是在这两种情况下，代码的运行速度都比仅使用 3 个线程且不使用 FFTW OpenMP 慢。

我使用的数组的大小是512x512x256。我使用的是配备 M1Pro 和 32GB 内存的 MacBookPro。

我是否需要

omp_set_nested(.true.)

？我该怎么做才能在执行 FFTW 时看到所有 6 个处理器都已连接（在此代码中，超过 95% 的时间都是如此！）？

关于尝试什么的任何建议。

Answer 1

从 OpenMP 的角度来看，您希望在外层执行两个具有 3 个线程的嵌套并行区域，然后在内层执行 n 个线程。从 OpenMP 5.0 开始，不推荐使用

nested-var

及其 API，现在可以通过设置

max-active-levels-var

或通过

nthreads-var

列表中的元素（在您的情况下为

3,n

）进行控制。

omp_get_max_threads()

定义为提供

nthreads-var

列表的头值。因此，从串行上下文来看，它将为您提供 3。

因此，在您的代码中，您要么告诉 fftw 计划使用 2 个线程执行 (

OMP_NUM_THREADS=6

)，但然后使用单个线程执行，或者告诉 fftw 计划使用 1 个线程执行 (

OMP_NUM_THREADS=3,2

)，然后执行有 2 个线程。

要正确规划 fftw 执行，棘手的部分是从列表中访问

n

。您可以生成一个并行区域来模拟 3 部分区域，并从

nthreads-var

列表中弹出 3 个区域。此时，您可以使用

omp_get_max_threads()

访问第二个值。

     if (fftw_init_threads().eq.0) then
       print *,'fftw has problem: fftw_init_threads'
     else
!$omp parallel
!$omp single
       call fftw_plan_with_nthreads(max(1,omp_get_max_threads()))
       print *,'each fftw will use ',max(1,omp_get_max_threads()),' threads'
!$omp end single
!$omp end parallel
     endif

GCC 至少从版本 9 开始就实现了基于

nthreads-var

列表的嵌套，所以你不需要另外调用

omp_set_nested