1: #include <petsc/private/deviceimpl.h>
  2: #include "objpool.hpp"

  4: const char *const PetscStreamTypes[] = {
  5:   "global_blocking",
  6:   "default_blocking",
  7:   "global_nonblocking",
  8:   "max",
  9:   "PetscStreamType",
 10:   "PETSC_STREAM_",
 11:   nullptr
 12: };

 14: const char *const PetscDeviceContextJoinModes[] = {
 15:   "destroy",
 16:   "sync",
 17:   "no_sync",
 18:   "PetscDeviceContextJoinMode",
 19:   "PETSC_DEVICE_CONTEXT_JOIN_",
 20:   nullptr
 21: };

 23: /* Define the allocator */
 24: struct PetscDeviceContextAllocator : Petsc::AllocatorBase<PetscDeviceContext>
 25: {
 26:   static PetscInt PetscDeviceContextID;

 28:   PETSC_NODISCARD static PetscErrorCode create(PetscDeviceContext *dctx) noexcept
 29:   {
 30:     PetscDeviceContext dc;

 32:     PetscNew(&dc);
 33:     dc->id         = PetscDeviceContextID++;
 34:     dc->streamType = PETSC_STREAM_DEFAULT_BLOCKING;
 35:     *dctx          = dc;
 36:     return 0;
 37:   }

 39:   PETSC_NODISCARD static PetscErrorCode destroy(PetscDeviceContext dctx) noexcept
 40:   {
 41:     PetscAssert(!dctx->numChildren,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"Device context still has %" PetscInt_FMT " un-joined children, must call PetscDeviceContextJoin() with all children before destroying",dctx->numChildren);
 42:     if (dctx->ops->destroy) (*dctx->ops->destroy)(dctx);
 43:     PetscDeviceDestroy(&dctx->device);
 44:     PetscFree(dctx->childIDs);
 45:     PetscFree(dctx);
 46:     return 0;
 47:   }

 49:   PETSC_NODISCARD static PetscErrorCode reset(PetscDeviceContext dctx) noexcept
 50:   {
 51:     /* don't deallocate the child array, rather just zero it out */
 52:     PetscArrayzero(dctx->childIDs,dctx->maxNumChildren);
 53:     dctx->setup       = PETSC_FALSE;
 54:     dctx->numChildren = 0;
 55:     dctx->streamType  = PETSC_STREAM_DEFAULT_BLOCKING;
 56:     return 0;
 57:   }

 59:   PETSC_NODISCARD static constexpr PetscErrorCode finalize() noexcept { return 0; }
 60: };
 61: /* an ID = 0 is invalid */
 62: PetscInt PetscDeviceContextAllocator::PetscDeviceContextID = 1;

 64: static Petsc::ObjectPool<PetscDeviceContext,PetscDeviceContextAllocator> contextPool;

 66: /*@C
 67:   PetscDeviceContextCreate - Creates a PetscDeviceContext

 69:   Not Collective, Asynchronous

 71:   Output Paramemter:
 72: . dctx - The PetscDeviceContext

 74:   Notes:
 75:   Unlike almost every other PETSc class it is advised that most users use
 76:   PetscDeviceContextDuplicate() rather than this routine to create new contexts. Contexts
 77:   of different types are incompatible with one another; using
 78:   PetscDeviceContextDuplicate() ensures compatible types.

 80:   Level: beginner

 82: .seealso: PetscDeviceContextDuplicate(), PetscDeviceContextSetDevice(),
 83: PetscDeviceContextSetStreamType(), PetscDeviceContextSetUp(),
 84: PetscDeviceContextSetFromOptions(), PetscDeviceContextDestroy()
 85: @*/
 86: PetscErrorCode PetscDeviceContextCreate(PetscDeviceContext *dctx)
 87: {
 89:   PetscDeviceInitializePackage();
 90:   contextPool.get(*dctx);
 91:   return 0;
 92: }

 94: /*@C
 95:   PetscDeviceContextDestroy - Frees a PetscDeviceContext

 97:   Not Collective, Asynchronous

 99:   Input Parameters:
100: . dctx - The PetscDeviceContext

102:   Notes:
103:   No implicit synchronization occurs due to this routine, all resources are released completely asynchronously
104:   w.r.t. the host. If one needs to guarantee access to the data produced on this contexts stream one should perform the
105:   appropriate synchronization before calling this routine.

107:   Developer Notes:
108:   The context is never actually "destroyed", only returned to an ever growing pool of
109:   contexts. There are currently no safeguards on the size of the pool, this should perhaps
110:   be implemented.

112:   Level: beginner

114: .seealso: PetscDeviceContextCreate(), PetscDeviceContextSetDevice(), PetscDeviceContextSetUp(), PetscDeviceContextSynchronize()
115: @*/
116: PetscErrorCode PetscDeviceContextDestroy(PetscDeviceContext *dctx)
117: {
118:   if (!*dctx) return 0;
119:   contextPool.reclaim(std::move(*dctx));
120:   *dctx = nullptr;
121:   return 0;
122: }

