xref: /petsc/src/sys/objects/device/interface/device.cxx (revision f4d061e980d13bc62f06124c58b76593bdf99e72)
10e6b6b59SJacob Faibussowitsch #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/
217f48955SJacob Faibussowitsch #include <petsc/private/petscadvancedmacros.h>
3030f984aSJacob Faibussowitsch 
40e6b6b59SJacob Faibussowitsch #include "../impls/host/hostdevice.hpp"
50e6b6b59SJacob Faibussowitsch #include "../impls/cupm/cupmdevice.hpp"
60e6b6b59SJacob Faibussowitsch #include "../impls/sycl/sycldevice.hpp"
70e6b6b59SJacob Faibussowitsch 
80e6b6b59SJacob Faibussowitsch #include <limits>  // std::numeric_limits
90e6b6b59SJacob Faibussowitsch #include <utility> // std::make_pair
100e6b6b59SJacob Faibussowitsch 
110e6b6b59SJacob Faibussowitsch using namespace Petsc::device;
12030f984aSJacob Faibussowitsch 
13cf3a2253SJacob Faibussowitsch /*
14cf3a2253SJacob Faibussowitsch   note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to
15cf3a2253SJacob Faibussowitsch   be picked up by the switch-case macros below
16cf3a2253SJacob Faibussowitsch */
170e6b6b59SJacob Faibussowitsch static host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
18030f984aSJacob Faibussowitsch #if PetscDefined(HAVE_CUDA)
190e6b6b59SJacob Faibussowitsch static cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
20030f984aSJacob Faibussowitsch #endif
21030f984aSJacob Faibussowitsch #if PetscDefined(HAVE_HIP)
220e6b6b59SJacob Faibussowitsch static cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
23030f984aSJacob Faibussowitsch #endif
24a2158755SJunchao Zhang #if PetscDefined(HAVE_SYCL)
250e6b6b59SJacob Faibussowitsch static sycl::Device SYCLDevice{PetscDeviceContextCreate_SYCL};
26a2158755SJunchao Zhang #endif
27030f984aSJacob Faibussowitsch 
2817f48955SJacob Faibussowitsch #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
2917f48955SJacob Faibussowitsch   case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
309566063dSJacob Faibussowitsch     PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
3117f48955SJacob Faibussowitsch   } break
32a4af0ceeSJacob Faibussowitsch 
33cf3a2253SJacob Faibussowitsch /*
34cf3a2253SJacob Faibussowitsch   Suppose you have:
35cf3a2253SJacob Faibussowitsch 
36cf3a2253SJacob Faibussowitsch   CUDADevice.myFunction(arg1,arg2)
37cf3a2253SJacob Faibussowitsch 
38cf3a2253SJacob Faibussowitsch   that you would like to conditionally define and call in a switch-case:
39cf3a2253SJacob Faibussowitsch 
40cf3a2253SJacob Faibussowitsch   switch(PetscDeviceType) {
41cf3a2253SJacob Faibussowitsch   #if PetscDefined(HAVE_CUDA)
42cf3a2253SJacob Faibussowitsch   case PETSC_DEVICE_CUDA: {
439566063dSJacob Faibussowitsch     PetscCall(CUDADevice.myFunction(arg1,arg2));
44cf3a2253SJacob Faibussowitsch   } break;
45cf3a2253SJacob Faibussowitsch   #endif
46cf3a2253SJacob Faibussowitsch   }
47cf3a2253SJacob Faibussowitsch 
48cf3a2253SJacob Faibussowitsch   then calling this macro:
49cf3a2253SJacob Faibussowitsch 
50cf3a2253SJacob Faibussowitsch   PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
51cf3a2253SJacob Faibussowitsch 
52cf3a2253SJacob Faibussowitsch   will expand to the following case statement:
53cf3a2253SJacob Faibussowitsch 
54cf3a2253SJacob Faibussowitsch   case PETSC_DEVICE_CUDA: {
559566063dSJacob Faibussowitsch     PetscCall(CUDADevice.myFunction(arg1,arg2));
56cf3a2253SJacob Faibussowitsch   } break
57cf3a2253SJacob Faibussowitsch 
58cf3a2253SJacob Faibussowitsch   if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise
5917f48955SJacob Faibussowitsch */
609371c9d4SSatish Balay #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(PetscConcat_(HAVE_, IMPLS), PETSC_DEVICE_CASE, PetscExpandToNothing)(IMPLS, func, __VA_ARGS__)
61030f984aSJacob Faibussowitsch 
62030f984aSJacob Faibussowitsch /*@C
63811af0c4SBarry Smith   PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
64030f984aSJacob Faibussowitsch 
650e6b6b59SJacob Faibussowitsch   Not Collective
66030f984aSJacob Faibussowitsch 
67f1a722f8SMatthew G. Knepley   Input Parameters:
68811af0c4SBarry Smith + type  - The type of `PetscDevice`
69811af0c4SBarry Smith - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically)
70030f984aSJacob Faibussowitsch 
71030f984aSJacob Faibussowitsch   Output Parameter:
72811af0c4SBarry Smith . device - The `PetscDevice`
73030f984aSJacob Faibussowitsch 
74030f984aSJacob Faibussowitsch   Notes:
750e6b6b59SJacob Faibussowitsch   This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
760e6b6b59SJacob Faibussowitsch   device synchronization.
77a4af0ceeSJacob Faibussowitsch 
78811af0c4SBarry Smith   `devid` is what you might pass to `cudaSetDevice()` for example.
79030f984aSJacob Faibussowitsch 
80030f984aSJacob Faibussowitsch   Level: beginner
81030f984aSJacob Faibussowitsch 
820e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`,
830e6b6b59SJacob Faibussowitsch `PetscDeviceInitialize()`,`PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
840e6b6b59SJacob Faibussowitsch `PetscDeviceView()`, `PetscDeviceDestroy()`
85030f984aSJacob Faibussowitsch @*/
869371c9d4SSatish Balay PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device) {
87030f984aSJacob Faibussowitsch   static PetscInt PetscDeviceCounter = 0;
88030f984aSJacob Faibussowitsch 
89030f984aSJacob Faibussowitsch   PetscFunctionBegin;
90a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
91a4af0ceeSJacob Faibussowitsch   PetscValidPointer(device, 3);
929566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializePackage());
930e6b6b59SJacob Faibussowitsch   PetscCall(PetscNew(device));
940e6b6b59SJacob Faibussowitsch   (*device)->id     = PetscDeviceCounter++;
950e6b6b59SJacob Faibussowitsch   (*device)->type   = type;
960e6b6b59SJacob Faibussowitsch   (*device)->refcnt = 1;
97cf3a2253SJacob Faibussowitsch   /*
98cf3a2253SJacob Faibussowitsch     if you are adding a device, you also need to add it's initialization in
99cf3a2253SJacob Faibussowitsch     PetscDeviceInitializeTypeFromOptions_Private() below
100cf3a2253SJacob Faibussowitsch   */
101a4af0ceeSJacob Faibussowitsch   switch (type) {
1020e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
1030e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
1040e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
1050e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
106030f984aSJacob Faibussowitsch   default:
10717f48955SJacob Faibussowitsch     /* in case the above macros expand to nothing this silences any unused variable warnings */
10817f48955SJacob Faibussowitsch     (void)(devid);
10998921bdaSJacob Faibussowitsch     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]);
110030f984aSJacob Faibussowitsch   }
111030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
112030f984aSJacob Faibussowitsch }
113030f984aSJacob Faibussowitsch 
114030f984aSJacob Faibussowitsch /*@C
115811af0c4SBarry Smith   PetscDeviceDestroy - Free a `PetscDevice`
116030f984aSJacob Faibussowitsch 
1170e6b6b59SJacob Faibussowitsch   Not Collective
118030f984aSJacob Faibussowitsch 
119030f984aSJacob Faibussowitsch   Input Parameter:
1200e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
121030f984aSJacob Faibussowitsch 
122030f984aSJacob Faibussowitsch   Level: beginner
123030f984aSJacob Faibussowitsch 
1240e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
1250e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
126030f984aSJacob Faibussowitsch @*/
1279371c9d4SSatish Balay PetscErrorCode PetscDeviceDestroy(PetscDevice *device) {
128a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
1290e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 1);
130a4af0ceeSJacob Faibussowitsch   if (!*device) PetscFunctionReturn(0);
131a4af0ceeSJacob Faibussowitsch   PetscValidDevice(*device, 1);
1329566063dSJacob Faibussowitsch   PetscCall(PetscDeviceDereference_Internal(*device));
133a4af0ceeSJacob Faibussowitsch   if ((*device)->refcnt) {
1340e6b6b59SJacob Faibussowitsch     *device = nullptr;
135a4af0ceeSJacob Faibussowitsch     PetscFunctionReturn(0);
136030f984aSJacob Faibussowitsch   }
1379566063dSJacob Faibussowitsch   PetscCall(PetscFree((*device)->data));
1389566063dSJacob Faibussowitsch   PetscCall(PetscFree(*device));
139030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
140030f984aSJacob Faibussowitsch }
141030f984aSJacob Faibussowitsch 
142a4af0ceeSJacob Faibussowitsch /*@C
143811af0c4SBarry Smith   PetscDeviceConfigure - Configure a particular `PetscDevice`
144030f984aSJacob Faibussowitsch 
1450e6b6b59SJacob Faibussowitsch   Not Collective
146a4af0ceeSJacob Faibussowitsch 
147a4af0ceeSJacob Faibussowitsch   Input Parameter:
148811af0c4SBarry Smith . device - The `PetscDevice` to configure
149a4af0ceeSJacob Faibussowitsch 
1500e6b6b59SJacob Faibussowitsch   Notes:
1510e6b6b59SJacob Faibussowitsch   The user should not assume that this is a cheap operation.
152a4af0ceeSJacob Faibussowitsch 
153a4af0ceeSJacob Faibussowitsch   Level: beginner
154a4af0ceeSJacob Faibussowitsch 
1550e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
1560e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
157a4af0ceeSJacob Faibussowitsch @*/
1589371c9d4SSatish Balay PetscErrorCode PetscDeviceConfigure(PetscDevice device) {
159030f984aSJacob Faibussowitsch   PetscFunctionBegin;
160a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
161cf3a2253SJacob Faibussowitsch   /*
162cf3a2253SJacob Faibussowitsch     if no available configuration is available, this cascades all the way down to default
163cf3a2253SJacob Faibussowitsch     and error
164cf3a2253SJacob Faibussowitsch   */
1650e6b6b59SJacob Faibussowitsch   switch (const auto dtype = device->type) {
1660e6b6b59SJacob Faibussowitsch   case PETSC_DEVICE_HOST:
1670e6b6b59SJacob Faibussowitsch     if (PetscDefined(HAVE_HOST)) break; // always true
1689371c9d4SSatish Balay   case PETSC_DEVICE_CUDA:
1699371c9d4SSatish Balay     if (PetscDefined(HAVE_CUDA)) break;
1700e6b6b59SJacob Faibussowitsch     goto error;
1719371c9d4SSatish Balay   case PETSC_DEVICE_HIP:
1729371c9d4SSatish Balay     if (PetscDefined(HAVE_HIP)) break;
1730e6b6b59SJacob Faibussowitsch     goto error;
1749371c9d4SSatish Balay   case PETSC_DEVICE_SYCL:
1759371c9d4SSatish Balay     if (PetscDefined(HAVE_SYCL)) break;
176*f4d061e9SPierre Jolivet     goto error;
1770e6b6b59SJacob Faibussowitsch   default:
1780e6b6b59SJacob Faibussowitsch   error:
1790e6b6b59SJacob Faibussowitsch     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
180a4af0ceeSJacob Faibussowitsch   }
181dbbe0bcdSBarry Smith   PetscUseTypeMethod(device, configure);
182a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
183a4af0ceeSJacob Faibussowitsch }
184a4af0ceeSJacob Faibussowitsch 
185a4af0ceeSJacob Faibussowitsch /*@C
186811af0c4SBarry Smith   PetscDeviceView - View a `PetscDevice`
187a4af0ceeSJacob Faibussowitsch 
1880e6b6b59SJacob Faibussowitsch   Collective on viewer
189a4af0ceeSJacob Faibussowitsch 
19091e63d38SStefano Zampini   Input Parameters:
191811af0c4SBarry Smith + device - The `PetscDevice` to view
1920e6b6b59SJacob Faibussowitsch - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
193a4af0ceeSJacob Faibussowitsch 
194a4af0ceeSJacob Faibussowitsch   Level: beginner
195a4af0ceeSJacob Faibussowitsch 
1960e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
1970e6b6b59SJacob Faibussowitsch `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
198a4af0ceeSJacob Faibussowitsch @*/
1999371c9d4SSatish Balay PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer) {
2000e6b6b59SJacob Faibussowitsch   auto      sub = viewer;
2010e6b6b59SJacob Faibussowitsch   PetscBool iascii;
2020e6b6b59SJacob Faibussowitsch 
203a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
204a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
2050e6b6b59SJacob Faibussowitsch   if (viewer) {
206a4af0ceeSJacob Faibussowitsch     PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2070e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
2080e6b6b59SJacob Faibussowitsch   } else {
2090e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
2100e6b6b59SJacob Faibussowitsch     iascii = PETSC_TRUE;
2110e6b6b59SJacob Faibussowitsch   }
2120e6b6b59SJacob Faibussowitsch 
2130e6b6b59SJacob Faibussowitsch   if (iascii) {
2140e6b6b59SJacob Faibussowitsch     auto        dtype = PETSC_DEVICE_HOST;
2150e6b6b59SJacob Faibussowitsch     MPI_Comm    comm;
2160e6b6b59SJacob Faibussowitsch     PetscMPIInt size;
2170e6b6b59SJacob Faibussowitsch     PetscInt    id = 0;
2180e6b6b59SJacob Faibussowitsch 
2190e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
2200e6b6b59SJacob Faibussowitsch     PetscCallMPI(MPI_Comm_size(comm, &size));
2210e6b6b59SJacob Faibussowitsch 
2220e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetDeviceId(device, &id));
2230e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetType(device, &dtype));
2240e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
2250e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
2260e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPushTab(sub));
2270e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
2280e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
2290e6b6b59SJacob Faibussowitsch   }
2300e6b6b59SJacob Faibussowitsch 
2310e6b6b59SJacob Faibussowitsch   // see if impls has extra viewer stuff
2320e6b6b59SJacob Faibussowitsch   PetscTryTypeMethod(device, view, sub);
2330e6b6b59SJacob Faibussowitsch 
2340e6b6b59SJacob Faibussowitsch   if (iascii) {
2350e6b6b59SJacob Faibussowitsch     // undo the ASCII specific stuff
2360e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPopTab(sub));
2370e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
2380e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerFlush(viewer));
2390e6b6b59SJacob Faibussowitsch   }
240a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
241a4af0ceeSJacob Faibussowitsch }
242a4af0ceeSJacob Faibussowitsch 
24391e63d38SStefano Zampini /*@C
2440e6b6b59SJacob Faibussowitsch   PetscDeviceGetType - Get the type of device
24591e63d38SStefano Zampini 
2460e6b6b59SJacob Faibussowitsch   Not Collective
24791e63d38SStefano Zampini 
24891e63d38SStefano Zampini   Input Parameter:
249811af0c4SBarry Smith . device - The `PetscDevice`
25091e63d38SStefano Zampini 
25191e63d38SStefano Zampini   Output Parameter:
2520e6b6b59SJacob Faibussowitsch . type - The `PetscDeviceType`
25391e63d38SStefano Zampini 
25491e63d38SStefano Zampini   Level: beginner
25591e63d38SStefano Zampini 
2560e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
2570e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
2580e6b6b59SJacob Faibussowitsch `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
2590e6b6b59SJacob Faibussowitsch @*/
2600e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type) {
2610e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
2620e6b6b59SJacob Faibussowitsch   PetscValidDevice(device, 1);
2630e6b6b59SJacob Faibussowitsch   PetscValidPointer(type, 2);
2640e6b6b59SJacob Faibussowitsch   *type = device->type;
2650e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
2660e6b6b59SJacob Faibussowitsch }
2670e6b6b59SJacob Faibussowitsch 
2680e6b6b59SJacob Faibussowitsch /*@C
2690e6b6b59SJacob Faibussowitsch   PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
2700e6b6b59SJacob Faibussowitsch 
2710e6b6b59SJacob Faibussowitsch   Not Collective
2720e6b6b59SJacob Faibussowitsch 
2730e6b6b59SJacob Faibussowitsch   Input Parameter:
2740e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
2750e6b6b59SJacob Faibussowitsch 
2760e6b6b59SJacob Faibussowitsch   Output Parameter:
2770e6b6b59SJacob Faibussowitsch . id - The id
2780e6b6b59SJacob Faibussowitsch 
2790e6b6b59SJacob Faibussowitsch   Notes:
2800e6b6b59SJacob Faibussowitsch   The returned ID may have been assigned by the underlying device backend. For example if the
2810e6b6b59SJacob Faibussowitsch   backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
2820e6b6b59SJacob Faibussowitsch   this device was configured.
2830e6b6b59SJacob Faibussowitsch 
2840e6b6b59SJacob Faibussowitsch   Level: beginner
2850e6b6b59SJacob Faibussowitsch 
2860e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
28791e63d38SStefano Zampini @*/
2889371c9d4SSatish Balay PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id) {
28991e63d38SStefano Zampini   PetscFunctionBegin;
29091e63d38SStefano Zampini   PetscValidDevice(device, 1);
29191e63d38SStefano Zampini   PetscValidIntPointer(id, 2);
29291e63d38SStefano Zampini   *id = device->deviceId;
29391e63d38SStefano Zampini   PetscFunctionReturn(0);
29491e63d38SStefano Zampini }
29591e63d38SStefano Zampini 
2960e6b6b59SJacob Faibussowitsch struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
2970e6b6b59SJacob Faibussowitsch   PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
2980e6b6b59SJacob Faibussowitsch 
2990e6b6b59SJacob Faibussowitsch   PETSC_NODISCARD PetscErrorCode finalize_() noexcept {
3000e6b6b59SJacob Faibussowitsch     PetscFunctionBegin;
3010e6b6b59SJacob Faibussowitsch     type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
3020e6b6b59SJacob Faibussowitsch     PetscFunctionReturn(0);
3030e6b6b59SJacob Faibussowitsch   }
3040e6b6b59SJacob Faibussowitsch };
3050e6b6b59SJacob Faibussowitsch 
3060e6b6b59SJacob Faibussowitsch static auto default_device_type = DefaultDeviceType();
3070e6b6b59SJacob Faibussowitsch 
3080e6b6b59SJacob Faibussowitsch /*@C
3090e6b6b59SJacob Faibussowitsch   PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
3100e6b6b59SJacob Faibussowitsch 
3110e6b6b59SJacob Faibussowitsch   Not Collective
3120e6b6b59SJacob Faibussowitsch 
3130e6b6b59SJacob Faibussowitsch   Notes:
3140e6b6b59SJacob Faibussowitsch   Unless selected by the user, the default device is selected in the following order\:
3150e6b6b59SJacob Faibussowitsch   `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
3160e6b6b59SJacob Faibussowitsch 
3170e6b6b59SJacob Faibussowitsch   Level: beginner
3180e6b6b59SJacob Faibussowitsch 
3190e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
3200e6b6b59SJacob Faibussowitsch @*/
3210e6b6b59SJacob Faibussowitsch PetscDeviceType PETSC_DEVICE_DEFAULT(void) {
3220e6b6b59SJacob Faibussowitsch   return default_device_type.type;
3230e6b6b59SJacob Faibussowitsch }
3240e6b6b59SJacob Faibussowitsch 
3250e6b6b59SJacob Faibussowitsch /*@C
3260e6b6b59SJacob Faibussowitsch   PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
3270e6b6b59SJacob Faibussowitsch 
3280e6b6b59SJacob Faibussowitsch   Not Collective
3290e6b6b59SJacob Faibussowitsch 
3300e6b6b59SJacob Faibussowitsch   Input Parameter:
3310e6b6b59SJacob Faibussowitsch . type - the new default device type
3320e6b6b59SJacob Faibussowitsch 
3330e6b6b59SJacob Faibussowitsch   Notes:
3340e6b6b59SJacob Faibussowitsch   This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
3350e6b6b59SJacob Faibussowitsch 
3360e6b6b59SJacob Faibussowitsch   Level: beginner
3370e6b6b59SJacob Faibussowitsch 
3380e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
3390e6b6b59SJacob Faibussowitsch @*/
3400e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type) {
3410e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
3420e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3430e6b6b59SJacob Faibussowitsch   if (default_device_type.type != type) {
3440e6b6b59SJacob Faibussowitsch     // no need to waster a PetscRegisterFinalize() slot if we don't change it
3450e6b6b59SJacob Faibussowitsch     default_device_type.type = type;
3460e6b6b59SJacob Faibussowitsch     PetscCall(default_device_type.register_finalize());
3470e6b6b59SJacob Faibussowitsch   }
3480e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
3490e6b6b59SJacob Faibussowitsch }
3500e6b6b59SJacob Faibussowitsch 
3510e6b6b59SJacob Faibussowitsch static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
3520e6b6b59SJacob Faibussowitsch 
3530e6b6b59SJacob Faibussowitsch /*
3540e6b6b59SJacob Faibussowitsch   Actual intialization function; any functions claiming to initialize PetscDevice or
3550e6b6b59SJacob Faibussowitsch   PetscDeviceContext will have to run through this one
3560e6b6b59SJacob Faibussowitsch */
3570e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId) {
3580e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
3590e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3600e6b6b59SJacob Faibussowitsch   if (PetscUnlikely(!PetscDeviceInitialized(type))) {
3610e6b6b59SJacob Faibussowitsch     auto &dev  = defaultDevices[type].first;
3620e6b6b59SJacob Faibussowitsch     auto &init = defaultDevices[type].second;
3630e6b6b59SJacob Faibussowitsch 
3640e6b6b59SJacob Faibussowitsch     PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
3650e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
3660e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceConfigure(dev));
3670e6b6b59SJacob Faibussowitsch     init = true;
3680e6b6b59SJacob Faibussowitsch   }
3690e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
3700e6b6b59SJacob Faibussowitsch }
371a4af0ceeSJacob Faibussowitsch 
372a4af0ceeSJacob Faibussowitsch /*@C
373811af0c4SBarry Smith   PetscDeviceInitialize - Initialize `PetscDevice`
374a4af0ceeSJacob Faibussowitsch 
3750e6b6b59SJacob Faibussowitsch   Not Collective
376a4af0ceeSJacob Faibussowitsch 
377a4af0ceeSJacob Faibussowitsch   Input Parameter:
378811af0c4SBarry Smith . type - The `PetscDeviceType` to initialize
379a4af0ceeSJacob Faibussowitsch 
3800e6b6b59SJacob Faibussowitsch   Notes:
3810e6b6b59SJacob Faibussowitsch   Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
3820e6b6b59SJacob Faibussowitsch   result in device synchronization.
383a4af0ceeSJacob Faibussowitsch 
384a4af0ceeSJacob Faibussowitsch   Level: beginner
385a4af0ceeSJacob Faibussowitsch 
3860e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
3870e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
388a4af0ceeSJacob Faibussowitsch @*/
3899371c9d4SSatish Balay PetscErrorCode PetscDeviceInitialize(PetscDeviceType type) {
390a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
391a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3929566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
393a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
394a4af0ceeSJacob Faibussowitsch }
395a4af0ceeSJacob Faibussowitsch 
396a4af0ceeSJacob Faibussowitsch /*@C
397811af0c4SBarry Smith   PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
398811af0c4SBarry Smith   `PetscDeviceType`
399a4af0ceeSJacob Faibussowitsch 
4000e6b6b59SJacob Faibussowitsch   Not Collective
401a4af0ceeSJacob Faibussowitsch 
402a4af0ceeSJacob Faibussowitsch   Input Parameter:
403811af0c4SBarry Smith . type - The `PetscDeviceType` to check
404a4af0ceeSJacob Faibussowitsch 
4050e6b6b59SJacob Faibussowitsch   Notes:
4060e6b6b59SJacob Faibussowitsch   Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
407a4af0ceeSJacob Faibussowitsch 
408811af0c4SBarry Smith   If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
409811af0c4SBarry Smith   return `PETSC_FALSE` for that `PetscDeviceType`.
410a4af0ceeSJacob Faibussowitsch 
411a4af0ceeSJacob Faibussowitsch   Level: beginner
412a4af0ceeSJacob Faibussowitsch 
4130e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
4140e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
415a4af0ceeSJacob Faibussowitsch @*/
4169371c9d4SSatish Balay PetscBool PetscDeviceInitialized(PetscDeviceType type) {
4170e6b6b59SJacob Faibussowitsch   return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
4180e6b6b59SJacob Faibussowitsch }
4190e6b6b59SJacob Faibussowitsch 
4200e6b6b59SJacob Faibussowitsch /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
4210e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device) {
4220e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
4230e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 2);
4240e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitialize(type));
4250e6b6b59SJacob Faibussowitsch   *device = defaultDevices[type].first;
4260e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
427a4af0ceeSJacob Faibussowitsch }
428a4af0ceeSJacob Faibussowitsch 
429a16fd2c9SJacob Faibussowitsch /*@C
430a16fd2c9SJacob Faibussowitsch   PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
431a16fd2c9SJacob Faibussowitsch 
4320e6b6b59SJacob Faibussowitsch   Not Collective
433a16fd2c9SJacob Faibussowitsch 
434a16fd2c9SJacob Faibussowitsch   Input Parameters:
435a16fd2c9SJacob Faibussowitsch + device - The `PetscDevice`
436a16fd2c9SJacob Faibussowitsch - attr   - The attribute
437a16fd2c9SJacob Faibussowitsch 
438a16fd2c9SJacob Faibussowitsch   Output Parameter:
439a16fd2c9SJacob Faibussowitsch . value - The value of the attribute
440a16fd2c9SJacob Faibussowitsch 
441a16fd2c9SJacob Faibussowitsch   Notes:
442a16fd2c9SJacob Faibussowitsch   Since different attributes are often different types `value` is a `void *` to accommodate
443a16fd2c9SJacob Faibussowitsch   them all. The underlying type of the attribute is therefore included in the name of the
444a16fd2c9SJacob Faibussowitsch   `PetscDeviceAttribute` reponsible for querying it. For example,
445a16fd2c9SJacob Faibussowitsch   `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
446a16fd2c9SJacob Faibussowitsch 
4470e6b6b59SJacob Faibussowitsch   Level: intermediate
4480e6b6b59SJacob Faibussowitsch 
449a16fd2c9SJacob Faibussowitsch .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
450a16fd2c9SJacob Faibussowitsch @*/
451a16fd2c9SJacob Faibussowitsch PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value) {
452a16fd2c9SJacob Faibussowitsch   PetscFunctionBegin;
453a16fd2c9SJacob Faibussowitsch   PetscValidDevice(device, 1);
454a16fd2c9SJacob Faibussowitsch   PetscValidDeviceAttribute(attr, 2);
455a16fd2c9SJacob Faibussowitsch   PetscValidPointer(value, 3);
456a16fd2c9SJacob Faibussowitsch   PetscUseTypeMethod(device, getattribute, attr, value);
457a16fd2c9SJacob Faibussowitsch   PetscFunctionReturn(0);
458a16fd2c9SJacob Faibussowitsch }
459a16fd2c9SJacob Faibussowitsch 
4609371c9d4SSatish Balay static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType) {
461a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
462a4af0ceeSJacob Faibussowitsch   if (!PetscDeviceConfiguredFor_Internal(type)) {
4630e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
4640e6b6b59SJacob Faibussowitsch     defaultDevices[type].first = nullptr;
465a4af0ceeSJacob Faibussowitsch     PetscFunctionReturn(0);
466a4af0ceeSJacob Faibussowitsch   }
4670e6b6b59SJacob Faibussowitsch   PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
468a4af0ceeSJacob Faibussowitsch   /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
469a4af0ceeSJacob Faibussowitsch   switch (type) {
4700e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4710e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4720e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4730e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4749371c9d4SSatish Balay   default: SETERRQ(comm, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
475a4af0ceeSJacob Faibussowitsch   }
4760e6b6b59SJacob Faibussowitsch   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)]));
477cf3a2253SJacob Faibussowitsch   /*
4780e6b6b59SJacob Faibussowitsch     defaultInitType, defaultView  and defaultDeviceId now represent what the individual TYPES
4790e6b6b59SJacob Faibussowitsch     have decided to initialize as
480cf3a2253SJacob Faibussowitsch   */
4810e6b6b59SJacob Faibussowitsch   if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
4820e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
4839566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
4840e6b6b59SJacob Faibussowitsch     if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
4850e6b6b59SJacob Faibussowitsch   }
4860e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
4870e6b6b59SJacob Faibussowitsch }
488a4af0ceeSJacob Faibussowitsch 
4890e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView) {
4900e6b6b59SJacob Faibussowitsch   PetscInt initIdx       = PETSC_DEVICE_INIT_LAZY;
4910e6b6b59SJacob Faibussowitsch   auto     initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
4920e6b6b59SJacob Faibussowitsch   auto     flg           = PETSC_FALSE;
4930e6b6b59SJacob Faibussowitsch 
4940e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
4950e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
4960e6b6b59SJacob Faibussowitsch   if (flg) PetscCall(PetscLogGpuTime());
4970e6b6b59SJacob Faibussowitsch 
4980e6b6b59SJacob Faibussowitsch   PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
4990e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
5000e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
5010e6b6b59SJacob Faibussowitsch   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));
5020e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
5030e6b6b59SJacob Faibussowitsch   PetscOptionsEnd();
5040e6b6b59SJacob Faibussowitsch 
5050e6b6b59SJacob Faibussowitsch   if (initIdx == PETSC_DEVICE_INIT_NONE) {
5060e6b6b59SJacob Faibussowitsch     /* disabled all device initialization if devices are globally disabled */
5070e6b6b59SJacob Faibussowitsch     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");
5080e6b6b59SJacob Faibussowitsch     *defaultView  = PETSC_FALSE;
5090e6b6b59SJacob Faibussowitsch     initDeviceIdx = PETSC_DEVICE_HOST;
5100e6b6b59SJacob Faibussowitsch   } else {
5110e6b6b59SJacob Faibussowitsch     *defaultView = static_cast<PetscBool>(*defaultView && flg);
5120e6b6b59SJacob Faibussowitsch     if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
513a4af0ceeSJacob Faibussowitsch   }
5140e6b6b59SJacob Faibussowitsch   *defaultInitType         = PetscDeviceInitTypeCast(initIdx);
5150e6b6b59SJacob Faibussowitsch   *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
516030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
517030f984aSJacob Faibussowitsch }
518030f984aSJacob Faibussowitsch 
519030f984aSJacob Faibussowitsch /* called from PetscFinalize() do not call yourself! */
5200e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceFinalize_Private() {
521030f984aSJacob Faibussowitsch   PetscFunctionBegin;
522a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
523bf025ffbSJacob Faibussowitsch     const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] {
524a4af0ceeSJacob Faibussowitsch       PetscFunctionBegin;
5250e6b6b59SJacob Faibussowitsch       for (auto &&device : defaultDevices) {
5260e6b6b59SJacob Faibussowitsch         const auto dev = device.first;
5270e6b6b59SJacob Faibussowitsch 
5280e6b6b59SJacob Faibussowitsch         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);
5290e6b6b59SJacob Faibussowitsch       }
530a4af0ceeSJacob Faibussowitsch       PetscFunctionReturn(0);
531a4af0ceeSJacob Faibussowitsch     };
532bf025ffbSJacob Faibussowitsch     /*
533bf025ffbSJacob Faibussowitsch       you might be thinking, why on earth are you registered yet another finalizer in a
534bf025ffbSJacob Faibussowitsch       function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
535bf025ffbSJacob Faibussowitsch       because it is.
536bf025ffbSJacob Faibussowitsch 
537bf025ffbSJacob Faibussowitsch       The crux of the problem is that the initializer (and therefore the ~finalizer~) of
538bf025ffbSJacob Faibussowitsch       PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
539bf025ffbSJacob Faibussowitsch       a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
540bf025ffbSJacob Faibussowitsch       won't be destroyed yet. So we need to repeat the check that all devices have been
541bf025ffbSJacob Faibussowitsch       destroyed again ~after~ the global context is destroyed. In summary:
542bf025ffbSJacob Faibussowitsch 
543bf025ffbSJacob Faibussowitsch       1. This finalizer runs and destroys all devices, except it may not because the global
544bf025ffbSJacob Faibussowitsch          context may still hold a reference!
545bf025ffbSJacob Faibussowitsch       2. The global context finalizer runs and does the final reference count decrement
546bf025ffbSJacob Faibussowitsch          required, which actually destroys the held device.
547bf025ffbSJacob Faibussowitsch       3. Our newly added finalizer runs and checks that all is well.
548a4af0ceeSJacob Faibussowitsch     */
5490e6b6b59SJacob Faibussowitsch     PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize)));
550a4af0ceeSJacob Faibussowitsch   }
5510e6b6b59SJacob Faibussowitsch   for (auto &&device : defaultDevices) {
5520e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceDestroy(&device.first));
5530e6b6b59SJacob Faibussowitsch     device.second = false;
5540e6b6b59SJacob Faibussowitsch   }
555030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
556030f984aSJacob Faibussowitsch }
557030f984aSJacob Faibussowitsch 
558cf3a2253SJacob Faibussowitsch /*
559cf3a2253SJacob Faibussowitsch   Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
560cf3a2253SJacob Faibussowitsch   initialization types:
561cf3a2253SJacob Faibussowitsch 
562a4af0ceeSJacob Faibussowitsch   1. defaultInitType - how does PetscDevice as a whole expect to initialize?
563a4af0ceeSJacob Faibussowitsch   2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
564a4af0ceeSJacob Faibussowitsch      e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
565a4af0ceeSJacob Faibussowitsch      have all CUDA devices still initialize.
566a4af0ceeSJacob Faibussowitsch 
567a4af0ceeSJacob Faibussowitsch   All told the following happens:
568cf3a2253SJacob Faibussowitsch 
569a4af0ceeSJacob Faibussowitsch   0. defaultInitType -> LAZY
570a4af0ceeSJacob Faibussowitsch   1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
571a4af0ceeSJacob Faibussowitsch   2. PetscDevice initializes each sub type with deviceDefaultInitType.
572a4af0ceeSJacob Faibussowitsch   2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
573a4af0ceeSJacob Faibussowitsch       to checking for specific device init. if view or specific device init
574a4af0ceeSJacob Faibussowitsch       subTypeDefaultInitType -> EAGER. disabled once again overrides all.
575a4af0ceeSJacob Faibussowitsch */
5760e6b6b59SJacob Faibussowitsch 
5779371c9d4SSatish Balay PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm) {
5787a101e5eSJacob Faibussowitsch   auto defaultView                    = PETSC_FALSE;
5797a101e5eSJacob Faibussowitsch   auto initializeDeviceContextEagerly = PETSC_FALSE;
5800e6b6b59SJacob Faibussowitsch   auto defaultDeviceSet               = PETSC_FALSE;
5817a101e5eSJacob Faibussowitsch   auto defaultDevice                  = PetscInt{PETSC_DECIDE};
5820e6b6b59SJacob Faibussowitsch   auto deviceContextInitDevice        = PETSC_DEVICE_DEFAULT();
5830e6b6b59SJacob Faibussowitsch   auto defaultInitType                = PETSC_DEVICE_INIT_LAZY;
584a4af0ceeSJacob Faibussowitsch 
585a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
586a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
587a4af0ceeSJacob Faibussowitsch     int result;
588a4af0ceeSJacob Faibussowitsch 
5899566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
590a4af0ceeSJacob Faibussowitsch     /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
591a4af0ceeSJacob Faibussowitsch      * global space */
592a4af0ceeSJacob Faibussowitsch     if (PetscUnlikely(result != MPI_IDENT)) {
593a4af0ceeSJacob Faibussowitsch       char name[MPI_MAX_OBJECT_NAME] = {};
594a4af0ceeSJacob Faibussowitsch       int  len; /* unused */
595a4af0ceeSJacob Faibussowitsch 
5969566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
59798921bdaSJacob Faibussowitsch       SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
598a4af0ceeSJacob Faibussowitsch     }
599a4af0ceeSJacob Faibussowitsch   }
600a4af0ceeSJacob Faibussowitsch   comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
6019566063dSJacob Faibussowitsch   PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
602a4af0ceeSJacob Faibussowitsch 
6030e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
6047a101e5eSJacob Faibussowitsch 
6050e6b6b59SJacob Faibussowitsch   // the precise values don't matter here, so long as they are sequential
6060e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, "");
6070e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, "");
6080e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, "");
6090e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, "");
6100e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_MAX) == 4, "");
6110e6b6b59SJacob Faibussowitsch   for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
6120e6b6b59SJacob Faibussowitsch     const auto deviceType = PetscDeviceTypeCast(i);
613a4af0ceeSJacob Faibussowitsch     auto       initType   = defaultInitType;
614a4af0ceeSJacob Faibussowitsch 
6159566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
6160e6b6b59SJacob Faibussowitsch     if (PetscDeviceConfiguredFor_Internal(deviceType)) {
6170e6b6b59SJacob Faibussowitsch       if (initType == PETSC_DEVICE_INIT_EAGER) {
618a4af0ceeSJacob Faibussowitsch         initializeDeviceContextEagerly = PETSC_TRUE;
6190e6b6b59SJacob Faibussowitsch         // only update the default device if the user hasn't set it previously
6200e6b6b59SJacob Faibussowitsch         if (!defaultDeviceSet) {
621a4af0ceeSJacob Faibussowitsch           deviceContextInitDevice = deviceType;
6220e6b6b59SJacob Faibussowitsch           PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
6230e6b6b59SJacob Faibussowitsch         }
6240e6b6b59SJacob Faibussowitsch       } else if (initType == PETSC_DEVICE_INIT_NONE) {
6250e6b6b59SJacob Faibussowitsch         if (deviceType != PETSC_DEVICE_HOST) PetscCheck(deviceType != deviceContextInitDevice, comm, PETSC_ERR_USER_INPUT, "Cannot explicitly disable the device set as default device type (%s)", PetscDeviceTypes[deviceType]);
626a4af0ceeSJacob Faibussowitsch       }
627a4af0ceeSJacob Faibussowitsch     }
6280e6b6b59SJacob Faibussowitsch   }
6290e6b6b59SJacob Faibussowitsch 
6300e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
6310e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
6320e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
6330e6b6b59SJacob Faibussowitsch   /*                       PetscDevice is now fully initialized                          */
6340e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
6350e6b6b59SJacob Faibussowitsch   {
6360e6b6b59SJacob Faibussowitsch     /*
6370e6b6b59SJacob Faibussowitsch       query the options db to get the root settings from the user (if any).
6380e6b6b59SJacob Faibussowitsch 
6390e6b6b59SJacob Faibussowitsch       This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
6400e6b6b59SJacob Faibussowitsch       PetscDeviceContextSetFromOptions() before we even have one, then set a few static
6410e6b6b59SJacob Faibussowitsch       variables in that file with the results.
6420e6b6b59SJacob Faibussowitsch     */
6430e6b6b59SJacob Faibussowitsch     auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
6440e6b6b59SJacob Faibussowitsch     auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
6450e6b6b59SJacob Faibussowitsch 
6460e6b6b59SJacob Faibussowitsch     PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
6470e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
6480e6b6b59SJacob Faibussowitsch     PetscOptionsEnd();
6490e6b6b59SJacob Faibussowitsch 
6500e6b6b59SJacob Faibussowitsch     if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
6510e6b6b59SJacob Faibussowitsch     if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
6520e6b6b59SJacob Faibussowitsch   }
6530e6b6b59SJacob Faibussowitsch 
654a4af0ceeSJacob Faibussowitsch   if (initializeDeviceContextEagerly) {
655a4af0ceeSJacob Faibussowitsch     PetscDeviceContext dctx;
656a4af0ceeSJacob Faibussowitsch 
6570e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
6580e6b6b59SJacob Faibussowitsch     /* instantiates the device context */
6599566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
6609566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextSetUp(dctx));
661a4af0ceeSJacob Faibussowitsch   }
662a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
663a4af0ceeSJacob Faibussowitsch }
664