1 #ifndef __SEQAIJKOKKOSIMPL_HPP 2 #define __SEQAIJKOKKOSIMPL_HPP 3 4 #include <petsc/private/vecimpl_kokkos.hpp> 5 #include <../src/mat/impls/aij/seq/aij.h> 6 #include <KokkosSparse_CrsMatrix.hpp> 7 #include <KokkosSparse_spiluk.hpp> 8 #include <string> 9 10 namespace 11 { 12 PETSC_NODISCARD inline decltype(auto) NoInit(std::string label) 13 { 14 return Kokkos::view_alloc(Kokkos::WithoutInitializing, std::move(label)); 15 } 16 } // namespace 17 18 using MatRowMapType = PetscInt; 19 using MatColIdxType = PetscInt; 20 using MatScalarType = PetscScalar; 21 22 template <class MemorySpace> 23 using KokkosCsrMatrixType = typename KokkosSparse::CrsMatrix<MatScalarType, MatColIdxType, MemorySpace, void /* MemoryTraits */, MatRowMapType>; 24 template <class MemorySpace> 25 using KokkosCsrGraphType = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type; 26 27 using KokkosCsrGraph = KokkosCsrGraphType<DefaultMemorySpace>; 28 using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>; 29 30 using KokkosCsrMatrix = KokkosCsrMatrixType<DefaultMemorySpace>; 31 using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>; 32 33 using MatRowMapKokkosView = KokkosCsrGraph::row_map_type::non_const_type; 34 using MatColIdxKokkosView = KokkosCsrGraph::entries_type::non_const_type; 35 using MatScalarKokkosView = KokkosCsrMatrix::values_type::non_const_type; 36 37 using MatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::non_const_type; 38 using MatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::non_const_type; 39 using MatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::non_const_type; 40 41 using ConstMatRowMapKokkosView = KokkosCsrGraph::row_map_type::const_type; 42 using ConstMatColIdxKokkosView = KokkosCsrGraph::entries_type::const_type; 43 using ConstMatScalarKokkosView = KokkosCsrMatrix::values_type::const_type; 44 45 using ConstMatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::const_type; 46 using ConstMatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::const_type; 47 using ConstMatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::const_type; 48 49 using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType *>; 50 using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType *>; 51 using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType *>; 52 53 using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType, MatColIdxType, MatScalarType, DefaultExecutionSpace, DefaultMemorySpace, DefaultMemorySpace>; 54 55 using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type; 56 57 /* For mat->spptr of a factorized matrix */ 58 struct Mat_SeqAIJKokkosTriFactors { 59 MatRowMapKokkosView iL_d, iU_d, iLt_d, iUt_d; /* rowmap for L, U, L^t, U^t of A=LU */ 60 MatColIdxKokkosView jL_d, jU_d, jLt_d, jUt_d; /* column ids */ 61 MatScalarKokkosView aL_d, aU_d, aLt_d, aUt_d; /* matrix values */ 62 KernelHandle kh, khL, khU, khLt, khUt; /* Kernel handles for A, L, U, L^t, U^t */ 63 PetscBool transpose_updated; /* Are L^T, U^T updated wrt L, U*/ 64 PetscBool sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */ 65 PetscScalarKokkosView workVector; 66 67 Mat_SeqAIJKokkosTriFactors(PetscInt n) : transpose_updated(PETSC_FALSE), sptrsv_symbolic_completed(PETSC_FALSE), workVector("workVector", n) { } 68 69 ~Mat_SeqAIJKokkosTriFactors() { Destroy(); } 70 71 void Destroy() 72 { 73 kh.destroy_spiluk_handle(); 74 khL.destroy_sptrsv_handle(); 75 khU.destroy_sptrsv_handle(); 76 khLt.destroy_sptrsv_handle(); 77 khUt.destroy_sptrsv_handle(); 78 transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE; 79 } 80 }; 81 82 /* For mat->spptr of a regular matrix */ 83 struct Mat_SeqAIJKokkos { 84 MatRowMapKokkosDualView i_dual; 85 MatColIdxKokkosDualView j_dual; 86 MatScalarKokkosDualView a_dual; 87 88 MatRowMapKokkosDualView diag_dual; /* Diagonal pointer, built on demand */ 89 90 KokkosCsrMatrix csrmat; /* The CSR matrix, used to call KK functions */ 91 PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */ 92 93 KokkosCsrMatrix csrmatT, csrmatH; /* Transpose and Hermitian of the matrix (built on demand) */ 94 PetscBool transpose_updated, hermitian_updated; /* Are At, Ah updated wrt the matrix? */ 95 MatRowMapKokkosView transpose_perm; // A permutation array making Ta(i) = Aa(perm(i)), where T = A^t 96 97 /* COO stuff */ 98 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 */ 99 PetscCountKokkosView perm_d; /* The permutation array in sorting (i,j) by row and then by col */ 100 101 /* Construct a nrows by ncols matrix with nnz nonzeros from the given (i,j,a) on host. Caller also specifies a nonzero state */ 102 Mat_SeqAIJKokkos(PetscInt nrows, PetscInt ncols, PetscInt nnz, const MatRowMapType *i, MatColIdxType *j, MatScalarType *a, PetscObjectState nzstate, PetscBool copyValues = PETSC_TRUE) 103 { 104 MatScalarKokkosViewHost a_h(a, nnz); 105 MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType *>(i), nrows + 1); 106 MatColIdxKokkosViewHost j_h(j, nnz); 107 108 auto a_d = Kokkos::create_mirror_view(DefaultMemorySpace(), a_h); 109 auto i_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), i_h); 110 auto j_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), j_h); 111 112 a_dual = MatScalarKokkosDualView(a_d, a_h); 113 i_dual = MatRowMapKokkosDualView(i_d, i_h); 114 j_dual = MatColIdxKokkosDualView(j_d, j_h); 115 116 a_dual.modify_host(); /* Since caller provided values on host */ 117 if (copyValues) a_dual.sync_device(); 118 119 csrmat = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d)); 120 nonzerostate = nzstate; 121 transpose_updated = hermitian_updated = PETSC_FALSE; 122 } 123 124 /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */ 125 Mat_SeqAIJKokkos(const KokkosCsrMatrix &csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */ 126 { 127 auto a_d = csr.values; 128 /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */ 129 MatRowMapKokkosView i_d(const_cast<MatRowMapType *>(csr.graph.row_map.data()), csr.graph.row_map.extent(0)); 130 auto j_d = csr.graph.entries; 131 auto a_h = Kokkos::create_mirror_view(Kokkos::HostSpace(), a_d); 132 auto i_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), i_d); 133 auto j_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(), j_d); 134 135 a_dual = MatScalarKokkosDualView(a_d, a_h); 136 a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */ 137 i_dual = MatRowMapKokkosDualView(i_d, i_h); 138 j_dual = MatColIdxKokkosDualView(j_d, j_h); 139 Init(); 140 } 141 142 Mat_SeqAIJKokkos(PetscInt nrows, PetscInt ncols, PetscInt nnz, MatRowMapKokkosDualView &i, MatColIdxKokkosDualView &j, MatScalarKokkosDualView a) : i_dual(i), j_dual(j), a_dual(a) 143 { 144 csrmat = KokkosCsrMatrix("csrmat", nrows, ncols, nnz, a.view_device(), i.view_device(), j.view_device()); 145 Init(); 146 } 147 148 MatScalarType *a_host_data() { return a_dual.view_host().data(); } 149 MatRowMapType *i_host_data() { return i_dual.view_host().data(); } 150 MatColIdxType *j_host_data() { return j_dual.view_host().data(); } 151 152 MatScalarType *a_device_data() { return a_dual.view_device().data(); } 153 MatRowMapType *i_device_data() { return i_dual.view_device().data(); } 154 MatColIdxType *j_device_data() { return j_dual.view_device().data(); } 155 156 MatColIdxType nrows() { return csrmat.numRows(); } 157 MatColIdxType ncols() { return csrmat.numCols(); } 158 MatRowMapType nnz() { return csrmat.nnz(); } 159 160 /* Change the csrmat size to n */ 161 void SetColSize(MatColIdxType n) { csrmat = KokkosCsrMatrix("csrmat", n, a_dual.view_device(), csrmat.graph); } 162 163 void SetUpCOO(const Mat_SeqAIJ *aij) 164 { 165 jmap_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(aij->jmap, aij->nz + 1)); 166 perm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(aij->perm, aij->Atot)); 167 } 168 169 void SetDiagonal(const MatRowMapType *diag) 170 { 171 MatRowMapKokkosViewHost diag_h(const_cast<MatRowMapType *>(diag), nrows()); 172 auto diag_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), diag_h); 173 diag_dual = MatRowMapKokkosDualView(diag_d, diag_h); 174 } 175 176 /* Shared init stuff */ 177 void Init(void) 178 { 179 transpose_updated = hermitian_updated = PETSC_FALSE; 180 nonzerostate = 0; 181 } 182 183 PetscErrorCode DestroyMatTranspose(void) 184 { 185 PetscFunctionBegin; 186 csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */ 187 csrmatH = KokkosCsrMatrix(); 188 PetscFunctionReturn(PETSC_SUCCESS); 189 } 190 }; 191 192 struct MatProductData_SeqAIJKokkos { 193 KernelHandle kh; 194 PetscBool reusesym; 195 MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE) { } 196 }; 197 198 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat, Mat_SeqAIJKokkos *); 199 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm, Mat_SeqAIJKokkos *, Mat *); 200 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat, Mat, MatReuse, Mat *); 201 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat); 202 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGetKokkosCsrMatrix(Mat, KokkosCsrMatrix *); 203 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm, KokkosCsrMatrix, Mat *); 204 PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix &csrmat); 205 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat, MatType, MatReuse, Mat *); 206 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat); 207 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGenerateTranspose_Private(Mat, KokkosCsrMatrix *); 208 209 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat, MatScalarKokkosView *); 210 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, MatScalarKokkosView *); 211 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, MatScalarKokkosView *); 212 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, MatScalarKokkosView *); 213 #endif 214