Actual source code: mpihipsparsematimpl.h
1: /* Portions of this code are under:
2: Copyright (c) 2022 Advanced Micro Devices, Inc. All rights reserved.
3: */
4: #ifndef __MPIHIPSPARSEMATIMPL
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>
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 */
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;
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: };
43: #endif