Actual source code: device.cxx
1: #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/
2: #include <petsc/private/petscadvancedmacros.h>
4: #include <petsc/private/cpp/register_finalize.hpp>
6: #include "../impls/host/hostdevice.hpp"
7: #if PetscDefined(HAVE_CUPM)
8: #include "../impls/cupm/cupmdevice.hpp"
9: #endif
10: #if PetscDefined(HAVE_SYCL)
11: #include "../impls/sycl/sycldevice.hpp"
12: #endif
14: #include <utility> // std::make_pair
16: using namespace Petsc::device;
18: #if PetscDefined(HAVE_CUPM)
19: int PetscDeviceCUPMRuntimeArch = 0;
20: #endif
22: namespace
23: {
25: /*
26: note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to
27: be picked up by the switch-case macros below
28: */
29: host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
30: #if PetscDefined(HAVE_CUDA)
31: cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
32: #endif
33: #if PetscDefined(HAVE_HIP)
34: cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
35: #endif
36: #if PetscDefined(HAVE_SYCL)
37: sypm::Device SYCLDevice{PetscDeviceContextCreate_SYCL};
38: #endif
40: } // namespace
42: #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
43: case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
44: PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
45: } break
47: #define PETSC_VOID_0(...) ((void)0)
49: /*
50: Suppose you have:
52: CUDADevice.myFunction(arg1,arg2)
54: that you would like to conditionally define and call in a switch-case:
56: switch(PetscDeviceType) {
57: #if PetscDefined(HAVE_CUDA)
58: case PETSC_DEVICE_CUDA: {
59: PetscCall(CUDADevice.myFunction(arg1,arg2));
60: } break;
61: #endif
62: }
64: then calling this macro:
66: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
68: will expand to the following case statement:
70: case PETSC_DEVICE_CUDA: {
71: PetscCall(CUDADevice.myFunction(arg1,arg2));
72: } break
74: if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise
75: */
76: #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(HAVE_##IMPLS, PETSC_DEVICE_CASE, PETSC_VOID_0)(IMPLS, func, __VA_ARGS__)
78: /*@
79: PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
81: Not Collective
83: Input Parameters:
84: + type - The type of `PetscDevice`
85: - devid - The device selection strategy or the numeric ID# of the device
87: Output Parameter:
88: . device - The `PetscDevice`
90: Options Database Key:
91: . -device_select N - The device selection strategy or the numeric ID# of the device
93: Level: beginner
95: Notes:
96: This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
97: device synchronization.
99: Available device selection strategies\:
101: | Value | Behavior |
102: |-------|---------------------------------------------------------------|
103: | -1 | Use the configured default device (`PETSC_DECIDE`) |
104: | -2 | Topology-aware allocation (CUDA/HIP with hwloc 2.0 or later) |
105: | >=0 | Select a specific GPU device by index |
107: SYCL gives -2 its own meaning, the host device, so in a CUPM and SYCL build
108: `-device_select -2` is topology-aware for CUDA or HIP but the host device for SYCL.
110: When `PETSC_DECIDE` (-1) is passed to `-device_select` and to this function, either
111: a round-robin allocation across all available GPUs is used, or if `torchrun` was used to
112: launch the application, the value of the `LOCAL_RANK` environment variable set by
113: `torchrun` is selected. When the "Topology-aware allocation" strategy is selected, PETSc
114: uses `hwloc` to inspect the hardware layout of the current system and determine which GPU
115: is closest to the CPU core(s) the current process has been bound to. Each process must
116: actually be bound (e.g. `mpiexec -bind-to core`, `srun --cpu-bind=cores`); an unbound process
117: has no locality to use and falls back to `PETSC_DECIDE`. For example, the LUMI, Frontier and
118: Setonix HPC systems all have nodes comprised of 64-core AMD EPYC "Trento" CPUs connected
119: to 4 AMD MI250X GPUs, each of which is comprised of two Graphics Compute Die (GCD) that
120: appears to the operating system as a distinct GPU device. The CPU is divided into four NUMA
121: nodes, each of which is divided into two chiplets with a direct connection to each GCD. Each
122: GCD is assigned an index from 0 to 7 by the GPU runtime. The diagram below shows how `hwloc`
123: sees one of these nodes.
125: 
127: In order to construct a mapping between the GCD index and the devices found by `hwloc`, we
128: inspect the PCIe Bus ID for each GCD. The results of this are shown in the table below
130: | PCIe Bus ID | GCD Index |
131: |-------------|----------:|
132: | d1:00.0 | 4 |
133: | d6:00.0 | 5 |
134: | c9:00.0 | 2 |
135: | ce:00.0 | 3 |
136: | d9:00.0 | 6 |
137: | de:00.0 | 7 |
138: | c1:00.0 | 0 |
139: | c6:00.0 | 1 |
141: By cross-referencing the figure and the table, we can see that the 'closest' GPU to e.g. core 0 has the PCIe
142: Bus ID `d1:00.0`, which corresponds to GCD 4.
144: `hwloc` uses the notion of "depth" of devices to describe the different layers of the topology it
145: discovers. The top layer is always Machine and is assigned a value 0. Beneath that is the Package
146: layer, which is assigned depth 1 and generally corresponds to CPU sockets. In order to determine which GCD
147: is topologically closest, PETSc uses `hwloc` to determine the depth value of the nearest common ancestor
148: for the CPU the current process is bound to and each device on the system. The higher the depth value,
149: the closer the hardware controller for the device is to CPU core 0, and therefore the closer the GCD
150: is to CPU core 0. The table below shows the nearest common ancestor between CPU core 0 and each of the
151: GCDs identified on the system
153: | PCIe Bus ID | Nearest Common Ancestor to Core 0 | Depth |
154: |-------------|----------------------------------:|------:|
155: | d1:00.0 | L3 (chiplet) | 3 |
156: | d6:00.0 | Group (NUMA Node) | 2 |
157: | c9:00.0 | Package (CPU socket) | 1 |
158: | ce:00.0 | Package (CPU socket) | 1 |
159: | d9:00.0 | Package (CPU socket) | 1 |
160: | de:00.0 | Package (CPU socket) | 1 |
161: | c1:00.0 | Package (CPU socket) | 1 |
162: | c6:00.0 | Package (CPU socket) | 1 |
164: Choosing the device with the highest depth value gives the intuitive result that CPU
165: core 0 is closest to the GCD on PCIe Bus ID d1:00.0, or GCD 4. In the case where multiple devices
166: are found at the same depth, the logical ordering of the CPU cores under the device at that depth
167: is used to determine which device to select. For instance, a system with 2 GPUs at the same depth
168: and 16 cores will assign device 0 to processes on cores 0-7 and device 1 to processes on cores 8-15.
169: A summary of the topologically aware device selection method is as follows.
171: - If there is a single GPU on the system, select it and return
172: - Enumerate all devices on the system and map PCIe Bus ID to device ID
173: - Discover system topology with hwloc
174: - Determine current process's CPU binding
175: - Find GPU devices by PCIe Bus ID in discovered topology
176: - Find the depth of nearest common ancestor between current CPU core and every GPU device
177: - If there is a single device at a maximum depth, select that device and return
178: - If there are multiple devices at the same depth, select between those devices based on logical CPU core ordering
180: .seealso: `PetscDevice`, `PetscDeviceInitType`,
181: `PetscDeviceInitialize()`, `PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
182: `PetscDeviceView()`, `PetscDeviceDestroy()`
183: @*/
184: PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device)
185: {
186: static PetscInt PetscDeviceCounter = 0;
188: PetscFunctionBegin;
190: PetscAssertPointer(device, 3);
191: PetscCall(PetscDeviceInitializePackage());
193: PetscCall(PetscNew(device));
194: (*device)->id = PetscDeviceCounter++;
195: (*device)->type = type;
196: (*device)->refcnt = 1;
197: /*
198: if you are adding a device, you also need to add its initialization in
199: PetscDeviceInitializeTypeFromOptions_Private() below
200: */
201: switch (type) {
202: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
203: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
204: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
205: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
206: default:
207: /* in case the above macros expand to nothing this silences any unused variable warnings */
208: (void)devid;
209: 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]);
210: }
211: PetscFunctionReturn(PETSC_SUCCESS);
212: }
214: /*@
215: PetscDeviceDestroy - Free a `PetscDevice`
217: Not Collective
219: Input Parameter:
220: . device - The `PetscDevice`
222: Level: beginner
224: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
225: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
226: @*/
227: PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
228: {
229: PetscFunctionBegin;
230: PetscAssertPointer(device, 1);
231: if (!*device) PetscFunctionReturn(PETSC_SUCCESS);
233: PetscCall(PetscDeviceDereference_Internal(*device));
234: if ((*device)->refcnt) {
235: *device = nullptr;
236: PetscFunctionReturn(PETSC_SUCCESS);
237: }
238: PetscCall(PetscFree((*device)->data));
239: PetscCall(PetscFree(*device));
240: PetscFunctionReturn(PETSC_SUCCESS);
241: }
243: /*@
244: PetscDeviceConfigure - Configure a particular `PetscDevice`
246: Not Collective
248: Input Parameter:
249: . device - The `PetscDevice` to configure
251: Level: beginner
253: Notes:
254: The user should not assume that this is a cheap operation.
256: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
257: `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
258: @*/
259: PetscErrorCode PetscDeviceConfigure(PetscDevice device)
260: {
261: PetscFunctionBegin;
263: /*
264: if no available configuration is available, this cascades all the way down to default
265: and error
266: */
267: switch (const auto dtype = device->type) {
268: case PETSC_DEVICE_HOST:
269: if (PetscDefined(HAVE_HOST)) break; // always true
270: case PETSC_DEVICE_CUDA:
271: if (PetscDefined(HAVE_CUDA)) break;
272: goto error;
273: case PETSC_DEVICE_HIP:
274: if (PetscDefined(HAVE_HIP)) break;
275: goto error;
276: case PETSC_DEVICE_SYCL:
277: if (PetscDefined(HAVE_SYCL)) break;
278: goto error;
279: default:
280: error:
281: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
282: }
283: PetscUseTypeMethod(device, configure);
284: PetscFunctionReturn(PETSC_SUCCESS);
285: }
287: /*@
288: PetscDeviceView - View a `PetscDevice`
290: Collective on viewer
292: Input Parameters:
293: + device - The `PetscDevice` to view
294: - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
296: Level: beginner
298: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
299: `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
300: @*/
301: PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer)
302: {
303: auto sub = viewer;
304: PetscBool isascii;
306: PetscFunctionBegin;
308: if (viewer) {
310: PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &isascii));
311: } else {
312: PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
313: isascii = PETSC_TRUE;
314: }
316: if (isascii) {
317: auto dtype = PETSC_DEVICE_HOST;
318: MPI_Comm comm;
319: PetscMPIInt size;
320: PetscInt id = 0;
322: PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
323: PetscCallMPI(MPI_Comm_size(comm, &size));
325: PetscCall(PetscDeviceGetDeviceId(device, &id));
326: PetscCall(PetscDeviceGetType(device, &dtype));
327: PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
328: PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
329: PetscCall(PetscViewerASCIIPushTab(sub));
330: PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
331: PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
332: }
334: // see if impls has extra viewer stuff
335: PetscTryTypeMethod(device, view, sub);
337: if (isascii) {
338: // undo the ASCII specific stuff
339: PetscCall(PetscViewerASCIIPopTab(sub));
340: PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
341: }
342: PetscFunctionReturn(PETSC_SUCCESS);
343: }
345: /*@
346: PetscDeviceGetType - Get the type of device
348: Not Collective
350: Input Parameter:
351: . device - The `PetscDevice`
353: Output Parameter:
354: . type - The `PetscDeviceType`
356: Level: beginner
358: .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
359: `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
360: `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
361: @*/
362: PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type)
363: {
364: PetscFunctionBegin;
366: PetscAssertPointer(type, 2);
367: *type = device->type;
368: PetscFunctionReturn(PETSC_SUCCESS);
369: }
371: /*@
372: PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
374: Not Collective
376: Input Parameter:
377: . device - The `PetscDevice`
379: Output Parameter:
380: . id - The id
382: Level: beginner
384: Notes:
385: The returned ID may have been assigned by the underlying device backend. For example if the
386: backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
387: this device was configured.
389: .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
390: @*/
391: PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
392: {
393: PetscFunctionBegin;
395: PetscAssertPointer(id, 2);
396: *id = device->deviceId;
397: PetscFunctionReturn(PETSC_SUCCESS);
398: }
400: namespace
401: {
403: struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
404: PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
406: PetscErrorCode finalize_() noexcept
407: {
408: PetscFunctionBegin;
409: type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
410: PetscFunctionReturn(PETSC_SUCCESS);
411: }
412: };
414: auto default_device_type = DefaultDeviceType();
416: } // namespace
418: /*@
419: PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
421: Not Collective
423: Level: beginner
425: Notes:
426: Unless selected by the user, the default device is selected in the following order\:
427: `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
429: .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
430: @*/
431: PetscDeviceType PETSC_DEVICE_DEFAULT(void)
432: {
433: return default_device_type.type;
434: }
436: /*@
437: PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
439: Not Collective
441: Input Parameter:
442: . type - the new default device type
444: Level: beginner
446: Notes:
447: This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
449: .seealso: `PetscDeviceType`, `PetscDeviceGetType`
450: @*/
451: PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
452: {
453: PetscFunctionBegin;
455: if (default_device_type.type != type) {
456: // no need to waste a PetscRegisterFinalize() slot if we don't change it
457: default_device_type.type = type;
458: PetscCall(default_device_type.register_finalize());
459: }
460: PetscFunctionReturn(PETSC_SUCCESS);
461: }
463: namespace
464: {
466: std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
468: /*
469: Actual initialization function; any functions claiming to initialize PetscDevice or
470: PetscDeviceContext will have to run through this one
471: */
472: PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId)
473: {
474: PetscFunctionBegin;
476: if (PetscUnlikely(!PetscDeviceInitialized(type))) {
477: auto &dev = defaultDevices[type].first;
478: auto &init = defaultDevices[type].second;
480: PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
481: PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
482: PetscCall(PetscDeviceConfigure(dev));
483: init = true;
484: }
485: PetscFunctionReturn(PETSC_SUCCESS);
486: }
488: } // namespace
490: /*@
491: PetscDeviceInitialize - Initialize `PetscDevice`
493: Not Collective
495: Input Parameter:
496: . type - The `PetscDeviceType` to initialize
498: Level: beginner
500: Notes:
501: Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
502: result in device synchronization.
504: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
505: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
506: @*/
507: PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
508: {
509: PetscFunctionBegin;
511: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
512: PetscFunctionReturn(PETSC_SUCCESS);
513: }
515: /*@
516: PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
517: `PetscDeviceType`
519: Not Collective
521: Input Parameter:
522: . type - The `PetscDeviceType` to check
524: Level: beginner
526: Notes:
527: Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
529: If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
530: return `PETSC_FALSE` for that `PetscDeviceType`.
532: .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
533: `PetscDeviceCreate()`, `PetscDeviceDestroy()`
534: @*/
535: PetscBool PetscDeviceInitialized(PetscDeviceType type)
536: {
537: return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
538: }
540: /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
541: PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device)
542: {
543: PetscFunctionBegin;
544: PetscAssertPointer(device, 2);
545: PetscCall(PetscDeviceInitialize(type));
546: *device = defaultDevices[type].first;
547: PetscFunctionReturn(PETSC_SUCCESS);
548: }
550: /*@
551: PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
553: Not Collective
555: Input Parameters:
556: + device - The `PetscDevice`
557: - attr - The attribute
559: Output Parameter:
560: . value - The value of the attribute
562: Level: intermediate
564: Notes:
565: Since different attributes are often different types `value` is a `void *` to accommodate
566: them all. The underlying type of the attribute is therefore included in the name of the
567: `PetscDeviceAttribute` responsible for querying it. For example,
568: `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
570: .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
571: @*/
572: PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
573: {
574: PetscFunctionBegin;
577: PetscAssertPointer(value, 3);
578: PetscUseTypeMethod(device, getattribute, attr, value);
579: PetscFunctionReturn(PETSC_SUCCESS);
580: }
582: namespace
583: {
585: PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType)
586: {
587: PetscFunctionBegin;
588: if (!PetscDeviceConfiguredFor_Internal(type)) {
589: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
590: defaultDevices[type].first = nullptr;
591: PetscFunctionReturn(PETSC_SUCCESS);
592: }
593: PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
594: /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
595: switch (type) {
596: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
597: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
598: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
599: PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
600: default:
601: SETERRQ(comm, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
602: }
603: 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)]));
604: /*
605: defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES
606: have decided to initialize as
607: */
608: if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
609: PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
610: PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
611: if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
612: }
613: PetscFunctionReturn(PETSC_SUCCESS);
614: }
616: PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDeviceId, PetscBool *defaultDeviceIdSet, PetscBool *defaultView)
617: {
618: PetscInt initIdx = PETSC_DEVICE_INIT_LAZY;
619: auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
620: auto flg = PETSC_FALSE;
622: PetscFunctionBegin;
623: PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
624: if (flg) PetscCall(PetscLogGpuTime());
625: PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_energy_meter", &flg));
626: if (flg) PetscCall(PetscLogGpuEnergyMeter());
627: PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_energy", &flg));
628: if (flg) PetscCall(PetscLogGpuEnergy());
630: PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
631: PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
632: PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceIdSet));
633: PetscCall(PetscOptionsRangeInt("-device_select", PETSC_DEVICE_SELECT_HELP, "PetscDeviceCreate", *defaultDeviceId, defaultDeviceId, nullptr, PETSC_DEVICE_SELECT_LOWER_BOUND, PETSC_DEVICE_MAX_DEVICES));
634: PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView", *defaultView, defaultView, &flg));
635: PetscOptionsEnd();
637: if (initIdx == PETSC_DEVICE_INIT_NONE) {
638: /* disabled all device initialization if devices are globally disabled */
639: PetscCheck(*defaultDeviceId == PETSC_DECIDE, comm, PETSC_ERR_USER_INPUT, "You have disabled devices but also specified a particular device to use, these options are mutually exclusive");
640: *defaultView = PETSC_FALSE;
641: initDeviceIdx = PETSC_DEVICE_HOST;
642: } else {
643: *defaultView = static_cast<PetscBool>(*defaultView && flg);
644: if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
645: }
646: *defaultInitType = PetscDeviceInitTypeCast(initIdx);
647: *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
648: PetscFunctionReturn(PETSC_SUCCESS);
649: }
651: /* called from PetscFinalize() do not call yourself! */
652: PetscErrorCode PetscDeviceFinalize_Private()
653: {
654: PetscFunctionBegin;
655: if (PetscDefined(USE_DEBUG)) {
656: /*
657: you might be thinking, why on earth are you registered yet another finalizer in a
658: function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
659: because it is.
661: The crux of the problem is that the initializer (and therefore the ~finalizer~) of
662: PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
663: a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
664: won't be destroyed yet. So we need to repeat the check that all devices have been
665: destroyed again ~after~ the global context is destroyed. In summary:
667: 1. This finalizer runs and destroys all devices, except it may not because the global
668: context may still hold a reference!
669: 2. The global context finalizer runs and does the final reference count decrement
670: required, which actually destroys the held device.
671: 3. Our newly added finalizer runs and checks that all is well.
672: */
673: PetscCall(PetscRegisterFinalize([] {
674: PetscFunctionBegin;
675: for (auto &&device : defaultDevices) {
676: const auto dev = device.first;
678: 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);
679: }
680: PetscFunctionReturn(PETSC_SUCCESS);
681: }));
682: }
683: for (auto &&device : defaultDevices) {
684: PetscCall(PetscDeviceDestroy(&device.first));
685: device.second = false;
686: }
687: PetscFunctionReturn(PETSC_SUCCESS);
688: }
690: } // namespace
692: /*
693: Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
694: initialization types:
696: 1. defaultInitType - how does PetscDevice as a whole expect to initialize?
697: 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
698: e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
699: have all CUDA devices still initialize.
701: All told the following happens:
703: 0. defaultInitType -> LAZY
704: 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
705: 2. PetscDevice initializes each sub type with deviceDefaultInitType.
706: 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
707: to checking for specific device init. if view or specific device init
708: subTypeDefaultInitType -> EAGER. disabled once again overrides all.
709: */
711: PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm)
712: {
713: auto defaultView = PETSC_FALSE;
714: auto initializeDeviceContextEagerly = PETSC_FALSE;
715: auto defaultDeviceIdSet = PETSC_FALSE;
716: auto defaultDeviceId = PetscInt{PETSC_DECIDE};
717: auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT();
718: auto defaultInitType = PETSC_DEVICE_INIT_LAZY;
720: PetscFunctionBegin;
721: if (PetscDefined(USE_DEBUG)) {
722: int result;
724: PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
725: /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
726: * global space */
727: if (PetscUnlikely(result != MPI_IDENT)) {
728: char name[MPI_MAX_OBJECT_NAME] = {};
729: int len; /* unused */
731: PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
732: SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
733: }
734: }
735: comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
736: PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
738: PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDeviceId, &defaultDeviceIdSet, &defaultView));
740: // the precise values don't matter here, so long as they are sequential
741: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HOST) == 0, "");
742: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_CUDA) == 1, "");
743: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HIP) == 2, "");
744: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_SYCL) == 3, "");
745: static_assert(Petsc::util::to_underlying(PETSC_DEVICE_MAX) == 4, "");
746: for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
747: const auto deviceType = PetscDeviceTypeCast(i);
748: auto initType = defaultInitType;
750: PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDeviceId, defaultView, &initType));
751: if (PetscDeviceConfiguredFor_Internal(deviceType)) {
752: if (initType == PETSC_DEVICE_INIT_EAGER) {
753: initializeDeviceContextEagerly = PETSC_TRUE;
754: // only update the default device if the user hasn't set it previously
755: if (!defaultDeviceIdSet) {
756: deviceContextInitDevice = deviceType;
757: PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
758: }
759: } else if (initType == PETSC_DEVICE_INIT_NONE) {
760: if (deviceType != PETSC_DEVICE_HOST) PetscCheck(!defaultDeviceIdSet || (deviceType != deviceContextInitDevice), comm, PETSC_ERR_USER_INPUT, "Cannot explicitly disable the device set as default device type (%s)", PetscDeviceTypes[deviceType]);
761: }
762: }
763: }
765: PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
766: PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
767: /* PetscDevice is now fully initialized */
768: {
769: /*
770: query the options db to get the root settings from the user (if any).
772: This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
773: PetscDeviceContextSetFromOptions() before we even have one, then set a few static
774: variables in that file with the results.
775: */
776: auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
777: auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
779: PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
780: PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
781: PetscOptionsEnd();
783: if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
784: if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
785: }
787: if (initializeDeviceContextEagerly) {
788: PetscDeviceContext dctx;
790: PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
791: /* instantiates the device context */
792: PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
793: PetscCall(PetscDeviceContextSetUp(dctx));
794: }
795: PetscFunctionReturn(PETSC_SUCCESS);
796: }