124: /*@C
125:   PetscDeviceContextSetStreamType - Set the implementation type of the underlying stream for a PetscDeviceContext

127:   Not Collective, Asynchronous

129:   Input Parameters:
130: + dctx - The PetscDeviceContext
131: - type - The PetscStreamType

133:   Notes:
134:   See PetscStreamType in include/petscdevicetypes.h for more information on the available
135:   types and their interactions. If the PetscDeviceContext was previously set up and stream
136:   type was changed, you must call PetscDeviceContextSetUp() again after this routine.

138:   Level: intermediate

140: .seealso: PetscStreamType, PetscDeviceContextGetStreamType(), PetscDeviceContextCreate(), PetscDeviceContextSetUp(), PetscDeviceContextSetFromOptions()
141: @*/
142: PetscErrorCode PetscDeviceContextSetStreamType(PetscDeviceContext dctx, PetscStreamType type)
143: {
146:   /* only need to do complex swapping if the object has already been setup */
147:   if (dctx->setup && (dctx->streamType != type)) {
148:     (*dctx->ops->changestreamtype)(dctx,type);
149:     dctx->setup = PETSC_FALSE;
150:   }
151:   dctx->streamType = type;
152:   return 0;
153: }

155: /*@C
156:   PetscDeviceContextGetStreamType - Get the implementation type of the underlying stream for a PetscDeviceContext

158:   Not Collective, Asynchronous

160:   Input Parameter:
161: . dctx - The PetscDeviceContext

163:   Output Parameter:
164: . type - The PetscStreamType

166:   Notes:
167:   See PetscStreamType in include/petscdevicetypes.h for more information on the available types and their interactions

169:   Level: intermediate

171: .seealso: PetscDeviceContextSetStreamType(), PetscDeviceContextCreate(), PetscDeviceContextSetFromOptions()
172: @*/
173: PetscErrorCode PetscDeviceContextGetStreamType(PetscDeviceContext dctx, PetscStreamType *type)
174: {
177:   *type = dctx->streamType;
178:   return 0;
179: }

181: /*@C
182:   PetscDeviceContextSetDevice - Set the underlying device for the PetscDeviceContext

184:   Not Collective, Possibly Synchronous

186:   Input Parameters:
187: + dctx   - The PetscDeviceContext
188: - device - The PetscDevice

190:   Notes:
191:   This routine is effectively PetscDeviceContext's "set-type" (so every PetscDeviceContext
192:   must also have an attached PetscDevice). Unlike the usual set-type semantics, it is
193:   not stricly necessary to set a contexts device to enable usage, any created device
194:   contexts will always come equipped with the "default" device.

196:   This routine is a no-op if dctx is already attached to device.

198:   This routine may initialize the backend device and incur synchronization.

200:   Level: intermediate

202: .seealso: PetscDeviceCreate(), PetscDeviceConfigure(), PetscDeviceContextGetDevice()
203: @*/
204: PetscErrorCode PetscDeviceContextSetDevice(PetscDeviceContext dctx, PetscDevice device)
205: {
208:   if (dctx->device) {
209:     /* can't do a strict pointer equality check since PetscDevice's are reused */
210:     if (dctx->device->ops->createcontext == device->ops->createcontext) return 0;
211:   }
212:   PetscDeviceDestroy(&dctx->device);
213:   if (dctx->ops->destroy) (*dctx->ops->destroy)(dctx);
214:   PetscMemzero(dctx->ops,sizeof(*dctx->ops));
215:   (*device->ops->createcontext)(dctx);
216:   PetscDeviceReference_Internal(device);
217:   dctx->device = device;
218:   dctx->setup  = PETSC_FALSE;
219:   return 0;
220: }

