xref: /petsc/src/mat/impls/aij/mpi/mpicusparse/mpicusparsematimpl.h (revision d6685f554fbda8d96c6a5d73ab0e7a4e21a05c51)
1 #if !defined(__MPICUSPARSEMATIMPL)
2 #define __MPICUSPARSEMATIMPL
3 
4 #include <cusparse_v2.h>
5 #include <petsc/private/cudavecimpl.h>
6 
7 struct Mat_MPIAIJCUSPARSE {
8   /* The following are used by GPU capabilities to store matrix storage formats on the device */
9   MatCUSPARSEStorageFormat   diagGPUMatFormat;
10   MatCUSPARSEStorageFormat   offdiagGPUMatFormat;
11   cudaStream_t               stream;
12   cusparseHandle_t           handle;
13   PetscSplitCSRDataStructure deviceMat;
14   PetscInt                   coo_nd,coo_no; /* number of nonzero entries in coo for the diag/offdiag part */
15   THRUSTINTARRAY             *coo_p; /* the permutation array that partitions the coo array into diag/offdiag parts */
16   THRUSTARRAY                *coo_pw; /* the work array that stores the partitioned coo scalar values */
17 
18   /* Extended COO stuff */
19   PetscCount  *Aimap1_d,*Ajmap1_d,*Aperm1_d; /* Local entries to diag */
20   PetscCount  *Bimap1_d,*Bjmap1_d,*Bperm1_d; /* Local entries to offdiag */
21   PetscCount  *Aimap2_d,*Ajmap2_d,*Aperm2_d; /* Remote entries to diag */
22   PetscCount  *Bimap2_d,*Bjmap2_d,*Bperm2_d; /* Remote entries to offdiag */
23   PetscCount  *Cperm1_d; /* Permutation to fill send buffer. 'C' for communication */
24   PetscScalar *sendbuf_d,*recvbuf_d; /* Buffers for remote values in MatSetValuesCOO() */
25   PetscBool   use_extended_coo;
26 
27   Mat_MPIAIJCUSPARSE() {
28     diagGPUMatFormat    = MAT_CUSPARSE_CSR;
29     offdiagGPUMatFormat = MAT_CUSPARSE_CSR;
30     coo_p               = NULL;
31     coo_pw              = NULL;
32     stream              = 0;
33     deviceMat           = NULL;
34     use_extended_coo    = PETSC_FALSE;
35   }
36 };
37 
38 PETSC_INTERN PetscErrorCode MatCUSPARSESetStream(Mat, const cudaStream_t stream);
39 PETSC_INTERN PetscErrorCode MatCUSPARSESetHandle(Mat, const cusparseHandle_t handle);
40 PETSC_INTERN PetscErrorCode MatCUSPARSEClearHandle(Mat);
41 
42 #endif
43