Actual source code: mpikok.kokkos.cxx
1: /*
2: This file contains routines for Parallel vector operations.
3: */
4: #include <petsc_kokkos.hpp>
5: #include <petscvec_kokkos.hpp>
6: #include <petsc/private/deviceimpl.h>
7: #include <petsc/private/vecimpl.h>
8: #include <../src/vec/vec/impls/mpi/pvecimpl.h>
9: #include <../src/vec/vec/impls/seq/kokkos/veckokkosimpl.hpp>
10: #include <petscsf.h>
12: static PetscErrorCode VecDestroy_MPIKokkos(Vec v)
13: {
14: PetscFunctionBegin;
15: delete static_cast<Vec_Kokkos *>(v->spptr);
16: v->spptr = NULL;
17: PetscCall(VecDestroy_MPI(v));
18: PetscFunctionReturn(PETSC_SUCCESS);
19: }
21: static PetscErrorCode VecNorm_MPIKokkos(Vec xin, NormType type, PetscReal *z)
22: {
23: PetscFunctionBegin;
24: PetscCall(VecNorm_MPI_Default(xin, type, z, VecNorm_SeqKokkos));
25: PetscFunctionReturn(PETSC_SUCCESS);
26: }
28: static PetscErrorCode VecErrorWeightedNorms_MPIKokkos(Vec U, Vec Y, Vec E, NormType wnormtype, PetscReal atol, Vec vatol, PetscReal rtol, Vec vrtol, PetscReal ignore_max, PetscReal *norm, PetscInt *norm_loc, PetscReal *norma, PetscInt *norma_loc, PetscReal *normr, PetscInt *normr_loc)
29: {
30: PetscFunctionBegin;
31: PetscCall(VecErrorWeightedNorms_MPI_Default(U, Y, E, wnormtype, atol, vatol, rtol, vrtol, ignore_max, norm, norm_loc, norma, norma_loc, normr, normr_loc, VecErrorWeightedNorms_SeqKokkos));
32: PetscFunctionReturn(PETSC_SUCCESS);
33: }
35: /* z = y^H x */
36: static PetscErrorCode VecDot_MPIKokkos(Vec xin, Vec yin, PetscScalar *z)
37: {
38: PetscFunctionBegin;
39: PetscCall(VecXDot_MPI_Default(xin, yin, z, VecDot_SeqKokkos));
40: PetscFunctionReturn(PETSC_SUCCESS);
41: }
43: /* z = y^T x */
44: static PetscErrorCode VecTDot_MPIKokkos(Vec xin, Vec yin, PetscScalar *z)
45: {
46: PetscFunctionBegin;
47: PetscCall(VecXDot_MPI_Default(xin, yin, z, VecTDot_SeqKokkos));
48: PetscFunctionReturn(PETSC_SUCCESS);
49: }
51: static PetscErrorCode VecMDot_MPIKokkos(Vec xin, PetscInt nv, const Vec y[], PetscScalar *z)
52: {
53: PetscFunctionBegin;
54: PetscCall(VecMXDot_MPI_Default(xin, nv, y, z, VecMDot_SeqKokkos));
55: PetscFunctionReturn(PETSC_SUCCESS);
56: }
58: static PetscErrorCode VecMTDot_MPIKokkos(Vec xin, PetscInt nv, const Vec y[], PetscScalar *z)
59: {
60: PetscFunctionBegin;
61: PetscCall(VecMXDot_MPI_Default(xin, nv, y, z, VecMTDot_SeqKokkos));
62: PetscFunctionReturn(PETSC_SUCCESS);
63: }
65: static PetscErrorCode VecMDot_MPIKokkos_GEMV(Vec xin, PetscInt nv, const Vec y[], PetscScalar *z)
66: {
67: PetscFunctionBegin;
68: PetscCall(VecMXDot_MPI_Default(xin, nv, y, z, VecMDot_SeqKokkos_GEMV));
69: PetscFunctionReturn(PETSC_SUCCESS);
70: }
72: static PetscErrorCode VecMTDot_MPIKokkos_GEMV(Vec xin, PetscInt nv, const Vec y[], PetscScalar *z)
73: {
74: PetscFunctionBegin;
75: PetscCall(VecMXDot_MPI_Default(xin, nv, y, z, VecMTDot_SeqKokkos_GEMV));
76: PetscFunctionReturn(PETSC_SUCCESS);
77: }
79: static PetscErrorCode VecMax_MPIKokkos(Vec xin, PetscInt *idx, PetscReal *z)
80: {
81: const MPI_Op ops[] = {MPIU_MAXLOC, MPIU_MAX};
83: PetscFunctionBegin;
84: PetscCall(VecMinMax_MPI_Default(xin, idx, z, VecMax_SeqKokkos, ops));
85: PetscFunctionReturn(PETSC_SUCCESS);
86: }
88: static PetscErrorCode VecMin_MPIKokkos(Vec xin, PetscInt *idx, PetscReal *z)
89: {
90: const MPI_Op ops[] = {MPIU_MINLOC, MPIU_MIN};
92: PetscFunctionBegin;
93: PetscCall(VecMinMax_MPI_Default(xin, idx, z, VecMin_SeqKokkos, ops));
94: PetscFunctionReturn(PETSC_SUCCESS);
95: }
97: static PetscErrorCode VecCreate_MPIKokkos_Common(Vec); // forward declaration
99: static PetscErrorCode VecDuplicate_MPIKokkos(Vec win, Vec *vv)
100: {
101: Vec v;
102: Vec_Kokkos *veckok;
103: Vec_MPI *wdata = (Vec_MPI *)win->data, *vdata;
104: PetscScalarKokkosDualView w_dual;
106: PetscFunctionBegin;
107: PetscCallCXX(w_dual = PetscScalarKokkosDualView("w_dual", win->map->n + wdata->nghost)); // Kokkos init's v_dual to zero
109: /* Reuse VecDuplicate_MPI, which contains a lot of stuff */
110: PetscCall(VecDuplicateWithArray_MPI(win, w_dual.view_host().data(), &v)); /* after the call, v is a VECMPI */
111: PetscCall(PetscObjectChangeTypeName((PetscObject)v, VECMPIKOKKOS));
112: // In case win is a ghost vector, we also need to convert its localrep to VECKOKKOS. We provide the device array, so allocation in w_dual is not wasted
113: vdata = static_cast<Vec_MPI *>(v->data);
114: if (vdata->localrep) PetscCall(VecConvert_Seq_SeqKokkos_inplace(vdata->localrep, w_dual.view_device().data()));
115: PetscCall(VecCreate_MPIKokkos_Common(v));
116: v->ops[0] = win->ops[0]; // always follow ops[] in win
118: /* Build the Vec_Kokkos struct */
119: veckok = new Vec_Kokkos(v->map->n, w_dual.view_host().data(), w_dual.view_device().data());
120: veckok->w_dual = w_dual;
121: v->spptr = veckok;
122: *vv = v;
123: PetscFunctionReturn(PETSC_SUCCESS);
124: }
126: static PetscErrorCode VecDotNorm2_MPIKokkos(Vec s, Vec t, PetscScalar *dp, PetscScalar *nm)
127: {
128: PetscFunctionBegin;
129: PetscCall(VecDotNorm2_MPI_Default(s, t, dp, nm, VecDotNorm2_SeqKokkos));
130: PetscFunctionReturn(PETSC_SUCCESS);
131: }
133: static PetscErrorCode VecGetSubVector_MPIKokkos(Vec x, IS is, Vec *y)
134: {
135: PetscFunctionBegin;
136: PetscCall(VecGetSubVector_Kokkos_Private(x, PETSC_TRUE, is, y));
137: PetscFunctionReturn(PETSC_SUCCESS);
138: }
140: static PetscErrorCode VecSetPreallocationCOO_MPIKokkos(Vec x, PetscCount ncoo, const PetscInt coo_i[])
141: {
142: const auto vecmpi = static_cast<Vec_MPI *>(x->data);
143: const auto veckok = static_cast<Vec_Kokkos *>(x->spptr);
144: PetscInt m;
146: PetscFunctionBegin;
147: PetscCall(VecGetLocalSize(x, &m));
148: PetscCall(VecSetPreallocationCOO_MPI(x, ncoo, coo_i));
149: PetscCall(veckok->SetUpCOO(vecmpi, m));
150: PetscFunctionReturn(PETSC_SUCCESS);
151: }
153: static PetscErrorCode VecSetValuesCOO_MPIKokkos(Vec x, const PetscScalar v[], InsertMode imode)
154: {
155: const auto vecmpi = static_cast<Vec_MPI *>(x->data);
156: const auto veckok = static_cast<Vec_Kokkos *>(x->spptr);
157: const PetscCountKokkosView &jmap1 = veckok->jmap1_d;
158: const PetscCountKokkosView &perm1 = veckok->perm1_d;
159: const PetscCountKokkosView &imap2 = veckok->imap2_d;
160: const PetscCountKokkosView &jmap2 = veckok->jmap2_d;
161: const PetscCountKokkosView &perm2 = veckok->perm2_d;
162: const PetscCountKokkosView &Cperm = veckok->Cperm_d;
163: PetscScalarKokkosView &sendbuf = veckok->sendbuf_d;
164: PetscScalarKokkosView &recvbuf = veckok->recvbuf_d;
165: PetscScalarKokkosView xv;
166: ConstPetscScalarKokkosView vv;
167: PetscMemType memtype;
168: PetscInt m;
170: PetscFunctionBegin;
171: PetscCall(VecGetLocalSize(x, &m));
172: PetscCall(PetscGetMemType(v, &memtype));
173: if (PetscMemTypeHost(memtype)) { /* If user gave v[] in host, we might need to copy it to device if any */
174: vv = Kokkos::create_mirror_view_and_copy(DefaultMemorySpace(), PetscScalarKokkosViewHost(const_cast<PetscScalar *>(v), vecmpi->coo_n));
175: PetscCall(PetscLogCpuToGpu(vecmpi->coo_n * sizeof(PetscScalar)));
176: } else {
177: vv = ConstPetscScalarKokkosView(v, vecmpi->coo_n); /* Directly use v[]'s memory */
178: }
180: /* Pack entries to be sent to remote */
181: Kokkos::parallel_for(Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, vecmpi->sendlen), KOKKOS_LAMBDA(const PetscCount i) { sendbuf(i) = vv(Cperm(i)); });
182: PetscCall(PetscSFReduceWithMemTypeBegin(vecmpi->coo_sf, MPIU_SCALAR, PETSC_MEMTYPE_KOKKOS, sendbuf.data(), PETSC_MEMTYPE_KOKKOS, recvbuf.data(), MPI_REPLACE));
184: if (imode == INSERT_VALUES) PetscCall(VecGetKokkosViewWrite(x, &xv)); /* write vector */
185: else PetscCall(VecGetKokkosView(x, &xv)); /* read & write vector */
187: Kokkos::parallel_for(
188: Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, m), KOKKOS_LAMBDA(const PetscCount i) {
189: PetscScalar sum = 0.0;
190: for (PetscCount k = jmap1(i); k < jmap1(i + 1); k++) sum += vv(perm1(k));
191: xv(i) = (imode == INSERT_VALUES ? 0.0 : xv(i)) + sum;
192: });
194: PetscCall(PetscSFReduceEnd(vecmpi->coo_sf, MPIU_SCALAR, sendbuf.data(), recvbuf.data(), MPI_REPLACE));
196: /* Add received remote entries */
197: Kokkos::parallel_for(
198: Kokkos::RangePolicy<>(PetscGetKokkosExecutionSpace(), 0, vecmpi->nnz2), KOKKOS_LAMBDA(PetscCount i) {
199: for (PetscCount k = jmap2(i); k < jmap2(i + 1); k++) xv(imap2(i)) += recvbuf(perm2(k));
200: });
202: if (imode == INSERT_VALUES) PetscCall(VecRestoreKokkosViewWrite(x, &xv));
203: else PetscCall(VecRestoreKokkosView(x, &xv));
204: PetscFunctionReturn(PETSC_SUCCESS);
205: }
207: // Shared by all VecCreate/Duplicate routines for VecMPIKokkos
208: static PetscErrorCode VecCreate_MPIKokkos_Common(Vec v)
209: {
210: PetscFunctionBegin;
211: v->boundtocpu = PetscDefined(HAVE_KOKKOS_WITHOUT_GPU) ? PETSC_TRUE : PETSC_FALSE; // VECKOKKOS has yet to support CPU binding. But in this case, we deem it is bound to CPU.
212: v->ops->bindtocpu = VecBindToCPU_SeqKokkos;
213: v->ops->abs = VecAbs_SeqKokkos;
214: v->ops->reciprocal = VecReciprocal_SeqKokkos;
215: v->ops->pointwisemult = VecPointwiseMult_SeqKokkos;
216: v->ops->setrandom = VecSetRandom_SeqKokkos;
217: v->ops->dotnorm2 = VecDotNorm2_MPIKokkos;
218: v->ops->waxpy = VecWAXPY_SeqKokkos;
219: v->ops->norm = VecNorm_MPIKokkos;
220: v->ops->min = VecMin_MPIKokkos;
221: v->ops->max = VecMax_MPIKokkos;
222: v->ops->sum = VecSum_SeqKokkos;
223: v->ops->shift = VecShift_SeqKokkos;
224: v->ops->scale = VecScale_SeqKokkos;
225: v->ops->copy = VecCopy_SeqKokkos;
226: v->ops->set = VecSet_SeqKokkos;
227: v->ops->swap = VecSwap_SeqKokkos;
228: v->ops->axpy = VecAXPY_SeqKokkos;
229: v->ops->axpby = VecAXPBY_SeqKokkos;
230: v->ops->maxpy = VecMAXPY_SeqKokkos;
231: v->ops->aypx = VecAYPX_SeqKokkos;
232: v->ops->axpbypcz = VecAXPBYPCZ_SeqKokkos;
233: v->ops->pointwisedivide = VecPointwiseDivide_SeqKokkos;
234: v->ops->placearray = VecPlaceArray_SeqKokkos;
235: v->ops->replacearray = VecReplaceArray_SeqKokkos;
236: v->ops->resetarray = VecResetArray_SeqKokkos;
238: v->ops->dot = VecDot_MPIKokkos;
239: v->ops->tdot = VecTDot_MPIKokkos;
240: v->ops->mdot = VecMDot_MPIKokkos;
241: v->ops->mtdot = VecMTDot_MPIKokkos;
243: v->ops->dot_local = VecDot_SeqKokkos;
244: v->ops->tdot_local = VecTDot_SeqKokkos;
245: v->ops->mdot_local = VecMDot_SeqKokkos;
246: v->ops->mtdot_local = VecMTDot_SeqKokkos;
248: v->ops->norm_local = VecNorm_SeqKokkos;
249: v->ops->duplicate = VecDuplicate_MPIKokkos;
250: v->ops->destroy = VecDestroy_MPIKokkos;
251: v->ops->getlocalvector = VecGetLocalVector_SeqKokkos;
252: v->ops->restorelocalvector = VecRestoreLocalVector_SeqKokkos;
253: v->ops->getlocalvectorread = VecGetLocalVectorRead_SeqKokkos;
254: v->ops->restorelocalvectorread = VecRestoreLocalVectorRead_SeqKokkos;
255: v->ops->getarraywrite = VecGetArrayWrite_SeqKokkos;
256: v->ops->getarray = VecGetArray_SeqKokkos;
257: v->ops->restorearray = VecRestoreArray_SeqKokkos;
258: v->ops->getarrayandmemtype = VecGetArrayAndMemType_SeqKokkos;
259: v->ops->restorearrayandmemtype = VecRestoreArrayAndMemType_SeqKokkos;
260: v->ops->getarraywriteandmemtype = VecGetArrayWriteAndMemType_SeqKokkos;
261: v->ops->getsubvector = VecGetSubVector_MPIKokkos;
262: v->ops->restoresubvector = VecRestoreSubVector_SeqKokkos;
264: v->ops->setpreallocationcoo = VecSetPreallocationCOO_MPIKokkos;
265: v->ops->setvaluescoo = VecSetValuesCOO_MPIKokkos;
267: v->ops->errorwnorm = VecErrorWeightedNorms_MPIKokkos;
269: v->offloadmask = PETSC_OFFLOAD_KOKKOS; // Mark this is a VECKOKKOS; We use this flag for cheap VECKOKKOS test.
270: PetscFunctionReturn(PETSC_SUCCESS);
271: }
273: PETSC_INTERN PetscErrorCode VecConvert_MPI_MPIKokkos_inplace(Vec v)
274: {
275: Vec_MPI *vecmpi;
277: PetscFunctionBegin;
278: PetscCall(PetscKokkosInitializeCheck());
279: PetscCall(PetscLayoutSetUp(v->map));
280: PetscCall(PetscObjectChangeTypeName((PetscObject)v, VECMPIKOKKOS));
281: PetscCall(VecCreate_MPIKokkos_Common(v));
282: PetscCheck(!v->spptr, PETSC_COMM_SELF, PETSC_ERR_PLIB, "v->spptr not NULL");
283: vecmpi = static_cast<Vec_MPI *>(v->data);
284: if (vecmpi->localrep) { // It is a ghost vector
285: Vec local = vecmpi->localrep;
286: Vec_Kokkos *veckok;
288: PetscCall(VecConvert_Seq_SeqKokkos_inplace(local, NULL));
289: veckok = static_cast<Vec_Kokkos *>(local->spptr);
290: // TODO: can we subview on veckok->v_dual?
291: PetscCallCXX(v->spptr = new Vec_Kokkos(v->map->n, veckok->v_dual.view_host().data(), veckok->v_dual.view_device().data()));
292: } else PetscCallCXX(v->spptr = new Vec_Kokkos(v->map->n, vecmpi->array, NULL));
293: PetscFunctionReturn(PETSC_SUCCESS);
294: }
296: // Duplicate a VECMPIKOKKOS
297: static PetscErrorCode VecDuplicateVecs_MPIKokkos_GEMV(Vec w, PetscInt m, Vec *V[])
298: {
299: PetscInt64 lda; // use 64-bit as we will do "m * lda"
300: PetscScalar *array_h, *array_d;
301: PetscLayout map;
302: Vec_MPI *wmpi = (Vec_MPI *)w->data;
303: PetscScalarKokkosDualView w_dual;
305: PetscFunctionBegin;
306: PetscCall(PetscKokkosInitializeCheck()); // as we'll call kokkos_malloc()
307: if (wmpi->nghost) { // currently only do GEMV optimization for vectors without ghosts
308: w->ops->duplicatevecs = VecDuplicateVecs_Default;
309: PetscCall(VecDuplicateVecs(w, m, V));
310: } else {
311: PetscCall(PetscMalloc1(m, V));
312: PetscCall(VecGetLayout(w, &map));
313: VecGetLocalSizeAligned(w, 64, &lda); // get in lda the 64-bytes aligned local size
315: // See comments in VecCreate_SeqKokkos() on why we use DualView to allocate the memory
316: PetscCallCXX(w_dual = PetscScalarKokkosDualView("VecDuplicateVecs", m * lda)); // Kokkos init's w_dual to zero
318: // create the m vectors with raw arrays
319: array_h = w_dual.view_host().data();
320: array_d = w_dual.view_device().data();
321: for (PetscInt i = 0; i < m; i++) {
322: Vec v;
323: PetscCall(VecCreateMPIKokkosWithLayoutAndArrays_Private(map, &array_h[i * lda], &array_d[i * lda], &v));
324: PetscCallCXX(static_cast<Vec_Kokkos *>(v->spptr)->v_dual.modify_host()); // as we only init'ed array_h
325: PetscCall(PetscObjectListDuplicate(((PetscObject)w)->olist, &((PetscObject)v)->olist));
326: PetscCall(PetscFunctionListDuplicate(((PetscObject)w)->qlist, &((PetscObject)v)->qlist));
327: v->ops[0] = w->ops[0];
328: v->stash.donotstash = w->stash.donotstash;
329: v->stash.ignorenegidx = w->stash.ignorenegidx;
330: v->stash.bs = w->stash.bs;
331: v->bstash.bs = w->bstash.bs;
332: (*V)[i] = v;
333: }
335: // let the first vector own the raw arrays, so when it is destroyed it will free the arrays
336: if (m) {
337: Vec v = (*V)[0];
339: static_cast<Vec_Kokkos *>(v->spptr)->w_dual = w_dual; // stash the memory
340: // disable replacearray of the first vector, as freeing its memory also frees others in the group.
341: // But replacearray of others is ok, as they don't own their array.
342: if (m > 1) v->ops->replacearray = VecReplaceArray_Default_GEMV_Error;
343: }
344: }
345: PetscFunctionReturn(PETSC_SUCCESS);
346: }
348: /*MC
349: VECMPIKOKKOS - VECMPIKOKKOS = "mpikokkos" - The basic parallel vector, modified to use Kokkos
351: Options Database Keys:
352: . -vec_type mpikokkos - sets the vector type to VECMPIKOKKOS during a call to VecSetFromOptions()
354: Level: beginner
356: .seealso: `VecCreate()`, `VecSetType()`, `VecSetFromOptions()`, `VecCreateMPIKokkosWithArray()`, `VECMPI`, `VecType`, `VecCreateMPI()`
357: M*/
358: PetscErrorCode VecCreate_MPIKokkos(Vec v)
359: {
360: PetscBool mdot_use_gemv = PETSC_TRUE;
361: PetscBool maxpy_use_gemv = PETSC_FALSE; // default is false as we saw bad performance with vendors' GEMV with tall skinny matrices.
362: PetscScalarKokkosDualView v_dual;
364: PetscFunctionBegin;
365: PetscCall(PetscKokkosInitializeCheck());
366: PetscCall(PetscLayoutSetUp(v->map));
368: PetscCallCXX(v_dual = PetscScalarKokkosDualView("v_dual", v->map->n)); // Kokkos init's v_dual to zero
369: PetscCall(VecCreate_MPI_Private(v, PETSC_FALSE, 0, v_dual.view_host().data()));
371: PetscCall(PetscObjectChangeTypeName((PetscObject)v, VECMPIKOKKOS));
372: PetscCall(VecCreate_MPIKokkos_Common(v));
373: PetscCheck(!v->spptr, PETSC_COMM_SELF, PETSC_ERR_PLIB, "v->spptr not NULL");
374: PetscCallCXX(v->spptr = new Vec_Kokkos(v_dual));
375: PetscCall(PetscOptionsGetBool(NULL, NULL, "-vec_mdot_use_gemv", &mdot_use_gemv, NULL));
376: PetscCall(PetscOptionsGetBool(NULL, NULL, "-vec_maxpy_use_gemv", &maxpy_use_gemv, NULL));
378: // allocate multiple vectors together
379: if (mdot_use_gemv || maxpy_use_gemv) v->ops[0].duplicatevecs = VecDuplicateVecs_MPIKokkos_GEMV;
381: if (mdot_use_gemv) {
382: v->ops[0].mdot = VecMDot_MPIKokkos_GEMV;
383: v->ops[0].mtdot = VecMTDot_MPIKokkos_GEMV;
384: v->ops[0].mdot_local = VecMDot_SeqKokkos_GEMV;
385: v->ops[0].mtdot_local = VecMTDot_SeqKokkos_GEMV;
386: }
388: if (maxpy_use_gemv) v->ops[0].maxpy = VecMAXPY_SeqKokkos_GEMV;
389: PetscFunctionReturn(PETSC_SUCCESS);
390: }
392: // Create a VECMPIKOKKOS with layout and arrays
393: PetscErrorCode VecCreateMPIKokkosWithLayoutAndArrays_Private(PetscLayout map, const PetscScalar harray[], const PetscScalar darray[], Vec *v)
394: {
395: Vec w;
397: PetscFunctionBegin;
398: if (map->n > 0) PetscCheck(darray, map->comm, PETSC_ERR_ARG_WRONG, "darray cannot be NULL");
399: #if defined(KOKKOS_ENABLE_UNIFIED_MEMORY)
400: PetscCheck(harray == darray, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "harray and darray must be the same");
401: #endif
402: PetscCall(VecCreateMPIWithLayoutAndArray_Private(map, harray, &w));
403: PetscCall(PetscObjectChangeTypeName((PetscObject)w, VECMPIKOKKOS)); // Change it to VECKOKKOS
404: PetscCall(VecCreate_MPIKokkos_Common(w));
405: PetscCallCXX(w->spptr = new Vec_Kokkos(map->n, const_cast<PetscScalar *>(harray), const_cast<PetscScalar *>(darray)));
406: *v = w;
407: PetscFunctionReturn(PETSC_SUCCESS);
408: }
410: /*@
411: VecCreateMPIKokkosWithArray - Creates a parallel, array-style vector,
412: where the user provides the GPU array space to store the vector values.
414: Collective
416: Input Parameters:
417: + comm - the MPI communicator to use
418: . bs - block size, same meaning as VecSetBlockSize()
419: . n - local vector length, cannot be PETSC_DECIDE
420: . N - global vector length (or PETSC_DECIDE to have calculated)
421: - darray - the user provided GPU array to store the vector values
423: Output Parameter:
424: . v - the vector
426: Notes:
427: Use VecDuplicate() or VecDuplicateVecs() to form additional vectors of the
428: same type as an existing vector.
430: If the user-provided array is NULL, then VecKokkosPlaceArray() can be used
431: at a later stage to SET the array for storing the vector values.
433: PETSc does NOT free the array when the vector is destroyed via VecDestroy().
434: The user should not free the array until the vector is destroyed.
436: Level: intermediate
438: .seealso: `VecCreateSeqKokkosWithArray()`, `VecCreateMPIWithArray()`, `VecCreateSeqWithArray()`,
439: `VecCreate()`, `VecDuplicate()`, `VecDuplicateVecs()`, `VecCreateGhost()`,
440: `VecCreateMPI()`, `VecCreateGhostWithArray()`, `VecPlaceArray()`
441: @*/
442: PetscErrorCode VecCreateMPIKokkosWithArray(MPI_Comm comm, PetscInt bs, PetscInt n, PetscInt N, const PetscScalar darray[], Vec *v)
443: {
444: Vec w;
445: Vec_Kokkos *veckok;
446: Vec_MPI *vecmpi;
447: PetscScalar *harray;
449: PetscFunctionBegin;
450: PetscCheck(n != PETSC_DECIDE, PETSC_COMM_SELF, PETSC_ERR_ARG_OUTOFRANGE, "Must set local size of vector");
451: PetscCall(PetscKokkosInitializeCheck());
452: PetscCall(PetscSplitOwnership(comm, &n, &N));
453: PetscCall(VecCreate(comm, &w));
454: PetscCall(VecSetSizes(w, n, N));
455: PetscCall(VecSetBlockSize(w, bs));
456: PetscCall(PetscLayoutSetUp(w->map));
458: if (std::is_same<DefaultMemorySpace, HostMirrorMemorySpace>::value) {
459: harray = const_cast<PetscScalar *>(darray);
460: } else PetscCall(PetscMalloc1(w->map->n, &harray)); /* If device is not the same as host, allocate the host array ourselves */
462: PetscCall(VecCreate_MPI_Private(w, PETSC_FALSE /*alloc*/, 0 /*nghost*/, harray)); /* Build a sequential vector with provided data */
463: vecmpi = static_cast<Vec_MPI *>(w->data);
465: if (!std::is_same<DefaultMemorySpace, HostMirrorMemorySpace>::value) vecmpi->array_allocated = harray; /* The host array was allocated by PETSc */
467: PetscCall(PetscObjectChangeTypeName((PetscObject)w, VECMPIKOKKOS));
468: PetscCall(VecCreate_MPIKokkos_Common(w));
469: veckok = new Vec_Kokkos(n, harray, const_cast<PetscScalar *>(darray));
470: veckok->v_dual.modify_device(); /* Mark the device is modified */
471: w->spptr = static_cast<void *>(veckok);
472: *v = w;
473: PetscFunctionReturn(PETSC_SUCCESS);
474: }
476: /*
477: VecCreateMPIKokkosWithArrays_Private - Creates a Kokkos parallel, array-style vector
478: with user-provided arrays on host and device.
480: Collective
482: Input Parameter:
483: + comm - the communicator
484: . bs - the block size
485: . n - the local vector length
486: . N - the global vector length
487: - harray - host memory where the vector elements are to be stored.
488: - darray - device memory where the vector elements are to be stored.
490: Output Parameter:
491: . v - the vector
493: Notes:
494: If there is no device, then harray and darray must be the same.
495: If n is not zero, then harray and darray must be allocated.
496: After the call, the created vector is supposed to be in a synchronized state, i.e.,
497: we suppose harray and darray have the same data.
499: PETSc does NOT free the array when the vector is destroyed via VecDestroy().
500: The user should not free the array until the vector is destroyed.
501: */
502: PetscErrorCode VecCreateMPIKokkosWithArrays_Private(MPI_Comm comm, PetscInt bs, PetscInt n, PetscInt N, const PetscScalar harray[], const PetscScalar darray[], Vec *v)
503: {
504: Vec w;
506: PetscFunctionBegin;
507: PetscCall(PetscKokkosInitializeCheck());
508: if (n) {
509: PetscAssertPointer(harray, 5);
510: PetscCheck(darray, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "darray cannot be NULL");
511: }
512: if (std::is_same<DefaultMemorySpace, HostMirrorMemorySpace>::value) PetscCheck(harray == darray, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "harray and darray must be the same");
513: PetscCall(VecCreateMPIWithArray(comm, bs, n, N, harray, &w));
514: PetscCall(PetscObjectChangeTypeName((PetscObject)w, VECMPIKOKKOS)); /* Change it to Kokkos */
515: PetscCall(VecCreate_MPIKokkos_Common(w));
516: PetscCallCXX(w->spptr = new Vec_Kokkos(n, const_cast<PetscScalar *>(harray), const_cast<PetscScalar *>(darray)));
517: *v = w;
518: PetscFunctionReturn(PETSC_SUCCESS);
519: }
521: /*MC
522: VECKOKKOS - VECKOKKOS = "kokkos" - The basic vector, modified to use Kokkos
524: Options Database Keys:
525: . -vec_type kokkos - sets the vector type to VECKOKKOS during a call to VecSetFromOptions()
527: Level: beginner
529: .seealso: `VecCreate()`, `VecSetType()`, `VecSetFromOptions()`, `VecCreateMPIKokkosWithArray()`, `VECMPI`, `VecType`, `VecCreateMPI()`
530: M*/
531: PetscErrorCode VecCreate_Kokkos(Vec v)
532: {
533: PetscMPIInt size;
535: PetscFunctionBegin;
536: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)v), &size));
537: if (size == 1) PetscCall(VecSetType(v, VECSEQKOKKOS));
538: else PetscCall(VecSetType(v, VECMPIKOKKOS));
539: PetscFunctionReturn(PETSC_SUCCESS);
540: }