222: /*@C
223:   PetscDeviceContextGetDevice - Get the underlying PetscDevice for a PetscDeviceContext

225:   Not Collective, Asynchronous

227:   Input Parameter:
228: . dctx - the PetscDeviceContext

230:   Output Parameter:
231: . device - The PetscDevice

233:   Notes:
234:   This is a borrowed reference, the user should not destroy the device.

236:   Level: intermediate

238: .seealso: PetscDeviceContextSetDevice(), PetscDevice
239: @*/
240: PetscErrorCode PetscDeviceContextGetDevice(PetscDeviceContext dctx, PetscDevice *device)
241: {
244:   PetscAssert(dctx->device,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"PetscDeviceContext %" PetscInt_FMT " has no attached PetscDevice to get",dctx->id);
245:   *device = dctx->device;
246:   return 0;
247: }

249: /*@C
250:   PetscDeviceContextSetUp - Prepares a PetscDeviceContext for use

252:   Not Collective, Asynchronous

254:   Input Parameter:
255: . dctx - The PetscDeviceContext

257:   Developer Notes:
258:   This routine is usually the stage where a PetscDeviceContext acquires device-side data structures such as streams,
259:   events, and (possibly) handles.

261:   Level: beginner

263: .seealso: PetscDeviceContextCreate(), PetscDeviceContextSetDevice(), PetscDeviceContextDestroy(), PetscDeviceContextSetFromOptions()
264: @*/
265: PetscErrorCode PetscDeviceContextSetUp(PetscDeviceContext dctx)
266: {
268:   if (!dctx->device) {
269:     PetscInfo(nullptr,"PetscDeviceContext %" PetscInt_FMT " did not have an explicitly attached PetscDevice, using default with type %s\n",dctx->id,PetscDeviceTypes[PETSC_DEVICE_DEFAULT]);
270:     PetscDeviceContextSetDefaultDevice_Internal(dctx);
271:   }
272:   if (dctx->setup) return 0;
273:   (*dctx->ops->setup)(dctx);
274:   dctx->setup = PETSC_TRUE;
275:   return 0;
276: }

278: /*@C
279:   PetscDeviceContextDuplicate - Duplicates a PetscDeviceContext object

281:   Not Collective, Asynchronous

283:   Input Parameter:
284: . dctx - The PetscDeviceContext to duplicate

286:   Output Parameter:
287: . dctxdup - The duplicated PetscDeviceContext

289:   Notes:
290:   This is a shorthand method for creating a PetscDeviceContext with the exact same
291:   settings as another. Note however that the duplicated PetscDeviceContext does not "share"
292:   any of the underlying data with the original, (including its current stream-state) they
293:   are completely separate objects.

295:   Level: beginner

297: .seealso: PetscDeviceContextCreate(), PetscDeviceContextSetDevice(), PetscDeviceContextSetStreamType()
298: @*/
299: PetscErrorCode PetscDeviceContextDuplicate(PetscDeviceContext dctx, PetscDeviceContext *dctxdup)
300: {
301:   PetscDeviceContext dup;

305:   PetscDeviceContextCreate(&dup);
306:   PetscDeviceContextSetStreamType(dup,dctx->streamType);
307:   if (dctx->device) PetscDeviceContextSetDevice(dup,dctx->device);
308:   PetscDeviceContextSetUp(dup);
309:   *dctxdup = dup;
310:   return 0;
311: }

