xref: /petsc/src/mat/impls/aij/mpi/mpihipsparse/mpihipsparsematimpl.h (revision a8d3b578d34eb8d5be1b3eaad4a98fb37a558ece)
1 /* Portions of this code are under:
2    Copyright (c) 2022 Advanced Micro Devices, Inc. All rights reserved.
3 */
4 #ifndef __MPIHIPSPARSEMATIMPL
5 #define __MPIHIPSPARSEMATIMPL
6 
7 #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
8   #include <hipsparse/hipsparse.h>
9 #else
10   #include <hipsparse.h>
11 #endif
12 #include <petsc/private/hipvecimpl.h>
13 
14 struct Mat_MPIAIJHIPSPARSE {
15   /* The following are used by GPU capabilities to store matrix storage formats on the device */
16   MatHIPSPARSEStorageFormat  diagGPUMatFormat;
17   MatHIPSPARSEStorageFormat  offdiagGPUMatFormat;
18   PetscSplitCSRDataStructure deviceMat;
19   PetscInt                   coo_nd, coo_no; /* number of nonzero entries in coo for the diag/offdiag part */
20   THRUSTINTARRAY            *coo_p;          /* the permutation array that partitions the coo array into diag/offdiag parts */
21   THRUSTARRAY               *coo_pw;         /* the work array that stores the partitioned coo scalar values */
22 
23   /* Extended COO stuff */
24   PetscCount  *Ajmap1_d, *Aperm1_d;            /* Local entries to diag */
25   PetscCount  *Bjmap1_d, *Bperm1_d;            /* Local entries to offdiag */
26   PetscCount  *Aimap2_d, *Ajmap2_d, *Aperm2_d; /* Remote entries to diag */
27   PetscCount  *Bimap2_d, *Bjmap2_d, *Bperm2_d; /* Remote entries to offdiag */
28   PetscCount  *Cperm1_d;                       /* Permutation to fill send buffer. 'C' for communication */
29   PetscScalar *sendbuf_d, *recvbuf_d;          /* Buffers for remote values in MatSetValuesCOO() */
30   PetscBool    use_extended_coo;
31 
32   Mat_MPIAIJHIPSPARSE()
33   {
34     diagGPUMatFormat    = MAT_HIPSPARSE_CSR;
35     offdiagGPUMatFormat = MAT_HIPSPARSE_CSR;
36     coo_p               = NULL;
37     coo_pw              = NULL;
38     deviceMat           = NULL;
39     use_extended_coo    = PETSC_FALSE;
40   }
41 };
42 
43 #endif
44