为何多线程pyFFTW运行更慢？求排查方案与优化建议

问题背景 #

近期需要对二维矩阵（图像）执行并行FFT操作，选用pyFFTW库后发现多线程模式下运行速度反而更慢，调整矩阵尺寸与threads参数后仍无明显加速效果。测试代码如下：

def pyfftw64(inp):
    u0 = pyfftw.empty_aligned(inp.shape, dtype='float64')
    u1 = pyfftw.empty_aligned((inp.shape[0], inp.shape[-1]//2 + 1), dtype='complex128')
    
    fft_plan = pyfftw.FFTW(u0, u1, axes=(-2,-1), direction='FFTW_FORWARD', flags=('FFTW_MEASURE',), threads=4)
    ifft_plan = pyfftw.FFTW(u1, u0, axes=(-2,-1), direction='FFTW_BACKWARD', flags=('FFTW_MEASURE',), threads=4)
    
    u0[:] = inp

    start = time.time()
    for i in range(0,1000):
        fft_plan()
        ifft_plan()
    
    #fft_plan()
    #ifft_plan()

    end = time.time()
    
    return end-start

排查方法与优化建议 #

1. 拆分计划创建与执行的时间统计 #

当前代码把计划创建时间和FFT执行时间混在一起了，FFTW_MEASURE会花费大量时间去寻找最优执行路径，这部分时间可能掩盖了多线程的加速效果。修改代码单独统计两段时间：

def pyfftw64(inp):
    u0 = pyfftw.empty_aligned(inp.shape, dtype='float64')
    u1 = pyfftw.empty_aligned((inp.shape[0], inp.shape[-1]//2 + 1), dtype='complex128')
    
    # 单独统计计划创建耗时
    plan_start = time.time()
    fft_plan = pyfftw.FFTW(u0, u1, axes=(-2,-1), direction='FFTW_FORWARD', flags=('FFTW_MEASURE',), threads=4)
    ifft_plan = pyfftw.FFTW(u1, u0, axes=(-2,-1), direction='FFTW_BACKWARD', flags=('FFTW_MEASURE',), threads=4)
    plan_end = time.time()
    print(f"计划创建耗时: {plan_end - plan_start:.2f}s")
    
    u0[:] = inp

    start = time.time()
    for i in range(0,1000):
        fft_plan()
        ifft_plan()
    end = time.time()
    
    print(f"执行耗时: {end - start:.2f}s")
    return end-start

若计划创建时间占比过高，测试阶段可改用FFTW_ESTIMATE快速生成计划，生产环境则提前创建计划并复用。

2. 检查矩阵尺寸的并行友好性 #

FFTW对2的幂次、小质数乘积的尺寸优化更充分，多线程加速效果更明显。如果矩阵尺寸是质数或大质数乘积，多线程调度开销可能超过并行收益。可测试将矩阵补全到最近的2的幂次尺寸对比性能：

import numpy as np
# 补全到2的幂次
new_shape = (2**int(np.ceil(np.log2(inp.shape[0]))), 2**int(np.ceil(np.log2(inp.shape[1]))))
padded_inp = np.zeros(new_shape, dtype=np.float64)
padded_inp[:inp.shape[0], :inp.shape[1]] = inp

3. 控制线程数与CPU核心匹配 #

不要盲目设置threads=4，需与CPU物理核心数匹配。比如4核心8线程的CPU，设置threads=4（物理核心）可能比threads=8（超线程）效果更好——超线程在计算密集型任务中加速有限，反而会增加调度开销。可测试threads=1,2,4,8等不同值，找到最优解。

4. 消除额外内存与预热开销 #

• u0[:] = inp的内存拷贝有开销，若允许直接修改输入矩阵，可直接用输入矩阵创建对齐内存：

  u0 = pyfftw.byte_align(inp, dtype='float64')
  

• 首次执行计划会有缓存加载开销，建议先预热再统计时间：

  # 预热
  fft_plan()
  ifft_plan()
  # 正式统计
  start = time.time()
  for i in range(0,1000):
      fft_plan()
      ifft_plan()
  end = time.time()
  

5. 检查pyFFTW的并行编译配置 #

pyFFTW是否启用了OpenMP并行后端？可通过pyfftw.config.NUM_THREADS查看默认线程数，若未启用，可重新安装带并行支持的版本：

pip install pyfftw --no-binary pyfftw

确保编译时链接了FFTW的并行库（fftw3_threads）。

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