Actual source code: veckokkosimpl.hpp
1: #pragma once
3: #include <../src/vec/vec/impls/mpi/pvecimpl.h>
4: #include <petsc/private/kokkosimpl.hpp>
6: #if PetscDefined(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
20: /* Stuff related to Vec_Kokkos */
22: struct Vec_Kokkos {
23: PetscScalarKokkosDualView v_dual;
24: PetscScalarKokkosView unplaced_d; // Unplaced device array in VecKokkosPlaceArray()
25: PetscScalarKokkosDualView unplaced_dual; // Unplaced v_dual, used in VecGetLocalVector_SeqKokkos()
27: /* COO stuff */
28: PetscCountKokkosView jmap1_d; /* [m+1]: i-th entry of the vector has jmap1[i+1]-jmap1[i] repeats in COO arrays */
29: PetscCountKokkosView perm1_d; /* [tot1]: permutation array for local entries */
31: PetscCountKokkosView imap2_d; /* [nnz2]: i-th unique entry in recvbuf is imap2[i]-th entry in the vector */
32: PetscCountKokkosView jmap2_d; /* [nnz2+1] */
33: PetscCountKokkosView perm2_d; /* [recvlen] */
34: PetscCountKokkosView Cperm_d; /* [sendlen]: permutation array to fill sendbuf[]. 'C' for communication */
35: PetscScalarKokkosView sendbuf_d, recvbuf_d; /* Buffers for remote values in VecSetValuesCOO() */
37: // (internal use only) sometimes we need to allocate multiple vectors from a contiguous memory block.
38: // We stash the memory in w_dual, which has the same lifespan as this vector. See VecDuplicateVecs_SeqKokkos_GEMV.
39: PetscScalarKokkosDualView w_dual;
41: /* Construct Vec_Kokkos with the given array(s). n is the length of the array.
42: If device array (array_d) is NULL but host array (array_h) is not, 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: If both arrays are NULL, the vector is array-less: no memory is allocated, the views wrap NULL with extent n and
46: nothing is marked as modified, until an array is placed with VecPlaceArray() or VecKokkosPlaceArray(), see UpdateArray()
47: */
48: Vec_Kokkos(PetscInt n, PetscScalar *array_h, PetscScalar *array_d = NULL)
49: {
50: PetscScalarKokkosViewHost v_h(array_h, n);
51: PetscScalarKokkosView v_d;
53: if (array_d || !array_h) {
54: v_d = PetscScalarKokkosView(array_d, n); /* Use the given device array, or none for an array-less vector */
55: } else {
56: v_d = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, DefaultMemorySpace(), v_h); /* Create a mirror in DefaultMemorySpace but do not copy values */
57: }
58: v_dual = PetscScalarKokkosDualView(v_d, v_h);
59: if (array_h && !array_d) v_dual.modify_host();
60: }
62: // Construct Vec_Kokkos with the given DualView. Use the sync state as is. With reference counting, Kokkos manages its lifespan.
63: Vec_Kokkos(PetscScalarKokkosDualView dual) : v_dual(dual) { }
65: /* SFINAE: Update the object with an array in the given memory space,
66: assuming the given array contains the latest value for this vector.
67: A NULL array takes the vector back to the array-less state in that memory space (see the constructor),
68: so nothing is marked as modified, as a view wrapping NULL can never be the source or the target of a sync.
69: */
70: 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>
71: PetscErrorCode UpdateArray(PetscScalar *array)
72: {
73: PetscScalarKokkosView v_d(array, v_dual.extent(0));
74: PetscScalarKokkosViewHost v_h(array, v_dual.extent(0));
76: PetscFunctionBegin;
77: /* Kokkos said they would add error-checking so that users won't accidentally pass two different Views in this case */
78: PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_d, v_h));
79: PetscFunctionReturn(PETSC_SUCCESS);
80: }
82: 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>
83: PetscErrorCode UpdateArray(PetscScalar *array)
84: {
85: PetscScalarKokkosViewHost v_h(array, v_dual.extent(0));
86: PetscScalarKokkosView v_d = v_dual.view<DefaultMemorySpace>();
88: PetscFunctionBegin;
89: /* An array-less vector gets its device mirror with the first host array placed in it */
90: if (array && !v_d.data()) PetscCallCXX(v_d = Kokkos::create_mirror_view(Kokkos::WithoutInitializing, DefaultMemorySpace(), v_h));
91: PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_d, v_h));
92: if (array) PetscCallCXX(v_dual.modify_host());
93: PetscFunctionReturn(PETSC_SUCCESS);
94: }
96: 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>
97: PetscErrorCode UpdateArray(PetscScalar *array)
98: {
99: PetscScalarKokkosView v_d(array, v_dual.extent(0));
101: PetscFunctionBegin;
102: PetscCallCXX(v_dual = PetscScalarKokkosDualView(v_d, v_dual.view_host()));
103: if (array) PetscCallCXX(v_dual.modify_device());
104: PetscFunctionReturn(PETSC_SUCCESS);
105: }
107: PetscErrorCode SetUpCOO(const Vec_Seq *vecseq, PetscInt m)
108: {
109: PetscFunctionBegin;
110: PetscCallCXX(jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->jmap1, m + 1)));
111: PetscCallCXX(perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecseq->perm1, vecseq->tot1)));
112: PetscFunctionReturn(PETSC_SUCCESS);
113: }
115: PetscErrorCode SetUpCOO(const Vec_MPI *vecmpi, PetscInt m)
116: {
117: PetscFunctionBegin;
118: PetscCallCXX(jmap1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap1, m + 1)));
119: PetscCallCXX(perm1_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm1, vecmpi->tot1)));
120: PetscCallCXX(imap2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->imap2, vecmpi->nnz2)));
121: PetscCallCXX(jmap2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->jmap2, vecmpi->nnz2 + 1)));
122: PetscCallCXX(perm2_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->perm2, vecmpi->recvlen)));
123: PetscCallCXX(Cperm_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscCountKokkosViewHost(vecmpi->Cperm, vecmpi->sendlen)));
124: PetscCallCXX(sendbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->sendbuf, vecmpi->sendlen)));
125: PetscCallCXX(recvbuf_d = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(vecmpi->recvbuf, vecmpi->recvlen)));
126: PetscFunctionReturn(PETSC_SUCCESS);
127: }
128: };
130: PETSC_INTERN PetscErrorCode VecAbs_SeqKokkos(Vec);
131: PETSC_INTERN PetscErrorCode VecReciprocal_SeqKokkos(Vec);
132: PETSC_INTERN PetscErrorCode VecDotNorm2_SeqKokkos(Vec, Vec, PetscScalar *, PetscScalar *);
133: PETSC_INTERN PetscErrorCode VecPointwiseDivide_SeqKokkos(Vec, Vec, Vec);
134: PETSC_INTERN PetscErrorCode VecWAXPY_SeqKokkos(Vec, PetscScalar, Vec, Vec);
135: PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *);
136: PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos(Vec, PetscInt, const Vec[], PetscScalar *);
137: PETSC_INTERN PetscErrorCode VecSet_SeqKokkos(Vec, PetscScalar);
138: PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos(Vec, PetscInt, const PetscScalar *, Vec *);
139: PETSC_INTERN PetscErrorCode VecAXPBYPCZ_SeqKokkos(Vec, PetscScalar, PetscScalar, PetscScalar, Vec, Vec);
140: PETSC_INTERN PetscErrorCode VecPointwiseMult_SeqKokkos(Vec, Vec, Vec);
141: PETSC_INTERN PetscErrorCode VecPlaceArray_SeqKokkos(Vec, const PetscScalar *);
142: PETSC_INTERN PetscErrorCode VecResetArray_SeqKokkos(Vec);
143: PETSC_INTERN PetscErrorCode VecReplaceArray_SeqKokkos(Vec, const PetscScalar *);
144: PETSC_INTERN PetscErrorCode VecDot_SeqKokkos(Vec, Vec, PetscScalar *);
145: PETSC_INTERN PetscErrorCode VecTDot_SeqKokkos(Vec, Vec, PetscScalar *);
146: PETSC_INTERN PetscErrorCode VecScale_SeqKokkos(Vec, PetscScalar);
147: PETSC_INTERN PetscErrorCode VecCopy_SeqKokkos(Vec, Vec);
148: PETSC_INTERN PetscErrorCode VecSwap_SeqKokkos(Vec, Vec);
149: PETSC_INTERN PetscErrorCode VecAXPY_SeqKokkos(Vec, PetscScalar, Vec);
150: PETSC_INTERN PetscErrorCode VecAXPBY_SeqKokkos(Vec, PetscScalar, PetscScalar, Vec);
151: PETSC_INTERN PetscErrorCode VecConjugate_SeqKokkos(Vec xin);
152: PETSC_INTERN PetscErrorCode VecNorm_SeqKokkos(Vec, NormType, PetscReal *);
153: PETSC_INTERN PetscErrorCode VecErrorWeightedNorms_SeqKokkos(Vec, Vec, Vec, NormType, PetscReal, Vec, PetscReal, Vec, PetscReal, PetscReal *, PetscInt *, PetscReal *, PetscInt *, PetscReal *, PetscInt *);
154: PETSC_INTERN PetscErrorCode VecCreate_SeqKokkos(Vec);
155: PETSC_INTERN PetscErrorCode VecCreate_MPIKokkos(Vec);
156: PETSC_INTERN PetscErrorCode VecCreate_Kokkos(Vec);
157: PETSC_INTERN PetscErrorCode VecAYPX_SeqKokkos(Vec, PetscScalar, Vec);
158: PETSC_INTERN PetscErrorCode VecSetRandom_SeqKokkos(Vec, PetscRandom);
159: PETSC_INTERN PetscErrorCode VecGetLocalVector_SeqKokkos(Vec, Vec);
160: PETSC_INTERN PetscErrorCode VecGetLocalVectorRead_SeqKokkos(Vec, Vec);
161: PETSC_INTERN PetscErrorCode VecRestoreLocalVector_SeqKokkos(Vec, Vec);
162: PETSC_INTERN PetscErrorCode VecRestoreLocalVectorRead_SeqKokkos(Vec, Vec);
163: PETSC_INTERN PetscErrorCode VecGetArrayWrite_SeqKokkos(Vec, PetscScalar **);
164: PETSC_INTERN PetscErrorCode VecMin_SeqKokkos(Vec, PetscInt *, PetscReal *);
165: PETSC_INTERN PetscErrorCode VecMax_SeqKokkos(Vec, PetscInt *, PetscReal *);
166: PETSC_INTERN PetscErrorCode VecSum_SeqKokkos(Vec, PetscScalar *);
167: PETSC_INTERN PetscErrorCode VecShift_SeqKokkos(Vec, PetscScalar);
168: PETSC_INTERN PetscErrorCode VecGetArray_SeqKokkos(Vec, PetscScalar **);
169: PETSC_INTERN PetscErrorCode VecRestoreArray_SeqKokkos(Vec, PetscScalar **);
171: PETSC_INTERN PetscErrorCode VecGetArrayAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *);
172: PETSC_INTERN PetscErrorCode VecRestoreArrayAndMemType_SeqKokkos(Vec, PetscScalar **);
173: PETSC_INTERN PetscErrorCode VecGetArrayWriteAndMemType_SeqKokkos(Vec, PetscScalar **, PetscMemType *);
174: PETSC_INTERN PetscErrorCode VecGetSubVector_Kokkos_Private(Vec, PetscBool, IS, Vec *);
175: PETSC_INTERN PetscErrorCode VecRestoreSubVector_SeqKokkos(Vec, IS, Vec *);
177: PETSC_INTERN PetscErrorCode VecMDot_SeqKokkos_GEMV(Vec, PetscInt, const Vec *, PetscScalar *);
178: PETSC_INTERN PetscErrorCode VecMTDot_SeqKokkos_GEMV(Vec, PetscInt, const Vec *, PetscScalar *);
179: PETSC_INTERN PetscErrorCode VecMAXPY_SeqKokkos_GEMV(Vec, PetscInt, const PetscScalar *, Vec *);
181: PETSC_INTERN PetscErrorCode VecCreateMPIKokkosWithLayoutAndArrays_Private(PetscLayout map, const PetscScalar *, const PetscScalar *, Vec *);
182: PETSC_INTERN PetscErrorCode VecBindToCPU_SeqKokkos(Vec, PetscBool);