CUDA 中 3D 矩阵的列和行的 1D FFT

Question

我正在尝试使用 cufftPlanMany 计算批量 1D FFT。数据集来自存储在 1D 数组中的 3D 字段，我想在其中计算 x 和 y 方向的 1D FFT。数据存储如下图所示；连续在x 然后y 然后z。

在x 方向进行批量 FFT（我相信）是直截了当的；使用输入 stride=1、distance=nx 和 batch=ny * nz，它计算元素 {0,1,2,3}、{4,5,6,7}、...、{28,29,30,31} 上的 FFT。但是，我想不出一种方法来为y 方向的 FFT 实现相同的目标。每个xy 平面的批次再次简单明了（输入stride=nx、dist=1、batch=nx 导致在{0,4,8,12}、{1,5,9,13} 上的 FFT 等）。但是对于batch=nx * nz，从{3,7,11,15} 到{16,20,24,28}，距离大于1。这可以通过 cufftPlanMany 以某种方式完成吗？

Answer 1

我认为您的问题的简短回答（使用单个 cufftPlanMany 对 3D 矩阵的列执行 1D FFT 的可能性）是否。

确实，根据cufftPlanMany 执行的转换，你称之为

cufftPlanMany(&handle, rank, n, 
              inembed, istride, idist,
              onembed, ostride, odist, CUFFT_C2C, batch);

必须遵守Advanced Data Layout。具体来说，一维 FFT 是根据以下布局计算出来的

input[b * idist + x * istride]

其中b 处理b-th 信号，istride 是同一信号中两个连续项目之间的距离。如果 3D 矩阵的维度为M * N * Q，并且如果您想沿列执行 1D 变换，则两个连续元素之间的距离将为M，而两个连续信号之间的距离将为1。此外，批处理执行的数量必须设置为等于M。使用这些参数，您只能覆盖 3D 矩阵的一个切片。实际上，如果您尝试增加M，那么 cuFFT 将开始尝试从第二行开始计算新的按列 FFT。解决此问题的唯一方法是迭代调用 cufftExecC2C 以覆盖所有 Q 切片。

作为记录，以下代码提供了一个完整的示例，说明如何对 3D 矩阵的列执行 1D FFT。

#include <thrust/device_vector.h>
#include <cufft.h>

/********************/
/* CUDA ERROR CHECK */
/********************/
#define gpuErrchk(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
   if (code != cudaSuccess) 
   {
      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
      if (abort) exit(code);
   }
}

int main() {

    const int M = 3;
    const int N = 4;
    const int Q = 2;

    thrust::host_vector<float2> h_matrix(M * N * Q);

    for (int k=0; k<Q; k++) 
        for (int j=0; j<N; j++) 
            for (int i=0; i<M; i++) {
                float2 temp;
                temp.x = (float)(j + k * M); 
                //temp.x = 1.f; 
                temp.y = 0.f;
                h_matrix[k*M*N+j*M+i] = temp;
                printf("%i %i %i %f %f\n", i, j, k, temp.x, temp.y); 
            }
    printf("\n");

    thrust::device_vector<float2> d_matrix(h_matrix);

    thrust::device_vector<float2> d_matrix_out(M * N * Q);

    // --- Advanced data layout
    //     input[b * idist + x * istride]
    //     output[b * odist + x * ostride]
    //     b = signal number
    //     x = element of the b-th signal

    cufftHandle handle;
    int rank = 1;                           // --- 1D FFTs
    int n[] = { N };                        // --- Size of the Fourier transform
    int istride = M, ostride = M;           // --- Distance between two successive input/output elements
    int idist = 1, odist = 1;               // --- Distance between batches
    int inembed[] = { 0 };                  // --- Input size with pitch (ignored for 1D transforms)
    int onembed[] = { 0 };                  // --- Output size with pitch (ignored for 1D transforms)
    int batch = M;                          // --- Number of batched executions
    cufftPlanMany(&handle, rank, n, 
                  inembed, istride, idist,
                  onembed, ostride, odist, CUFFT_C2C, batch);

    for (int k=0; k<Q; k++)
        cufftExecC2C(handle, (cufftComplex*)(thrust::raw_pointer_cast(d_matrix.data()) + k * M * N), (cufftComplex*)(thrust::raw_pointer_cast(d_matrix_out.data()) + k * M * N), CUFFT_FORWARD);
    cufftDestroy(handle);

    for (int k=0; k<Q; k++) 
        for (int j=0; j<N; j++) 
            for (int i=0; i<M; i++) { 
                float2 temp = d_matrix_out[k*M*N+j*M+i];
                printf("%i %i %i %f %f\n", i, j, k, temp.x, temp.y); 
            }

}

