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