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> using KokkosCsrMatrixType = typename KokkosSparse::CrsMatrix<MatScalarType,MatColIdxType,MemorySpace,void/* MemoryTraits */,MatRowMapType>;
21: template<class MemorySpace> using KokkosCsrGraphType = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type;
23: using KokkosCsrGraph = KokkosCsrGraphType<DefaultMemorySpace>;
24: using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>;
26: using KokkosCsrMatrix = KokkosCsrMatrixType<DefaultMemorySpace>;
27: using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>;
29: using MatRowMapKokkosView = KokkosCsrGraph::row_map_type::non_const_type;
30: using MatColIdxKokkosView = KokkosCsrGraph::entries_type::non_const_type;
31: using MatScalarKokkosView = KokkosCsrMatrix::values_type::non_const_type;
33: using MatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::non_const_type;
34: using MatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::non_const_type;
35: using MatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::non_const_type;
37: using ConstMatRowMapKokkosView = KokkosCsrGraph::row_map_type::const_type;
38: using ConstMatColIdxKokkosView = KokkosCsrGraph::entries_type::const_type;
39: using ConstMatScalarKokkosView = KokkosCsrMatrix::values_type::const_type;
41: using ConstMatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::const_type;
42: using ConstMatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::const_type;
43: using ConstMatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::const_type;
45: using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType*>;
46: using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType*>;
47: using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType*>;
49: using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType,MatColIdxType,MatScalarType,DefaultExecutionSpace,DefaultMemorySpace,DefaultMemorySpace>;
51: using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type;
53: using PetscCountKokkosView = Kokkos::View<PetscCount*,DefaultMemorySpace>;
54: using PetscCountKokkosViewHost = Kokkos::View<PetscCount*,Kokkos::HostSpace>;
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)
67: : transpose_updated(PETSC_FALSE),sptrsv_symbolic_completed(PETSC_FALSE),workVector("workVector",n) {}
69: ~Mat_SeqAIJKokkosTriFactors() {Destroy();}
71: void Destroy() {
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: KokkosCsrMatrix csrmat; /* The CSR matrix, used to call KK functions */
88: PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */
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? */
93: /* COO stuff */
94: 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 */
95: PetscCountKokkosView perm_d; /* The permutation array in sorting (i,j) by row and then by col */
97: Kokkos::View<PetscInt*> i_uncompressed_d;
98: Kokkos::View<PetscInt*> colmap_d; // ugh, this is a parallel construct
99: Kokkos::View<SplitCSRMat,DefaultMemorySpace> device_mat_d;
100: Kokkos::View<PetscInt*> diag_d; // factorizations
102: /* Construct a nrows by ncols matrix with nnz nonzeros from the given (i,j,a) on host. Caller also specifies a nonzero state */
103: Mat_SeqAIJKokkos(PetscInt nrows,PetscInt ncols,PetscInt nnz,const MatRowMapType *i,MatColIdxType *j,MatScalarType *a,PetscObjectState nzstate,PetscBool copyValues=PETSC_TRUE)
104: {
105: MatScalarKokkosViewHost a_h(a,nnz);
106: MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType*>(i),nrows+1);
107: MatColIdxKokkosViewHost j_h(j,nnz);
109: auto a_d = Kokkos::create_mirror_view(DefaultMemorySpace(),a_h);
110: auto i_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),i_h);
111: auto j_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),j_h);
113: a_dual = MatScalarKokkosDualView(a_d,a_h);
114: i_dual = MatRowMapKokkosDualView(i_d,i_h);
115: j_dual = MatColIdxKokkosDualView(j_d,j_h);
117: a_dual.modify_host(); /* Since caller provided values on host */
118: if (copyValues) a_dual.sync_device();
120: csrmat = KokkosCsrMatrix("csrmat",ncols,a_d,KokkosCsrGraph(j_d,i_d));
121: nonzerostate = nzstate;
122: transpose_updated = hermitian_updated = PETSC_FALSE;
123: }
125: /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */
126: Mat_SeqAIJKokkos(const KokkosCsrMatrix& csr) : csrmat(csr) /* Shallow-copy csr's views to csrmat */
127: {
128: auto a_d = csr.values;
129: /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */
130: MatRowMapKokkosView i_d(const_cast<MatRowMapType*>(csr.graph.row_map.data()),csr.graph.row_map.extent(0));
131: auto j_d = csr.graph.entries;
132: auto a_h = Kokkos::create_mirror_view(Kokkos::HostSpace(),a_d);
133: auto i_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),i_d);
134: auto j_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),j_d);
136: a_dual = MatScalarKokkosDualView(a_d,a_h);
137: a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */
138: i_dual = MatRowMapKokkosDualView(i_d,i_h);
139: j_dual = MatColIdxKokkosDualView(j_d,j_h);
140: Init();
141: }
143: Mat_SeqAIJKokkos(PetscInt nrows,PetscInt ncols,PetscInt nnz,
144: MatRowMapKokkosDualView& i,MatColIdxKokkosDualView& j,MatScalarKokkosDualView a)
145: :i_dual(i),j_dual(j),a_dual(a)
146: {
147: csrmat = KokkosCsrMatrix("csrmat",nrows,ncols,nnz,a.view_device(),i.view_device(),j.view_device());
148: Init();
149: }
151: MatScalarType* a_host_data() {return a_dual.view_host().data();}
152: MatRowMapType* i_host_data() {return i_dual.view_host().data();}
153: MatColIdxType* j_host_data() {return j_dual.view_host().data();}
155: MatScalarType* a_device_data() {return a_dual.view_device().data();}
156: MatRowMapType* i_device_data() {return i_dual.view_device().data();}
157: MatColIdxType* j_device_data() {return j_dual.view_device().data();}
159: MatColIdxType nrows() {return csrmat.numRows();}
160: MatColIdxType ncols() {return csrmat.numCols();}
161: MatRowMapType nnz() {return csrmat.nnz();}
163: /* Change the csrmat size to n */
164: void SetColSize(MatColIdxType n) {csrmat = KokkosCsrMatrix("csrmat",n,a_dual.view_device(),csrmat.graph);}
166: void SetUpCOO(const Mat_SeqAIJ *aij) {
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: /* Shared init stuff */
172: void Init(void)
173: {
174: transpose_updated = hermitian_updated = PETSC_FALSE;
175: nonzerostate = 0;
176: }
178: PetscErrorCode DestroyMatTranspose(void)
179: {
180: csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */
181: csrmatH = KokkosCsrMatrix();
182: return 0;
183: }
184: };
186: struct MatProductData_SeqAIJKokkos {
187: KernelHandle kh;
188: PetscBool reusesym;
189: MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE){}
190: };
192: PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat,Mat_SeqAIJKokkos*);
193: PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm,Mat_SeqAIJKokkos*,Mat*);
194: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat,Mat,MatReuse,Mat*);
195: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat);
196: PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix& csrmat);
197: PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat,MatType,MatReuse,Mat*);
198: PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat);
200: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat,MatScalarKokkosView*);
201: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat,MatScalarKokkosView*);
202: PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat,MatScalarKokkosView*);
203: PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat,MatScalarKokkosView*);
204: #endif