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: #include <petscpkg_version.h>
8: #if PETSC_PKG_HIP_VERSION_GE(5, 2, 0)
9: #include <hipsparse/hipsparse.h>
10: #else
11: #include <hipsparse.h>
12: #endif
13: #include <petsc/private/veccupmimpl.h>
15: struct Mat_MPIAIJHIPSPARSE {
16: /* The following are used by GPU capabilities to store matrix storage formats on the device */
17: MatHIPSPARSEStorageFormat diagGPUMatFormat;
18: MatHIPSPARSEStorageFormat offdiagGPUMatFormat;
19: PetscSplitCSRDataStructure deviceMat;
20: PetscInt coo_nd, coo_no; /* number of nonzero entries in coo for the diag/offdiag part */
21: THRUSTINTARRAY *coo_p; /* the permutation array that partitions the coo array into diag/offdiag parts */
22: THRUSTARRAY *coo_pw; /* the work array that stores the partitioned coo scalar values */
24: /* Extended COO stuff */
25: PetscCount *Ajmap1_d, *Aperm1_d; /* Local entries to diag */
26: PetscCount *Bjmap1_d, *Bperm1_d; /* Local entries to offdiag */
27: PetscCount *Aimap2_d, *Ajmap2_d, *Aperm2_d; /* Remote entries to diag */
28: PetscCount *Bimap2_d, *Bjmap2_d, *Bperm2_d; /* Remote entries to offdiag */
29: PetscCount *Cperm1_d; /* Permutation to fill send buffer. 'C' for communication */
30: PetscScalar *sendbuf_d, *recvbuf_d; /* Buffers for remote values in MatSetValuesCOO() */
31: PetscBool use_extended_coo;
33: Mat_MPIAIJHIPSPARSE()
34: {
35: diagGPUMatFormat = MAT_HIPSPARSE_CSR;
36: offdiagGPUMatFormat = MAT_HIPSPARSE_CSR;
37: coo_p = NULL;
38: coo_pw = NULL;
39: deviceMat = NULL;
40: use_extended_coo = PETSC_FALSE;
41: }
42: };
44: #endif