313: /*@C
314:   PetscDeviceContextQueryIdle - Returns whether or not a PetscDeviceContext is idle

316:   Not Collective, Asynchronous

318:   Input Parameter:
319: . dctx - The PetscDeviceContext object

321:   Output Parameter:
322: . idle - PETSC_TRUE if PetscDeviceContext has NO work, PETSC_FALSE if it has work

324:   Notes:
325:   This routine only refers a singular context and does NOT take any of its children into
326:   account. That is, if dctx is idle but has dependents who do have work, this routine still
327:   returns PETSC_TRUE.

329:   Level: intermediate

331: .seealso: PetscDeviceContextCreate(), PetscDeviceContextWaitForContext(), PetscDeviceContextFork()
332: @*/
333: PetscErrorCode PetscDeviceContextQueryIdle(PetscDeviceContext dctx, PetscBool *idle)
334: {
337:   (*dctx->ops->query)(dctx,idle);
338:   PetscInfo(nullptr,"PetscDeviceContext id %" PetscInt_FMT " %s idle\n",dctx->id,*idle ? "was" : "was not");
339:   return 0;
340: }

342: /*@C
343:   PetscDeviceContextWaitForContext - Make one context wait for another context to finish

345:   Not Collective, Asynchronous

347:   Input Parameters:
348: + dctxa - The PetscDeviceContext object that is waiting
349: - dctxb - The PetscDeviceContext object that is being waited on

351:   Notes:
352:   Serializes two PetscDeviceContexts. This routine uses only the state of dctxb at the moment this routine was
353:   called, so any future work queued will not affect dctxa. It is safe to pass the same context to both arguments.

355:   Level: beginner

357: .seealso: PetscDeviceContextCreate(), PetscDeviceContextQueryIdle(), PetscDeviceContextJoin()
358: @*/
359: PetscErrorCode PetscDeviceContextWaitForContext(PetscDeviceContext dctxa, PetscDeviceContext dctxb)
360: {
362:   if (dctxa == dctxb) return 0;
363:   (*dctxa->ops->waitforcontext)(dctxa,dctxb);
364:   return 0;
365: }

