1 #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/ 2 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; 6 7 /* when PetsDevice initializes PetscDeviceContext eagerly the type of device created should 8 * match whatever device is eagerly intialized */ 9 PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType type) { 10 PetscFunctionBegin; 11 PetscValidDeviceType(type, 1); 12 rootDeviceType = type; 13 PetscFunctionReturn(0); 14 } 15 16 PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType type) { 17 PetscFunctionBegin; 18 PetscValidStreamType(type, 1); 19 rootStreamType = type; 20 PetscFunctionReturn(0); 21 } 22 23 static PetscErrorCode PetscDeviceContextSetupGlobalContext_Private() noexcept { 24 PetscFunctionBegin; 25 if (PetscUnlikely(!globalContext)) { 26 PetscObject pobj; 27 const auto dtype = rootDeviceType; 28 const auto finalizer = [] { 29 PetscDeviceType dtype; 30 31 PetscFunctionBegin; 32 PetscCall(PetscDeviceContextGetDeviceType(globalContext, &dtype)); 33 PetscCall(PetscInfo(globalContext, "Destroying global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype])); 34 PetscCall(PetscDeviceContextDestroy(&globalContext)); 35 rootDeviceType = PETSC_DEVICE_CONTEXT_DEFAULT_DEVICE_TYPE; 36 rootStreamType = PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE; 37 PetscFunctionReturn(0); 38 }; 39 40 /* this exists purely as a valid device check. */ 41 PetscCall(PetscDeviceInitializePackage()); 42 PetscCall(PetscRegisterFinalize(std::move(finalizer))); 43 PetscCall(PetscDeviceContextCreate(&globalContext)); 44 PetscCall(PetscInfo(globalContext, "Initializing global PetscDeviceContext with device type %s\n", PetscDeviceTypes[dtype])); 45 pobj = PetscObjectCast(globalContext); 46 PetscCall(PetscObjectSetName(pobj, "global root")); 47 PetscCall(PetscObjectSetOptionsPrefix(pobj, "root_")); 48 PetscCall(PetscDeviceContextSetStreamType(globalContext, rootStreamType)); 49 PetscCall(PetscDeviceContextSetDefaultDeviceForType_Internal(globalContext, dtype)); 50 PetscCall(PetscDeviceContextSetUp(globalContext)); 51 } 52 PetscFunctionReturn(0); 53 } 54 55 /*@C 56 PetscDeviceContextGetCurrentContext - Get the current active `PetscDeviceContext` 57 58 Not Collective 59 60 Output Parameter: 61 . dctx - The `PetscDeviceContext` 62 63 Notes: 64 The user generally should not destroy contexts retrieved with this routine unless they 65 themselves have created them. There exists no protection against destroying the root 66 context. 67 68 Developer Notes: 69 Unless the user has set their own, this routine creates the "root" context the first time it 70 is called, registering its destructor to `PetscFinalize()`. 71 72 Level: beginner 73 74 .seealso: `PetscDeviceContextSetCurrentContext()`, `PetscDeviceContextFork()`, 75 `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()` 76 @*/ 77 PetscErrorCode PetscDeviceContextGetCurrentContext(PetscDeviceContext *dctx) { 78 PetscFunctionBegin; 79 PetscValidPointer(dctx, 1); 80 PetscCall(PetscDeviceContextSetupGlobalContext_Private()); 81 /* while the static analyzer can find global variables, it will throw a warning about not 82 * being able to connect this back to the function arguments */ 83 PetscDisableStaticAnalyzerForExpressionUnderstandingThatThisIsDangerousAndBugprone(PetscValidDeviceContext(globalContext, -1)); 84 *dctx = globalContext; 85 PetscFunctionReturn(0); 86 } 87 88 /*@C 89 PetscDeviceContextSetCurrentContext - Set the current active `PetscDeviceContext` 90 91 Not Collective 92 93 Input Parameter: 94 . dctx - The `PetscDeviceContext` 95 96 Notes: 97 This routine can be used to set the defacto "root" `PetscDeviceContext` to a user-defined 98 implementation by calling this routine immediately after `PetscInitialize()` and ensuring that 99 `PetscDevice` is not greedily intialized. In this case the user is responsible for destroying 100 their `PetscDeviceContext` before `PetscFinalize()` returns. 101 102 The old context is not stored in any way by this routine; if one is overriding a context that 103 they themselves do not control, one should take care to temporarily store it by calling 104 `PetscDeviceContextGetCurrentContext()` before calling this routine. 105 106 Level: beginner 107 108 .seealso: `PetscDeviceContextGetCurrentContext()`, `PetscDeviceContextFork()`, 109 `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()` 110 @*/ 111 PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext dctx) { 112 PetscDeviceType dtype; 113 114 PetscFunctionBegin; 115 PetscCall(PetscDeviceContextGetOptionalNullContext_Internal(&dctx)); 116 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); 117 PetscCall(PetscDeviceContextGetDeviceType(dctx, &dtype)); 118 PetscCall(PetscDeviceSetDefaultDeviceType(dtype)); 119 globalContext = dctx; 120 PetscCall(PetscInfo(dctx, "Set global PetscDeviceContext id %" PetscInt64_FMT "\n", PetscObjectCast(dctx)->id)); 121 PetscFunctionReturn(0); 122 } 123