Actual source code: cufft.cu
petsc-3.7.1 2016-05-15
Report Typos and Errors

      2: /*
      3:     Provides an interface to the CUFFT package.
      4:     Testing examples can be found in ~src/mat/examples/tests
      5: */

      7: #include <petsc/private/matimpl.h>          /*I "petscmat.h" I*/
      8: EXTERN_C_BEGIN
      9: #include <cuda.h>
     10: #include <cuda_runtime.h>
     11: #include <cufft.h>
     12: EXTERN_C_END

     14: typedef struct {
     15:   PetscInt     ndim;
     16:   PetscInt     *dim;
     17:   cufftHandle  p_forward, p_backward;
     18:   cufftComplex *devArray;
     19: } Mat_CUFFT;

     23: PetscErrorCode MatMult_SeqCUFFT(Mat A, Vec x, Vec y)
     24: {
     25:   Mat_CUFFT      *cufft    = (Mat_CUFFT*) A->data;
     26:   cufftComplex   *devArray = cufft->devArray;
     27:   PetscInt       ndim      = cufft->ndim, *dim = cufft->dim;
     28:   PetscScalar    *x_array, *y_array;
     29:   cufftResult    result;

     33:   VecGetArray(x, &x_array);
     34:   VecGetArray(y, &y_array);
     35:   if (!cufft->p_forward) {
     36:     cufftResult result;
     37:     /* create a plan, then execute it */
     38:     switch (ndim) {
     39:     case 1:
     40:       result = cufftPlan1d(&cufft->p_forward, dim[0], CUFFT_C2C, 1);CHKERRQ(result != CUFFT_SUCCESS);
     41:       break;
     42:     case 2:
     43:       result = cufftPlan2d(&cufft->p_forward, dim[0], dim[1], CUFFT_C2C);CHKERRQ(result != CUFFT_SUCCESS);
     44:       break;
     45:     case 3:
     46:       result = cufftPlan3d(&cufft->p_forward, dim[0], dim[1], dim[2], CUFFT_C2C);CHKERRQ(result != CUFFT_SUCCESS);
     47:       break;
     48:     default:
     49:       SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_USER, "Cannot create plan for %d-dimensional transform", ndim);
     50:     }
     51:   }
     52:   /* transfer to GPU memory */
     53:   cudaMemcpy(devArray, x_array, sizeof(cufftComplex)*dim[ndim], cudaMemcpyHostToDevice);
     54:   /* execute transform */
     55:   result = cufftExecC2C(cufft->p_forward, devArray, devArray, CUFFT_FORWARD);CHKERRQ(result != CUFFT_SUCCESS);
     56:   /* transfer from GPU memory */
     57:   cudaMemcpy(y_array, devArray, sizeof(cufftComplex)*dim[ndim], cudaMemcpyDeviceToHost);
     58:   VecRestoreArray(y, &y_array);
     59:   VecRestoreArray(x, &x_array);
     60:   return(0);
     61: }

     65: PetscErrorCode MatMultTranspose_SeqCUFFT(Mat A, Vec x, Vec y)
     66: {
     67:   Mat_CUFFT      *cufft    = (Mat_CUFFT*) A->data;
     68:   cufftComplex   *devArray = cufft->devArray;
     69:   PetscInt       ndim      = cufft->ndim, *dim = cufft->dim;
     70:   PetscScalar    *x_array, *y_array;
     71:   cufftResult    result;

     75:   VecGetArray(x, &x_array);
     76:   VecGetArray(y, &y_array);
     77:   if (!cufft->p_backward) {
     78:     /* create a plan, then execute it */
     79:     switch (ndim) {
     80:     case 1:
     81:       result = cufftPlan1d(&cufft->p_backward, dim[0], CUFFT_C2C, 1);CHKERRQ(result != CUFFT_SUCCESS);
     82:       break;
     83:     case 2:
     84:       result = cufftPlan2d(&cufft->p_backward, dim[0], dim[1], CUFFT_C2C);CHKERRQ(result != CUFFT_SUCCESS);
     85:       break;
     86:     case 3:
     87:       result = cufftPlan3d(&cufft->p_backward, dim[0], dim[1], dim[2], CUFFT_C2C);CHKERRQ(result != CUFFT_SUCCESS);
     88:       break;
     89:     default:
     90:       SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_USER, "Cannot create plan for %d-dimensional transform", ndim);
     91:     }
     92:   }
     93:   /* transfer to GPU memory */
     94:   cudaMemcpy(devArray, x_array, sizeof(cufftComplex)*dim[ndim], cudaMemcpyHostToDevice);
     95:   /* execute transform */
     96:   result = cufftExecC2C(cufft->p_forward, devArray, devArray, CUFFT_INVERSE);CHKERRQ(result != CUFFT_SUCCESS);
     97:   /* transfer from GPU memory */
     98:   cudaMemcpy(y_array, devArray, sizeof(cufftComplex)*dim[ndim], cudaMemcpyDeviceToHost);
     99:   VecRestoreArray(y, &y_array);
    100:   VecRestoreArray(x, &x_array);
    101:   return(0);
    102: }

    106: PetscErrorCode MatDestroy_SeqCUFFT(Mat A)
    107: {
    108:   Mat_CUFFT      *cufft = (Mat_CUFFT*) A->data;
    109:   cufftResult    result;

    113:   PetscFree(cufft->dim);
    114:   if (cufft->p_forward)  {result = cufftDestroy(cufft->p_forward);CHKERRQ(result != CUFFT_SUCCESS);}
    115:   if (cufft->p_backward) {result = cufftDestroy(cufft->p_backward);CHKERRQ(result != CUFFT_SUCCESS);}
    116:   cudaFree(cufft->devArray);
    117:   PetscFree(A->data);
    118:   PetscObjectChangeTypeName((PetscObject)A,0);
    119:   return(0);
    120: }

    124: /*@
    125:   MatCreateSeqCUFFT - Creates a matrix object that provides sequential FFT via the external package CUFFT

    127:   Collective on MPI_Comm

    129:   Input Parameters:
    130: + comm - MPI communicator, set to PETSC_COMM_SELF
    131: . ndim - the ndim-dimensional transform
    132: - dim  - array of size ndim, dim[i] contains the vector length in the i-dimension

    134:   Output Parameter:
    135: . A - the matrix

    137:   Options Database Keys:
    138: . -mat_cufft_plannerflags - set CUFFT planner flags

    140:   Level: intermediate
    141: @*/
    142: PetscErrorCode  MatCreateSeqCUFFT(MPI_Comm comm, PetscInt ndim, const PetscInt dim[], Mat *A)
    143: {
    144:   Mat_CUFFT      *cufft;
    145:   PetscInt       m, d;

    149:   if (ndim < 0) SETERRQ1(PETSC_COMM_SELF, PETSC_ERR_USER, "ndim %d must be > 0", ndim);
    150:   MatCreate(comm, A);
    151:   m    = 1;
    152:   for (d = 0; d < ndim; ++d) {
    153:     if (dim[d] < 0) SETERRQ2(PETSC_COMM_SELF, PETSC_ERR_USER, "dim[%d]=%d must be > 0", d, dim[d]);
    154:     m *= dim[d];
    155:   }
    156:   MatSetSizes(*A, m, m, m, m);
    157:   PetscObjectChangeTypeName((PetscObject)*A, MATSEQCUFFT);

    159:   PetscNewLog(*A,&cufft);
    160:   (*A)->data = (void*) cufft;
    161:   PetscMalloc1(ndim+1, &cufft->dim);
    162:   PetscMemcpy(cufft->dim, dim, ndim*sizeof(PetscInt));

    164:   cufft->ndim       = ndim;
    165:   cufft->p_forward  = 0;
    166:   cufft->p_backward = 0;
    167:   cufft->dim[ndim]  = m;

    169:   /* GPU memory allocation */
    170:   cudaMalloc((void**) &cufft->devArray, sizeof(cufftComplex)*m);

    172:   (*A)->ops->mult          = MatMult_SeqCUFFT;
    173:   (*A)->ops->multtranspose = MatMultTranspose_SeqCUFFT;
    174:   (*A)->assembled          = PETSC_TRUE;
    175:   (*A)->ops->destroy       = MatDestroy_SeqCUFFT;

    177:   /* get runtime options */
    178:   PetscOptionsBegin(comm, ((PetscObject)(*A))->prefix, "CUFFT Options", "Mat");
    179:   PetscOptionsEnd();
    180:   return(0);
    181: }