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 initialized */ 9 PetscErrorCode PetscDeviceContextSetRootDeviceType_Internal(PetscDeviceType type) 10 { 11 PetscFunctionBegin; 12 PetscValidDeviceType(type, 1); 13 rootDeviceType = type; 14 PetscFunctionReturn(PETSC_SUCCESS); 15 } 16 17 PetscErrorCode PetscDeviceContextSetRootStreamType_Internal(PetscStreamType type) 18 { 19 PetscFunctionBegin; 20 PetscValidStreamType(type, 1); 21 rootStreamType = type; 22 PetscFunctionReturn(PETSC_SUCCESS); 23 } 24 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; 33 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 }; 42 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 } 57 58 /*@C 59 PetscDeviceContextGetCurrentContext - Get the current active `PetscDeviceContext` 60 61 Not Collective 62 63 Output Parameter: 64 . dctx - The `PetscDeviceContext` 65 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. 70 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()`. 74 75 Level: beginner 76 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 */ 87 PetscDisableStaticAnalyzerForExpressionUnderstandingThatThisIsDangerousAndBugprone(PetscValidDeviceContext(globalContext, -1)); 88 *dctx = globalContext; 89 PetscFunctionReturn(PETSC_SUCCESS); 90 } 91 92 /*@C 93 PetscDeviceContextSetCurrentContext - Set the current active `PetscDeviceContext` 94 95 Not Collective 96 97 Input Parameter: 98 . dctx - The `PetscDeviceContext` 99 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. 105 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. 109 110 Level: beginner 111 112 .seealso: `PetscDeviceContextGetCurrentContext()`, `PetscDeviceContextFork()`, 113 `PetscDeviceContextJoin()`, `PetscDeviceContextCreate()` 114 @*/ 115 PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext dctx) 116 { 117 PetscDeviceType dtype; 118 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 } 128