1 #ifndef __SEQAIJKOKKOSIMPL_HPP 2 #define __SEQAIJKOKKOSIMPL_HPP 3 4 #include <petscaijdevice.h> 5 #include <petsc/private/vecimpl_kokkos.hpp> 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> using KokkosCsrMatrixType = typename KokkosSparse::CrsMatrix<MatScalarType,MatColIdxType,MemorySpace,void/* MemoryTraits */,MatRowMapType>; 20 template<class MemorySpace> using KokkosCsrGraphType = typename KokkosCsrMatrixType<MemorySpace>::staticcrsgraph_type; 21 22 using KokkosCsrGraph = KokkosCsrGraphType<DefaultMemorySpace>; 23 using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>; 24 25 using KokkosCsrMatrix = KokkosCsrMatrixType<DefaultMemorySpace>; 26 using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>; 27 28 using MatRowMapKokkosView = KokkosCsrGraph::row_map_type::non_const_type; 29 using MatColIdxKokkosView = KokkosCsrGraph::entries_type::non_const_type; 30 using MatScalarKokkosView = KokkosCsrMatrix::values_type::non_const_type; 31 32 using MatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::non_const_type; 33 using MatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::non_const_type; 34 using MatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::non_const_type; 35 36 using ConstMatRowMapKokkosView = KokkosCsrGraph::row_map_type::const_type; 37 using ConstMatColIdxKokkosView = KokkosCsrGraph::entries_type::const_type; 38 using ConstMatScalarKokkosView = KokkosCsrMatrix::values_type::const_type; 39 40 using ConstMatRowMapKokkosViewHost = KokkosCsrGraphHost::row_map_type::const_type; 41 using ConstMatColIdxKokkosViewHost = KokkosCsrGraphHost::entries_type::const_type; 42 using ConstMatScalarKokkosViewHost = KokkosCsrMatrixHost::values_type::const_type; 43 44 using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType*>; 45 using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType*>; 46 using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType*>; 47 48 using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType,MatColIdxType,MatScalarType,DefaultExecutionSpace,DefaultMemorySpace,DefaultMemorySpace>; 49 50 using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type; 51 52 /* For mat->spptr of a factorized matrix */ 53 struct Mat_SeqAIJKokkosTriFactors { 54 MatRowMapKokkosView iL_d,iU_d,iLt_d,iUt_d; /* rowmap for L, U, L^t, U^t of A=LU */ 55 MatColIdxKokkosView jL_d,jU_d,jLt_d,jUt_d; /* column ids */ 56 MatScalarKokkosView aL_d,aU_d,aLt_d,aUt_d; /* matrix values */ 57 KernelHandle kh,khL,khU,khLt,khUt; /* Kernel handles for A, L, U, L^t, U^t */ 58 PetscBool transpose_updated; /* Are L^T, U^T updated wrt L, U*/ 59 PetscBool sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */ 60 PetscScalarKokkosView workVector; 61 62 Mat_SeqAIJKokkosTriFactors(PetscInt n) 63 : transpose_updated(PETSC_FALSE),sptrsv_symbolic_completed(PETSC_FALSE),workVector("workVector",n) {} 64 65 ~Mat_SeqAIJKokkosTriFactors() {Destroy();} 66 67 void Destroy() { 68 kh.destroy_spiluk_handle(); 69 khL.destroy_sptrsv_handle(); 70 khU.destroy_sptrsv_handle(); 71 khLt.destroy_sptrsv_handle(); 72 khUt.destroy_sptrsv_handle(); 73 transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE; 74 } 75 }; 76 77 /* For mat->spptr of a regular matrix */ 78 struct Mat_SeqAIJKokkos { 79 MatRowMapKokkosDualView i_dual; 80 MatColIdxKokkosDualView j_dual; 81 MatScalarKokkosDualView a_dual; 82 83 KokkosCsrMatrix csrmat; /* The CSR matrix, used to call KK functions */ 84 PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */ 85 86 KokkosCsrMatrix csrmatT,csrmatH; /* Transpose and Hermitian of the matrix (built on demand) */ 87 PetscBool transpose_updated,hermitian_updated; /* Are At, Ah updated wrt the matrix? */ 88 89 PetscInt coo_n; /* Number of entries in MatSetPreallocationCOO() */ 90 PetscBool coo_has_repeats; /* Are there replicated (i,j) pairs? (excluding those ignored) */ 91 MatRowMapKokkosView jmap_d; /* perm[disp+jmap[i]..disp+jmap[i+1]) gives indices of entries in v[] associated with i-th nonzero of the matrix */ 92 MatRowMapKokkosView perm_d; /* The permutation array in sorting (i,j) by row and then by col */ 93 94 Kokkos::View<PetscInt*> *i_uncompressed_d; 95 Kokkos::View<PetscInt*> *colmap_d; // ugh, this is a parallel construct 96 Kokkos::View<SplitCSRMat,DefaultMemorySpace> device_mat_d; 97 Kokkos::View<PetscInt*> *diag_d; // factorizations 98 99 /* Construct a nrows by ncols matrix with nnz nonzeros from the given (i,j,a) on host. Caller also specifies a nonzero state */ 100 Mat_SeqAIJKokkos(PetscInt nrows,PetscInt ncols,PetscInt nnz,const MatRowMapType *i,MatColIdxType *j,MatScalarType *a,PetscObjectState nzstate,PetscBool copyValues=PETSC_TRUE) 101 { 102 MatScalarKokkosViewHost a_h(a,nnz); 103 MatRowMapKokkosViewHost i_h(const_cast<MatRowMapType*>(i),nrows+1); 104 MatColIdxKokkosViewHost j_h(j,nnz); 105 106 auto a_d = Kokkos::create_mirror_view(DefaultMemorySpace(),a_h); 107 auto i_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),i_h); 108 auto j_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),j_h); 109 110 a_dual = MatScalarKokkosDualView(a_d,a_h); 111 i_dual = MatRowMapKokkosDualView(i_d,i_h); 112 j_dual = MatColIdxKokkosDualView(j_d,j_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 i_uncompressed_d = colmap_d = diag_d = NULL; 121 } 122 123 /* Construct with a KokkosCsrMatrix. For performance, only i, j are copied to host, but not the matrix values. */ 124 Mat_SeqAIJKokkos(const KokkosCsrMatrix& csr) : csrmat(csr) 125 { 126 auto a_d = csr.values; 127 /* Get a non-const version since I don't want to deal with DualView<const T*>, which is not well defined */ 128 MatRowMapKokkosView i_d(const_cast<MatRowMapType*>(csr.graph.row_map.data()),csr.graph.row_map.extent(0)); 129 auto j_d = csr.graph.entries; 130 auto a_h = Kokkos::create_mirror_view(Kokkos::HostSpace(),a_d); 131 auto i_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),i_d); 132 auto j_h = Kokkos::create_mirror_view_and_copy(Kokkos::HostSpace(),j_d); 133 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, 142 MatRowMapKokkosDualView& i,MatColIdxKokkosDualView& j,MatScalarKokkosDualView a) 143 :i_dual(i),j_dual(j),a_dual(a) 144 { 145 csrmat = KokkosCsrMatrix("csrmat",nrows,ncols,nnz,a.view_device(),i.view_device(),j.view_device()); 146 Init(); 147 } 148 149 MatScalarType* a_host_data() {return a_dual.view_host().data();} 150 MatRowMapType* i_host_data() {return i_dual.view_host().data();} 151 MatColIdxType* j_host_data() {return j_dual.view_host().data();} 152 153 MatScalarType* a_device_data() {return a_dual.view_device().data();} 154 MatRowMapType* i_device_data() {return i_dual.view_device().data();} 155 MatColIdxType* j_device_data() {return j_dual.view_device().data();} 156 157 MatColIdxType nrows() {return csrmat.numRows();} 158 MatColIdxType ncols() {return csrmat.numCols();} 159 MatRowMapType nnz() {return csrmat.nnz();} 160 161 /* Change the csrmat size to n */ 162 void SetColSize(MatColIdxType n) {csrmat = KokkosCsrMatrix("csrmat",n,a_dual.view_device(),csrmat.graph);} 163 164 void SetUpCOO(PetscInt n,PetscBool has_repeats,MatRowMapKokkosViewHost& jmap_h,MatRowMapKokkosViewHost& perm_h) { 165 coo_n = n; 166 coo_has_repeats = has_repeats; 167 if (coo_has_repeats) {jmap_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),jmap_h);} 168 perm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),perm_h); 169 } 170 171 /* Shared init stuff */ 172 void Init(void) 173 { 174 transpose_updated = hermitian_updated = PETSC_FALSE; 175 i_uncompressed_d = colmap_d = diag_d = NULL; 176 nonzerostate = 0; 177 178 coo_n = 0; 179 coo_has_repeats = PETSC_FALSE; 180 } 181 182 PetscErrorCode DestroyMatTranspose(void) 183 { 184 PetscFunctionBegin; 185 csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */ 186 csrmatH = KokkosCsrMatrix(); 187 PetscFunctionReturn(0); 188 } 189 }; 190 191 struct MatProductData_SeqAIJKokkos { 192 KernelHandle kh; 193 PetscBool reusesym; 194 MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE){} 195 }; 196 197 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat,Mat_SeqAIJKokkos*); 198 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm,Mat_SeqAIJKokkos*,Mat*); 199 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat,Mat,MatReuse,Mat*); 200 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat); 201 PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix& csrmat); 202 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat,MatType,MatReuse,Mat*); 203 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat); 204 PetscErrorCode MatSeqAIJGetKokkosView(Mat,MatScalarKokkosView*); 205 PetscErrorCode MatSeqAIJRestoreKokkosView(Mat,MatScalarKokkosView*); 206 #endif 207