当您想要对行执行一维变换时，情况会有所不同。在这种情况下，两个连续元素之间的距离为1，而两个连续信号之间的距离为M。这允许您设置多个N * Q 转换，然后只调用一次cufftExecC2C。作为记录，下面的代码提供了 3D 矩阵行的 1D 转换的完整示例。

#include <thrust/device_vector.h>
#include <cufft.h>

/********************/
/* CUDA ERROR CHECK */
/********************/
#define gpuErrchk(ans) { gpuAssert((ans), __FILE__, __LINE__); }
inline void gpuAssert(cudaError_t code, const char *file, int line, bool abort=true)
{
   if (code != cudaSuccess) 
   {
      fprintf(stderr,"GPUassert: %s %s %d\n", cudaGetErrorString(code), file, line);
      if (abort) exit(code);
   }
}

int main() {

    const int M = 3;
    const int N = 4;
    const int Q = 2;

    thrust::host_vector<float2> h_matrix(M * N * Q);

    for (int k=0; k<Q; k++) 
        for (int j=0; j<N; j++) 
            for (int i=0; i<M; i++) {
                float2 temp;
                temp.x = (float)(j + k * M); 
                //temp.x = 1.f; 
                temp.y = 0.f;
                h_matrix[k*M*N+j*M+i] = temp;
                printf("%i %i %i %f %f\n", i, j, k, temp.x, temp.y); 
            }
    printf("\n");

    thrust::device_vector<float2> d_matrix(h_matrix);

    thrust::device_vector<float2> d_matrix_out(M * N * Q);

    // --- Advanced data layout
    //     input[b * idist + x * istride]
    //     output[b * odist + x * ostride]
    //     b = signal number
    //     x = element of the b-th signal

    cufftHandle handle;
    int rank = 1;                           // --- 1D FFTs
    int n[] = { M };                        // --- Size of the Fourier transform
    int istride = 1, ostride = 1;           // --- Distance between two successive input/output elements
    int idist = M, odist = M;               // --- Distance between batches
    int inembed[] = { 0 };                  // --- Input size with pitch (ignored for 1D transforms)
    int onembed[] = { 0 };                  // --- Output size with pitch (ignored for 1D transforms)
    int batch = N * Q;                      // --- Number of batched executions
    cufftPlanMany(&handle, rank, n, 
                  inembed, istride, idist,
                  onembed, ostride, odist, CUFFT_C2C, batch);

    cufftExecC2C(handle, (cufftComplex*)(thrust::raw_pointer_cast(d_matrix.data())), (cufftComplex*)(thrust::raw_pointer_cast(d_matrix_out.data())), CUFFT_FORWARD);
    cufftDestroy(handle);

    for (int k=0; k<Q; k++) 
        for (int j=0; j<N; j++) 
            for (int i=0; i<M; i++) { 
                float2 temp = d_matrix_out[k*M*N+j*M+i];
                printf("%i %i %i %f %f\n", i, j, k, temp.x, temp.y); 
            }

}

Answer 2

我想， idist=nx*nz 也可以跳过整个平面，而 batch=nz 将覆盖一个 yx 平面。应根据 nx 或 nz 是否较大来决定。