1 #pragma once 2 3 #include <../src/vec/vec/impls/mpi/pvecimpl.h> 4 #include <petsc/private/kokkosimpl.hpp> 5 6 #if defined(PETSC_USE_DEBUG) 7 #define VecErrorIfNotKokkos(v) \ 8 do { \ 9 PetscBool isKokkos = PETSC_FALSE; \ 10 PetscCall(PetscObjectTypeCompareAny((PetscObject)(v), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, VECKOKKOS, "")); \ 11 PetscCheck(isKokkos, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Calling VECKOKKOS methods on a non-VECKOKKOS object"); \ 12 } while (0) 13 #else 14 #define VecErrorIfNotKokkos(v) \ 15 do { \ 16 (void)(v); \ 17 } while (0) 18 #endif 19 20 /* Stuff related to Vec_Kokkos */ 21 22 struct Vec_Kokkos { 23 PetscScalarKokkosDualView v_dual; 24 PetscScalarKokkosView unplaced_d; /* Unplaced device array in VecKokkosPlaceArray() */ 25 26 /* COO stuff */ 27 PetscCountKokkosView jmap1_d; /* [m+1]: i-th entry of the vector has jmap1[i+1]-jmap1[i] repeats in COO arrays */ 28 PetscCountKokkosView perm1_d; /* [tot1]: permutation array for local entries */ 29 30 PetscCountKokkosView imap2_d; /* [nnz2]: i-th unique entry in recvbuf is imap2[i]-th entry in the vector */ 31 PetscCountKokkosView jmap2_d; /* [nnz2+1] */ 32 PetscCountKokkosView perm2_d; /* [recvlen] */ 33 PetscCountKokkosView Cperm_d; /* [sendlen]: permutation array to fill sendbuf[]. 'C' for communication */ 34 PetscScalarKokkosView sendbuf_d, recvbuf_d; /* Buffers for remote values in VecSetValuesCOO() */ 35 36 // (internal use only) sometimes we need to allocate multiple vectors from a contiguous memory block. 37 // We stash the memory in w_dual, which has the same lifespan as this vector. See VecDuplicateVecs_SeqKokkos_GEMV. 38 PetscScalarKokkosDualView w_dual; 39 40 /* Construct Vec_Kokkos with the given array(s). n is the length of the array. 41 If n != 0, host array (array_h) must not be NULL. 42 If device array (array_d) is NULL, then a proper device mirror will be allocated. 43 Otherwise, the mirror will be created using the given array_d. 44 If both arrays are given, we assume they contain the same value (i.e., sync'ed) 45 */ Vec_KokkosVec_Kokkos46 Vec_Kokkos(PetscInt n, PetscScalar *array_h, PetscScalar *array_d = NULL) 47 { 48 PetscScalarKokkosViewHost v_h(array_h, n); 49 PetscScalarKokkosView v_d; 50 51 if (array_d) { 52 v_d = PetscScalarKokkosView(array_d, n); /* Use the given device array */ 53 } else { 54 v_d = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, DefaultMemorySpace(), v_h); /* Create a mirror in DefaultMemorySpace but do not copy values */ 55 } 56 v_dual = PetscScalarKokkosDualView(v_d, v_h); 57 if (!array_d) v_dual.modify_host(); 58 } 59 60 // Construct Vec_Kokkos with the given DualView. Use the sync state as is. With reference counting, Kokkos manages its lifespan. Vec_KokkosVec_Kokkos61 Vec_Kokkos(PetscScalarKokkosDualView dual) : v_dual(dual) { } 62 63 /* SFINAE: Update the object with an array in the given memory space, 64 assuming the given array contains the latest value for this vector. 65 */ 66 template <typename MemorySpace, std::enable_if_t<std::is_same<MemorySpace, HostMirrorMemorySpace>::value, bool> = true, std::enable_if_t<std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true> UpdateArrayVec_Kokkos67 PetscErrorCode UpdateArray(PetscScalar *array) 68 { 69 PetscScalarKokkosView v_d(array, v_dual.extent(0)); 70 PetscScalarKokkosViewHost v_h(array, v_dual.extent(0)); 71 72 PetscFunctionBegin; 73 /* Kokkos said they would add error-checking so that users won't accidentally pass two different Views in this case */ 74 PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_d, v_h)); 75 PetscFunctionReturn(PETSC_SUCCESS); 76 } 77 78 template <typename MemorySpace, std::enable_if_t<std::is_same<MemorySpace, HostMirrorMemorySpace>::value, bool> = true, std::enable_if_t<!std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true> UpdateArrayVec_Kokkos79 PetscErrorCode UpdateArray(PetscScalar *array) 80 { 81 PetscScalarKokkosViewHost v_h(array, v_dual.extent(0)); 82 83 PetscFunctionBegin; 84 PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_dual.view<DefaultMemorySpace>(), v_h)); 85 PetscCallCXX(v_dual.modify_host()); 86 PetscFunctionReturn(PETSC_SUCCESS); 87 } 88 89 template <typename MemorySpace, std::enable_if_t<!std::is_same<MemorySpace, HostMirrorMemorySpace>::value, bool> = true, std::enable_if_t<std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true> UpdateArrayVec_Kokkos90 PetscErrorCode UpdateArray(PetscScalar *array) 91 { 92 PetscScalarKokkosView v_d(array, v_dual.extent(0)); 93 94 PetscFunctionBegin; 95 PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_d, v_dual.view_host())); 96 PetscCallCXX(v_dual.modify_device()); 97 PetscFunctionReturn(PETSC_SUCCESS); 98 } 99 SetUpCOOVec_Kokkos100 PetscErrorCode SetUpCOO(const Vec_Seq *vecseq, PetscInt m) 101 { 102 PetscFunctionBegin; 103 PetscCallCXX(jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->jmap1, m + 1))); 104 PetscCallCXX(perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->perm1, vecseq->tot1))); 105 PetscFunctionReturn(PETSC_SUCCESS); 106 } 107 SetUpCOOVec_Kokkos108 PetscErrorCode SetUpCOO(const Vec_MPI *vecmpi, PetscInt m) 109 { 110 PetscFunctionBegin; 111 PetscCallCXX(jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap1, m + 1))); 112 PetscCallCXX(perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm1, vecmpi->tot1))); 113 PetscCallCXX(imap2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->imap2, vecmpi->nnz2))); 114 PetscCallCXX(jmap2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap2, vecmpi->nnz2 + 1))); 115 PetscCallCXX(perm2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm2, vecmpi->recvlen))); 116 PetscCallCXX(Cperm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->Cperm, vecmpi->sendlen))); 117 PetscCallCXX(sendbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->sendbuf, vecmpi->sendlen))); 118 PetscCallCXX(recvbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->recvbuf, vecmpi->recvlen))); 119 PetscFunctionReturn(PETSC_SUCCESS); 120 } 121 }; 122 123 PETSC_INTERN PetscErrorCode VecAbs_SeqKokkos(Vec); 124 PETSC_INTERN PetscErrorCode VecReciprocal_SeqKokkos(Vec); 125 PETSC_INTERN PetscErrorCode VecDotNorm2_SeqKokkos(Vec, Vec, PetscScalar *, PetscScalar *); 126 PETSC_INTERN PetscErrorCode VecPointwiseDivide_SeqKokkos(Vec, Vec, Vec); 127 PETSC_INTERN PetscErrorCode VecWAXPY_SeqKokkos(Vec, PetscScalar, Vec, Vec); 128 PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *); 129 PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *); 130 PETSC_INTERN PetscErrorCode VecSet_SeqKokkos(Vec, PetscScalar); 131 PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos(Vec, PetscInt, const PetscScalar *, Vec *); 132 PETSC_INTERN PetscErrorCode VecAXPBYPCZ_SeqKokkos(Vec, PetscScalar, PetscScalar, PetscScalar, Vec, Vec); 133 PETSC_INTERN PetscErrorCode VecPointwiseMult_SeqKokkos(Vec, Vec, Vec); 134 PETSC_INTERN PetscErrorCode VecPlaceArray_SeqKokkos(Vec, const PetscScalar *); 135 PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos(Vec); 136 PETSC_INTERN PetscErrorCode VecReplaceArray_SeqKokkos(Vec, const PetscScalar *); 137 PETSC_INTERN PetscErrorCode VecDot_SeqKokkos(Vec, Vec, PetscScalar *); 138 PETSC_INTERN PetscErrorCode VecTDot_SeqKokkos(Vec, Vec, PetscScalar *); 139 PETSC_INTERN PetscErrorCode VecScale_SeqKokkos(Vec, PetscScalar); 140 PETSC_INTERN PetscErrorCode VecCopy_SeqKokkos(Vec, Vec); 141 PETSC_INTERN PetscErrorCode VecSwap_SeqKokkos(Vec, Vec); 142 PETSC_INTERN PetscErrorCode VecAXPY_SeqKokkos(Vec, PetscScalar, Vec); 143 PETSC_INTERN PetscErrorCode VecAXPBY_SeqKokkos(Vec, PetscScalar, PetscScalar, Vec); 144 PETSC_INTERN PetscErrorCode VecConjugate_SeqKokkos(Vec xin); 145 PETSC_INTERN PetscErrorCode VecNorm_SeqKokkos(Vec, NormType, PetscReal *); 146 PETSC_INTERN PetscErrorCode VecErrorWeightedNorms_SeqKokkos(Vec, Vec, Vec, NormType, PetscReal, Vec, PetscReal, Vec, PetscReal, PetscReal *, PetscInt *, PetscReal *, PetscInt *, PetscReal *, PetscInt *); 147 PETSC_INTERN PetscErrorCode VecCreate_SeqKokkos(Vec); 148 PETSC_INTERN PetscErrorCode VecCreate_SeqKokkos_Private(Vec, const PetscScalar *); 149 PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos(Vec); 150 PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos_Private(Vec, PetscBool, PetscInt, const PetscScalar *); 151 PETSC_INTERN PetscErrorCode VecCreate_Kokkos(Vec); 152 PETSC_INTERN PetscErrorCode VecAYPX_SeqKokkos(Vec, PetscScalar, Vec); 153 PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos(Vec, PetscRandom); 154 PETSC_INTERN PetscErrorCode VecGetLocalVector_SeqKokkos(Vec, Vec); 155 PETSC_INTERN PetscErrorCode VecRestoreLocalVector_SeqKokkos(Vec, Vec); 156 PETSC_INTERN PetscErrorCode VecGetArrayWrite_SeqKokkos(Vec, PetscScalar **); 157 PETSC_INTERN PetscErrorCode VecCopy_SeqKokkos_Private(Vec, Vec); 158 PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos_Private(Vec, PetscRandom); 159 PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos_Private(Vec); 160 PETSC_INTERN PetscErrorCode VecMin_SeqKokkos(Vec, PetscInt *, PetscReal *); 161 PETSC_INTERN PetscErrorCode VecMax_SeqKokkos(Vec, PetscInt *, PetscReal *); 162 PETSC_INTERN PetscErrorCode VecSum_SeqKokkos(Vec, PetscScalar *); 163 PETSC_INTERN PetscErrorCode VecShift_SeqKokkos(Vec, PetscScalar); 164 PETSC_INTERN PetscErrorCode VecGetArray_SeqKokkos(Vec, PetscScalar **); 165 PETSC_INTERN PetscErrorCode VecRestoreArray_SeqKokkos(Vec, PetscScalar **); 166 167 PETSC_INTERN PetscErrorCode VecGetArrayAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *); 168 PETSC_INTERN PetscErrorCode VecRestoreArrayAndMemType_SeqKokkos(Vec, PetscScalar **); 169 PETSC_INTERN PetscErrorCode VecGetArrayWriteAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *); 170 PETSC_INTERN PetscErrorCode VecGetSubVector_Kokkos_Private(Vec, PetscBool, IS, Vec *); 171 PETSC_INTERN PetscErrorCode VecRestoreSubVector_SeqKokkos(Vec, IS, Vec *); 172 173 PETSC_INTERN PetscErrorCode VecDuplicateVecs_Kokkos_GEMV(Vec, PetscInt, Vec **); 174 PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos_GEMV(Vec, PetscInt, const Vec *, PetscScalar *); 175 PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos_GEMV(Vec, PetscInt, const Vec *, PetscScalar *); 176 PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos_GEMV(Vec, PetscInt, const PetscScalar *, Vec *); 177 178 PETSC_INTERN PetscErrorCode VecCreateMPIKokkosWithLayoutAndArrays_Private(PetscLayout map, const PetscScalar *, const PetscScalar *, Vec *); 179