fft2_cuda.c
上传用户:chinakey
上传日期:2019-06-01
资源大小:2k
文件大小:6k
源码类别:

并行计算

开发平台:

C/C++

  1. /* 
  2.         This is a mex file to offload  the FFT2 function in Matlab to CUDA:
  3. Y = fft2(X)
  4. or 
  5.         Y = fft2(X,M,N)
  6. Y = fft2(X) returns the two-dimensional FFT of X. The result Y is the same size as X.
  7. Y = fft2(X,M,N) truncates X, or pads X with zeros to create an M-by-N array before doing the transform. 
  8. The result is M-by-N.
  9.         On Windows:
  10. Setup the mex build from a Matlab shell:
  11. mex -setup
  12.         Compile the mex file:
  13. mex -IC:CUDAinclude fft2_cuda.c -LC:cudalib -lcufft -lcudart
  14.         On Linux:
  15.         Use the makefile
  16.  */
  17. #include "mex.h"
  18. #include "cufft.h"
  19. #include "cuda_runtime.h"
  20. /**************************************************************************/
  21. /* MATLAB stores complex numbers in separate arrays for the real and
  22.    imaginary parts.  The following functions take the data in
  23.    this format and pack it into a complex work array, or
  24.    unpack it, respectively.  
  25.    We are using cufftComplex defined in cufft.h to  handle complex on Windows and Linux
  26. */
  27. void pack_r2c(cufftComplex  *output_float, 
  28.               double *input_re, 
  29.               int Ntot)
  30. {
  31.     int i;
  32.     for (i = 0; i < Ntot; i++) 
  33.       {
  34.                output_float[i].x = input_re[i];
  35.                output_float[i].y = 0.0f;
  36.       }
  37. }
  38. void pack_c2c(cufftComplex  *output_float, 
  39.               double *input_re, 
  40.               double *input_im, 
  41.               int Ntot)
  42. {
  43.     int i;
  44.     for (i = 0; i < Ntot; i++) 
  45.      {
  46.                output_float[i].x = input_re[i];
  47.                output_float[i].y = input_im[i];
  48.      }
  49. }
  50. void pack_r2c_expand(cufftComplex  *output_float, 
  51.                      double *input_re, 
  52.                      int originalM, 
  53.                      int originalN, 
  54.                      int M, 
  55.                      int N)
  56. {
  57.      int i, j;
  58.     for (i = 0; i < originalM; i++) 
  59.      for (j = 0; j < originalN; j++) 
  60.      {
  61.                output_float[i+M*j].x = input_re[i+originalM*j];
  62.                output_float[i+M*j].y = 0.0f;
  63.      }
  64. }
  65. void pack_c2c_expand(cufftComplex  *output_float, 
  66.                      double *input_re, 
  67.                      double *input_im, 
  68.                      int originalM, 
  69.                      int originalN, 
  70.                      int M, 
  71.                      int N)
  72. {
  73.      int i, j;
  74.     for (i = 0; i < originalM; i++) 
  75.      for (j = 0; j < originalN; j++) 
  76.      {
  77.                output_float[i+M*j].x = input_re[i+originalM*j];
  78.                output_float[i+M*j].y = input_im[i+originalM*j];
  79.      }
  80. }
  81. void unpack_c2c(cufftComplex  *input_float, 
  82.                 double *output_re, 
  83.                 double *output_im,  
  84.                 int Ntot)
  85. {
  86.     int i;
  87.     for (i = 0; i < Ntot; i++) 
  88.     {
  89.                output_re[i] = input_float[i].x;
  90.                output_im[i] = input_float[i].y;
  91.     }
  92. }
  93. /**************************************************************************/
  94. void mexFunction( int nlhs, mxArray *plhs[],
  95.                   int nrhs, const mxArray *prhs[])
  96. {
  97.   int originalM, originalN, M, N;
  98.   double *ar,*ai;
  99.   cufftComplex *input_single ;
  100.   cufftHandle            plan;
  101.   cufftComplex *rhs_complex_d;
  102.   /* 
  103.      Find out the  dimension of the array we want to transform:
  104.      prhs(originalM,originalN) 
  105.      originalM = Number of rows    in the mxArray prhs 
  106.      originalN = Number of columns in the mxArray prhs 
  107.   */
  108.     originalM = mxGetM(prhs[0]);
  109.     originalN = mxGetN(prhs[0]);
  110.     M = originalM;
  111.     N = originalN;
  112.   /*
  113.     Find out if the result is of the same size of the input.
  114.     plhs(M,N)
  115.     M = The number of rows in the mxArray plhs
  116.     N = The number of columns in the mxArray plhs
  117.   */
  118.    if (nrhs == 3) 
  119.    {
  120.     M = mxGetScalar(prhs[1]);
  121.     N = mxGetScalar(prhs[2]);
  122.    }
  123.    
  124.   /* 
  125.     Find out if the input array was real or complex.
  126.     Matlab is passing two separate pointers for the real and imaginary part.
  127.     The current version of  CUDAFFT is expecting interleaved data.
  128.     We will need to pack and unpack the data.
  129.     The current version of CUDA supports only single precision, 
  130.     so the original double precision data need to be converted.
  131.   */
  132.   /* Allocating working array on host */
  133.     if(nrhs == 1) input_single  = (cufftComplex*) mxMalloc(sizeof(cufftComplex)*N*M);
  134.     if(nrhs == 3) input_single  = (cufftComplex*) mxCalloc(N*M,sizeof(cufftComplex));
  135.   /* Pointer for the real part of the input */
  136.     ar =  (double *) mxGetData(prhs[0]);
  137.   if(mxIsComplex(prhs[0])) 
  138.   {
  139.    /* The input array is complex */
  140.    ai =  (double *) mxGetImagData(prhs[0]); 
  141.    
  142.    /* Input and output have same dimensions */
  143.    if(nrhs ==1) pack_c2c(input_single, ar, ai, N*M); 
  144.    /* Input and output have different dimensions */
  145.    if(nrhs ==3) pack_c2c_expand(input_single, ar, ai, originalM, originalN, M, N); 
  146.   }
  147.   else
  148.   {
  149.    /* The input array is real */
  150.    /* Input and output have same dimensions */
  151.    if(nrhs ==1) pack_r2c(input_single, ar, N*M); 
  152.    /* Input and output have different dimensions */
  153.    if(nrhs ==3) pack_r2c_expand(input_single, ar, originalM, originalN, M, N); 
  154.   }
  155.  
  156.   /* Allocate array on device */
  157.   cudaMalloc( (void **) &rhs_complex_d,sizeof(cufftComplex)*N*M);
  158.   /* Copy input array in interleaved format to the device */
  159.   cudaMemcpy( rhs_complex_d, input_single, sizeof(cufftComplex)*N*M, cudaMemcpyHostToDevice);
  160.   /* Create plan for CUDA FFT */
  161.   cufftPlan2d(&plan, N, M, CUFFT_C2C) ;
  162.   /* Execute FFT on device */
  163.     cufftExecC2C(plan, rhs_complex_d, rhs_complex_d, CUFFT_FORWARD) ;
  164.   /* Destroy plan */
  165.     cufftDestroy(plan);
  166.   /* Copy result back to host */
  167.   cudaMemcpy( input_single, rhs_complex_d, sizeof(cufftComplex)*N*M, cudaMemcpyDeviceToHost);
  168.   plhs[0]=mxCreateDoubleMatrix(M,N,mxCOMPLEX);
  169.   ar = mxGetPr(plhs[0]); 
  170.   ai = mxGetPi(plhs[0]);
  171.   unpack_c2c(input_single, ar, ai, N*M); 
  172.   mxFree(input_single);
  173.   cudaFree(rhs_complex_d);
  174.   return;
  175. }