1 #pragma once 2 #include <petsc_kokkos.hpp> 3 #include <petscmat_kokkos.hpp> 4 #include <petsc/private/kokkosimpl.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<HostMirrorMemorySpace>; 29 30 using KokkosCsrMatrix = KokkosCsrMatrixType<DefaultMemorySpace>; 31 using KokkosCsrMatrixHost = KokkosCsrMatrixType<HostMirrorMemorySpace>; 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 ILU factorization of A, and TRSV of L, U, L^t, U^t */ 63 PetscScalarKokkosView workVector; 64 PetscBool transpose_updated; /* Are L^T, U^T updated wrt L, U*/ 65 PetscBool sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */ 66 67 MatRowMapKokkosViewHost iL_h, iU_h, iLt_h, iUt_h; // temp. buffers when we do factorization with PETSc on host. We copy L, U to these buffers and then copy to device 68 MatColIdxKokkosViewHost jL_h, jU_h, jLt_h, jUt_h; 69 MatScalarKokkosViewHost aL_h, aU_h, aLt_h, aUt_h, D_h; // D is for LDLT factorization 70 MatScalarKokkosView D_d; 71 Mat L, U, Lt, Ut; // MATSEQAIJ on host if needed. Their arrays are alias to (iL_h, jL_h, aL_h), (iU_h, jU_h, aU_h) and their transpose. 72 // MatTranspose() on host might be faster than KK's csr transpose on device. 73 74 PetscIntKokkosView rowperm, colperm; // row permutation and column permutation 75 76 Mat_SeqAIJKokkosTriFactors(PetscInt n) : workVector("workVector", n) 77 { 78 L = U = Lt = Ut = nullptr; 79 transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE; 80 } 81 82 ~Mat_SeqAIJKokkosTriFactors() { Destroy(); } 83 84 void Destroy() 85 { 86 PetscFunctionBeginUser; 87 kh.destroy_spiluk_handle(); 88 khL.destroy_sptrsv_handle(); 89 khU.destroy_sptrsv_handle(); 90 khLt.destroy_sptrsv_handle(); 91 khUt.destroy_sptrsv_handle(); 92 PetscCallVoid(MatDestroy(&L)); 93 PetscCallVoid(MatDestroy(&U)); 94 PetscCallVoid(MatDestroy(&Lt)); 95 PetscCallVoid(MatDestroy(&Ut)); 96 L = U = Lt = Ut = nullptr; 97 transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE; 98 PetscFunctionReturnVoid(); 99 } 100 }; 101 102 /* For mat->spptr of a regular matrix */ 103 struct Mat_SeqAIJKokkos { 104 MatRowMapKokkosDualView i_dual; 105 MatColIdxKokkosDualView j_dual; 106 MatScalarKokkosDualView a_dual; 107 PetscBool host_aij_allocated_by_kokkos = PETSC_FALSE; /* Are host views of a, i, j in the duals allocated by Kokkos? */ 108 109 MatRowMapKokkosDualView diag_dual; /* Diagonal pointer, built on demand */ 110 111 KokkosCsrMatrix csrmat; /* The CSR matrix, used to call KK functions */ 112 PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */ 113 114 KokkosCsrMatrix csrmatT, csrmatH; /* Transpose and Hermitian of the matrix (built on demand) */ 115 PetscBool transpose_updated, hermitian_updated; /* Are At, Ah updated wrt the matrix? */ 116 MatRowMapKokkosView transpose_perm; // A permutation array making Ta(i) = Aa(perm(i)), where T = A^t 117 118 /* Construct a nrows by ncols matrix with given aseq on host. Caller also specifies a nonzero state */ 119 Mat_SeqAIJKokkos(Mat A, PetscInt nrows, PetscInt ncols, Mat_SeqAIJ *aseq, PetscObjectState nzstate, PetscBool copyValues = PETSC_TRUE) 120 { 121 auto exec = PetscGetKokkosExecutionSpace(); 122 123 // a->diag must exist before new Mat_SeqAIJKokkos() 124 const PetscInt *adiag; 125 PetscCallVoid(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL)); 126 127 MatScalarKokkosViewHost a_h(aseq->a, aseq->nz); 128 MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType *>(aseq->i), nrows + 1); 129 MatColIdxKokkosViewHost j_h(aseq->j, aseq->nz); 130 MatRowMapKokkosViewHost diag_h(aseq->diag, nrows); 131 132 auto a_d = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, exec, a_h); 133 auto i_d = Kokkos::create_mirror_view_and_copy(exec, i_h); 134 auto j_d = Kokkos::create_mirror_view_and_copy(exec, j_h); 135 auto diag_d = Kokkos::create_mirror_view_and_copy(exec, diag_h); 136 a_dual = MatScalarKokkosDualView(a_d, a_h); 137 i_dual = MatRowMapKokkosDualView(i_d, i_h); 138 j_dual = MatColIdxKokkosDualView(j_d, j_h); 139 diag_dual = MatColIdxKokkosDualView(diag_d, diag_h); 140 141 a_dual.modify_host(); /* Since caller provided values on host */ 142 if (copyValues) (void)KokkosDualViewSyncDevice(a_dual, exec); 143 144 csrmat = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d)); 145 Init(nzstate); 146 } 147 148 /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */ 149 Mat_SeqAIJKokkos(const KokkosCsrMatrix &csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */ 150 { 151 auto a_d = csr.values; 152 /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */ 153 MatRowMapKokkosView i_d(const_cast<MatRowMapType *>(csr.graph.row_map.data()), csr.graph.row_map.extent(0)); 154 auto j_d = csr.graph.entries; 155 auto a_h = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, HostMirrorMemorySpace(), a_d); 156 auto i_h = Kokkos::create_mirror_view_and_copy(HostMirrorMemorySpace(), i_d); 157 auto j_h = Kokkos::create_mirror_view_and_copy(HostMirrorMemorySpace(), j_d); 158 159 // diag_dual is set until MatAssemblyEnd() where we copy diag from host to device 160 a_dual = MatScalarKokkosDualView(a_d, a_h); 161 a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */ 162 i_dual = MatRowMapKokkosDualView(i_d, i_h); 163 j_dual = MatColIdxKokkosDualView(j_d, j_h); 164 host_aij_allocated_by_kokkos = PETSC_TRUE; /* That means after deleting aijkok, one shouldn't access aijseq->{a,i,j} anymore! */ 165 Init(); 166 } 167 168 // Don't use DualView argument types as we want to be sure that a,i,j on host are allocated by Mat_SeqAIJKokkos itself (vs. by users) 169 Mat_SeqAIJKokkos(PetscInt nrows, PetscInt ncols, PetscInt nnz, const MatRowMapKokkosView &i_d, const MatColIdxKokkosView &j_d, const MatScalarKokkosView &a_d) : Mat_SeqAIJKokkos(KokkosCsrMatrix("csrmat", nrows, ncols, nnz, a_d, i_d, j_d)) { } 170 171 MatScalarType *a_host_data() { return a_dual.view_host().data(); } 172 MatRowMapType *i_host_data() { return i_dual.view_host().data(); } 173 MatColIdxType *j_host_data() { return j_dual.view_host().data(); } 174 175 MatScalarType *a_device_data() { return a_dual.view_device().data(); } 176 MatRowMapType *i_device_data() { return i_dual.view_device().data(); } 177 MatColIdxType *j_device_data() { return j_dual.view_device().data(); } 178 179 MatColIdxType nrows() { return csrmat.numRows(); } 180 MatColIdxType ncols() { return csrmat.numCols(); } 181 MatRowMapType nnz() { return csrmat.nnz(); } 182 183 /* Change the csrmat size to n */ 184 void SetColSize(MatColIdxType n) { csrmat = KokkosCsrMatrix("csrmat", n, a_dual.view_device(), csrmat.graph); } 185 186 void SetDiagonal(const MatRowMapType *diag) 187 { 188 MatRowMapKokkosViewHost diag_h(const_cast<MatRowMapType *>(diag), nrows()); 189 auto diag_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), diag_h); 190 diag_dual = MatRowMapKokkosDualView(diag_d, diag_h); 191 } 192 193 /* Shared init stuff */ 194 void Init(PetscObjectState nzstate = 0) 195 { 196 nonzerostate = nzstate; 197 transpose_updated = PETSC_FALSE; 198 hermitian_updated = PETSC_FALSE; 199 } 200 201 PetscErrorCode DestroyMatTranspose(void) 202 { 203 PetscFunctionBegin; 204 csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */ 205 csrmatH = KokkosCsrMatrix(); 206 PetscFunctionReturn(PETSC_SUCCESS); 207 } 208 }; 209 210 struct MatProductCtx_SeqAIJKokkos { 211 KernelHandle kh; 212 PetscBool reusesym; 213 MatProductCtx_SeqAIJKokkos() : reusesym(PETSC_FALSE) { } 214 }; 215 216 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat, Mat_SeqAIJKokkos *); 217 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm, Mat_SeqAIJKokkos *, Mat *); 218 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat, Mat, MatReuse, Mat *); 219 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat); 220 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGetKokkosCsrMatrix(Mat, KokkosCsrMatrix *); 221 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm, KokkosCsrMatrix, Mat *); 222 PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix &csrmat); 223 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat, MatType, MatReuse, Mat *); 224 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat); 225 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGenerateTranspose_Private(Mat, KokkosCsrMatrix *); 226 227 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat, MatScalarKokkosView *); 228 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, MatScalarKokkosView *); 229 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, MatScalarKokkosView *); 230 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, MatScalarKokkosView *); 231 PETSC_INTERN PetscErrorCode MatInvertVariableBlockDiagonal_SeqAIJKokkos(Mat, const PetscIntKokkosView &, const PetscIntKokkosView &, const PetscIntKokkosView &, PetscScalarKokkosView &, PetscScalarKokkosView &); 232