Actual source code: aijkokkosimpl.hpp


  4: #include <petsc/private/vecimpl_kokkos.hpp>
  5: #include <KokkosSparse_CrsMatrix.hpp>

  7: using MatRowOffsetType    = PetscInt;
  8: using MatColumnIndexType  = PetscInt;
  9: using MatValueType        = PetscScalar;

 11: template<class MemorySpace> using KokkosCsrMatrixType   = typename KokkosSparse::CrsMatrix<MatValueType,MatColumnIndexType,MemorySpace,void/* MemoryTraits */,MatRowOffsetType>;
 12: template<class MemorySpace> using KokkosCsrGraphType    = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type;

 14: using KokkosCsrGraph                      = KokkosCsrGraphType<DefaultMemorySpace>;
 15: using KokkosCsrMatrix                     = KokkosCsrMatrixType<DefaultMemorySpace>;

 17: using KokkosCsrGraphHost                  = KokkosCsrGraphType<DefaultMemorySpace>::HostMirror;

 19: using MatColumnIndexKokkosView            = KokkosCsrGraph::entries_type;
 20: using MatRowOffsetKokkosView              = KokkosCsrGraph::row_map_type;
 21: using MatValueKokkosView                  = KokkosCsrMatrix::values_type;

 23: using MatColumnIndexKokkosViewHost        = MatColumnIndexKokkosView::HostMirror;
 24: using MatRowOffsetKokkosViewHost          = MatRowOffsetKokkosView::HostMirror;
 25: using MatValueKokkosViewHost              = MatValueKokkosView::HostMirror;

 27: using MatValueKokkosDualView              = Kokkos::DualView<MatValueType*>;

 29: struct Mat_SeqAIJKokkos {
 30:   MatRowOffsetKokkosViewHost     i_h;
 31:   MatRowOffsetKokkosView         i_d;

 33:   MatColumnIndexKokkosViewHost   j_h;
 34:   MatColumnIndexKokkosView       j_d;

 36:   MatValueKokkosViewHost         a_h;
 37:   MatValueKokkosView             a_d;

 39:   MatValueKokkosDualView         a_dual;

 41:   KokkosCsrGraphHost             csrgraph_h;
 42:   KokkosCsrGraph                 csrgraph_d;

 44:   KokkosCsrMatrix                csrmat; /* The CSR matrix */
 45:   PetscObjectState               nonzerostate; /* State of the nonzero pattern (graph) on device */

 47:   Mat                            At,Ah; /* Transpose and Hermitian of the matrix in MATAIJKOKKOS type (built on demand) */
 48:   PetscBool                      need_sync_At_values,need_sync_Ah_values; /* A's values are updated but At, Ah are not */

 50:   Kokkos::View<PetscInt*>        *i_uncompressed_d;
 51:   Kokkos::View<PetscInt*>        *colmap_d; // ugh, this is a parallel construct
 52:   Kokkos::View<PetscSplitCSRDataStructure,DefaultMemorySpace> device_mat_d;
 53:   Kokkos::View<PetscInt*>        *diag_d; // factorizations

 55:   /* Construct a nrows by ncols matrix of nnz nonzeros with (i,j,a) for the CSR */
 56:   Mat_SeqAIJKokkos(MatColumnIndexType nrows,MatColumnIndexType ncols,MatRowOffsetType nnz,MatRowOffsetType *i,MatColumnIndexType *j,MatValueType *a)
 57:    : i_h(i,nrows+1),j_h(j,nnz),a_h(a,nnz),At(NULL),Ah(NULL),need_sync_At_values(PETSC_FALSE),need_sync_Ah_values(PETSC_FALSE),
 58:      i_uncompressed_d(NULL),colmap_d(NULL),device_mat_d(NULL),diag_d(NULL)
 59:   {
 60:      i_d        = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),i_h);
 61:      j_d        = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),j_h);
 62:      a_d        = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),a_h);
 63:      csrgraph_d = KokkosCsrGraph(j_d,i_d);
 64:      csrgraph_h = KokkosCsrGraphHost(j_h,i_h);
 65:      a_dual     = MatValueKokkosDualView(a_d,a_h);
 66:      csrmat     = KokkosCsrMatrix("csrmat",ncols,a_d,csrgraph_d);
 67:   }

 69:   ~Mat_SeqAIJKokkos()
 70:   {
 71:     DestroyMatTranspose();
 72:   }

 74:   PetscErrorCode DestroyMatTranspose(void)
 75:   {
 78:     MatDestroy(&At);
 79:     MatDestroy(&Ah);
 80:     return(0);
 81:   }
 82: };

 84: #endif