3

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

假设我有一个形状为 (2, 3, 4) 的 3D 数组。步幅为 (12, 4, 1)，这意味着为了沿最后一个轴移动一个单位，我们在平面阵列中移动 1 个单位，而要沿第一个轴移动一个单位，我们必须跨过3 * 4 = 12 个单位。（该数组是一个 numpy ndarray，当轴被转置时，它也可以有其他的步幅，但我很高兴得到一个答案，它只解决给定步幅的这个特定的 3D 案例）

现在假设我想沿中轴计算一维 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

参数说明对于一维 fft，批次 b 中位置 x 的元素将取自： input[b * idist + x * istride]

但是，位置 [b][x][z] 处的元素存储在：

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

因此尚不清楚如何使 CUFFT 在第三 (z) 轴上循环。

如果我设置：

  • idist 和 odist 为 3*4=12 （因此递增 b 使我们沿第一个轴移动）并且，
  • isstride 和 ostride 到 4（这样递增的 x 沿着第二个轴移动，这是我们想要 fft 的轴），
  • 批次 = 2
  • inembed 和 oneembed 为 3（但根据文档，这些对于 1D 变换被忽略）

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

这似乎是一个常见的用例，但我似乎无法弄清楚如何在不进行多次调用（这在 GPU 上很昂贵）或使用不同的内存布局进行复制的情况下使用给定的参数来做到这一点。

4

0 回答 0