在多维数组的给定轴上的FFTW / CUFFT

Question

是否有一种有效的方法可以使用FFTW / CUFFT（它们具有类似的API）在多维数组的给定轴上执行fft？

我们说我有一个3D阵列形状（2,3,4）。步幅是（12,4,1），这意味着为了沿着最后一个轴移动一个单位，我们在平面阵列中移动1个单位，而沿着第一个轴移动一个单位，我们必须跨过3 * 4 = 12个单位。 （数组是一个numpy ndarray，当转换轴时，它也可以有其他步幅，但我对满足这个特定3D情况的答案感到满意，并且给定了步幅）

现在让我们说我想沿着中轴计算 1D fft 。 CUFFT公开以下功能：

cufftResult cufftPlanMany(
    cufftHandle *plan,        // Plan to be initialized
    int rank,                 // Rank = 1 for 1D fft
    int *n,                   // shape of the fft = 3
    int *inembed,
    int istride,
    int idist,
    int *onembed,
    int ostride,
    int odist,
    cufftType type,           // e.g. 64 bit float to 128 bit complex
    int batch                 // Could use batch = 2 for the first axis
);

我认为我们需要nembed，stride，dist参数来进行转换。它们记录在这里：  http://docs.nvidia.com/cuda/cufft/index.html#advanced-data-layout

dumentation指出，对于1D fft，位置x处的批次b中的元素将取自： input[b * idist + x * istride]

然而，位置[b] [x] [z]的元素存储在：

input[b * 12 + x * 4 + z]

因此，不清楚如何在第三（z）轴上进行CUFFT循环。

如果我设置：

• idist and odist to 3 * 4 = 12（所以递增b使我们沿第一轴移动），
• istride and ostride to 4（这样增量x沿第二轴移动，这是我们想要fft的轴），
• batch = 2
• 不合标准，并且标记为3（但根据文档，这些将被忽略1D转换）

然后它为最后一个轴索引为0的2个批次中的每个批次计算正确的fft，但保留最后一个索引为1,2或3的子数组不受影响。

这似乎是一个常见的用例，但我似乎无法弄清楚如何使用给定的参数执行此操作而无需进行多次调用（这在GPU上很昂贵）或制作具有不同内存的副本布局。