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 /* For mat->spptr of a factorized matrix */ 54 struct Mat_SeqAIJKokkosTriFactors { 55 MatRowMapKokkosView iL_d,iU_d,iLt_d,iUt_d; /* rowmap for L, U, L^t, U^t of A=LU */ 56 MatColIdxKokkosView jL_d,jU_d,jLt_d,jUt_d; /* column ids */ 57 MatScalarKokkosView aL_d,aU_d,aLt_d,aUt_d; /* matrix values */ 58 KernelHandle kh,khL,khU,khLt,khUt; /* Kernel handles for A, L, U, L^t, U^t */ 59 PetscBool transpose_updated; /* Are L^T, U^T updated wrt L, U*/ 60 PetscBool sptrsv_symbolic_completed; /* Have we completed the symbolic solve for L and U */ 61 PetscScalarKokkosView workVector; 62 63 Mat_SeqAIJKokkosTriFactors(PetscInt n) 64 : transpose_updated(PETSC_FALSE),sptrsv_symbolic_completed(PETSC_FALSE),workVector("workVector",n) {} 65 66 ~Mat_SeqAIJKokkosTriFactors() {Destroy();} 67 68 void Destroy() { 69 kh.destroy_spiluk_handle(); 70 khL.destroy_sptrsv_handle(); 71 khU.destroy_sptrsv_handle(); 72 khLt.destroy_sptrsv_handle(); 73 khUt.destroy_sptrsv_handle(); 74 transpose_updated = sptrsv_symbolic_completed = PETSC_FALSE; 75 } 76 }; 77 78 /* For mat->spptr of a regular matrix */ 79 struct Mat_SeqAIJKokkos { 80 MatRowMapKokkosDualView i_dual; 81 MatColIdxKokkosDualView j_dual; 82 MatScalarKokkosDualView a_dual; 83 84 KokkosCsrMatrix csrmat; /* The CSR matrix, used to call KK functions */ 85 PetscObjectState nonzerostate; /* State of the nonzero pattern (graph) on device */ 86 87 KokkosCsrMatrix csrmatT,csrmatH; /* Transpose and Hermitian of the matrix (built on demand) */ 88 PetscBool transpose_updated,hermitian_updated; /* Are At, Ah updated wrt the matrix? */ 89 90 /* COO stuff */ 91 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 */ 92 PetscCountKokkosView 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(const Mat_SeqAIJ *aij) { 164 jmap_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),PetscCountKokkosViewHost(aij->jmap,aij->nz+1)); 165 perm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(),PetscCountKokkosViewHost(aij->perm,aij->Atot)); 166 } 167 168 /* Shared init stuff */ 169 void Init(void) 170 { 171 transpose_updated = hermitian_updated = PETSC_FALSE; 172 nonzerostate = 0; 173 } 174 175 PetscErrorCode DestroyMatTranspose(void) 176 { 177 PetscFunctionBegin; 178 csrmatT = KokkosCsrMatrix(); /* Overwrite with empty matrices */ 179 csrmatH = KokkosCsrMatrix(); 180 PetscFunctionReturn(0); 181 } 182 }; 183 184 struct MatProductData_SeqAIJKokkos { 185 KernelHandle kh; 186 PetscBool reusesym; 187 MatProductData_SeqAIJKokkos() : reusesym(PETSC_FALSE){} 188 }; 189 190 PETSC_INTERN PetscErrorCode MatSetSeqAIJKokkosWithCSRMatrix(Mat,Mat_SeqAIJKokkos*); 191 PETSC_INTERN PetscErrorCode MatCreateSeqAIJKokkosWithCSRMatrix(MPI_Comm,Mat_SeqAIJKokkos*,Mat*); 192 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosMergeMats(Mat,Mat,MatReuse,Mat*); 193 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosSyncDevice(Mat); 194 PETSC_INTERN PetscErrorCode PrintCsrMatrix(const KokkosCsrMatrix& csrmat); 195 PETSC_INTERN PetscErrorCode MatConvert_SeqAIJ_SeqAIJKokkos(Mat,MatType,MatReuse,Mat*); 196 PETSC_INTERN PetscErrorCode MatSeqAIJKokkosModifyDevice(Mat); 197 198 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosView(Mat,MatScalarKokkosView*); 199 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosView(Mat,MatScalarKokkosView*); 200 PETSC_INTERN PetscErrorCode MatSeqAIJGetKokkosViewWrite(Mat,MatScalarKokkosView*); 201 PETSC_INTERN PetscErrorCode MatSeqAIJRestoreKokkosViewWrite(Mat,MatScalarKokkosView*); 202 #endif 203