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: }