Actual source code: veckokkosimpl.hpp

  1: #ifndef __VECKOKKOSIMPL_HPP

  4: #include <../src/vec/vec/impls/mpi/pvecimpl.h>
  5: #include <petsc/private/vecimpl_kokkos.hpp>

  7: #if defined(PETSC_USE_DEBUG)
  8:   #define VecErrorIfNotKokkos(v) \
  9:     do { \
 10:       PetscBool isKokkos = PETSC_FALSE; \
 11:       PetscObjectTypeCompareAny((PetscObject)(v), &isKokkos, VECSEQKOKKOS, VECMPIKOKKOS, VECKOKKOS, ""); \
 13:     } while (0)
 14: #else
 15:   #define VecErrorIfNotKokkos(v) \
 16:     do { \
 17:       (void)(v); \
 18:     } while (0)
 19: #endif

 21: /* Stuff related to Vec_Kokkos */

 23: struct Vec_Kokkos {
 24:   PetscScalarKokkosDualView v_dual;

 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 */

 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() */

 36:   /* Construct Vec_Kokkos with the given array(s). n is the length of the array.
 37:     If n != 0, host array (array_h) must not be NULL.
 38:     If device array (array_d) is NULL, then a proper device mirror will be allocated.
 39:     Otherwise, the mirror will be created using the given array_d.
 40:   */
 41:   Vec_Kokkos(PetscInt n, PetscScalar *array_h, PetscScalar *array_d = NULL)
 42:   {
 43:     PetscScalarKokkosViewHost v_h(array_h, n);
 44:     PetscScalarKokkosView     v_d;

 46:     if (array_d) {
 47:       v_d = PetscScalarKokkosView(array_d, n); /* Use the given device array */
 48:     } else {
 49:       v_d = Kokkos::create_mirror_view(DefaultMemorySpace(), v_h); /* Create a mirror in DefaultMemorySpace but do not copy values */
 50:     }
 51:     v_dual = PetscScalarKokkosDualView(v_d, v_h);
 52:     if (!array_d) v_dual.modify_host();
 53:   }

 55:   /* SFINAE: Update the object with an array in the given memory space,
 56:      assuming the given array contains the latest value for this vector.
 57:    */
 58:   template <typename MemorySpace, std::enable_if_t<std::is_same<MemorySpace, Kokkos::HostSpace>::value, bool> = true, std::enable_if_t<std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true>
 59:   void UpdateArray(PetscScalar *array)
 60:   {
 61:     PetscScalarKokkosViewHost v_h(array, v_dual.extent(0));
 62:     /* Kokkos said they would add error-checking so that users won't accidentally pass two different Views in this case */
 63:     v_dual = PetscScalarKokkosDualView(v_h, v_h);
 64:   }

 66:   template <typename MemorySpace, std::enable_if_t<std::is_same<MemorySpace, Kokkos::HostSpace>::value, bool> = true, std::enable_if_t<!std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true>
 67:   void UpdateArray(PetscScalar *array)
 68:   {
 69:     PetscScalarKokkosViewHost v_h(array, v_dual.extent(0));
 70:     v_dual = PetscScalarKokkosDualView(v_dual.view<DefaultMemorySpace>(), v_h);
 71:     v_dual.modify_host();
 72:   }

 74:   template <typename MemorySpace, std::enable_if_t<!std::is_same<MemorySpace, Kokkos::HostSpace>::value, bool> = true, std::enable_if_t<std::is_same<MemorySpace, DefaultMemorySpace>::value, bool> = true>
 75:   void UpdateArray(PetscScalar *array)
 76:   {
 77:     PetscScalarKokkosView v_d(array, v_dual.extent(0));
 78:     v_dual = PetscScalarKokkosDualView(v_d, v_dual.view<Kokkos::HostSpace>());
 79:     v_dual.modify_device();
 80:   }

 82:   void SetUpCOO(const Vec_Seq *vecseq, PetscInt m)
 83:   {
 84:     jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->jmap1, m + 1));
 85:     perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->perm1, vecseq->tot1));
 86:   }

 88:   void SetUpCOO(const Vec_MPI *vecmpi, PetscInt m)
 89:   {
 90:     jmap1_d   = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap1, m + 1));
 91:     perm1_d   = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm1, vecmpi->tot1));
 92:     imap2_d   = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->imap2, vecmpi->nnz2));
 93:     jmap2_d   = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap2, vecmpi->nnz2 + 1));
 94:     perm2_d   = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm2, vecmpi->recvlen));
 95:     Cperm_d   = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->Cperm, vecmpi->sendlen));
 96:     sendbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->sendbuf, vecmpi->sendlen));
 97:     recvbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->recvbuf, vecmpi->recvlen));
 98:   }
 99: };