367: #define PETSC_USE_DEBUG_AND_INFO (PetscDefined(USE_DEBUG) && PetscDefined(USE_INFO))
368: #if PETSC_USE_DEBUG_AND_INFO
369: #include <string>
370: #endif
371: /*@C
372:   PetscDeviceContextFork - Create a set of dependent child contexts from a parent context

374:   Not Collective, Asynchronous

376:   Input Parameters:
377: + dctx - The parent PetscDeviceContext
378: - n    - The number of children to create

380:   Output Parameter:
381: . dsub - The created child context(s)

383:   Notes:
384:   This routine creates n edges of a DAG from a source node which are causally dependent on the source node, meaning
385:   that work queued on child contexts will not start until the parent context finishes its work. This accounts for work
386:   queued on the parent up until calling this function, any subsequent work enqueued on the parent has no effect on the children.

388:   Any children created with this routine have their lifetimes bounded by the parent. That is, the parent context expects
389:   to free all of it's children (and ONLY its children) before itself is freed.

391:   DAG representation:
392: .vb
393:   time ->

395:   -> dctx \----> dctx ------>
396:            \---> dsub[0] --->
397:             \--> ... ------->
398:              \-> dsub[n-1] ->
399: .ve

401:   Level: intermediate

403: .seealso: PetscDeviceContextJoin(), PetscDeviceContextSynchronize(), PetscDeviceContextQueryIdle()
404: @*/
405: PetscErrorCode PetscDeviceContextFork(PetscDeviceContext dctx, PetscInt n, PetscDeviceContext **dsub)
406: {
407: #if PETSC_USE_DEBUG_AND_INFO
408:   const PetscInt      nBefore = n;
409:   static std::string  idList;
410: #endif
411:   PetscDeviceContext *dsubTmp = nullptr;
412:   PetscInt            i = 0;

416:   PetscAssert(n >= 0,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Number of contexts requested %" PetscInt_FMT " < 0",n);
417: #if PETSC_USE_DEBUG_AND_INFO
418:   /* reserve 4 chars per id, 2 for number and 2 for ', ' separator */
419:   idList.reserve(4*n);
420: #endif
421:   /* update child totals */
422:   dctx->numChildren += n;
423:   /* now to find out if we have room */
424:   if (dctx->numChildren > dctx->maxNumChildren) {
425:     /* no room, either from having too many kids or not having any */
426:     if (dctx->childIDs) {
427:       /* have existing children, must reallocate them */
428:       PetscRealloc(dctx->numChildren*sizeof(*dctx->childIDs),&dctx->childIDs);
429:       /* clear the extra memory since realloc doesn't do it for us */
430:       PetscArrayzero((dctx->childIDs)+(dctx->maxNumChildren),(dctx->numChildren)-(dctx->maxNumChildren));
431:     } else {
432:       /* have no children */
433:       PetscCalloc1(dctx->numChildren,&dctx->childIDs);
434:     }
435:     /* update total number of children */
436:     dctx->maxNumChildren = dctx->numChildren;
437:   }
438:   PetscMalloc1(n,&dsubTmp);
439:   while (n) {
440:     /* empty child slot */
441:     if (!(dctx->childIDs[i])) {
442:       /* create the child context in the image of its parent */
443:       PetscDeviceContextDuplicate(dctx,dsubTmp+i);
444:       PetscDeviceContextWaitForContext(dsubTmp[i],dctx);
445:       /* register the child with its parent */
446:       dctx->childIDs[i] = dsubTmp[i]->id;
447: #if PETSC_USE_DEBUG_AND_INFO
448:       idList += std::to_string(dsubTmp[i]->id);
449:       if (n != 1) idList += ", ";
450: #endif
451:       --n;
452:     }
453:     ++i;
454:   }
455: #if PETSC_USE_DEBUG_AND_INFO
456:   PetscInfo(nullptr,"Forked %" PetscInt_FMT " children from parent %" PetscInt_FMT " with IDs: %s\n",nBefore,dctx->id,idList.c_str());
457:   /* resets the size but doesn't deallocate the memory */
458:   idList.clear();
459: #endif
460:   /* pass the children back to caller */
461:   *dsub = dsubTmp;
462:   return 0;
463: }

465: /*@C
466:   PetscDeviceContextJoin - Converge a set of child contexts

468:   Not Collective, Asynchronous

470:   Input Parameters:
471: + dctx         - A PetscDeviceContext to converge on
472: . n            - The number of sub contexts to converge
473: . joinMode     - The type of join to perform
474: - dsub         - The sub contexts to converge

476:   Notes:
477:   If PetscDeviceContextFork() creates n edges from a source node which all depend on the
478:   source node, then this routine is the exact mirror. That is, it creates a node
479:   (represented in dctx) which recieves n edges (and optionally destroys them) which is
480:   dependent on the completion of all incoming edges.

482:   If joinMode is PETSC_DEVICE_CONTEXT_JOIN_DESTROY all contexts in dsub will be destroyed
483:   by this routine. Thus all sub contexts must have been created with the dctx passed to
484:   this routine.

486:   if joinMode is PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC dctx waits for all sub contexts but the
487:   sub contexts do not wait for one another afterwards.

489:   If joinMode is PETSC_DEVICE_CONTEXT_JOIN_SYNC all sub contexts will additionally
490:   wait on dctx after converging. This has the effect of "synchronizing" the outgoing
491:   edges.

493:   DAG representations:
494:   If joinMode is PETSC_DEVICE_CONTEXT_JOIN_DESTROY
495: .vb
496:   time ->

498:   -> dctx ---------/- dctx ->
499:   -> dsub[0] -----/
500:   ->  ... -------/
501:   -> dsub[n-1] -/
502: .ve
503:   If joinMode is PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC
504: .vb
505:   time ->

507:   -> dctx ---------/- dctx ->
508:   -> dsub[0] -----/--------->
509:   ->  ... -------/---------->
510:   -> dsub[n-1] -/----------->
511: .ve
512:   If joinMode is PETSC_DEVICE_CONTEXT_JOIN_SYNC
513: .vb
514:   time ->

516:   -> dctx ---------/- dctx -\----> dctx ------>
517:   -> dsub[0] -----/          \---> dsub[0] --->
518:   ->  ... -------/            \--> ... ------->
519:   -> dsub[n-1] -/              \-> dsub[n-1] ->
520: .ve

522:   Level: intermediate

524: .seealso: PetscDeviceContextFork(), PetscDeviceContextSynchronize(), PetscDeviceContextJoinMode
525: @*/
526: PetscErrorCode PetscDeviceContextJoin(PetscDeviceContext dctx, PetscInt n, PetscDeviceContextJoinMode joinMode, PetscDeviceContext **dsub)
527: {
528: #if defined(PETSC_USE_DEBUG) && defined(PETSC_USE_INFO)
529:   static std::string idList;
530: #endif

532:   /* validity of dctx is checked in the wait-for loop */
534:   PetscAssert(n >= 0,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Number of contexts merged %" PetscInt_FMT " < 0",n);
535: #if defined(PETSC_USE_DEBUG) && defined(PETSC_USE_INFO)
536:   /* reserve 4 chars per id, 2 for number and 2 for ', ' separator */
537:   idList.reserve(4*n);
538: #endif
539:   /* first dctx waits on all the incoming edges */
540:   for (PetscInt i = 0; i < n; ++i) {
542:     PetscDeviceContextWaitForContext(dctx,(*dsub)[i]);
543: #if defined(PETSC_USE_DEBUG) && defined(PETSC_USE_INFO)
544:     idList += std::to_string((*dsub)[i]->id);
545:     if (i+1 < n) idList += ", ";
546: #endif
547:   }

549:   /* now we handle the aftermath */
550:   switch (joinMode) {
551:   case PETSC_DEVICE_CONTEXT_JOIN_DESTROY:
552:     {
553:       PetscInt j = 0;

555:       PetscAssert(n <= dctx->numChildren,PETSC_COMM_SELF,PETSC_ERR_ARG_OUTOFRANGE,"Trying to destroy %" PetscInt_FMT " children of a parent context that only has %" PetscInt_FMT " children, likely trying to restore to wrong parent",n,dctx->numChildren);
556:       /* update child count while it's still fresh in memory */
557:       dctx->numChildren -= n;
558:       for (PetscInt i = 0; i < dctx->maxNumChildren; ++i) {
559:         if (dctx->childIDs[i] && (dctx->childIDs[i] == (*dsub)[j]->id)) {
560:           /* child is one of ours, can destroy it */
561:           PetscDeviceContextDestroy((*dsub)+j);
562:           /* reset the child slot */
563:           dctx->childIDs[i] = 0;
564:           if (++j == n) break;
565:         }
566:       }
567:       /* gone through the loop but did not find every child, if this triggers (or well, doesn't) on perf-builds we leak the remaining contexts memory */
568:       PetscAssert(j == n,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"%" PetscInt_FMT " contexts still remain after destroy, this may be because you are trying to restore to the wrong parent context, or the device contexts are not in the same order as they were checked out out in.",n-j);
569:       PetscFree(*dsub);
570:     }
571:     break;
572:   case PETSC_DEVICE_CONTEXT_JOIN_SYNC:
573:     for (PetscInt i = 0; i < n; ++i) PetscDeviceContextWaitForContext((*dsub)[i],dctx);
574:   case PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC:
575:     break;
576:   default:
577:     SETERRQ(PETSC_COMM_SELF,PETSC_ERR_ARG_WRONG,"Unknown PetscDeviceContextJoinMode given");
578:   }

580: #if defined(PETSC_USE_DEBUG) && defined(PETSC_USE_INFO)
581:   PetscInfo(nullptr,"Joined %" PetscInt_FMT " ctxs to ctx %" PetscInt_FMT ", mode %s with IDs: %s\n",n,dctx->id,PetscDeviceContextJoinModes[joinMode],idList.c_str());
582:   idList.clear();
583: #endif
584:   return 0;
585: }

587: /*@C
588:   PetscDeviceContextSynchronize - Block the host until all work queued on or associated with a PetscDeviceContext has finished

590:   Not Collective, Synchronous

592:   Input Parameters:
593: . dctx - The PetscDeviceContext to synchronize

595:   Level: beginner

597: .seealso: PetscDeviceContextFork(), PetscDeviceContextJoin(), PetscDeviceContextQueryIdle()
598: @*/
599: PetscErrorCode PetscDeviceContextSynchronize(PetscDeviceContext dctx)
600: {
602:   /* if it isn't setup there is nothing to sync on */
603:   if (dctx->setup) (*dctx->ops->synchronize)(dctx);
604:   return 0;
605: }

607: #define PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE PETSC_DEVICE_DEFAULT
608: // REMOVE ME (change)
609: #define PETSC_DEVICE_CONTEXT_DEFAULT_STREAM PETSC_STREAM_GLOBAL_BLOCKING

611: static PetscDeviceType    rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE;
612: static PetscStreamType    rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM;
613: static PetscDeviceContext globalContext  = nullptr;

615: /* when PetsDevice initializes PetscDeviceContext eagerly the type of device created should
616:  * match whatever device is eagerly intialized */
617: PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType type)
618: {
620:   rootDeviceType = type;
621:   return 0;
622: }

624: #if 0
625: /* currently unused */
626: PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType type)
627: {
629:   rootStreamType = type;
630:   return 0;
631: }
632: #endif

