1 #ifndef __SEQAIJKOKKOSIMPL_HPP 2 #define __SEQAIJKOKKOSIMPL_HPP 3 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> 9 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; 15 16 using MatRowMapType = PetscInt; 17 using MatColIdxType = PetscInt; 18 using MatScalarType = PetscScalar; 19 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; 22 23 using KokkosCsrGraph = KokkosCsrGraphType<DefaultMemorySpace>; 24 using KokkosCsrGraphHost = KokkosCsrGraphType<Kokkos::HostSpace>; 25 26 using KokkosCsrMatrix = KokkosCsrMatrixType<DefaultMemorySpace>; 27 using KokkosCsrMatrixHost = KokkosCsrMatrixType<Kokkos::HostSpace>; 28 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; 32 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; 36 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; 40 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; 44 45 using MatRowMapKokkosDualView = Kokkos::DualView<MatRowMapType*>; 46 using MatColIdxKokkosDualView = Kokkos::DualView<MatColIdxType*>; 47 using MatScalarKokkosDualView = Kokkos::DualView<MatScalarType*>; 48 49 using KernelHandle = KokkosKernels::Experimental::KokkosKernelsHandle<MatRowMapType,MatColIdxType,MatScalarType,DefaultExecutionSpace,DefaultMemorySpace,DefaultMemorySpace>; 50 51 using KokkosTeamMemberType = Kokkos::TeamPolicy<DefaultExecutionSpace>::member_type; 52 53 using PetscCountKokkosView = Kokkos::View<PetscCount*,DefaultMemorySpace>; 54 using PetscCountKokkosViewHost = Kokkos::View<PetscCount*,Kokkos::HostSpace>; 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) 67 : transpose_updated(PETSC_FALSE),sptrsv_symbolic_completed(PETSC_FALSE),workVector("workVector",n) {} 68 69 ~Mat_SeqAIJKokkosTriFactors() {Destroy();} 70 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 }; 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 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 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 */ 96 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 101 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); 108 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); 112 113 a_dual = MatScalarKokkosDualView(a_d,a_h); 114 i_dual = MatRowMapKokkosDualView(i_d,i_h); 115 j_dual = MatColIdxKokkosDualView(j_d,j_h); 116 117 a_dual.modify_host(); /* Since caller provided values on host */ 118 if (copyValues) a_dual.sync_device(); 119 120 csrmat = KokkosCsrMatrix("csrmat",ncols,a_d,KokkosCsrGraph(j_d,i_d)); 121 nonzerostate = nzstate; 122 transpose_updated = hermitian_updated = PETSC_FALSE; 123 } 124 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); 135 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 } 142 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 } 150 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();} 154 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();} 158 159 MatColIdxType nrows() {return csrmat.numRows();} 160 MatColIdxType ncols() {return csrmat.numCols();} 161 MatRowMapType nnz() {return csrmat.nnz();} 162 163 /* Change the csrmat size to n */ 164 void SetColSize(MatColIdxType n) {csrmat = KokkosCsrMatrix("csrmat",n,a_dual.view_device(),csrmat.graph);} 165 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 } 170 171 /* Shared init stuff */ 172 void Init(void) 173 { 174 transpose_updated = hermitian_updated = PETSC_FALSE; 175 nonzerostate = 0; 176 } 177 178 PetscErrorCode DestroyMatTranspose(void) 179 { 180 PetscFunctionBegin; 181 csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */ 182 csrmatH = KokkosCsrMatrix(); 183 PetscFunctionReturn(0); 184 } 185 }; 186 187 struct MatProductData_SeqAIJKokkos { 188 KernelHandle kh; 189 PetscBool reusesym; 190 MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE){} 191 }; 192 193 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat,Mat_SeqAIJKokkos*); 194 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm,Mat_SeqAIJKokkos*,Mat*); 195 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat,Mat,MatReuse,Mat*); 196 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(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 201 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat,MatScalarKokkosView*); 202 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat,MatScalarKokkosView*); 203 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat,MatScalarKokkosView*); 204 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat,MatScalarKokkosView*); 205 #endif 206