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