634: static PetscErrorCode PetscDeviceContextSetupGlobalContext_Private(void)
635: {
636:   static const auto PetscDeviceContextFinalizer = []() -> PetscErrorCode {

638:     PetscDeviceContextDestroy(&globalContext);
639:     rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE;
640:     rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM;
641:     return 0;
642:   };

644:   if (globalContext) return 0;
645:   /* this exists purely as a valid device check. */
646:   PetscDeviceInitializePackage();
647:   PetscRegisterFinalize(PetscDeviceContextFinalizer);
648:   PetscInfo(nullptr,"Initializing global PetscDeviceContext\n");
649:   /* we call the allocator directly here since the ObjectPool creates a PetscContainer which
650:    * eventually tries to call logging functions. However, this routine may be purposefully
651:    * called __before__ logging is initialized, so the logging function would PETSCABORT */
652:   contextPool.allocator().create(&globalContext);
653:   PetscDeviceContextSetStreamType(globalContext,rootStreamType);
654:   PetscDeviceContextSetDefaultDeviceForType_Internal(globalContext,rootDeviceType);
655:   PetscDeviceContextSetUp(globalContext);
656:   return 0;
657: }

659: /*@C
660:   PetscDeviceContextGetCurrentContext - Get the current active PetscDeviceContext

662:   Not Collective, Asynchronous

664:   Output Parameter:
665: . dctx - The PetscDeviceContext

667:   Notes:
668:   The user generally should not destroy contexts retrieved with this routine unless they
669:   themselves have created them. There exists no protection against destroying the root
670:   context.

672:   Developer Notes:
673:   Unless the user has set their own, this routine creates the "root" context the first time it
674:   is called, registering its destructor to PetscFinalize().

676:   Level: beginner

678: .seealso: PetscDeviceContextSetCurrentContext(), PetscDeviceContextFork(),
679: PetscDeviceContextJoin(), PetscDeviceContextCreate()
680: @*/
681: PetscErrorCode PetscDeviceContextGetCurrentContext(PetscDeviceContext *dctx)
682: {
684:   PetscDeviceContextSetupGlobalContext_Private();
685:   /* while the static analyzer can find global variables, it will throw a warning about not
686:    * being able to connect this back to the function arguments */
688:   *dctx = globalContext;
689:   return 0;
690: }

