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 #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> 14 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 */ 23 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; 32 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 }; 43 44 #endif 45