101: PETSC_INTERN PetscErrorCode VecAbs_SeqKokkos(Vec);
102: PETSC_INTERN PetscErrorCode VecReciprocal_SeqKokkos(Vec);
103: PETSC_INTERN PetscErrorCode VecDotNorm2_SeqKokkos(Vec, Vec, PetscScalar *, PetscScalar *);
104: PETSC_INTERN PetscErrorCode VecPointwiseDivide_SeqKokkos(Vec, Vec, Vec);
105: PETSC_INTERN PetscErrorCode VecWAXPY_SeqKokkos(Vec, PetscScalar, Vec, Vec);
106: PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *);
107: PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *);
108: PETSC_EXTERN PetscErrorCode VecSet_SeqKokkos(Vec, PetscScalar);
109: PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos(Vec, PetscInt, const PetscScalar *, Vec *);
110: PETSC_INTERN PetscErrorCode VecAXPBYPCZ_SeqKokkos(Vec, PetscScalar, PetscScalar, PetscScalar, Vec, Vec);
111: PETSC_INTERN PetscErrorCode VecPointwiseMult_SeqKokkos(Vec, Vec, Vec);
112: PETSC_INTERN PetscErrorCode VecPlaceArray_SeqKokkos(Vec, const PetscScalar *);
113: PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos(Vec);
114: PETSC_INTERN PetscErrorCode VecReplaceArray_SeqKokkos(Vec, const PetscScalar *);
115: PETSC_INTERN PetscErrorCode VecDot_SeqKokkos(Vec, Vec, PetscScalar *);
116: PETSC_INTERN PetscErrorCode VecTDot_SeqKokkos(Vec, Vec, PetscScalar *);
117: PETSC_INTERN PetscErrorCode VecScale_SeqKokkos(Vec, PetscScalar);
118: PETSC_EXTERN PetscErrorCode VecCopy_SeqKokkos(Vec, Vec);
119: PETSC_INTERN PetscErrorCode VecSwap_SeqKokkos(Vec, Vec);
120: PETSC_EXTERN PetscErrorCode VecAXPY_SeqKokkos(Vec, PetscScalar, Vec);
121: PETSC_INTERN PetscErrorCode VecAXPBY_SeqKokkos(Vec, PetscScalar, PetscScalar, Vec);
122: PETSC_INTERN PetscErrorCode VecDuplicate_SeqKokkos(Vec, Vec *);
123: PETSC_INTERN PetscErrorCode VecConjugate_SeqKokkos(Vec xin);
124: PETSC_INTERN PetscErrorCode VecNorm_SeqKokkos(Vec, NormType, PetscReal *);
125: PETSC_EXTERN PetscErrorCode VecCreate_SeqKokkos(Vec);
126: PETSC_INTERN PetscErrorCode VecCreate_SeqKokkos_Private(Vec, const PetscScalar *);
127: PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos(Vec);
128: PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos_Private(Vec, PetscBool, PetscInt, const PetscScalar *);
129: PETSC_INTERN PetscErrorCode VecCreate_Kokkos(Vec);
130: PETSC_INTERN PetscErrorCode VecDestroy_SeqKokkos(Vec);
131: PETSC_INTERN PetscErrorCode VecDestroy_MPIKokkos(Vec);
132: PETSC_INTERN PetscErrorCode VecAYPX_SeqKokkos(Vec, PetscScalar, Vec);
133: PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos(Vec, PetscRandom);
134: PETSC_INTERN PetscErrorCode VecGetLocalVector_SeqKokkos(Vec, Vec);
135: PETSC_INTERN PetscErrorCode VecRestoreLocalVector_SeqKokkos(Vec, Vec);
136: PETSC_INTERN PetscErrorCode VecGetArrayWrite_SeqKokkos(Vec, PetscScalar **);
137: PETSC_INTERN PetscErrorCode VecCopy_SeqKokkos_Private(Vec, Vec);
138: PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos_Private(Vec, PetscRandom);
139: PETSC_INTERN PetscErrorCode VecDestroy_SeqKokkos_Private(Vec);
140: PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos_Private(Vec);
141: PETSC_INTERN PetscErrorCode VecMin_SeqKokkos(Vec, PetscInt *, PetscReal *);
142: PETSC_INTERN PetscErrorCode VecMax_SeqKokkos(Vec, PetscInt *, PetscReal *);
143: PETSC_INTERN PetscErrorCode VecSum_SeqKokkos(Vec, PetscScalar *);
144: PETSC_INTERN PetscErrorCode VecShift_SeqKokkos(Vec, PetscScalar);
145: PETSC_INTERN PetscErrorCode VecGetArray_SeqKokkos(Vec, PetscScalar **);
146: PETSC_INTERN PetscErrorCode VecRestoreArray_SeqKokkos(Vec, PetscScalar **);

148: PETSC_INTERN PetscErrorCode VecGetArrayAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *);
149: PETSC_INTERN PetscErrorCode VecRestoreArrayAndMemType_SeqKokkos(Vec, PetscScalar **);
150: PETSC_INTERN PetscErrorCode VecGetArrayWriteAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *);
151: PETSC_INTERN PetscErrorCode VecGetSubVector_Kokkos_Private(Vec, PetscBool, IS, Vec *);
152: PETSC_INTERN PetscErrorCode VecRestoreSubVector_SeqKokkos(Vec, IS, Vec *);
153: #endif