如何用cuFFT对2D图像分块（含重叠）执行批量2D FFT？

将尺寸为[32,32]的2D图像划分为[8,8]的非重叠子块并执行批量2D FFT，无法仅通过cufftPlanMany()的参数配置实现，必须自定义CUDA核完成格式重排，原因如下：

常规格式转换逻辑 #

要运行标准批量2D FFT，需将原始[32,32]图像重塑为[8,8,16]格式（16为4×4的子块总数），核心是完成索引置换：先把图像拆分为[8,4,8,4]的四维结构，再交换中间两个维度得到[8,8,4×4]的批量格式。以下是实现该转换的CUDA核代码：

// Output tile coordinates
int x = threadIdx.x;
int y = threadIdx.y;

// Output tile size, e.g. [8,8]
int KX = blockDim.x;
int KY = blockDim.y;

// Full dataset size, e.g. [32,32]
int NX = blockDim.x * gridDim.x;
int NY = blockDim.y * gridDim.y;

// Batch ID
int batch = blockIdx.x + blockIdx.y * gridDim.x;

// Input data coordinates
int nx = threadIdx.x + blockIdx.x * blockDim.x;
int ny = threadIdx.y + blockIdx.y * blockDim.y;

// Copy-paste to a batched format, e.g. [8,8,4*4=16]
output[x + y*KX + batch*KX*KY] = input[nx + ny*NX];

cuFFT的数据布局限制 #

cuFFT通过以下线性索引公式访问输入输出数据：

input[ batch * idist + (y * inembed[1] + x) * istride ]
output[ batch * odist + (y * onembed[1] + x) * ostride ]

该公式仅支持线性的步长与偏移变换，无法实现上述格式转换中的维度置换逻辑（即交换[8,4,8,4]中的中间两个维度）。尝试映射参数时会发现始终存在无法抵消的索引偏移项，因此必须先通过自定义核完成格式重排，再调用cuFFT执行批量FFT。

类似1D场景中用cuFFT实现重叠批量FFT的需求，2D场景可通过cufftPlanMany()配合参数配置直接实现，无需自定义核，核心步骤如下：

1. 定义核心参数：

   • 子块尺寸：示例值Kx=8, Ky=8
   • 重叠量：设x方向重叠Sx，y方向重叠Sy，则子块在原始图像中的步长为step_x = Kx - Sx，step_y = Ky - Sy
   • 原始图像尺寸：NX=32, NY=32
   • 批量数：num_batches = ((NX - Kx) // step_x + 1) * ((NY - Ky) // step_y + 1)

2. 配置cuFFT关键参数：

   • inembed：设置为原始图像尺寸[NX, NY]，指定输入数据的逻辑形状
   • istride：设为1，按行连续访问子块内的元素
   • idist：设为step_x + step_y * NX，表示相邻批量子块在原始输入中的起始位置间隔
   • 基础参数：rank=2（2D FFT），n=[Kx, Ky]（子块尺寸），batch=num_batches

3. 边界处理：
   若(NX - Kx)无法被step_x整除，或(NY - Ky)无法被step_y整除，需对原始图像进行边界填充（如补零），或截断超出范围的子块，确保所有子块都能被正确覆盖。

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