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: namespace
19: {
21: /*
22: note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to
23: be picked up by the switch-case macros below
24: */
25: host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
26: #if PetscDefined(HAVE_CUDA)
27: cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
28: #endif
29: #if PetscDefined(HAVE_HIP)
30: cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
31: #endif
32: #if PetscDefined(HAVE_SYCL)
33: sycl::Device SYCLDevice{PetscDeviceContextCreate_SYCL};
34: #endif
36: } // namespace
38: #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
39: case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
40: PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
41: } break
43: #define PETSC_VOID_0(...) ((void)0)
45: /*
46: Suppose you have:
48: CUDADevice.myFunction(arg1,arg2)
50: that you would like to conditionally define and call in a switch-case:
52: switch(PetscDeviceType) {
53: #if PetscDefined(HAVE_CUDA)
54: case PETSC_DEVICE_CUDA: {
55: PetscCall(CUDADevice.myFunction(arg1,arg2));
56: } break;
57: #endif
58: }
60: then calling this macro:
62: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
64: will expand to the following case statement:
66: case PETSC_DEVICE_CUDA: {
67: PetscCall(CUDADevice.myFunction(arg1,arg2));
68: } break
70: if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise
71: */
72: #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(PetscConcat_(HAVE_, IMPLS), PETSC_DEVICE_CASE, PETSC_VOID_0)(IMPLS, func, __VA_ARGS__)
74: /*@C
75: PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
77: Not Collective
79: Input Parameters:
80: + type - The type of `PetscDevice`
81: - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically)
83: Output Parameter:
84: . device - The `PetscDevice`
86: Level: beginner
88: Notes:
89: This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
90: device synchronization.
92: `devid` is what you might pass to `cudaSetDevice()` for example.
94: .seealso: `PetscDevice`, `PetscDeviceInitType`,
95: `PetscDeviceInitialize()`, `PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
96: `PetscDeviceView()`, `PetscDeviceDestroy()`
97: @*/
98: PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device)
99: {
100: static PetscInt PetscDeviceCounter = 0;
102: PetscFunctionBegin;
104: PetscAssertPointer(device, 3);
105: PetscCall(PetscDeviceInitializePackage());
106: PetscCall(PetscNew(device));
107: (*device)->id = PetscDeviceCounter++;
108: (*device)->type = type;
109: (*device)->refcnt = 1;
110: /*
111: if you are adding a device, you also need to add its initialization in
112: PetscDeviceInitializeTypeFromOptions_Private() below
113: */
114: switch (type) {
115: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
116: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
117: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
118: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
119: default:
120: /* in case the above macros expand to nothing this silences any unused variable warnings */
121: (void)(devid);
122: 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]);
123: }
124: PetscFunctionReturn(PETSC_SUCCESS);
125: }
127: /*@C
128: PetscDeviceDestroy - Free a `PetscDevice`
130: Not Collective
132: Input Parameter:
133: . device - The `PetscDevice`
135: Level: beginner
137: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
138: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
139: @*/
140: PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
141: {
142: PetscFunctionBegin;
143: PetscAssertPointer(device, 1);
144: if (!*device) PetscFunctionReturn(PETSC_SUCCESS);
146: PetscCall(PetscDeviceDereference_Internal(*device));
147: if ((*device)->refcnt) {
148: *device = nullptr;
149: PetscFunctionReturn(PETSC_SUCCESS);
150: }
151: PetscCall(PetscFree((*device)->data));
152: PetscCall(PetscFree(*device));
153: PetscFunctionReturn(PETSC_SUCCESS);
154: }
156: /*@C
157: PetscDeviceConfigure - Configure a particular `PetscDevice`
159: Not Collective
161: Input Parameter:
162: . device - The `PetscDevice` to configure
164: Level: beginner
166: Notes:
167: The user should not assume that this is a cheap operation.
169: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
170: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
171: @*/
172: PetscErrorCode PetscDeviceConfigure(PetscDevice device)
173: {
174: PetscFunctionBegin;
176: /*
177: if no available configuration is available, this cascades all the way down to default
178: and error
179: */
180: switch (const auto dtype = device->type) {
181: case PETSC_DEVICE_HOST:
182: if (PetscDefined(HAVE_HOST)) break; // always true
183: case PETSC_DEVICE_CUDA:
184: if (PetscDefined(HAVE_CUDA)) break;
185: goto error;
186: case PETSC_DEVICE_HIP:
187: if (PetscDefined(HAVE_HIP)) break;
188: goto error;
189: case PETSC_DEVICE_SYCL:
190: if (PetscDefined(HAVE_SYCL)) break;
191: goto error;
192: default:
193: error:
194: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
195: }
196: PetscUseTypeMethod(device, configure);
197: PetscFunctionReturn(PETSC_SUCCESS);
198: }
200: /*@C
201: PetscDeviceView - View a `PetscDevice`
203: Collective on viewer
205: Input Parameters:
206: + device - The `PetscDevice` to view
207: - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
209: Level: beginner
211: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
212: `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
213: @*/
214: PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer)
215: {
216: auto sub = viewer;
217: PetscBool iascii;
219: PetscFunctionBegin;
221: if (viewer) {
223: PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
224: } else {
225: PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
226: iascii = PETSC_TRUE;
227: }
229: if (iascii) {
230: auto dtype = PETSC_DEVICE_HOST;
231: MPI_Comm comm;
232: PetscMPIInt size;
233: PetscInt id = 0;
235: PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
236: PetscCallMPI(MPI_Comm_size(comm, &size));
238: PetscCall(PetscDeviceGetDeviceId(device, &id));
239: PetscCall(PetscDeviceGetType(device, &dtype));
240: PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
241: PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
242: PetscCall(PetscViewerASCIIPushTab(sub));
243: PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
244: PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
245: }
247: // see if impls has extra viewer stuff
248: PetscTryTypeMethod(device, view, sub);
250: if (iascii) {
251: // undo the ASCII specific stuff
252: PetscCall(PetscViewerASCIIPopTab(sub));
253: PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
254: PetscCall(PetscViewerFlush(viewer));
255: }
256: PetscFunctionReturn(PETSC_SUCCESS);
257: }
259: /*@C
260: PetscDeviceGetType - Get the type of device
262: Not Collective
264: Input Parameter:
265: . device - The `PetscDevice`
267: Output Parameter:
268: . type - The `PetscDeviceType`
270: Level: beginner
272: .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
273: `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
274: `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
275: @*/
276: PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type)
277: {
278: PetscFunctionBegin;
280: PetscAssertPointer(type, 2);
281: *type = device->type;
282: PetscFunctionReturn(PETSC_SUCCESS);
283: }
285: /*@C
286: PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
288: Not Collective
290: Input Parameter:
291: . device - The `PetscDevice`
293: Output Parameter:
294: . id - The id
296: Level: beginner
298: Notes:
299: The returned ID may have been assigned by the underlying device backend. For example if the
300: backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
301: this device was configured.
303: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
304: @*/
305: PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
306: {
307: PetscFunctionBegin;
309: PetscAssertPointer(id, 2);
310: *id = device->deviceId;
311: PetscFunctionReturn(PETSC_SUCCESS);
312: }
314: namespace
315: {
317: struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
318: PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
320: PetscErrorCode finalize_() noexcept
321: {
322: PetscFunctionBegin;
323: type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
324: PetscFunctionReturn(PETSC_SUCCESS);
325: }
326: };
328: auto default_device_type = DefaultDeviceType();
330: } // namespace
332: /*@C
333: PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
335: Not Collective
337: Level: beginner
339: Notes:
340: Unless selected by the user, the default device is selected in the following order\:
341: `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
343: .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
344: @*/
345: PetscDeviceType PETSC_DEVICE_DEFAULT(void)
346: {
347: return default_device_type.type;
348: }
350: /*@C
351: PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
353: Not Collective
355: Input Parameter:
356: . type - the new default device type
358: Level: beginner
360: Notes:
361: This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
363: .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
364: @*/
365: PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
366: {
367: PetscFunctionBegin;
369: if (default_device_type.type != type) {
370: // no need to waster a PetscRegisterFinalize() slot if we don't change it
371: default_device_type.type = type;
372: PetscCall(default_device_type.register_finalize());
373: }
374: PetscFunctionReturn(PETSC_SUCCESS);
375: }
377: namespace
378: {
380: std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
382: /*
383: Actual initialization function; any functions claiming to initialize PetscDevice or
384: PetscDeviceContext will have to run through this one
385: */
386: PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId)
387: {
388: PetscFunctionBegin;
390: if (PetscUnlikely(!PetscDeviceInitialized(type))) {
391: auto &dev = defaultDevices[type].first;
392: auto &init = defaultDevices[type].second;
394: PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
395: PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
396: PetscCall(PetscDeviceConfigure(dev));
397: init = true;
398: }
399: PetscFunctionReturn(PETSC_SUCCESS);
400: }
402: } // namespace
404: /*@C
405: PetscDeviceInitialize - Initialize `PetscDevice`
407: Not Collective
409: Input Parameter:
410: . type - The `PetscDeviceType` to initialize
412: Level: beginner
414: Notes:
415: Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
416: result in device synchronization.
418: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
419: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
420: @*/
421: PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
422: {
423: PetscFunctionBegin;
425: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
426: PetscFunctionReturn(PETSC_SUCCESS);
427: }
429: /*@C
430: PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
431: `PetscDeviceType`
433: Not Collective
435: Input Parameter:
436: . type - The `PetscDeviceType` to check
438: Level: beginner
440: Notes:
441: Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
443: If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
444: return `PETSC_FALSE` for that `PetscDeviceType`.
446: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
447: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
448: @*/
449: PetscBool PetscDeviceInitialized(PetscDeviceType type)
450: {
451: return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
452: }
454: /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
455: PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device)
456: {
457: PetscFunctionBegin;
458: PetscAssertPointer(device, 2);
459: PetscCall(PetscDeviceInitialize(type));
460: *device = defaultDevices[type].first;
461: PetscFunctionReturn(PETSC_SUCCESS);
462: }
464: /*@C
465: PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
467: Not Collective
469: Input Parameters:
470: + device - The `PetscDevice`
471: - attr - The attribute
473: Output Parameter:
474: . value - The value of the attribute
476: Level: intermediate
478: Notes:
479: Since different attributes are often different types `value` is a `void *` to accommodate
480: them all. The underlying type of the attribute is therefore included in the name of the
481: `PetscDeviceAttribute` responsible for querying it. For example,
482: `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
484: .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
485: @*/
486: PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
487: {
488: PetscFunctionBegin;
491: PetscAssertPointer(value, 3);
492: PetscUseTypeMethod(device, getattribute, attr, value);
493: PetscFunctionReturn(PETSC_SUCCESS);
494: }
496: namespace
497: {
499: PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType)
500: {
501: PetscFunctionBegin;
502: if (!PetscDeviceConfiguredFor_Internal(type)) {
503: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
504: defaultDevices[type].first = nullptr;
505: PetscFunctionReturn(PETSC_SUCCESS);
506: }
507: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
508: /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
509: switch (type) {
510: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
511: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
512: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
513: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
514: default:
515: SETERRQ(comm, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
516: }
517: 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)]));
518: /*
519: defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES
520: have decided to initialize as
521: */
522: if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
523: PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
524: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
525: if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
526: }
527: PetscFunctionReturn(PETSC_SUCCESS);
528: }
530: PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView)
531: {
532: PetscInt initIdx = PETSC_DEVICE_INIT_LAZY;
533: auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
534: auto flg = PETSC_FALSE;
536: PetscFunctionBegin;
537: PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
538: if (flg) PetscCall(PetscLogGpuTime());
540: PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
541: PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
542: PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
543: 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));
544: PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
545: PetscOptionsEnd();
547: if (initIdx == PETSC_DEVICE_INIT_NONE) {
548: /* disabled all device initialization if devices are globally disabled */
549: 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");
550: *defaultView = PETSC_FALSE;
551: initDeviceIdx = PETSC_DEVICE_HOST;
552: } else {
553: *defaultView = static_cast<PetscBool>(*defaultView && flg);
554: if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
555: }
556: *defaultInitType = PetscDeviceInitTypeCast(initIdx);
557: *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
558: PetscFunctionReturn(PETSC_SUCCESS);
559: }
561: /* called from PetscFinalize() do not call yourself! */
562: PetscErrorCode PetscDeviceFinalize_Private()
563: {
564: PetscFunctionBegin;
565: if (PetscDefined(USE_DEBUG)) {
566: /*
567: you might be thinking, why on earth are you registered yet another finalizer in a
568: function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
569: because it is.
571: The crux of the problem is that the initializer (and therefore the ~finalizer~) of
572: PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
573: a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
574: won't be destroyed yet. So we need to repeat the check that all devices have been
575: destroyed again ~after~ the global context is destroyed. In summary:
577: 1. This finalizer runs and destroys all devices, except it may not because the global
578: context may still hold a reference!
579: 2. The global context finalizer runs and does the final reference count decrement
580: required, which actually destroys the held device.
581: 3. Our newly added finalizer runs and checks that all is well.
582: */
583: PetscCall(PetscRegisterFinalize([] {
584: PetscFunctionBegin;
585: for (auto &&device : defaultDevices) {
586: const auto dev = device.first;
588: 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);
589: }
590: PetscFunctionReturn(PETSC_SUCCESS);
591: }));
592: }
593: for (auto &&device : defaultDevices) {
594: PetscCall(PetscDeviceDestroy(&device.first));
595: device.second = false;
596: }
597: PetscFunctionReturn(PETSC_SUCCESS);
598: }
600: } // namespace
602: /*
603: Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
604: initialization types:
606: 1. defaultInitType - how does PetscDevice as a whole expect to initialize?
607: 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
608: e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
609: have all CUDA devices still initialize.
611: All told the following happens:
613: 0. defaultInitType -> LAZY
614: 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
615: 2. PetscDevice initializes each sub type with deviceDefaultInitType.
616: 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
617: to checking for specific device init. if view or specific device init
618: subTypeDefaultInitType -> EAGER. disabled once again overrides all.
619: */
621: PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm)
622: {
623: auto defaultView = PETSC_FALSE;
624: auto initializeDeviceContextEagerly = PETSC_FALSE;
625: auto defaultDeviceSet = PETSC_FALSE;
626: auto defaultDevice = PetscInt{PETSC_DECIDE};
627: auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT();
628: auto defaultInitType = PETSC_DEVICE_INIT_LAZY;
630: PetscFunctionBegin;
631: if (PetscDefined(USE_DEBUG)) {
632: int result;
634: PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
635: /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
636: * global space */
637: if (PetscUnlikely(result != MPI_IDENT)) {
638: char name[MPI_MAX_OBJECT_NAME] = {};
639: int len; /* unused */
641: PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
642: SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
643: }
644: }
645: comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
646: PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
648: PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
650: // the precise values don't matter here, so long as they are sequential
651: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HOST) == 0, "");
652: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_CUDA) == 1, "");
653: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HIP) == 2, "");
654: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_SYCL) == 3, "");
655: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_MAX) == 4, "");
656: for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
657: const auto deviceType = PetscDeviceTypeCast(i);
658: auto initType = defaultInitType;
660: PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
661: if (PetscDeviceConfiguredFor_Internal(deviceType)) {
662: if (initType == PETSC_DEVICE_INIT_EAGER) {
663: initializeDeviceContextEagerly = PETSC_TRUE;
664: // only update the default device if the user hasn't set it previously
665: if (!defaultDeviceSet) {
666: deviceContextInitDevice = deviceType;
667: PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
668: }
669: } else if (initType == PETSC_DEVICE_INIT_NONE) {
670: 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]);
671: }
672: }
673: }
675: PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
676: PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
677: /* ----------------------------------------------------------------------------------- */
678: /* PetscDevice is now fully initialized */
679: /* ----------------------------------------------------------------------------------- */
680: {
681: /*
682: query the options db to get the root settings from the user (if any).
684: This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
685: PetscDeviceContextSetFromOptions() before we even have one, then set a few static
686: variables in that file with the results.
687: */
688: auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
689: auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
691: PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
692: PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
693: PetscOptionsEnd();
695: if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
696: if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
697: }
699: if (initializeDeviceContextEagerly) {
700: PetscDeviceContext dctx;
702: PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
703: /* instantiates the device context */
704: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
705: PetscCall(PetscDeviceContextSetUp(dctx));
706: }
707: PetscFunctionReturn(PETSC_SUCCESS);
708: }