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 } 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::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 a_dual = MatScalarKokkosDualView(a_d,a_h); 134 a_dual.modify_device(); /* since we did not copy a_d to a_h, we mark device has the latest data */ 135 i_dual = MatRowMapKokkosDualView(i_d,i_h); 136 j_dual = MatColIdxKokkosDualView(j_d,j_h); 137 Init(); 138 } 139 140 Mat_SeqAIJKokkos(PetscInt nrows,PetscInt ncols,PetscInt nnz, 141 MatRowMapKokkosDualView& i,MatColIdxKokkosDualView& j,MatScalarKokkosDualView a) 142 :i_dual(i),j_dual(j),a_dual(a) 143 { 144 csrmat = KokkosCsrMatrix("csrmat",nrows,ncols,nnz,a.view_device(),i.view_device(),j.view_device()); 145 Init(); 146 } 147 148 MatScalarType* a_host_data() {return a_dual.view_host().data();} 149 MatRowMapType* i_host_data() {return i_dual.view_host().data();} 150 MatColIdxType* j_host_data() {return j_dual.view_host().data();} 151 152 MatScalarType* a_device_data() {return a_dual.view_device().data();} 153 MatRowMapType* i_device_data() {return i_dual.view_device().data();} 154 MatColIdxType* j_device_data() {return j_dual.view_device().data();} 155 156 MatColIdxType nrows() {return csrmat.numRows();} 157 MatColIdxType ncols() {return csrmat.numCols();} 158 MatRowMapType nnz() {return csrmat.nnz();} 159 160 /* Change the csrmat size to n */ 161 void SetColSize(MatColIdxType n) {csrmat = KokkosCsrMatrix("csrmat",n,a_dual.view_device(),csrmat.graph);} 162 163 void SetUpCOO(PetscInt n,PetscBool has_repeats,MatRowMapKokkosViewHost& jmap_h,MatRowMapKokkosViewHost& perm_h) { 164 coo_n = n; 165 coo_has_repeats = has_repeats; 166 if (coo_has_repeats) {jmap_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),jmap_h);} 167 perm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),perm_h); 168 } 169 170 /* Shared init stuff */ 171 void Init(void) 172 { 173 transpose_updated = hermitian_updated = PETSC_FALSE; 174 nonzerostate = 0; 175 176 coo_n = 0; 177 coo_has_repeats = PETSC_FALSE; 178 } 179 180 PetscErrorCode DestroyMatTranspose(void) 181 { 182 PetscFunctionBegin; 183 csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */ 184 csrmatH = KokkosCsrMatrix(); 185 PetscFunctionReturn(0); 186 } 187 }; 188 189 struct MatProductData_SeqAIJKokkos { 190 KernelHandle kh; 191 PetscBool reusesym; 192 MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE){} 193 }; 194 195 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat,Mat_SeqAIJKokkos*); 196 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm,Mat_SeqAIJKokkos*,Mat*); 197 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat,Mat,MatReuse,Mat*); 198 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat); 199 PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix& csrmat); 200 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat,MatType,MatReuse,Mat*); 201 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat); 202 PetscErrorCode MatSeqAIJGetKokkosView(Mat,MatScalarKokkosView*); 203 PetscErrorCode MatSeqAIJRestoreKokkosView(Mat,MatScalarKokkosView*); 204 #endif 205