Actual source code: aijkok.hpp
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: PETSC_PRAGMA_DIAGNOSTIC_IGNORED_BEGIN("-Wundef")
8: #include <KokkosSparse_spiluk.hpp>
9: PETSC_PRAGMA_DIAGNOSTIC_IGNORED_END()
10: #include <string>
12: namespace
13: {
14: PETSC_NODISCARD inline decltype(auto) NoInit(std::string label)
15: {
16: return Kokkos::view_alloc(Kokkos::WithoutInitializing, std::move(label));
17: }
18: } // namespace
20: using MatRowMapType = PetscInt;
21: using MatColIdxType = PetscInt;
22: using MatScalarType = PetscScalar;
24: template <class MemorySpace>
25: using KokkosCsrMatrixType = typename KokkosSparse::CrsMatrix<MatScalarType, MatColIdxType, MemorySpace, void /* MemoryTraits */, MatRowMapType>;
26: template <class MemorySpace>
27: using KokkosCsrGraphType = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type;
29: using KokkosCsrGraph = KokkosCsrGraphType<DefaultMemorySpace>;
30: using KokkosCsrGraphHost = KokkosCsrGraphType<HostMirrorMemorySpace>;
32: using KokkosCsrMatrix = KokkosCsrMatrixType<DefaultMemorySpace>;
33: using KokkosCsrMatrixHost = KokkosCsrMatrixType<HostMirrorMemorySpace>;
35: using MatRowMapKokkosView = KokkosCsrGraph::row_map_type::non_const_type;
36: using MatColIdxKokkosView = KokkosCsrGraph::entries_type::non_const_type;
37: using MatScalarKokkosView = KokkosCsrMatrix::values_type::non_const_type;
39: using MatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::non_const_type;
40: using MatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::non_const_type;
41: using MatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::non_const_type;
43: using ConstMatRowMapKokkosView = KokkosCsrGraph::row_map_type::const_type;
44: using ConstMatColIdxKokkosView = KokkosCsrGraph::entries_type::const_type;
45: using ConstMatScalarKokkosView = KokkosCsrMatrix::values_type::const_type;
47: using ConstMatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::const_type;
48: using ConstMatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::const_type;
49: using ConstMatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::const_type;
51: using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType *>;
52: using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType *>;
53: using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType *>;
55: using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType, MatColIdxType, MatScalarType, DefaultExecutionSpace, DefaultMemorySpace, DefaultMemorySpace>;
57: using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type;
59: /* For mat->spptr of a factorized matrix */
60: struct Mat_SeqAIJKokkosTriFactors {
61: MatRowMapKokkosView iL_d, iU_d, iLt_d, iUt_d; /* rowmap for L, U, L^t, U^t of A=LU */
62: MatColIdxKokkosView jL_d, jU_d, jLt_d, jUt_d; /* column ids */
63: MatScalarKokkosView aL_d, aU_d, aLt_d, aUt_d; /* matrix values */
64: KernelHandle kh, khL, khU, khLt, khUt; /* Kernel handles for ILU factorization of A, and TRSV of L, U, L^t, U^t */
65: PetscScalarKokkosView workVector;
66: PetscBool transpose_updated; /* Are L^T, U^T updated wrt L, U*/
67: PetscBool sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */
69: 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
70: MatColIdxKokkosViewHost jL_h, jU_h, jLt_h, jUt_h;
71: MatScalarKokkosViewHost aL_h, aU_h, aLt_h, aUt_h, D_h; // D is for LDLT factorization
72: MatScalarKokkosView D_d;
73: 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.
74: // MatTranspose() on host might be faster than KK's csr transpose on device.
76: PetscIntKokkosView rowperm, colperm; // row permutation and column permutation
78: Mat_SeqAIJKokkosTriFactors(PetscInt n) : workVector("workVector", n)
79: {
80: L = U = Lt = Ut = nullptr;
81: transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
82: }
84: ~Mat_SeqAIJKokkosTriFactors() { Destroy(); }
86: void Destroy()
87: {
88: PetscFunctionBeginUser;
89: kh.destroy_spiluk_handle();
90: khL.destroy_sptrsv_handle();
91: khU.destroy_sptrsv_handle();
92: khLt.destroy_sptrsv_handle();
93: khUt.destroy_sptrsv_handle();
94: PetscCallVoid(MatDestroy(&L));
95: PetscCallVoid(MatDestroy(&U));
96: PetscCallVoid(MatDestroy(&Lt));
97: PetscCallVoid(MatDestroy(&Ut));
98: L = U = Lt = Ut = nullptr;
99: transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
100: PetscFunctionReturnVoid();
101: }
102: };
104: /* For mat->spptr of a regular matrix */
105: struct Mat_SeqAIJKokkos {
106: MatRowMapKokkosDualView i_dual;
107: MatColIdxKokkosDualView j_dual;
108: MatScalarKokkosDualView a_dual;
109: PetscBool host_aij_allocated_by_kokkos = PETSC_FALSE; /* Are host views of a, i, j in the duals allocated by Kokkos? */
111: MatRowMapKokkosDualView diag_dual; /* Diagonal pointer, built on demand */
113: KokkosCsrMatrix csrmat; /* The CSR matrix, used to call KK functions */
114: PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */
116: KokkosCsrMatrix csrmatT, csrmatH; /* Transpose and Hermitian of the matrix (built on demand) */
117: PetscBool transpose_updated, hermitian_updated; /* Are At, Ah updated wrt the matrix? */
118: MatRowMapKokkosView transpose_perm; // A permutation array making Ta(i) = Aa(perm(i)), where T = A^t
120: /* Construct a nrows by ncols matrix with given aseq on host. Caller also specifies a nonzero state */
121: Mat_SeqAIJKokkos(Mat A, PetscInt nrows, PetscInt ncols, Mat_SeqAIJ *aseq, PetscObjectState nzstate, PetscBool copyValues = PETSC_TRUE)
122: {
123: auto exec = PetscGetKokkosExecutionSpace();
125: // a->diag must exist before new Mat_SeqAIJKokkos()
126: const PetscInt *adiag;
127: PetscCallVoid(MatGetDiagonalMarkers_SeqAIJ(A, &adiag, NULL));
129: MatScalarKokkosViewHost a_h(aseq->a, aseq->nz);
130: MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType *>(aseq->i), nrows + 1);
131: MatColIdxKokkosViewHost j_h(aseq->j, aseq->nz);
132: MatRowMapKokkosViewHost diag_h(aseq->diag, nrows);
134: auto a_d = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, exec, a_h);
135: auto i_d = Kokkos::create_mirror_view_and_copy(exec, i_h);
136: auto j_d = Kokkos::create_mirror_view_and_copy(exec, j_h);
137: auto diag_d = Kokkos::create_mirror_view_and_copy(exec, diag_h);
138: a_dual = MatScalarKokkosDualView(a_d, a_h);
139: i_dual = MatRowMapKokkosDualView(i_d, i_h);
140: j_dual = MatColIdxKokkosDualView(j_d, j_h);
141: diag_dual = MatColIdxKokkosDualView(diag_d, diag_h);
143: a_dual.modify_host(); /* Since caller provided values on host */
144: if (copyValues) (void)KokkosDualViewSyncDevice(a_dual, exec);
146: csrmat = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d));
147: Init(nzstate);
148: }
150: /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */
151: Mat_SeqAIJKokkos(const KokkosCsrMatrix &csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */
152: {
153: auto a_d = csr.values;
154: /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */
155: MatRowMapKokkosView i_d(const_cast<MatRowMapType *>(csr.graph.row_map.data()), csr.graph.row_map.extent(0));
156: auto j_d = csr.graph.entries;
157: auto a_h = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, HostMirrorMemorySpace(), a_d);
158: auto i_h = Kokkos::create_mirror_view_and_copy(HostMirrorMemorySpace(), i_d);
159: auto j_h = Kokkos::create_mirror_view_and_copy(HostMirrorMemorySpace(), j_d);
161: // diag_dual is set until MatAssemblyEnd() where we copy diag from host to device
162: a_dual = MatScalarKokkosDualView(a_d, a_h);
163: a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */
164: i_dual = MatRowMapKokkosDualView(i_d, i_h);
165: j_dual = MatColIdxKokkosDualView(j_d, j_h);
166: host_aij_allocated_by_kokkos = PETSC_TRUE; /* That means after deleting aijkok, one shouldn't access aijseq->{a,i,j} anymore! */
167: Init();
168: }
170: // 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)
171: 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)) { }
173: MatScalarType *a_host_data() { return a_dual.view_host().data(); }
174: MatRowMapType *i_host_data() { return i_dual.view_host().data(); }
175: MatColIdxType *j_host_data() { return j_dual.view_host().data(); }
177: MatScalarType *a_device_data() { return a_dual.view_device().data(); }
178: MatRowMapType *i_device_data() { return i_dual.view_device().data(); }
179: MatColIdxType *j_device_data() { return j_dual.view_device().data(); }
181: MatColIdxType nrows() { return csrmat.numRows(); }
182: MatColIdxType ncols() { return csrmat.numCols(); }
183: MatRowMapType nnz() { return csrmat.nnz(); }
185: /* Change the csrmat size to n */
186: void SetColSize(MatColIdxType n) { csrmat = KokkosCsrMatrix("csrmat", n, a_dual.view_device(), csrmat.graph); }
188: void SetDiagonal(const MatRowMapType *diag)
189: {
190: MatRowMapKokkosViewHost diag_h(const_cast<MatRowMapType *>(diag), nrows());
191: auto diag_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), diag_h);
192: diag_dual = MatRowMapKokkosDualView(diag_d, diag_h);
193: }
195: /* Shared init stuff */
196: void Init(PetscObjectState nzstate = 0)
197: {
198: nonzerostate = nzstate;
199: transpose_updated = PETSC_FALSE;
200: hermitian_updated = PETSC_FALSE;
201: }
203: PetscErrorCode DestroyMatTranspose(void)
204: {
205: PetscFunctionBegin;
206: csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
207: csrmatH = KokkosCsrMatrix();
208: PetscFunctionReturn(PETSC_SUCCESS);
209: }
210: };
212: struct MatProductCtx_SeqAIJKokkos {
213: KernelHandle kh;
214: PetscBool reusesym;
215: MatProductCtx_SeqAIJKokkos() : reusesym(PETSC_FALSE) { }
216: };
218: PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat, Mat_SeqAIJKokkos *);
219: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm, Mat_SeqAIJKokkos *, Mat *);
220: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat, Mat, MatReuse, Mat *);
221: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat);
222: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGetKokkosCsrMatrix(Mat, KokkosCsrMatrix *);
223: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm, KokkosCsrMatrix, Mat *);
224: PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix &csrmat);
225: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat, MatType, MatReuse, Mat *);
226: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat);
227: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGenerateTranspose_Private(Mat, KokkosCsrMatrix *);
229: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat, MatScalarKokkosView *);
230: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, MatScalarKokkosView *);
231: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, MatScalarKokkosView *);
232: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, MatScalarKokkosView *);
233: PETSC_INTERN PetscErrorCode MatInvertVariableBlockDiagonal_SeqAIJKokkos(Mat, const PetscIntKokkosView &, const PetscIntKokkosView &, const PetscIntKokkosView &, PetscScalarKokkosView &, PetscScalarKokkosView &);
234: PETSC_INTERN PetscErrorCode MatBindToCPU_SeqAIJKokkos(Mat, PetscBool);