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