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