Actual source code: device.cxx
1: #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/
2: #include <petsc/private/petscadvancedmacros.h>
4: #include <petsc/private/cpp/register_finalize.hpp>
6: #include "../impls/host/hostdevice.hpp"
7: #if PetscDefined(HAVE_CUPM)
8: #include "../impls/cupm/cupmdevice.hpp"
9: #endif
10: #if PetscDefined(HAVE_SYCL)
11: #include "../impls/sycl/sycldevice.hpp"
12: #endif
14: #include <utility> // std::make_pair
16: using namespace Petsc::device;
18: #if defined(PETSC_HAVE_CUPM)
19: int PetscDeviceCUPMRuntimeArch = 0;
20: #endif
22: namespace
23: {
25: /*
26: note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to
27: be picked up by the switch-case macros below
28: */
29: host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
30: #if PetscDefined(HAVE_CUDA)
31: cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
32: #endif
33: #if PetscDefined(HAVE_HIP)
34: cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
35: #endif
36: #if PetscDefined(HAVE_SYCL)
37: sycl::Device SYCLDevice{PetscDeviceContextCreate_SYCL};
38: #endif
40: } // namespace
42: #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
43: case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
44: PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
45: } break
47: #define PETSC_VOID_0(...) ((void)0)
49: /*
50: Suppose you have:
52: CUDADevice.myFunction(arg1,arg2)
54: that you would like to conditionally define and call in a switch-case:
56: switch(PetscDeviceType) {
57: #if PetscDefined(HAVE_CUDA)
58: case PETSC_DEVICE_CUDA: {
59: PetscCall(CUDADevice.myFunction(arg1,arg2));
60: } break;
61: #endif
62: }
64: then calling this macro:
66: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
68: will expand to the following case statement:
70: case PETSC_DEVICE_CUDA: {
71: PetscCall(CUDADevice.myFunction(arg1,arg2));
72: } break
74: if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise
75: */
76: #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(PetscConcat_(HAVE_, IMPLS), PETSC_DEVICE_CASE, PETSC_VOID_0)(IMPLS, func, __VA_ARGS__)
78: /*@C
79: PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
81: Not Collective
83: Input Parameters:
84: + type - The type of `PetscDevice`
85: - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically)
87: Output Parameter:
88: . device - The `PetscDevice`
90: Level: beginner
92: Notes:
93: This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
94: device synchronization.
96: `devid` is what you might pass to `cudaSetDevice()` for example.
98: .seealso: `PetscDevice`, `PetscDeviceInitType`,
99: `PetscDeviceInitialize()`, `PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
100: `PetscDeviceView()`, `PetscDeviceDestroy()`
101: @*/
102: PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device)
103: {
104: static PetscInt PetscDeviceCounter = 0;
106: PetscFunctionBegin;
108: PetscAssertPointer(device, 3);
109: PetscCall(PetscDeviceInitializePackage());
110: PetscCall(PetscNew(device));
111: (*device)->id = PetscDeviceCounter++;
112: (*device)->type = type;
113: (*device)->refcnt = 1;
114: /*
115: if you are adding a device, you also need to add its initialization in
116: PetscDeviceInitializeTypeFromOptions_Private() below
117: */
118: switch (type) {
119: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
120: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
121: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
122: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
123: default:
124: /* in case the above macros expand to nothing this silences any unused variable warnings */
125: (void)(devid);
126: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
127: }
128: PetscFunctionReturn(PETSC_SUCCESS);
129: }
131: /*@C
132: PetscDeviceDestroy - Free a `PetscDevice`
134: Not Collective
136: Input Parameter:
137: . device - The `PetscDevice`
139: Level: beginner
141: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
142: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
143: @*/
144: PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
145: {
146: PetscFunctionBegin;
147: PetscAssertPointer(device, 1);
148: if (!*device) PetscFunctionReturn(PETSC_SUCCESS);
150: PetscCall(PetscDeviceDereference_Internal(*device));
151: if ((*device)->refcnt) {
152: *device = nullptr;
153: PetscFunctionReturn(PETSC_SUCCESS);
154: }
155: PetscCall(PetscFree((*device)->data));
156: PetscCall(PetscFree(*device));
157: PetscFunctionReturn(PETSC_SUCCESS);
158: }
160: /*@C
161: PetscDeviceConfigure - Configure a particular `PetscDevice`
163: Not Collective
165: Input Parameter:
166: . device - The `PetscDevice` to configure
168: Level: beginner
170: Notes:
171: The user should not assume that this is a cheap operation.
173: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
174: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
175: @*/
176: PetscErrorCode PetscDeviceConfigure(PetscDevice device)
177: {
178: PetscFunctionBegin;
180: /*
181: if no available configuration is available, this cascades all the way down to default
182: and error
183: */
184: switch (const auto dtype = device->type) {
185: case PETSC_DEVICE_HOST:
186: if (PetscDefined(HAVE_HOST)) break; // always true
187: case PETSC_DEVICE_CUDA:
188: if (PetscDefined(HAVE_CUDA)) break;
189: goto error;
190: case PETSC_DEVICE_HIP:
191: if (PetscDefined(HAVE_HIP)) break;
192: goto error;
193: case PETSC_DEVICE_SYCL:
194: if (PetscDefined(HAVE_SYCL)) break;
195: goto error;
196: default:
197: error:
198: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
199: }
200: PetscUseTypeMethod(device, configure);
201: PetscFunctionReturn(PETSC_SUCCESS);
202: }
204: /*@C
205: PetscDeviceView - View a `PetscDevice`
207: Collective on viewer
209: Input Parameters:
210: + device - The `PetscDevice` to view
211: - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
213: Level: beginner
215: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
216: `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
217: @*/
218: PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer)
219: {
220: auto sub = viewer;
221: PetscBool iascii;
223: PetscFunctionBegin;
225: if (viewer) {
227: PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
228: } else {
229: PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
230: iascii = PETSC_TRUE;
231: }
233: if (iascii) {
234: auto dtype = PETSC_DEVICE_HOST;
235: MPI_Comm comm;
236: PetscMPIInt size;
237: PetscInt id = 0;
239: PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
240: PetscCallMPI(MPI_Comm_size(comm, &size));
242: PetscCall(PetscDeviceGetDeviceId(device, &id));
243: PetscCall(PetscDeviceGetType(device, &dtype));
244: PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
245: PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
246: PetscCall(PetscViewerASCIIPushTab(sub));
247: PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
248: PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
249: }
251: // see if impls has extra viewer stuff
252: PetscTryTypeMethod(device, view, sub);
254: if (iascii) {
255: // undo the ASCII specific stuff
256: PetscCall(PetscViewerASCIIPopTab(sub));
257: PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
258: }
259: PetscFunctionReturn(PETSC_SUCCESS);
260: }
262: /*@
263: PetscDeviceGetType - Get the type of device
265: Not Collective
267: Input Parameter:
268: . device - The `PetscDevice`
270: Output Parameter:
271: . type - The `PetscDeviceType`
273: Level: beginner
275: .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
276: `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
277: `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
278: @*/
279: PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type)
280: {
281: PetscFunctionBegin;
283: PetscAssertPointer(type, 2);
284: *type = device->type;
285: PetscFunctionReturn(PETSC_SUCCESS);
286: }
288: /*@C
289: PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
291: Not Collective
293: Input Parameter:
294: . device - The `PetscDevice`
296: Output Parameter:
297: . id - The id
299: Level: beginner
301: Notes:
302: The returned ID may have been assigned by the underlying device backend. For example if the
303: backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
304: this device was configured.
306: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
307: @*/
308: PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
309: {
310: PetscFunctionBegin;
312: PetscAssertPointer(id, 2);
313: *id = device->deviceId;
314: PetscFunctionReturn(PETSC_SUCCESS);
315: }
317: namespace
318: {
320: struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
321: PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
323: PetscErrorCode finalize_() noexcept
324: {
325: PetscFunctionBegin;
326: type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
327: PetscFunctionReturn(PETSC_SUCCESS);
328: }
329: };
331: auto default_device_type = DefaultDeviceType();
333: } // namespace
335: /*@C
336: PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
338: Not Collective
340: Level: beginner
342: Notes:
343: Unless selected by the user, the default device is selected in the following order\:
344: `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
346: .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
347: @*/
348: PetscDeviceType PETSC_DEVICE_DEFAULT(void)
349: {
350: return default_device_type.type;
351: }
353: /*@C
354: PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
356: Not Collective
358: Input Parameter:
359: . type - the new default device type
361: Level: beginner
363: Notes:
364: This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
366: .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
367: @*/
368: PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
369: {
370: PetscFunctionBegin;
372: if (default_device_type.type != type) {
373: // no need to waster a PetscRegisterFinalize() slot if we don't change it
374: default_device_type.type = type;
375: PetscCall(default_device_type.register_finalize());
376: }
377: PetscFunctionReturn(PETSC_SUCCESS);
378: }
380: namespace
381: {
383: std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
385: /*
386: Actual initialization function; any functions claiming to initialize PetscDevice or
387: PetscDeviceContext will have to run through this one
388: */
389: PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId)
390: {
391: PetscFunctionBegin;
393: if (PetscUnlikely(!PetscDeviceInitialized(type))) {
394: auto &dev = defaultDevices[type].first;
395: auto &init = defaultDevices[type].second;
397: PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
398: PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
399: PetscCall(PetscDeviceConfigure(dev));
400: init = true;
401: }
402: PetscFunctionReturn(PETSC_SUCCESS);
403: }
405: } // namespace
407: /*@C
408: PetscDeviceInitialize - Initialize `PetscDevice`
410: Not Collective
412: Input Parameter:
413: . type - The `PetscDeviceType` to initialize
415: Level: beginner
417: Notes:
418: Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
419: result in device synchronization.
421: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
422: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
423: @*/
424: PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
425: {
426: PetscFunctionBegin;
428: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
429: PetscFunctionReturn(PETSC_SUCCESS);
430: }
432: /*@C
433: PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
434: `PetscDeviceType`
436: Not Collective
438: Input Parameter:
439: . type - The `PetscDeviceType` to check
441: Level: beginner
443: Notes:
444: Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
446: If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
447: return `PETSC_FALSE` for that `PetscDeviceType`.
449: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
450: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
451: @*/
452: PetscBool PetscDeviceInitialized(PetscDeviceType type)
453: {
454: return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
455: }
457: /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
458: PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device)
459: {
460: PetscFunctionBegin;
461: PetscAssertPointer(device, 2);
462: PetscCall(PetscDeviceInitialize(type));
463: *device = defaultDevices[type].first;
464: PetscFunctionReturn(PETSC_SUCCESS);
465: }
467: /*@C
468: PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
470: Not Collective
472: Input Parameters:
473: + device - The `PetscDevice`
474: - attr - The attribute
476: Output Parameter:
477: . value - The value of the attribute
479: Level: intermediate
481: Notes:
482: Since different attributes are often different types `value` is a `void *` to accommodate
483: them all. The underlying type of the attribute is therefore included in the name of the
484: `PetscDeviceAttribute` responsible for querying it. For example,
485: `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
487: .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
488: @*/
489: PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
490: {
491: PetscFunctionBegin;
494: PetscAssertPointer(value, 3);
495: PetscUseTypeMethod(device, getattribute, attr, value);
496: PetscFunctionReturn(PETSC_SUCCESS);
497: }
499: namespace
500: {
502: PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType)
503: {
504: PetscFunctionBegin;
505: if (!PetscDeviceConfiguredFor_Internal(type)) {
506: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
507: defaultDevices[type].first = nullptr;
508: PetscFunctionReturn(PETSC_SUCCESS);
509: }
510: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
511: /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
512: switch (type) {
513: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
514: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
515: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
516: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
517: default:
518: SETERRQ(comm, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
519: }
520: PetscCall(PetscInfo(nullptr, "PetscDevice %s initialized, default device id %" PetscInt_FMT ", view %s, init type %s\n", PetscDeviceTypes[type], defaultDeviceId, PetscBools[defaultView], PetscDeviceInitTypes[Petsc::util::to_underlying(*defaultInitType)]));
521: /*
522: defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES
523: have decided to initialize as
524: */
525: if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
526: PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
527: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
528: if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
529: }
530: PetscFunctionReturn(PETSC_SUCCESS);
531: }
533: PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView)
534: {
535: PetscInt initIdx = PETSC_DEVICE_INIT_LAZY;
536: auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
537: auto flg = PETSC_FALSE;
539: PetscFunctionBegin;
540: PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
541: if (flg) PetscCall(PetscLogGpuTime());
543: PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
544: PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
545: PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
546: PetscCall(PetscOptionsRangeInt("-device_select", "Which device to use. Pass " PetscStringize(PETSC_DECIDE) " to have PETSc decide or (given they exist) [0-" PetscStringize(PETSC_DEVICE_MAX_DEVICES) ") for a specific device", "PetscDeviceCreate()", *defaultDevice, defaultDevice, nullptr, PETSC_DECIDE, PETSC_DEVICE_MAX_DEVICES));
547: PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
548: PetscOptionsEnd();
550: if (initIdx == PETSC_DEVICE_INIT_NONE) {
551: /* disabled all device initialization if devices are globally disabled */
552: PetscCheck(*defaultDevice == PETSC_DECIDE, comm, PETSC_ERR_USER_INPUT, "You have disabled devices but also specified a particular device to use, these options are mutually exclusive");
553: *defaultView = PETSC_FALSE;
554: initDeviceIdx = PETSC_DEVICE_HOST;
555: } else {
556: *defaultView = static_cast<PetscBool>(*defaultView && flg);
557: if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
558: }
559: *defaultInitType = PetscDeviceInitTypeCast(initIdx);
560: *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
561: PetscFunctionReturn(PETSC_SUCCESS);
562: }
564: /* called from PetscFinalize() do not call yourself! */
565: PetscErrorCode PetscDeviceFinalize_Private()
566: {
567: PetscFunctionBegin;
568: if (PetscDefined(USE_DEBUG)) {
569: /*
570: you might be thinking, why on earth are you registered yet another finalizer in a
571: function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
572: because it is.
574: The crux of the problem is that the initializer (and therefore the ~finalizer~) of
575: PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
576: a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
577: won't be destroyed yet. So we need to repeat the check that all devices have been
578: destroyed again ~after~ the global context is destroyed. In summary:
580: 1. This finalizer runs and destroys all devices, except it may not because the global
581: context may still hold a reference!
582: 2. The global context finalizer runs and does the final reference count decrement
583: required, which actually destroys the held device.
584: 3. Our newly added finalizer runs and checks that all is well.
585: */
586: PetscCall(PetscRegisterFinalize([] {
587: PetscFunctionBegin;
588: for (auto &&device : defaultDevices) {
589: const auto dev = device.first;
591: PetscCheck(!dev, PETSC_COMM_WORLD, PETSC_ERR_COR, "Device of type '%s' had reference count %" PetscInt_FMT " and was not fully destroyed during PetscFinalize()", PetscDeviceTypes[dev->type], dev->refcnt);
592: }
593: PetscFunctionReturn(PETSC_SUCCESS);
594: }));
595: }
596: for (auto &&device : defaultDevices) {
597: PetscCall(PetscDeviceDestroy(&device.first));
598: device.second = false;
599: }
600: PetscFunctionReturn(PETSC_SUCCESS);
601: }
603: } // namespace
605: /*
606: Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
607: initialization types:
609: 1. defaultInitType - how does PetscDevice as a whole expect to initialize?
610: 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
611: e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
612: have all CUDA devices still initialize.
614: All told the following happens:
616: 0. defaultInitType -> LAZY
617: 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
618: 2. PetscDevice initializes each sub type with deviceDefaultInitType.
619: 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
620: to checking for specific device init. if view or specific device init
621: subTypeDefaultInitType -> EAGER. disabled once again overrides all.
622: */
624: PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm)
625: {
626: auto defaultView = PETSC_FALSE;
627: auto initializeDeviceContextEagerly = PETSC_FALSE;
628: auto defaultDeviceSet = PETSC_FALSE;
629: auto defaultDevice = PetscInt{PETSC_DECIDE};
630: auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT();
631: auto defaultInitType = PETSC_DEVICE_INIT_LAZY;
633: PetscFunctionBegin;
634: if (PetscDefined(USE_DEBUG)) {
635: int result;
637: PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
638: /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
639: * global space */
640: if (PetscUnlikely(result != MPI_IDENT)) {
641: char name[MPI_MAX_OBJECT_NAME] = {};
642: int len; /* unused */
644: PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
645: SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
646: }
647: }
648: comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
649: PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
651: PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
653: // the precise values don't matter here, so long as they are sequential
654: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HOST) == 0, "");
655: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_CUDA) == 1, "");
656: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HIP) == 2, "");
657: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_SYCL) == 3, "");
658: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_MAX) == 4, "");
659: for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
660: const auto deviceType = PetscDeviceTypeCast(i);
661: auto initType = defaultInitType;
663: PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
664: if (PetscDeviceConfiguredFor_Internal(deviceType)) {
665: if (initType == PETSC_DEVICE_INIT_EAGER) {
666: initializeDeviceContextEagerly = PETSC_TRUE;
667: // only update the default device if the user hasn't set it previously
668: if (!defaultDeviceSet) {
669: deviceContextInitDevice = deviceType;
670: PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
671: }
672: } else if (initType == PETSC_DEVICE_INIT_NONE) {
673: if (deviceType != PETSC_DEVICE_HOST) PetscCheck(!defaultDeviceSet || (deviceType != deviceContextInitDevice), comm, PETSC_ERR_USER_INPUT, "Cannot explicitly disable the device set as default device type (%s)", PetscDeviceTypes[deviceType]);
674: }
675: }
676: }
678: PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
679: PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
680: /* ----------------------------------------------------------------------------------- */
681: /* PetscDevice is now fully initialized */
682: /* ----------------------------------------------------------------------------------- */
683: {
684: /*
685: query the options db to get the root settings from the user (if any).
687: This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
688: PetscDeviceContextSetFromOptions() before we even have one, then set a few static
689: variables in that file with the results.
690: */
691: auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
692: auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
694: PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
695: PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
696: PetscOptionsEnd();
698: if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
699: if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
700: }
702: if (initializeDeviceContextEagerly) {
703: PetscDeviceContext dctx;
705: PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
706: /* instantiates the device context */
707: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
708: PetscCall(PetscDeviceContextSetUp(dctx));
709: }
710: PetscFunctionReturn(PETSC_SUCCESS);
711: }