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

  4: /* rootContext holds the default context created by PETSc, for cleanup during PetscFinalize().
  5:    globalContext selects the current context, which may instead be supplied by the application.
  6:    PetscDeviceContextSetCurrentContext() changes this selection without transferring ownership,
  7:    so rootContext must be kept separately for PETSc to clean up the context it created. */
  8: static auto               rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE_TYPE;
  9: static auto               rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE;
 10: static PetscDeviceContext globalContext  = nullptr;
 11: static PetscDeviceContext rootContext    = nullptr;

 13: /* when PetscDevice initializes PetscDeviceContext eagerly the type of device created should
 14:  * match whatever device is eagerly initialized */
 15: PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType type)
 16: {
 17:   PetscFunctionBegin;
 19:   rootDeviceType = type;
 20:   PetscFunctionReturn(PETSC_SUCCESS);
 21: }

 23: PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType type)
 24: {
 25:   PetscFunctionBegin;
 27:   rootStreamType = type;
 28:   PetscFunctionReturn(PETSC_SUCCESS);
 29: }

 31: static inline PetscErrorCode PetscSetDefaultCUPMStreamFromDeviceContext(PetscDeviceContext dctx, PetscDeviceType dtype)
 32: {
 33:   PetscFunctionBegin;
 34: #if PetscDefined(HAVE_CUDA)
 35:   if (dtype == PETSC_DEVICE_CUDA) {
 36:     void *handle;

 38:     PetscCall(PetscDeviceContextGetStreamHandle_Internal(dctx, &handle));
 39:     PetscDefaultCudaStream = *static_cast<cudaStream_t *>(handle);
 40:   }
 41: #endif
 42: #if PetscDefined(HAVE_HIP)
 43:   if (dtype == PETSC_DEVICE_HIP) {
 44:     void *handle;

 46:     PetscCall(PetscDeviceContextGetStreamHandle_Internal(dctx, &handle));
 47:     PetscDefaultHipStream = *static_cast<hipStream_t *>(handle);
 48:   }
 49: #endif
 50: #if !PetscDefined(HAVE_CUDA) && !PetscDefined(HAVE_HIP)
 51:   (void)dctx, (void)dtype;
 52: #endif
 53:   PetscFunctionReturn(PETSC_SUCCESS);
 54: }

 56: static PetscErrorCode PetscDeviceContextSetupGlobalContext_Private() noexcept
 57: {
 58:   PetscFunctionBegin;
 59:   if (PetscUnlikely(!globalContext)) {
 60:     PetscObject pobj;
 61:     const auto  dtype     = rootDeviceType;
 62:     const auto  finalizer = [] {
 63:       PetscDeviceType dtype;

 65:       PetscFunctionBegin;
 66:       PetscCall(PetscDeviceContextGetDeviceType(rootContext, &dtype));
 67:       PetscCall(PetscInfo(rootContext, "Destroying global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype]));
 68:       PetscCall(PetscDeviceContextDestroy(&rootContext));
 69:       globalContext  = nullptr;
 70:       rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE_TYPE;
 71:       rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE;
 72:       PetscFunctionReturn(PETSC_SUCCESS);
 73:     };

 75:     /* this exists purely as a valid device check. */
 76:     PetscCall(PetscDeviceInitializePackage());
 77:     PetscCall(PetscRegisterFinalize(std::move(finalizer)));
 78:     PetscCall(PetscDeviceContextCreate(&globalContext));
 79:     rootContext = globalContext;
 80:     PetscCall(PetscInfo(globalContext, "Initializing global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype]));
 81:     pobj = PetscObjectCast(globalContext);
 82:     PetscCall(PetscObjectSetName(pobj, "global root"));
 83:     PetscCall(PetscObjectSetOptionsPrefix(pobj, "root_"));
 84:     PetscCall(PetscDeviceContextSetStreamType(globalContext, rootStreamType));
 85:     PetscCall(PetscDeviceContextSetDefaultDeviceForType_Internal(globalContext, dtype));
 86:     PetscCall(PetscDeviceContextSetUp(globalContext));
 87:     PetscCall(PetscSetDefaultCUPMStreamFromDeviceContext(globalContext, dtype));
 88:   }
 89:   PetscFunctionReturn(PETSC_SUCCESS);
 90: }

 92: /*@
 93:   PetscDeviceContextGetCurrentContext - Get the current active `PetscDeviceContext`

 95:   Not Collective

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

100:   Notes:
101:   The user generally should not destroy contexts retrieved with this routine unless they
102:   themselves have created them. There exists no protection against destroying the root
103:   context.

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

109:   Level: beginner

111: .seealso: `PetscDeviceContextSetCurrentContext()`, `PetscDeviceContextFork()`,
112:           `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()`
113: @*/
114: PetscErrorCode PetscDeviceContextGetCurrentContext(PetscDeviceContext *dctx)
115: {
116:   PetscFunctionBegin;
117:   PetscAssertPointer(dctx, 1);
118:   PetscCall(PetscDeviceContextSetupGlobalContext_Private());
119:   /* while the static analyzer can find global variables, it will throw a warning about not
120:    * being able to connect this back to the function arguments */
122:   *dctx = globalContext;
123:   PetscFunctionReturn(PETSC_SUCCESS);
124: }

126: /*@
127:   PetscDeviceContextSetCurrentContext - Set the current active `PetscDeviceContext`

129:   Not Collective

131:   Input Parameter:
132: . dctx - The `PetscDeviceContext`

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

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

144:   Level: beginner

146: .seealso: `PetscDeviceContextGetCurrentContext()`, `PetscDeviceContextFork()`,
147:           `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()`
148: @*/
149: PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext dctx)
150: {
151:   PetscDeviceType dtype;

153:   PetscFunctionBegin;
154:   PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx));
155:   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);
156:   PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype));
157:   PetscCall(PetscDeviceSetDefaultDeviceType(dtype));
158:   globalContext = dctx;
159:   PetscCall(PetscInfo(dctx, "Set global PetscDeviceContext id %" PetscInt64_FMT "\n", PetscObjectCast(dctx)->id));
160:   PetscCall(PetscSetDefaultCUPMStreamFromDeviceContext(globalContext, dtype));
161:   PetscFunctionReturn(PETSC_SUCCESS);
162: }