xref: /petsc/src/sys/objects/device/interface/device.cxx (revision 66c9fbdd036b1e887ebf0d2bef6dbdcafd086d45)
1 #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/
2 #include <petsc/private/petscadvancedmacros.h>
3 
4 #include "../impls/host/hostdevice.hpp"
5 #include "../impls/cupm/cupmdevice.hpp"
6 #include "../impls/sycl/sycldevice.hpp"
7 
8 #include <limits>  // std::numeric_limits
9 #include <utility> // std::make_pair
10 
11 using namespace Petsc::device;
12 
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
27 
28 #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
29   case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
30     PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
31   } break
32 
33 /*
34   Suppose you have:
35 
36   CUDADevice.myFunction(arg1,arg2)
37 
38   that you would like to conditionally define and call in a switch-case:
39 
40   switch(PetscDeviceType) {
41   #if PetscDefined(HAVE_CUDA)
42   case PETSC_DEVICE_CUDA: {
43     PetscCall(CUDADevice.myFunction(arg1,arg2));
44   } break;
45   #endif
46   }
47 
48   then calling this macro:
49 
50   PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
51 
52   will expand to the following case statement:
53 
54   case PETSC_DEVICE_CUDA: {
55     PetscCall(CUDADevice.myFunction(arg1,arg2));
56   } break
57 
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__)
61 
62 /*@C
63   PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
64 
65   Not Collective
66 
67   Input Parameters:
68 + type  - The type of `PetscDevice`
69 - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically)
70 
71   Output Parameter:
72 . device - The `PetscDevice`
73 
74   Notes:
75   This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
76   device synchronization.
77 
78   `devid` is what you might pass to `cudaSetDevice()` for example.
79 
80   Level: beginner
81 
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;
89 
90   PetscFunctionBegin;
91   PetscValidDeviceType(type, 1);
92   PetscValidPointer(device, 3);
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(0);
113 }
114 
115 /*@C
116   PetscDeviceDestroy - Free a `PetscDevice`
117 
118   Not Collective
119 
120   Input Parameter:
121 . device - The `PetscDevice`
122 
123   Level: beginner
124 
125 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
126 `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
127 @*/
128 PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
129 {
130   PetscFunctionBegin;
131   PetscValidPointer(device, 1);
132   if (!*device) PetscFunctionReturn(0);
133   PetscValidDevice(*device, 1);
134   PetscCall(PetscDeviceDereference_Internal(*device));
135   if ((*device)->refcnt) {
136     *device = nullptr;
137     PetscFunctionReturn(0);
138   }
139   PetscCall(PetscFree((*device)->data));
140   PetscCall(PetscFree(*device));
141   PetscFunctionReturn(0);
142 }
143 
144 /*@C
145   PetscDeviceConfigure - Configure a particular `PetscDevice`
146 
147   Not Collective
148 
149   Input Parameter:
150 . device - The `PetscDevice` to configure
151 
152   Notes:
153   The user should not assume that this is a cheap operation.
154 
155   Level: beginner
156 
157 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
158 `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
159 @*/
160 PetscErrorCode PetscDeviceConfigure(PetscDevice device)
161 {
162   PetscFunctionBegin;
163   PetscValidDevice(device, 1);
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(0);
186 }
187 
188 /*@C
189   PetscDeviceView - View a `PetscDevice`
190 
191   Collective on viewer
192 
193   Input Parameters:
194 + device - The `PetscDevice` to view
195 - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
196 
197   Level: beginner
198 
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;
206 
207   PetscFunctionBegin;
208   PetscValidDevice(device, 1);
209   if (viewer) {
210     PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
211     PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
212   } else {
213     PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
214     iascii = PETSC_TRUE;
215   }
216 
217   if (iascii) {
218     auto        dtype = PETSC_DEVICE_HOST;
219     MPI_Comm    comm;
220     PetscMPIInt size;
221     PetscInt    id = 0;
222 
223     PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
224     PetscCallMPI(MPI_Comm_size(comm, &size));
225 
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   }
234 
235   // see if impls has extra viewer stuff
236   PetscTryTypeMethod(device, view, sub);
237 
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(0);
245 }
246 
247 /*@C
248   PetscDeviceGetType - Get the type of device
249 
250   Not Collective
251 
252   Input Parameter:
253 . device - The `PetscDevice`
254 
255   Output Parameter:
256 . type - The `PetscDeviceType`
257 
258   Level: beginner
259 
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;
267   PetscValidDevice(device, 1);
268   PetscValidPointer(type, 2);
269   *type = device->type;
270   PetscFunctionReturn(0);
271 }
272 
273 /*@C
274   PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
275 
276   Not Collective
277 
278   Input Parameter:
279 . device - The `PetscDevice`
280 
281   Output Parameter:
282 . id - The id
283 
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.
288 
289   Level: beginner
290 
291 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
292 @*/
293 PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
294 {
295   PetscFunctionBegin;
296   PetscValidDevice(device, 1);
297   PetscValidIntPointer(id, 2);
298   *id = device->deviceId;
299   PetscFunctionReturn(0);
300 }
301 
302 struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
303   PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
304 
305   PETSC_NODISCARD PetscErrorCode finalize_() noexcept
306   {
307     PetscFunctionBegin;
308     type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
309     PetscFunctionReturn(0);
310   }
311 };
312 
313 static auto default_device_type = DefaultDeviceType();
314 
315 /*@C
316   PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
317 
318   Not Collective
319 
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`.
323 
324   Level: beginner
325 
326 .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
327 @*/
328 PetscDeviceType PETSC_DEVICE_DEFAULT(void)
329 {
330   return default_device_type.type;
331 }
332 
333 /*@C
334   PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
335 
336   Not Collective
337 
338   Input Parameter:
339 . type - the new default device type
340 
341   Notes:
342   This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
343 
344   Level: beginner
345 
346 .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
347 @*/
348 PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
349 {
350   PetscFunctionBegin;
351   PetscValidDeviceType(type, 1);
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(0);
358 }
359 
360 static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
361 
362 /*
363   Actual intialization 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;
369   PetscValidDeviceType(type, 1);
370   if (PetscUnlikely(!PetscDeviceInitialized(type))) {
371     auto &dev  = defaultDevices[type].first;
372     auto &init = defaultDevices[type].second;
373 
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(0);
380 }
381 
382 /*@C
383   PetscDeviceInitialize - Initialize `PetscDevice`
384 
385   Not Collective
386 
387   Input Parameter:
388 . type - The `PetscDeviceType` to initialize
389 
390   Notes:
391   Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
392   result in device synchronization.
393 
394   Level: beginner
395 
396 .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
397 `PetscDeviceCreate()`, `PetscDeviceDestroy()`
398 @*/
399 PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
400 {
401   PetscFunctionBegin;
402   PetscValidDeviceType(type, 1);
403   PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
404   PetscFunctionReturn(0);
405 }
406 
407 /*@C
408   PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
409   `PetscDeviceType`
410 
411   Not Collective
412 
413   Input Parameter:
414 . type - The `PetscDeviceType` to check
415 
416   Notes:
417   Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
418 
419   If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
420   return `PETSC_FALSE` for that `PetscDeviceType`.
421 
422   Level: beginner
423 
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 }
431 
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;
436   PetscValidPointer(device, 2);
437   PetscCall(PetscDeviceInitialize(type));
438   *device = defaultDevices[type].first;
439   PetscFunctionReturn(0);
440 }
441 
442 /*@C
443   PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
444 
445   Not Collective
446 
447   Input Parameters:
448 + device - The `PetscDevice`
449 - attr   - The attribute
450 
451   Output Parameter:
452 . value - The value of the attribute
453 
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` reponsible for querying it. For example,
458   `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
459 
460   Level: intermediate
461 
462 .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
463 @*/
464 PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
465 {
466   PetscFunctionBegin;
467   PetscValidDevice(device, 1);
468   PetscValidDeviceAttribute(attr, 2);
469   PetscValidPointer(value, 3);
470   PetscUseTypeMethod(device, getattribute, attr, value);
471   PetscFunctionReturn(0);
472 }
473 
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(0);
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::integral_value(*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(0);
503 }
504 
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;
510 
511   PetscFunctionBegin;
512   PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
513   if (flg) PetscCall(PetscLogGpuTime());
514 
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();
521 
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 exlusive");
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(0);
534 }
535 
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;
545 
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(0);
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.
554 
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:
560 
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(0);
574 }
575 
576 /*
577   Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
578   initialization types:
579 
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.
584 
585   All told the following happens:
586 
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 */
594 
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;
603 
604   PetscFunctionBegin;
605   if (PetscDefined(USE_DEBUG)) {
606     int result;
607 
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 */
614 
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));
621 
622   PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
623 
624   // the precise values don't matter here, so long as they are sequential
625   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, "");
626   static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, "");
627   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, "");
628   static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, "");
629   static_assert(Petsc::util::integral_value(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;
633 
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   }
648 
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).
657 
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);
664 
665     PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
666     PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
667     PetscOptionsEnd();
668 
669     if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
670     if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
671   }
672 
673   if (initializeDeviceContextEagerly) {
674     PetscDeviceContext dctx;
675 
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(0);
682 }
683