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: {
 11:   PetscFunctionBegin;
 13:   rootDeviceType = type;
 14:   PetscFunctionReturn(PETSC_SUCCESS);
 15: }

 17: PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType type)
 18: {
 19:   PetscFunctionBegin;
 21:   rootStreamType = type;
 22:   PetscFunctionReturn(PETSC_SUCCESS);
 23: }

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

 34:       PetscFunctionBegin;
 35:       PetscCall(PetscDeviceContextGetDeviceType(globalContext, &dtype));
 36:       PetscCall(PetscInfo(globalContext, "Destroying global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype]));
 37:       PetscCall(PetscDeviceContextDestroy(&globalContext));
 38:       rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE_TYPE;
 39:       rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE;
 40:       PetscFunctionReturn(PETSC_SUCCESS);
 41:     };

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

 58: /*@C
 59:   PetscDeviceContextGetCurrentContext - Get the current active `PetscDeviceContext`

 61:   Not Collective

 63:   Output Parameter:
 64: . dctx - The `PetscDeviceContext`

 66:   Notes:
 67:   The user generally should not destroy contexts retrieved with this routine unless they
 68:   themselves have created them. There exists no protection against destroying the root
 69:   context.

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

 75:   Level: beginner

 77: .seealso: `PetscDeviceContextSetCurrentContext()`, `PetscDeviceContextFork()`,
 78:           `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()`
 79: @*/
 80: PetscErrorCode PetscDeviceContextGetCurrentContext(PetscDeviceContext *dctx)
 81: {
 82:   PetscFunctionBegin;
 83:   PetscAssertPointer(dctx, 1);
 84:   PetscCall(PetscDeviceContextSetupGlobalContext_Private());
 85:   /* while the static analyzer can find global variables, it will throw a warning about not
 86:    * being able to connect this back to the function arguments */
 88:   *dctx = globalContext;
 89:   PetscFunctionReturn(PETSC_SUCCESS);
 90: }

 92: /*@C
 93:   PetscDeviceContextSetCurrentContext - Set the current active `PetscDeviceContext`

 95:   Not Collective

 97:   Input Parameter:
 98: . dctx - The `PetscDeviceContext`

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

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

110:   Level: beginner

112: .seealso: `PetscDeviceContextGetCurrentContext()`, `PetscDeviceContextFork()`,
113:           `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()`
114: @*/
115: PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext dctx)
116: {
117:   PetscDeviceType dtype;

119:   PetscFunctionBegin;
120:   PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
121:   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);
122:   PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype));
123:   PetscCall(PetscDeviceSetDefaultDeviceType(dtype));
124:   globalContext = dctx;
125:   PetscCall(PetscInfo(dctx, "Set global PetscDeviceContext id %" PetscInt64_FMT "\n", PetscObjectCast(dctx)->id));
126:   PetscFunctionReturn(PETSC_SUCCESS);
127: }