如何快速获得 cufftcomplex 幅度和相位

Question

我有一个 cufftcomplex 数据块，它是 cuda fft(R2C) 的结果。我知道数据保存为一个实数后跟图像编号的结构。现在我想通过快速的方式（不是循环）获得每个复杂元素的幅度=sqrt（R*R+I*I）和相位=arctan（I/R）。有什么好办法吗？或者任何图书馆都可以做到这一点？

Answer 1

由于cufftExecR2C 对 GPU 上的数据进行操作，因此结果已经在 GPU 上，（在您将它们复制回主机之前，如果您正在这样做。）

编写您自己的 cuda 内核来完成此操作应该很简单。您描述的幅度是cuCabs 或cuCabsf 在cuComplex.h 头文件中返回的值。通过查看该头文件中的函数，您应该能够弄清楚如何编写自己的函数来计算相位角。你会注意到cufftComplex 是just a typedef of cuComplex。

假设您的 cufftExecR2C 调用在大小为 sz 的数组 data 中留下了一些 cufftComplex 类型的结果。您的内核可能如下所示：

#include <math.h>
#include <cuComplex.h>
#include <cufft.h>
#define nTPB 256    // threads per block for kernel
#define sz 100000   // or whatever your output data size is from the FFT
...

__host__ __device__ float carg(const cuComplex& z) {return atan2(cuCimagf(z), cuCrealf(z));} // polar angle

__global__ void magphase(cufftComplex *data, float *mag, float *phase, int dsz){
  int idx = threadIdx.x + blockDim.x*blockIdx.x;
  if (idx < dsz){
    mag[idx]   = cuCabsf(data[idx]);
    phase[idx] = carg(data[idx]);
  }
}

...
int main(){
...
    /* Use the CUFFT plan to transform the signal in place. */
    /* Your code might be something like this already:      */
    if (cufftExecR2C(plan, (cufftReal*)data, data) != CUFFT_SUCCESS){
      fprintf(stderr, "CUFFT error: ExecR2C Forward failed");
      return;   
    }
    /* then you might add:                                  */
    float *h_mag, *h_phase, *d_mag, *d_phase;
    // malloc your h_ arrays using host malloc first, then...
    cudaMalloc((void **)&d_mag, sz*sizeof(float));
    cudaMalloc((void **)&d_phase, sz*sizeof(float));
    magphase<<<(sz+nTPB-1)/nTPB, nTPB>>>(data, d_mag, d_phase, sz);
    cudaMemcpy(h_mag, d_mag, sz*sizeof(float), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_phase, d_phase, sz*sizeof(float), cudaMemcpyDeviceToHost);

您也可以使用thrust 执行此操作，方法是为幅度和相位函数创建函子，并将这些函子与data、mag 和phase 一起传递给thrust::transform。

我相信您也可以使用CUBLAS 来实现，结合使用向量加法和向量乘法运算。

这个question/answer 可能也很有趣。我从那里提升了我的相位函数carg。