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