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