FFTW MPI并行大尺寸三维FFT双进程通信错误问题问询

AI 小北

2026-10-31

问题

使用FFTW开展MPI并行的原地三维快速傅里叶变换计算，具体调用fftw_mpi_plan_dft_r2c_3d()接口，整体运行正常。近期尝试将大尺寸数据集与少量CPU核心/MPI进程结合开展计算时出现异常：使用2个进程时，全局尺寸为gridsize*gridsize*gridsize的三维FFT计算会触发错误，报错内容如下：

[s81n11:31843] *** An error occurred in MPI_Sendrecv
[s81n11:31843] *** reported by process [1582235649,0]
[s81n11:31843] *** on communicator MPI COMMUNICATOR 6 DUP FROM 5
[s81n11:31843] *** MPI_ERR_COUNT: invalid count argument
[s81n11:31843] *** MPI_ERRORS_ARE_FATAL (processes in this communicator will now abort,
[s81n11:31843] *** and potentially your MPI job)

该问题仅在gridsize >= 2048时出现，使用1个或3个进程时计算可正常运行。推测问题来源于MPI_Sendrecv()调用中的MPI消息存在大小限制：双进程场景下全局网格沿第一维做分布，最后一维存在实数据FFT要求的少量padding，单进程持有本地网格的物理元素数量为size = (gridsize/2)*gridsize*(2*(gridsize/2 + 1))；当gridsize = 2048时，该值约为1.001*2³²，而gridsize = 2047时对应值约为0.999*2³²，此时计算无异常。

问题的精确根因是什么？希望能明确具体是哪个整型类型（如unsigned int）的长度限制触发了该异常。
如果由我直接控制MPI调用，可通过分块多次通信的方式完成传输，但由于MPI调用由FFTW内部封装，无法直接修改，应如何解决该问题？如果需要对FFTW/MPI源码做轻量适配，我也愿意尝试修改。

此前应该已有其他用户遇到过该问题，但尚未检索到FFTW相关传输大小限制的公开资料，希望能获得相关主题的参考内容。已在64位Linux环境、FFTW 3.3.9版本、同时搭配MPICH 3.3.2和OpenMPI 4.0.3的配置下复现了该问题。

示例代码

下方为可复现问题的示例C代码，文件名为fft.c：

#include <assert.h>
#include <stdio.h>
#include <stdlib.h>
#include <fftw3-mpi.h>

int main(int argc, char* argv[]) {
  /* Set grid size */
  ptrdiff_t gridsize = 32;
  if (argc > 1)
    gridsize = atoi(argv[1]);
  ptrdiff_t gridsize_padding = 2*(gridsize/2 + 1);

  /* Initialize MPI */
  MPI_Init(NULL, NULL);
  int rank;
  MPI_Comm_rank(MPI_COMM_WORLD, &rank);
  fftw_mpi_init();

  /* Allocate grid */
  ptrdiff_t gridsize_local;
  ptrdiff_t gridsize_local_i, gridstart_local_i;
  ptrdiff_t gridsize_local_j, gridstart_local_j;
  double* grid = fftw_alloc_real(fftw_mpi_local_size_3d_transposed(
    gridsize,
    gridsize,
    gridsize_padding,
    MPI_COMM_WORLD,
    &gridsize_local_i,
    &gridstart_local_i,
    &gridsize_local_j,
    &gridstart_local_j
  ));
  assert(gridsize_local_i == gridsize_local_j);
  gridsize_local = gridsize_local_i;
  if (rank == 0)
    printf("gridsize %td\n", gridsize);
  printf("rank %d got gridsize_local = %td\n", rank, gridsize_local);

  /* Nullify grid */
  ptrdiff_t i, j, k;
  for (i = 0; i < gridsize_local; i++) {
    for (j = 0; j < gridsize; j++) {
      for (k = 0; k < gridsize_padding; k++) {
        grid[(i*gridsize + j)*gridsize_padding + k] = 0.0;
      }
    }
  }

  /* Carry out FFT */
  fftw_plan plan = fftw_mpi_plan_dft_r2c_3d(
      gridsize,
      gridsize,
      gridsize,
      grid,
      (fftw_complex*) grid,
      MPI_COMM_WORLD,
      FFTW_ESTIMATE | FFTW_MPI_TRANSPOSED_OUT
  );
  fftw_mpi_gather_wisdom(MPI_COMM_WORLD);
  fftw_mpi_broadcast_wisdom(MPI_COMM_WORLD);
  fftw_execute(plan);

  /* Clean up */
  fftw_free(grid);
  fftw_destroy_plan(plan);
  fftw_mpi_cleanup();
  MPI_Finalize();
  return 0;
}

假设已完成MPI与FFTW的安装且环境配置正确，可使用如下命令编译上述fft.c：

mpicc -I/path/to/fftw/include -I/path/to/mpi/include -c -o fft.o fft.c
mpicc fft.o -o fft -L/path/to/fftw/lib -Wl,-rpath=/path/to/fftw/lib -lfftw3_mpi -lfftw3 -L/path/to/mpi/lib -Wl,-rpath=/path/to/mpi/lib -lmpi

使用2个进程执行例如全局网格尺寸gridsize=64的三维FFT时，可运行如下命令：

mpiexec -n 2 ./fft 64

程序运行后不会有额外输出，因为代码未开展结果校验或打印。当gridsize=2048时，程序运行约1分钟后会触发前述崩溃。注意：gridsize=2048的运行场景将占用64GB内存，可能无法直接开展测试。

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

FFTW MPI并行大尺寸三维FFT双进程通信错误问题问询

示例代码 #

示例代码