Actual source code: device.cxx
1: #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/
2: #include <petsc/private/petscadvancedmacros.h>
4: #include "../impls/host/hostdevice.hpp"
5: #include "../impls/cupm/cupmdevice.hpp"
6: #include "../impls/sycl/sycldevice.hpp"
8: #include <limits> // std::numeric_limits
9: #include <utility> // std::make_pair
11: using namespace Petsc::device;
13: /*
14: note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to
15: be picked up by the switch-case macros below
16: */
17: static host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
18: #if PetscDefined(HAVE_CUDA)
19: static cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
20: #endif
21: #if PetscDefined(HAVE_HIP)
22: static cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
23: #endif
24: #if PetscDefined(HAVE_SYCL)
25: static sycl::Device SYCLDevice{PetscDeviceContextCreate_SYCL};
26: #endif
28: #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
29: case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
30: PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
31: } break
33: /*
34: Suppose you have:
36: CUDADevice.myFunction(arg1,arg2)
38: that you would like to conditionally define and call in a switch-case:
40: switch(PetscDeviceType) {
41: #if PetscDefined(HAVE_CUDA)
42: case PETSC_DEVICE_CUDA: {
43: PetscCall(CUDADevice.myFunction(arg1,arg2));
44: } break;
45: #endif
46: }
48: then calling this macro:
50: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
52: will expand to the following case statement:
54: case PETSC_DEVICE_CUDA: {
55: PetscCall(CUDADevice.myFunction(arg1,arg2));
56: } break
58: if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise
59: */
60: #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(PetscConcat_(HAVE_, IMPLS), PETSC_DEVICE_CASE, PetscExpandToNothing)(IMPLS, func, __VA_ARGS__)
62: /*@C
63: PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
65: Not Collective
67: Input Parameters:
68: + type - The type of `PetscDevice`
69: - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically)
71: Output Parameter:
72: . device - The `PetscDevice`
74: Notes:
75: This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
76: device synchronization.
78: `devid` is what you might pass to `cudaSetDevice()` for example.
80: Level: beginner
82: .seealso: `PetscDevice`, `PetscDeviceInitType`,
83: `PetscDeviceInitialize()`, `PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
84: `PetscDeviceView()`, `PetscDeviceDestroy()`
85: @*/
86: PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device)
87: {
88: static PetscInt PetscDeviceCounter = 0;
90: PetscFunctionBegin;
93: PetscCall(PetscDeviceInitializePackage());
94: PetscCall(PetscNew(device));
95: (*device)->id = PetscDeviceCounter++;
96: (*device)->type = type;
97: (*device)->refcnt = 1;
98: /*
99: if you are adding a device, you also need to add it's initialization in
100: PetscDeviceInitializeTypeFromOptions_Private() below
101: */
102: switch (type) {
103: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
104: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
105: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
106: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
107: default:
108: /* in case the above macros expand to nothing this silences any unused variable warnings */
109: (void)(devid);
110: 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]);
111: }
112: PetscFunctionReturn(PETSC_SUCCESS);
113: }
115: /*@C
116: PetscDeviceDestroy - Free a `PetscDevice`
118: Not Collective
120: Input Parameter:
121: . device - The `PetscDevice`
123: Level: beginner
125: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
126: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
127: @*/
128: PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
129: {
130: PetscFunctionBegin;
132: if (!*device) PetscFunctionReturn(PETSC_SUCCESS);
134: PetscCall(PetscDeviceDereference_Internal(*device));
135: if ((*device)->refcnt) {
136: *device = nullptr;
137: PetscFunctionReturn(PETSC_SUCCESS);
138: }
139: PetscCall(PetscFree((*device)->data));
140: PetscCall(PetscFree(*device));
141: PetscFunctionReturn(PETSC_SUCCESS);
142: }
144: /*@C
145: PetscDeviceConfigure - Configure a particular `PetscDevice`
147: Not Collective
149: Input Parameter:
150: . device - The `PetscDevice` to configure
152: Notes:
153: The user should not assume that this is a cheap operation.
155: Level: beginner
157: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
158: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
159: @*/
160: PetscErrorCode PetscDeviceConfigure(PetscDevice device)
161: {
162: PetscFunctionBegin;
164: /*
165: if no available configuration is available, this cascades all the way down to default
166: and error
167: */
168: switch (const auto dtype = device->type) {
169: case PETSC_DEVICE_HOST:
170: if (PetscDefined(HAVE_HOST)) break; // always true
171: case PETSC_DEVICE_CUDA:
172: if (PetscDefined(HAVE_CUDA)) break;
173: goto error;
174: case PETSC_DEVICE_HIP:
175: if (PetscDefined(HAVE_HIP)) break;
176: goto error;
177: case PETSC_DEVICE_SYCL:
178: if (PetscDefined(HAVE_SYCL)) break;
179: goto error;
180: default:
181: error:
182: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
183: }
184: PetscUseTypeMethod(device, configure);
185: PetscFunctionReturn(PETSC_SUCCESS);
186: }
188: /*@C
189: PetscDeviceView - View a `PetscDevice`
191: Collective on viewer
193: Input Parameters:
194: + device - The `PetscDevice` to view
195: - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
197: Level: beginner
199: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
200: `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
201: @*/
202: PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer)
203: {
204: auto sub = viewer;
205: PetscBool iascii;
207: PetscFunctionBegin;
209: if (viewer) {
211: PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
212: } else {
213: PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
214: iascii = PETSC_TRUE;
215: }
217: if (iascii) {
218: auto dtype = PETSC_DEVICE_HOST;
219: MPI_Comm comm;
220: PetscMPIInt size;
221: PetscInt id = 0;
223: PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
224: PetscCallMPI(MPI_Comm_size(comm, &size));
226: PetscCall(PetscDeviceGetDeviceId(device, &id));
227: PetscCall(PetscDeviceGetType(device, &dtype));
228: PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
229: PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
230: PetscCall(PetscViewerASCIIPushTab(sub));
231: PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
232: PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
233: }
235: // see if impls has extra viewer stuff
236: PetscTryTypeMethod(device, view, sub);
238: if (iascii) {
239: // undo the ASCII specific stuff
240: PetscCall(PetscViewerASCIIPopTab(sub));
241: PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
242: PetscCall(PetscViewerFlush(viewer));
243: }
244: PetscFunctionReturn(PETSC_SUCCESS);
245: }
247: /*@C
248: PetscDeviceGetType - Get the type of device
250: Not Collective
252: Input Parameter:
253: . device - The `PetscDevice`
255: Output Parameter:
256: . type - The `PetscDeviceType`
258: Level: beginner
260: .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
261: `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
262: `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
263: @*/
264: PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type)
265: {
266: PetscFunctionBegin;
269: *type = device->type;
270: PetscFunctionReturn(PETSC_SUCCESS);
271: }
273: /*@C
274: PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
276: Not Collective
278: Input Parameter:
279: . device - The `PetscDevice`
281: Output Parameter:
282: . id - The id
284: Notes:
285: The returned ID may have been assigned by the underlying device backend. For example if the
286: backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
287: this device was configured.
289: Level: beginner
291: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
292: @*/
293: PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
294: {
295: PetscFunctionBegin;
298: *id = device->deviceId;
299: PetscFunctionReturn(PETSC_SUCCESS);
300: }
302: struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
303: PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
305: PetscErrorCode finalize_() noexcept
306: {
307: PetscFunctionBegin;
308: type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
309: PetscFunctionReturn(PETSC_SUCCESS);
310: }
311: };
313: static auto default_device_type = DefaultDeviceType();
315: /*@C
316: PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
318: Not Collective
320: Notes:
321: Unless selected by the user, the default device is selected in the following order\:
322: `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
324: Level: beginner
326: .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
327: @*/
328: PetscDeviceType PETSC_DEVICE_DEFAULT(void)
329: {
330: return default_device_type.type;
331: }
333: /*@C
334: PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
336: Not Collective
338: Input Parameter:
339: . type - the new default device type
341: Notes:
342: This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
344: Level: beginner
346: .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
347: @*/
348: PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
349: {
350: PetscFunctionBegin;
352: if (default_device_type.type != type) {
353: // no need to waster a PetscRegisterFinalize() slot if we don't change it
354: default_device_type.type = type;
355: PetscCall(default_device_type.register_finalize());
356: }
357: PetscFunctionReturn(PETSC_SUCCESS);
358: }
360: static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
362: /*
363: Actual initialization function; any functions claiming to initialize PetscDevice or
364: PetscDeviceContext will have to run through this one
365: */
366: static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId)
367: {
368: PetscFunctionBegin;
370: if (PetscUnlikely(!PetscDeviceInitialized(type))) {
371: auto &dev = defaultDevices[type].first;
372: auto &init = defaultDevices[type].second;
374: PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
375: PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
376: PetscCall(PetscDeviceConfigure(dev));
377: init = true;
378: }
379: PetscFunctionReturn(PETSC_SUCCESS);
380: }
382: /*@C
383: PetscDeviceInitialize - Initialize `PetscDevice`
385: Not Collective
387: Input Parameter:
388: . type - The `PetscDeviceType` to initialize
390: Notes:
391: Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
392: result in device synchronization.
394: Level: beginner
396: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
397: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
398: @*/
399: PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
400: {
401: PetscFunctionBegin;
403: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
404: PetscFunctionReturn(PETSC_SUCCESS);
405: }
407: /*@C
408: PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
409: `PetscDeviceType`
411: Not Collective
413: Input Parameter:
414: . type - The `PetscDeviceType` to check
416: Notes:
417: Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
419: If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
420: return `PETSC_FALSE` for that `PetscDeviceType`.
422: Level: beginner
424: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
425: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
426: @*/
427: PetscBool PetscDeviceInitialized(PetscDeviceType type)
428: {
429: return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
430: }
432: /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
433: PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device)
434: {
435: PetscFunctionBegin;
437: PetscCall(PetscDeviceInitialize(type));
438: *device = defaultDevices[type].first;
439: PetscFunctionReturn(PETSC_SUCCESS);
440: }
442: /*@C
443: PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
445: Not Collective
447: Input Parameters:
448: + device - The `PetscDevice`
449: - attr - The attribute
451: Output Parameter:
452: . value - The value of the attribute
454: Notes:
455: Since different attributes are often different types `value` is a `void *` to accommodate
456: them all. The underlying type of the attribute is therefore included in the name of the
457: `PetscDeviceAttribute` responsible for querying it. For example,
458: `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
460: Level: intermediate
462: .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
463: @*/
464: PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
465: {
466: PetscFunctionBegin;
470: PetscUseTypeMethod(device, getattribute, attr, value);
471: PetscFunctionReturn(PETSC_SUCCESS);
472: }
474: static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType)
475: {
476: PetscFunctionBegin;
477: if (!PetscDeviceConfiguredFor_Internal(type)) {
478: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
479: defaultDevices[type].first = nullptr;
480: PetscFunctionReturn(PETSC_SUCCESS);
481: }
482: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
483: /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
484: switch (type) {
485: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
486: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
487: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
488: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
489: default:
490: SETERRQ(comm, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
491: }
492: 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)]));
493: /*
494: defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES
495: have decided to initialize as
496: */
497: if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
498: PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
499: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
500: if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
501: }
502: PetscFunctionReturn(PETSC_SUCCESS);
503: }
505: static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView)
506: {
507: PetscInt initIdx = PETSC_DEVICE_INIT_LAZY;
508: auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
509: auto flg = PETSC_FALSE;
511: PetscFunctionBegin;
512: PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
513: if (flg) PetscCall(PetscLogGpuTime());
515: PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
516: PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
517: PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
518: 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));
519: PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
520: PetscOptionsEnd();
522: if (initIdx == PETSC_DEVICE_INIT_NONE) {
523: /* disabled all device initialization if devices are globally disabled */
524: 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");
525: *defaultView = PETSC_FALSE;
526: initDeviceIdx = PETSC_DEVICE_HOST;
527: } else {
528: *defaultView = static_cast<PetscBool>(*defaultView && flg);
529: if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
530: }
531: *defaultInitType = PetscDeviceInitTypeCast(initIdx);
532: *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
533: PetscFunctionReturn(PETSC_SUCCESS);
534: }
536: /* called from PetscFinalize() do not call yourself! */
537: static PetscErrorCode PetscDeviceFinalize_Private()
538: {
539: PetscFunctionBegin;
540: if (PetscDefined(USE_DEBUG)) {
541: const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] {
542: PetscFunctionBegin;
543: for (auto &&device : defaultDevices) {
544: const auto dev = device.first;
546: 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);
547: }
548: PetscFunctionReturn(PETSC_SUCCESS);
549: };
550: /*
551: you might be thinking, why on earth are you registered yet another finalizer in a
552: function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
553: because it is.
555: The crux of the problem is that the initializer (and therefore the ~finalizer~) of
556: PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
557: a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
558: won't be destroyed yet. So we need to repeat the check that all devices have been
559: destroyed again ~after~ the global context is destroyed. In summary:
561: 1. This finalizer runs and destroys all devices, except it may not because the global
562: context may still hold a reference!
563: 2. The global context finalizer runs and does the final reference count decrement
564: required, which actually destroys the held device.
565: 3. Our newly added finalizer runs and checks that all is well.
566: */
567: PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize)));
568: }
569: for (auto &&device : defaultDevices) {
570: PetscCall(PetscDeviceDestroy(&device.first));
571: device.second = false;
572: }
573: PetscFunctionReturn(PETSC_SUCCESS);
574: }
576: /*
577: Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
578: initialization types:
580: 1. defaultInitType - how does PetscDevice as a whole expect to initialize?
581: 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
582: e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
583: have all CUDA devices still initialize.
585: All told the following happens:
587: 0. defaultInitType -> LAZY
588: 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
589: 2. PetscDevice initializes each sub type with deviceDefaultInitType.
590: 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
591: to checking for specific device init. if view or specific device init
592: subTypeDefaultInitType -> EAGER. disabled once again overrides all.
593: */
595: PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm)
596: {
597: auto defaultView = PETSC_FALSE;
598: auto initializeDeviceContextEagerly = PETSC_FALSE;
599: auto defaultDeviceSet = PETSC_FALSE;
600: auto defaultDevice = PetscInt{PETSC_DECIDE};
601: auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT();
602: auto defaultInitType = PETSC_DEVICE_INIT_LAZY;
604: PetscFunctionBegin;
605: if (PetscDefined(USE_DEBUG)) {
606: int result;
608: PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
609: /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
610: * global space */
611: if (PetscUnlikely(result != MPI_IDENT)) {
612: char name[MPI_MAX_OBJECT_NAME] = {};
613: int len; /* unused */
615: PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
616: SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
617: }
618: }
619: comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
620: PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
622: PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
624: // the precise values don't matter here, so long as they are sequential
625: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HOST) == 0, "");
626: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_CUDA) == 1, "");
627: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HIP) == 2, "");
628: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_SYCL) == 3, "");
629: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_MAX) == 4, "");
630: for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
631: const auto deviceType = PetscDeviceTypeCast(i);
632: auto initType = defaultInitType;
634: PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
635: if (PetscDeviceConfiguredFor_Internal(deviceType)) {
636: if (initType == PETSC_DEVICE_INIT_EAGER) {
637: initializeDeviceContextEagerly = PETSC_TRUE;
638: // only update the default device if the user hasn't set it previously
639: if (!defaultDeviceSet) {
640: deviceContextInitDevice = deviceType;
641: PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
642: }
643: } else if (initType == PETSC_DEVICE_INIT_NONE) {
644: 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]);
645: }
646: }
647: }
649: PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
650: PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
651: /* ----------------------------------------------------------------------------------- */
652: /* PetscDevice is now fully initialized */
653: /* ----------------------------------------------------------------------------------- */
654: {
655: /*
656: query the options db to get the root settings from the user (if any).
658: This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
659: PetscDeviceContextSetFromOptions() before we even have one, then set a few static
660: variables in that file with the results.
661: */
662: auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
663: auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
665: PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
666: PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
667: PetscOptionsEnd();
669: if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
670: if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
671: }
673: if (initializeDeviceContextEagerly) {
674: PetscDeviceContext dctx;
676: PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
677: /* instantiates the device context */
678: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
679: PetscCall(PetscDeviceContextSetUp(dctx));
680: }
681: PetscFunctionReturn(PETSC_SUCCESS);
682: }