692: /*@C
693:   PetscDeviceContextSetCurrentContext - Set the current active PetscDeviceContext

695:   Not Collective, Asynchronous

697:   Input Parameter:
698: . dctx - The PetscDeviceContext

700:   Notes:
701:   This routine can be used to set the defacto "root" PetscDeviceContext to a user-defined
702:   implementation by calling this routine immediately after PetscInitialize() and ensuring that
703:   PetscDevice is not greedily intialized. In this case the user is responsible for destroying
704:   their PetscDeviceContext before PetscFinalize() returns.

706:   The old context is not stored in any way by this routine; if one is overriding a context that
707:   they themselves do not control, one should take care to temporarily store it by calling
708:   PetscDeviceContextGetCurrentContext() before calling this routine.

710:   Level: beginner

712: .seealso: PetscDeviceContextGetCurrentContext(), PetscDeviceContextFork(),
713: PetscDeviceContextJoin(), PetscDeviceContextCreate()
714: @*/
715: PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext dctx)
716: {
718:   PetscAssert(dctx->setup,PETSC_COMM_SELF,PETSC_ERR_ARG_WRONGSTATE,"PetscDeviceContext %" PetscInt_FMT " must be set up before being set as global context",dctx->id);
719:   globalContext = dctx;
720:   PetscInfo(nullptr,"Set global PetscDeviceContext id %" PetscInt_FMT "\n",dctx->id);
721:   return 0;
722: }

