Actual source code: global_dcontext.cxx

  1: #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/

  3: static auto               rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE_TYPE;
  4: static auto               rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE;
  5: static PetscDeviceContext globalContext  = nullptr;

  7: /* when PetsDevice initializes PetscDeviceContext eagerly the type of device created should
  8:  * match whatever device is eagerly initialized */
  9: PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType type)
 10: {
 12:   rootDeviceType = type;
 13:   return 0;
 14: }

 16: PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType type)
 17: {
 19:   rootStreamType = type;
 20:   return 0;
 21: }

 23: static PetscErrorCode PetscDeviceContextSetupGlobalContext_Private() noexcept
 24: {
 25:   if (PetscUnlikely(!globalContext)) {
 26:     PetscObject pobj;
 27:     const auto  dtype     = rootDeviceType;
 28:     const auto  finalizer = [] {
 29:       PetscDeviceType dtype;

 31:       PetscDeviceContextGetDeviceType(globalContext, &dtype);
 32:       PetscInfo(globalContext, "Destroying global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype]);
 33:       PetscDeviceContextDestroy(&globalContext);
 34:       rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE_TYPE;
 35:       rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE;
 36:       return 0;
 37:     };

 39:     /* this exists purely as a valid device check. */
 40:     PetscDeviceInitializePackage();
 41:     PetscRegisterFinalize(std::move(finalizer));
 42:     PetscDeviceContextCreate(&globalContext);
 43:     PetscInfo(globalContext, "Initializing global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype]);
 44:     pobj = PetscObjectCast(globalContext);
 45:     PetscObjectSetName(pobj, "global root");
 46:     PetscObjectSetOptionsPrefix(pobj, "root_");
 47:     PetscDeviceContextSetStreamType(globalContext, rootStreamType);
 48:     PetscDeviceContextSetDefaultDeviceForType_Internal(globalContext, dtype);
 49:     PetscDeviceContextSetUp(globalContext);
 50:   }
 51:   return 0;
 52: }

 54: /*@C
 55:   PetscDeviceContextGetCurrentContext - Get the current active `PetscDeviceContext`

 57:   Not Collective

 59:   Output Parameter:
 60: . dctx - The `PetscDeviceContext`

 62:   Notes:
 63:   The user generally should not destroy contexts retrieved with this routine unless they
 64:   themselves have created them. There exists no protection against destroying the root
 65:   context.

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

 71:   Level: beginner

 73: .seealso: `PetscDeviceContextSetCurrentContext()`, `PetscDeviceContextFork()`,
 74:           `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()`
 75: @*/
 76: PetscErrorCode PetscDeviceContextGetCurrentContext(PetscDeviceContext *dctx)
 77: {
 79:   PetscDeviceContextSetupGlobalContext_Private();
 80:   /* while the static analyzer can find global variables, it will throw a warning about not
 81:    * being able to connect this back to the function arguments */
 83:   *dctx = globalContext;
 84:   return 0;
 85: }

 87: /*@C
 88:   PetscDeviceContextSetCurrentContext - Set the current active `PetscDeviceContext`

 90:   Not Collective

 92:   Input Parameter:
 93: . dctx - The `PetscDeviceContext`

 95:   Notes:
 96:   This routine can be used to set the defacto "root" `PetscDeviceContext` to a user-defined
 97:   implementation by calling this routine immediately after `PetscInitialize()` and ensuring that
 98:   `PetscDevice` is not greedily initialized. In this case the user is responsible for destroying
 99:   their `PetscDeviceContext` before `PetscFinalize()` returns.

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

105:   Level: beginner

107: .seealso: `PetscDeviceContextGetCurrentContext()`, `PetscDeviceContextFork()`,
108:           `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()`
109: @*/
110: PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext dctx)
111: {
112:   PetscDeviceType dtype;

114:   PetscDeviceContextGetOptionalNullContext_Internal(&dctx);
115:   PetscAssert(dctx->setup, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONGSTATE, "PetscDeviceContext %" PetscInt64_FMT " must be set up before being set as global context", PetscObjectCast(dctx)->id);
116:   PetscDeviceContextGetDeviceType(dctx, &dtype);
117:   PetscDeviceSetDefaultDeviceType(dtype);
118:   globalContext = dctx;
119:   PetscInfo(dctx, "Set global PetscDeviceContext id %" PetscInt64_FMT "\n", PetscObjectCast(dctx)->id);
120:   return 0;
121: }