xref: /petsc/src/mat/impls/aij/seq/kokkos/aijkok.hpp (revision d9acb416d05abeed0a33bde3a81aeb2ea0364f6a) 
1 #pragma once
2 
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>
8 
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
16 
17 using MatRowMapType = PetscInt;
18 using MatColIdxType = PetscInt;
19 using MatScalarType = PetscScalar;
20 
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;
25 
26 using KokkosCsrGraph     = KokkosCsrGraphType<DefaultMemorySpace>;
27 using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>;
28 
29 using KokkosCsrMatrix     = KokkosCsrMatrixType<DefaultMemorySpace>;
30 using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>;
31 
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;
35 
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;
39 
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;
43 
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;
47 
48 using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType *>;
49 using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType *>;
50 using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType *>;
51 
52 using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType, MatColIdxType, MatScalarType, DefaultExecutionSpace, DefaultMemorySpace, DefaultMemorySpace>;
53 
54 using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type;
55 
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;
65 
66   Mat_SeqAIJKokkosTriFactors(PetscInt n) : transpose_updated(PETSC_FALSE), sptrsv_symbolic_completed(PETSC_FALSE), workVector("workVector", n) { }
67 
68   ~Mat_SeqAIJKokkosTriFactors() { Destroy(); }
69 
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 };
80 
81 /* For mat->spptr of a regular matrix */
82 struct Mat_SeqAIJKokkos {
83   MatRowMapKokkosDualView i_dual;
84   MatColIdxKokkosDualView j_dual;
85   MatScalarKokkosDualView a_dual;
86 
87   MatRowMapKokkosDualView diag_dual; /* Diagonal pointer, built on demand */
88 
89   KokkosCsrMatrix  csrmat;       /* The CSR matrix, used to call KK functions */
90   PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */
91 
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
95 
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);
103 
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);
108 
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);
113 
114     a_dual.modify_host(); /* Since caller provided values on host */
115     if (copyValues) a_dual.sync_device();
116 
117     csrmat            = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d));
118     nonzerostate      = nzstate;
119     transpose_updated = hermitian_updated = PETSC_FALSE;
120   }
121 
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);
132 
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   }
140 
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   }
146 
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(); }
150 
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(); }
154 
155   MatColIdxType nrows() { return csrmat.numRows(); }
156   MatColIdxType ncols() { return csrmat.numCols(); }
157   MatRowMapType nnz() { return csrmat.nnz(); }
158 
159   /* Change the csrmat size to n */
160   void SetColSize(MatColIdxType n) { csrmat = KokkosCsrMatrix("csrmat", n, a_dual.view_device(), csrmat.graph); }
161 
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   }
168 
169   /* Shared init stuff */
170   void Init(void)
171   {
172     transpose_updated = hermitian_updated = PETSC_FALSE;
173     nonzerostate                          = 0;
174   }
175 
176   PetscErrorCode DestroyMatTranspose(void)
177   {
178     PetscFunctionBegin;
179     csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
180     csrmatH = KokkosCsrMatrix();
181     PetscFunctionReturn(PETSC_SUCCESS);
182   }
183 };
184 
185 struct MatProductData_SeqAIJKokkos {
186   KernelHandle kh;
187   PetscBool    reusesym;
188   MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE) { }
189 };
190 
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 *);
201 
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 &);
207