724: /*@C
725:   PetscDeviceContextSetFromOptions - Configure a PetscDeviceContext from the options database

727:   Collective on comm, Asynchronous

729:   Input Parameters:
730: + comm   - MPI communicator on which to query the options database
731: . prefix - prefix to prepend to all options database queries, NULL if not needed
732: - dctx   - The PetscDeviceContext to configure

734:   Output Parameter:
735: . dctx - The PetscDeviceContext

737:   Options Database:
738: + -device_context_stream_type - type of stream to create inside the PetscDeviceContext -
739:    PetscDeviceContextSetStreamType()
740: - -device_context_device_type - the type of PetscDevice to attach by default - PetscDeviceType

742:   Level: beginner

744: .seealso: PetscDeviceContextSetStreamType(), PetscDeviceContextSetDevice()
745: @*/
746: PetscErrorCode PetscDeviceContextSetFromOptions(MPI_Comm comm, const char prefix[], PetscDeviceContext dctx)
747: {
748:   PetscBool      flag;
749:   PetscInt       stype,dtype;

754:   PetscOptionsBegin(comm,prefix,"PetscDeviceContext Options","Sys");
755:   PetscOptionsEList("-device_context_stream_type","PetscDeviceContext PetscStreamType","PetscDeviceContextSetStreamType",PetscStreamTypes,PETSC_STREAM_MAX,PetscStreamTypes[dctx->streamType],&stype,&flag);
756:   if (flag) PetscDeviceContextSetStreamType(dctx,static_cast<PetscStreamType>(stype));
757:   PetscOptionsEList("-device_context_device_type","Underlying PetscDevice","PetscDeviceContextSetDevice",PetscDeviceTypes+1,PETSC_DEVICE_MAX-1,dctx->device ? PetscDeviceTypes[dctx->device->type] : PetscDeviceTypes[PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE],&dtype,&flag);
758:   if (flag) {
759:     PetscDeviceContextSetDefaultDeviceForType_Internal(dctx,static_cast<PetscDeviceType>(dtype+1));
760:   }
761:   PetscOptionsEnd();
762:   return 0;
763: }