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