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