xref: /petsc/src/mat/impls/aij/seq/kokkos/aijkok.hpp (revision 2a28d964fe6fe359fc6e4247643cfa1e8e64ec79) 
1 #ifndef __SEQAIJKOKKOSIMPL_HPP
2 #define __SEQAIJKOKKOSIMPL_HPP
3 
4 #include <petsc/private/vecimpl_kokkos.hpp>
5 #include <../src/mat/impls/aij/seq/aij.h>
6 #include <KokkosSparse_CrsMatrix.hpp>
7 #include <KokkosSparse_spiluk.hpp>
8 
9 /*
10    Kokkos::View<struct _n_SplitCSRMat,DefaultMemorySpace> is not handled correctly so we define SplitCSRMat
11    for the singular purpose of working around this.
12 */
13 typedef struct _n_SplitCSRMat SplitCSRMat;
14 
15 using MatRowMapType = PetscInt;
16 using MatColIdxType = PetscInt;
17 using MatScalarType = PetscScalar;
18 
19 template <class MemorySpace>
20 using KokkosCsrMatrixType = typename KokkosSparse::CrsMatrix<MatScalarType, MatColIdxType, MemorySpace, void /* MemoryTraits */, MatRowMapType>;
21 template <class MemorySpace>
22 using KokkosCsrGraphType = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type;
23 
24 using KokkosCsrGraph     = KokkosCsrGraphType<DefaultMemorySpace>;
25 using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>;
26 
27 using KokkosCsrMatrix     = KokkosCsrMatrixType<DefaultMemorySpace>;
28 using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>;
29 
30 using MatRowMapKokkosView = KokkosCsrGraph::row_map_type::non_const_type;
31 using MatColIdxKokkosView = KokkosCsrGraph::entries_type::non_const_type;
32 using MatScalarKokkosView = KokkosCsrMatrix::values_type::non_const_type;
33 
34 using MatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::non_const_type;
35 using MatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::non_const_type;
36 using MatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::non_const_type;
37 
38 using ConstMatRowMapKokkosView = KokkosCsrGraph::row_map_type::const_type;
39 using ConstMatColIdxKokkosView = KokkosCsrGraph::entries_type::const_type;
40 using ConstMatScalarKokkosView = KokkosCsrMatrix::values_type::const_type;
41 
42 using ConstMatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::const_type;
43 using ConstMatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::const_type;
44 using ConstMatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::const_type;
45 
46 using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType *>;
47 using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType *>;
48 using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType *>;
49 
50 using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType, MatColIdxType, MatScalarType, DefaultExecutionSpace, DefaultMemorySpace, DefaultMemorySpace>;
51 
52 using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type;
53 
54 /* For mat->spptr of a factorized matrix */
55 struct Mat_SeqAIJKokkosTriFactors {
56   MatRowMapKokkosView   iL_d, iU_d, iLt_d, iUt_d;  /* rowmap for L, U, L^t, U^t of A=LU */
57   MatColIdxKokkosView   jL_d, jU_d, jLt_d, jUt_d;  /* column ids */
58   MatScalarKokkosView   aL_d, aU_d, aLt_d, aUt_d;  /* matrix values */
59   KernelHandle          kh, khL, khU, khLt, khUt;  /* Kernel handles for A, L, U, L^t, U^t */
60   PetscBool             transpose_updated;         /* Are L^T, U^T updated wrt L, U*/
61   PetscBool             sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */
62   PetscScalarKokkosView workVector;
63 
64   Mat_SeqAIJKokkosTriFactors(PetscInt n) : transpose_updated(PETSC_FALSE), sptrsv_symbolic_completed(PETSC_FALSE), workVector("workVector", n) { }
65 
66   ~Mat_SeqAIJKokkosTriFactors() { Destroy(); }
67 
68   void Destroy()
69   {
70     kh.destroy_spiluk_handle();
71     khL.destroy_sptrsv_handle();
72     khU.destroy_sptrsv_handle();
73     khLt.destroy_sptrsv_handle();
74     khUt.destroy_sptrsv_handle();
75     transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE;
76   }
77 };
78 
79 /* For mat->spptr of a regular matrix */
80 struct Mat_SeqAIJKokkos {
81   MatRowMapKokkosDualView i_dual;
82   MatColIdxKokkosDualView j_dual;
83   MatScalarKokkosDualView a_dual;
84 
85   MatRowMapKokkosDualView diag_dual; /* Diagonal pointer, built on demand */
86 
87   KokkosCsrMatrix  csrmat;       /* The CSR matrix, used to call KK functions */
88   PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */
89 
90   KokkosCsrMatrix     csrmatT, csrmatH;                     /* Transpose and Hermitian of the matrix (built on demand) */
91   PetscBool           transpose_updated, hermitian_updated; /* Are At, Ah updated wrt the matrix? */
92   MatRowMapKokkosView transpose_perm;                       // A permutation array making Ta(i) = Aa(perm(i)), where T = A^t
93 
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 */
97 
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
102 
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);
109 
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);
113 
114     a_dual = MatScalarKokkosDualView(a_d, a_h);
115     i_dual = MatRowMapKokkosDualView(i_d, i_h);
116     j_dual = MatColIdxKokkosDualView(j_d, j_h);
117 
118     a_dual.modify_host(); /* Since caller provided values on host */
119     if (copyValues) a_dual.sync_device();
120 
121     csrmat            = KokkosCsrMatrix("csrmat", ncols, a_d, KokkosCsrGraph(j_d, i_d));
122     nonzerostate      = nzstate;
123     transpose_updated = hermitian_updated = PETSC_FALSE;
124   }
125 
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);
136 
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   }
143 
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   }
149 
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(); }
153 
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(); }
157 
158   MatColIdxType nrows() { return csrmat.numRows(); }
159   MatColIdxType ncols() { return csrmat.numCols(); }
160   MatRowMapType nnz() { return csrmat.nnz(); }
161 
162   /* Change the csrmat size to n */
163   void SetColSize(MatColIdxType n) { csrmat = KokkosCsrMatrix("csrmat", n, a_dual.view_device(), csrmat.graph); }
164 
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   }
170 
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   }
177 
178   /* Shared init stuff */
179   void Init(void)
180   {
181     transpose_updated = hermitian_updated = PETSC_FALSE;
182     nonzerostate                          = 0;
183   }
184 
185   PetscErrorCode DestroyMatTranspose(void)
186   {
187     PetscFunctionBegin;
188     csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
189     csrmatH = KokkosCsrMatrix();
190     PetscFunctionReturn(PETSC_SUCCESS);
191   }
192 };
193 
194 struct MatProductData_SeqAIJKokkos {
195   KernelHandle kh;
196   PetscBool    reusesym;
197   MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE) { }
198 };
199 
200 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat, Mat_SeqAIJKokkos *);
201 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm, Mat_SeqAIJKokkos *, Mat *);
202 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat, Mat, MatReuse, Mat *);
203 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat);
204 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGetKokkosCsrMatrix(Mat, KokkosCsrMatrix *);
205 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithKokkosCsrMatrix(MPI_Comm, KokkosCsrMatrix, Mat *);
206 PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix &csrmat);
207 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat, MatType, MatReuse, Mat *);
208 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat);
209 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosGenerateTranspose_Private(Mat, KokkosCsrMatrix *);
210 
211 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat, MatScalarKokkosView *);
212 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat, MatScalarKokkosView *);
213 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat, MatScalarKokkosView *);
214 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat, MatScalarKokkosView *);
215 #endif
216