Actual source code: aijkok.hpp

      1: #ifndef __SEQAIJKOKKOSIMPL_HPP

      4: #include <petscaijdevice.h>
      5: #include <petsc/private/vecimpl_kokkos.hpp>
      6: #include <../src/mat/impls/aij/seq/aij.h>
      7: #include <KokkosSparse_CrsMatrix.hpp>
      8: #include <KokkosSparse_spiluk.hpp>

     10: /*
     11:    Kokkos::View<struct _n_SplitCSRMat,DefaultMemorySpace> is not handled correctly so we define SplitCSRMat
     12:    for the singular purpose of working around this.
     13: */
     14: typedef struct _n_SplitCSRMat SplitCSRMat;

     16: using MatRowMapType  = PetscInt;
     17: using MatColIdxType  = PetscInt;
     18: using MatScalarType  = PetscScalar;

     20: template<class MemorySpace> using KokkosCsrMatrixType   = typename KokkosSparse::CrsMatrix<MatScalarType,MatColIdxType,MemorySpace,void/* MemoryTraits */,MatRowMapType>;
     21: template<class MemorySpace> using KokkosCsrGraphType    = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type;

     23: using KokkosCsrGraph                 = KokkosCsrGraphType<DefaultMemorySpace>;
     24: using KokkosCsrGraphHost             = KokkosCsrGraphType<Kokkos::HostSpace>;

     26: using KokkosCsrMatrix                = KokkosCsrMatrixType<DefaultMemorySpace>;
     27: using KokkosCsrMatrixHost            = KokkosCsrMatrixType<Kokkos::HostSpace>;

     29: using MatRowMapKokkosView            = KokkosCsrGraph::row_map_type::non_const_type;
     30: using MatColIdxKokkosView            = KokkosCsrGraph::entries_type::non_const_type;
     31: using MatScalarKokkosView            = KokkosCsrMatrix::values_type::non_const_type;

     33: using MatRowMapKokkosViewHost        = KokkosCsrGraphHost::row_map_type::non_const_type;
     34: using MatColIdxKokkosViewHost        = KokkosCsrGraphHost::entries_type::non_const_type;
     35: using MatScalarKokkosViewHost        = KokkosCsrMatrixHost::values_type::non_const_type;

     37: using ConstMatRowMapKokkosView       = KokkosCsrGraph::row_map_type::const_type;
     38: using ConstMatColIdxKokkosView       = KokkosCsrGraph::entries_type::const_type;
     39: using ConstMatScalarKokkosView       = KokkosCsrMatrix::values_type::const_type;

     41: using ConstMatRowMapKokkosViewHost   = KokkosCsrGraphHost::row_map_type::const_type;
     42: using ConstMatColIdxKokkosViewHost   = KokkosCsrGraphHost::entries_type::const_type;
     43: using ConstMatScalarKokkosViewHost   = KokkosCsrMatrixHost::values_type::const_type;

     45: using MatRowMapKokkosDualView        = Kokkos::DualView<MatRowMapType*>;
     46: using MatColIdxKokkosDualView        = Kokkos::DualView<MatColIdxType*>;
     47: using MatScalarKokkosDualView        = Kokkos::DualView<MatScalarType*>;

     49: using KernelHandle                   = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType,MatColIdxType,MatScalarType,DefaultExecutionSpace,DefaultMemorySpace,DefaultMemorySpace>;

     51: using KokkosTeamMemberType           = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type;

     53: using PetscCountKokkosView           = Kokkos::View<PetscCount*,DefaultMemorySpace>;
     54: using PetscCountKokkosViewHost       = Kokkos::View<PetscCount*,Kokkos::HostSpace>;

     56: /* For mat->spptr of a factorized matrix */
     57: struct Mat_SeqAIJKokkosTriFactors {
     58:   MatRowMapKokkosView       iL_d,iU_d,iLt_d,iUt_d; /* rowmap for L, U, L^t, U^t of A=LU */
     59:   MatColIdxKokkosView       jL_d,jU_d,jLt_d,jUt_d; /* column ids */
     60:   MatScalarKokkosView       aL_d,aU_d,aLt_d,aUt_d; /* matrix values */
     61:   KernelHandle              kh,khL,khU,khLt,khUt;  /* Kernel handles for A, L, U, L^t, U^t */
     62:   PetscBool                 transpose_updated;     /* Are L^T, U^T updated wrt L, U*/
     63:   PetscBool                 sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */
     64:   PetscScalarKokkosView     workVector;

     66:   Mat_SeqAIJKokkosTriFactors(PetscInt n)
     67:     : transpose_updated(PETSC_FALSE),sptrsv_symbolic_completed(PETSC_FALSE),workVector("workVector",n) {}

     69:   ~Mat_SeqAIJKokkosTriFactors() {Destroy();}

     71:   void Destroy() {
     72:     kh.destroy_spiluk_handle();
     73:     khL.destroy_sptrsv_handle();
     74:     khU.destroy_sptrsv_handle();
     75:     khLt.destroy_sptrsv_handle();
     76:     khUt.destroy_sptrsv_handle();
     77:     transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
     78:   }
     79: };

     81: /* For mat->spptr of a regular matrix */
     82: struct Mat_SeqAIJKokkos {
     83:   MatRowMapKokkosDualView    i_dual;
     84:   MatColIdxKokkosDualView    j_dual;
     85:   MatScalarKokkosDualView    a_dual;

     87:   KokkosCsrMatrix            csrmat; /* The CSR matrix, used to call KK functions */
     88:   PetscObjectState           nonzerostate; /* State of the nonzero pattern (graph) on device */

     90:   KokkosCsrMatrix            csrmatT,csrmatH; /* Transpose and Hermitian of the matrix (built on demand) */
     91:   PetscBool                  transpose_updated,hermitian_updated; /* Are At, Ah updated wrt the matrix? */

     93:   /* COO stuff */
     94:   PetscCountKokkosView       jmap_d; /* perm[disp+jmap[i]..disp+jmap[i+1]) gives indices of entries in v[] associated with i-th nonzero of the matrix */
     95:   PetscCountKokkosView       perm_d; /* The permutation array in sorting (i,j) by row and then by col */

     97:   Kokkos::View<PetscInt*>         i_uncompressed_d;
     98:   Kokkos::View<PetscInt*>         colmap_d; // ugh, this is a parallel construct
     99:   Kokkos::View<SplitCSRMat,DefaultMemorySpace> device_mat_d;
    100:   Kokkos::View<PetscInt*>         diag_d; // factorizations

    102:   /* Construct a nrows by ncols matrix with nnz nonzeros from the given (i,j,a) on host. Caller also specifies a nonzero state */
    103:   Mat_SeqAIJKokkos(PetscInt nrows,PetscInt ncols,PetscInt nnz,const MatRowMapType *i,MatColIdxType *j,MatScalarType *a,PetscObjectState nzstate,PetscBool copyValues=PETSC_TRUE)
    104:   {
    105:     MatScalarKokkosViewHost    a_h(a,nnz);
    106:     MatRowMapKokkosViewHost    i_h(const_cast<MatRowMapType*>(i),nrows+1);
    107:     MatColIdxKokkosViewHost    j_h(j,nnz);

    109:     auto a_d = Kokkos::create_mirror_view(DefaultMemorySpace(),a_h);
    110:     auto i_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),i_h);
    111:     auto j_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),j_h);

    113:     a_dual   = MatScalarKokkosDualView(a_d,a_h);
    114:     i_dual   = MatRowMapKokkosDualView(i_d,i_h);
    115:     j_dual   = MatColIdxKokkosDualView(j_d,j_h);

    117:     a_dual.modify_host(); /* Since caller provided values on host */
    118:     if (copyValues) a_dual.sync_device();

    120:     csrmat       = KokkosCsrMatrix("csrmat",ncols,a_d,KokkosCsrGraph(j_d,i_d));
    121:     nonzerostate = nzstate;
    122:     transpose_updated = hermitian_updated = PETSC_FALSE;
    123:   }

    125:   /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */
    126:   Mat_SeqAIJKokkos(const KokkosCsrMatrix& csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */
    127:   {
    128:     auto a_d = csr.values;
    129:     /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */
    130:     MatRowMapKokkosView i_d(const_cast<MatRowMapType*>(csr.graph.row_map.data()),csr.graph.row_map.extent(0));
    131:     auto j_d = csr.graph.entries;
    132:     auto a_h = Kokkos::create_mirror_view(Kokkos::HostSpace(),a_d);
    133:     auto i_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),i_d);
    134:     auto j_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),j_d);

    136:     a_dual = MatScalarKokkosDualView(a_d,a_h);
    137:     a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */
    138:     i_dual = MatRowMapKokkosDualView(i_d,i_h);
    139:     j_dual = MatColIdxKokkosDualView(j_d,j_h);
    140:     Init();
    141:   }

    143:   Mat_SeqAIJKokkos(PetscInt nrows,PetscInt ncols,PetscInt nnz,
    144:                    MatRowMapKokkosDualView& i,MatColIdxKokkosDualView& j,MatScalarKokkosDualView a)
    145:     :i_dual(i),j_dual(j),a_dual(a)
    146:   {
    147:     csrmat = KokkosCsrMatrix("csrmat",nrows,ncols,nnz,a.view_device(),i.view_device(),j.view_device());
    148:     Init();
    149:   }

    151:   MatScalarType* a_host_data() {return a_dual.view_host().data();}
    152:   MatRowMapType* i_host_data() {return i_dual.view_host().data();}
    153:   MatColIdxType* j_host_data() {return j_dual.view_host().data();}

    155:   MatScalarType* a_device_data() {return a_dual.view_device().data();}
    156:   MatRowMapType* i_device_data() {return i_dual.view_device().data();}
    157:   MatColIdxType* j_device_data() {return j_dual.view_device().data();}

    159:   MatColIdxType  nrows() {return csrmat.numRows();}
    160:   MatColIdxType  ncols() {return csrmat.numCols();}
    161:   MatRowMapType  nnz()   {return csrmat.nnz();}

    163:   /* Change the csrmat size to n */
    164:   void SetColSize(MatColIdxType n) {csrmat = KokkosCsrMatrix("csrmat",n,a_dual.view_device(),csrmat.graph);}

    166:   void SetUpCOO(const Mat_SeqAIJ *aij) {
    167:     jmap_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),PetscCountKokkosViewHost(aij->jmap,aij->nz+1));
    168:     perm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),PetscCountKokkosViewHost(aij->perm,aij->Atot));
    169:   }

    171:   /* Shared init stuff */
    172:   void Init(void)
    173:   {
    174:     transpose_updated = hermitian_updated = PETSC_FALSE;
    175:     nonzerostate      = 0;
    176:   }

    178:   PetscErrorCode DestroyMatTranspose(void)
    179:   {
    180:     csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
    181:     csrmatH = KokkosCsrMatrix();
    182:     return 0;
    183:   }
    184: };

    186: struct MatProductData_SeqAIJKokkos {
    187:   KernelHandle kh;
    188:   PetscBool    reusesym;
    189:   MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE){}
    190: };

    192: PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat,Mat_SeqAIJKokkos*);
    193: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm,Mat_SeqAIJKokkos*,Mat*);
    194: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat,Mat,MatReuse,Mat*);
    195: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat);
    196: PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix& csrmat);
    197: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat,MatType,MatReuse,Mat*);
    198: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat);

    200: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat,MatScalarKokkosView*);
    201: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat,MatScalarKokkosView*);
    202: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat,MatScalarKokkosView*);
    203: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat,MatScalarKokkosView*);
    204: #endif