Actual source code: cufft.cu
1: /*
2: Provides an interface to the CUFFT package.
3: Testing examples can be found in ~src/mat/tests
4: */
6: #include <petscdevice_cuda.h>
7: #include <petsc/private/matimpl.h>
9: typedef struct {
10: PetscInt ndim;
11: PetscInt *dim;
12: cufftHandle p_forward, p_backward;
13: cufftComplex *devArray;
14: } Mat_CUFFT;
16: static PetscErrorCode MatMult_SeqCUFFT(Mat A, Vec x, Vec y)
17: {
18: Mat_CUFFT *cufft = (Mat_CUFFT *)A->data;
19: cufftComplex *devArray = cufft->devArray;
20: PetscInt ndim = cufft->ndim, *dim = cufft->dim;
21: PetscScalar *x_array, *y_array;
23: PetscFunctionBegin;
24: PetscCall(VecGetArray(x, &x_array));
25: PetscCall(VecGetArray(y, &y_array));
26: if (!cufft->p_forward) {
27: /* create a plan, then execute it */
28: switch (ndim) {
29: case 1:
30: PetscCallCUFFT(cufftPlan1d(&cufft->p_forward, dim[0], CUFFT_C2C, 1));
31: break;
32: case 2:
33: PetscCallCUFFT(cufftPlan2d(&cufft->p_forward, dim[0], dim[1], CUFFT_C2C));
34: break;
35: case 3:
36: PetscCallCUFFT(cufftPlan3d(&cufft->p_forward, dim[0], dim[1], dim[2], CUFFT_C2C));
37: break;
38: default:
39: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_USER, "Cannot create plan for %" PetscInt_FMT "-dimensional transform", ndim);
40: }
41: }
42: /* transfer to GPU memory */
43: PetscCallCUDA(cudaMemcpy(devArray, x_array, sizeof(cufftComplex) * dim[ndim], cudaMemcpyHostToDevice));
44: /* execute transform */
45: PetscCallCUFFT(cufftExecC2C(cufft->p_forward, devArray, devArray, CUFFT_FORWARD));
46: /* transfer from GPU memory */
47: PetscCallCUDA(cudaMemcpy(y_array, devArray, sizeof(cufftComplex) * dim[ndim], cudaMemcpyDeviceToHost));
48: PetscCall(VecRestoreArray(y, &y_array));
49: PetscCall(VecRestoreArray(x, &x_array));
50: PetscFunctionReturn(PETSC_SUCCESS);
51: }
53: static PetscErrorCode MatMultTranspose_SeqCUFFT(Mat A, Vec x, Vec y)
54: {
55: Mat_CUFFT *cufft = (Mat_CUFFT *)A->data;
56: cufftComplex *devArray = cufft->devArray;
57: PetscInt ndim = cufft->ndim, *dim = cufft->dim;
58: PetscScalar *x_array, *y_array;
60: PetscFunctionBegin;
61: PetscCall(VecGetArray(x, &x_array));
62: PetscCall(VecGetArray(y, &y_array));
63: if (!cufft->p_backward) {
64: /* create a plan, then execute it */
65: switch (ndim) {
66: case 1:
67: PetscCallCUFFT(cufftPlan1d(&cufft->p_backward, dim[0], CUFFT_C2C, 1));
68: break;
69: case 2:
70: PetscCallCUFFT(cufftPlan2d(&cufft->p_backward, dim[0], dim[1], CUFFT_C2C));
71: break;
72: case 3:
73: PetscCallCUFFT(cufftPlan3d(&cufft->p_backward, dim[0], dim[1], dim[2], CUFFT_C2C));
74: break;
75: default:
76: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_USER, "Cannot create plan for %" PetscInt_FMT "-dimensional transform", ndim);
77: }
78: }
79: /* transfer to GPU memory */
80: PetscCallCUDA(cudaMemcpy(devArray, x_array, sizeof(cufftComplex) * dim[ndim], cudaMemcpyHostToDevice));
81: /* execute transform */
82: PetscCallCUFFT(cufftExecC2C(cufft->p_forward, devArray, devArray, CUFFT_INVERSE));
83: /* transfer from GPU memory */
84: PetscCallCUDA(cudaMemcpy(y_array, devArray, sizeof(cufftComplex) * dim[ndim], cudaMemcpyDeviceToHost));
85: PetscCall(VecRestoreArray(y, &y_array));
86: PetscCall(VecRestoreArray(x, &x_array));
87: PetscFunctionReturn(PETSC_SUCCESS);
88: }
90: static PetscErrorCode MatDestroy_SeqCUFFT(Mat A)
91: {
92: Mat_CUFFT *cufft = (Mat_CUFFT *)A->data;
94: PetscFunctionBegin;
95: PetscCall(PetscFree(cufft->dim));
96: if (cufft->p_forward) PetscCallCUFFT(cufftDestroy(cufft->p_forward));
97: if (cufft->p_backward) PetscCallCUFFT(cufftDestroy(cufft->p_backward));
98: PetscCallCUDA(cudaFree(cufft->devArray));
99: PetscCall(PetscFree(A->data));
100: PetscCall(PetscObjectChangeTypeName((PetscObject)A, 0));
101: PetscFunctionReturn(PETSC_SUCCESS);
102: }
104: /*@
105: MatCreateSeqCUFFT - Creates a matrix object that provides `MATSEQCUFFT` via the NVIDIA package CuFFT
107: Collective
109: Input Parameters:
110: + comm - MPI communicator, set to `PETSC_COMM_SELF`
111: . ndim - the ndim-dimensional transform
112: - dim - array of size `ndim`, dim[i] contains the vector length in the i-dimension
114: Output Parameter:
115: . A - the matrix
117: Options Database Key:
118: . -mat_cufft_plannerflags - set CuFFT planner flags
120: Level: intermediate
122: .seealso: [](ch_matrices), `Mat`, `MATSEQCUFFT`
123: @*/
124: PetscErrorCode MatCreateSeqCUFFT(MPI_Comm comm, PetscInt ndim, const PetscInt dim[], Mat *A)
125: {
126: Mat_CUFFT *cufft;
127: PetscInt m = 1;
129: PetscFunctionBegin;
130: PetscCheck(ndim >= 0, PETSC_COMM_SELF, PETSC_ERR_USER, "ndim %" PetscInt_FMT " must be > 0", ndim);
131: if (ndim) PetscAssertPointer(dim, 3);
132: PetscAssertPointer(A, 4);
133: PetscCall(MatCreate(comm, A));
134: for (PetscInt d = 0; d < ndim; ++d) {
135: PetscCheck(dim[d] >= 0, PETSC_COMM_SELF, PETSC_ERR_USER, "dim[%" PetscInt_FMT "]=%" PetscInt_FMT " must be > 0", d, dim[d]);
136: m *= dim[d];
137: }
138: PetscCall(MatSetSizes(*A, m, m, m, m));
139: PetscCall(PetscObjectChangeTypeName((PetscObject)*A, MATSEQCUFFT));
141: PetscCall(PetscNew(&cufft));
142: (*A)->data = (void *)cufft;
143: PetscCall(PetscMalloc1(ndim + 1, &cufft->dim));
144: PetscCall(PetscArraycpy(cufft->dim, dim, ndim));
146: cufft->ndim = ndim;
147: cufft->p_forward = 0;
148: cufft->p_backward = 0;
149: cufft->dim[ndim] = m;
151: /* GPU memory allocation */
152: PetscCallCUDA(cudaMalloc((void **)&cufft->devArray, sizeof(cufftComplex) * m));
154: (*A)->ops->mult = MatMult_SeqCUFFT;
155: (*A)->ops->multtranspose = MatMultTranspose_SeqCUFFT;
156: (*A)->assembled = PETSC_TRUE;
157: (*A)->ops->destroy = MatDestroy_SeqCUFFT;
159: /* get runtime options ...what options????? */
160: PetscOptionsBegin(comm, ((PetscObject)*A)->prefix, "CUFFT Options", "Mat");
161: PetscOptionsEnd();
162: PetscFunctionReturn(PETSC_SUCCESS);
163: }