xref: /petsc/src/sys/objects/device/interface/device.cxx (revision 1015a2a477dd2695c1e05348307aec5fc5bfe5fb)
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 @*/
86d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device)
87d71ae5a4SJacob Faibussowitsch {
88030f984aSJacob Faibussowitsch   static PetscInt PetscDeviceCounter = 0;
89030f984aSJacob Faibussowitsch 
90030f984aSJacob Faibussowitsch   PetscFunctionBegin;
91a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
92a4af0ceeSJacob Faibussowitsch   PetscValidPointer(device, 3);
939566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializePackage());
940e6b6b59SJacob Faibussowitsch   PetscCall(PetscNew(device));
950e6b6b59SJacob Faibussowitsch   (*device)->id     = PetscDeviceCounter++;
960e6b6b59SJacob Faibussowitsch   (*device)->type   = type;
970e6b6b59SJacob Faibussowitsch   (*device)->refcnt = 1;
98cf3a2253SJacob Faibussowitsch   /*
99cf3a2253SJacob Faibussowitsch     if you are adding a device, you also need to add it's initialization in
100cf3a2253SJacob Faibussowitsch     PetscDeviceInitializeTypeFromOptions_Private() below
101cf3a2253SJacob Faibussowitsch   */
102a4af0ceeSJacob Faibussowitsch   switch (type) {
1030e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
1040e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
1050e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
1060e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
107030f984aSJacob Faibussowitsch   default:
10817f48955SJacob Faibussowitsch     /* in case the above macros expand to nothing this silences any unused variable warnings */
10917f48955SJacob Faibussowitsch     (void)(devid);
11098921bdaSJacob 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]);
111030f984aSJacob Faibussowitsch   }
112030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
113030f984aSJacob Faibussowitsch }
114030f984aSJacob Faibussowitsch 
115030f984aSJacob Faibussowitsch /*@C
116811af0c4SBarry Smith   PetscDeviceDestroy - Free a `PetscDevice`
117030f984aSJacob Faibussowitsch 
1180e6b6b59SJacob Faibussowitsch   Not Collective
119030f984aSJacob Faibussowitsch 
120030f984aSJacob Faibussowitsch   Input Parameter:
1210e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
122030f984aSJacob Faibussowitsch 
123030f984aSJacob Faibussowitsch   Level: beginner
124030f984aSJacob Faibussowitsch 
1250e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
1260e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
127030f984aSJacob Faibussowitsch @*/
128d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
129d71ae5a4SJacob Faibussowitsch {
130a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
1310e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 1);
132a4af0ceeSJacob Faibussowitsch   if (!*device) PetscFunctionReturn(0);
133a4af0ceeSJacob Faibussowitsch   PetscValidDevice(*device, 1);
1349566063dSJacob Faibussowitsch   PetscCall(PetscDeviceDereference_Internal(*device));
135a4af0ceeSJacob Faibussowitsch   if ((*device)->refcnt) {
1360e6b6b59SJacob Faibussowitsch     *device = nullptr;
137a4af0ceeSJacob Faibussowitsch     PetscFunctionReturn(0);
138030f984aSJacob Faibussowitsch   }
1399566063dSJacob Faibussowitsch   PetscCall(PetscFree((*device)->data));
1409566063dSJacob Faibussowitsch   PetscCall(PetscFree(*device));
141030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
142030f984aSJacob Faibussowitsch }
143030f984aSJacob Faibussowitsch 
144a4af0ceeSJacob Faibussowitsch /*@C
145811af0c4SBarry Smith   PetscDeviceConfigure - Configure a particular `PetscDevice`
146030f984aSJacob Faibussowitsch 
1470e6b6b59SJacob Faibussowitsch   Not Collective
148a4af0ceeSJacob Faibussowitsch 
149a4af0ceeSJacob Faibussowitsch   Input Parameter:
150811af0c4SBarry Smith . device - The `PetscDevice` to configure
151a4af0ceeSJacob Faibussowitsch 
1520e6b6b59SJacob Faibussowitsch   Notes:
1530e6b6b59SJacob Faibussowitsch   The user should not assume that this is a cheap operation.
154a4af0ceeSJacob Faibussowitsch 
155a4af0ceeSJacob Faibussowitsch   Level: beginner
156a4af0ceeSJacob Faibussowitsch 
1570e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
1580e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
159a4af0ceeSJacob Faibussowitsch @*/
160d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceConfigure(PetscDevice device)
161d71ae5a4SJacob Faibussowitsch {
162030f984aSJacob Faibussowitsch   PetscFunctionBegin;
163a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
164cf3a2253SJacob Faibussowitsch   /*
165cf3a2253SJacob Faibussowitsch     if no available configuration is available, this cascades all the way down to default
166cf3a2253SJacob Faibussowitsch     and error
167cf3a2253SJacob Faibussowitsch   */
1680e6b6b59SJacob Faibussowitsch   switch (const auto dtype = device->type) {
1690e6b6b59SJacob Faibussowitsch   case PETSC_DEVICE_HOST:
1700e6b6b59SJacob Faibussowitsch     if (PetscDefined(HAVE_HOST)) break; // always true
1719371c9d4SSatish Balay   case PETSC_DEVICE_CUDA:
1729371c9d4SSatish Balay     if (PetscDefined(HAVE_CUDA)) break;
1730e6b6b59SJacob Faibussowitsch     goto error;
1749371c9d4SSatish Balay   case PETSC_DEVICE_HIP:
1759371c9d4SSatish Balay     if (PetscDefined(HAVE_HIP)) break;
1760e6b6b59SJacob Faibussowitsch     goto error;
1779371c9d4SSatish Balay   case PETSC_DEVICE_SYCL:
1789371c9d4SSatish Balay     if (PetscDefined(HAVE_SYCL)) break;
179f4d061e9SPierre Jolivet     goto error;
1800e6b6b59SJacob Faibussowitsch   default:
1810e6b6b59SJacob Faibussowitsch   error:
1820e6b6b59SJacob Faibussowitsch     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
183a4af0ceeSJacob Faibussowitsch   }
184dbbe0bcdSBarry Smith   PetscUseTypeMethod(device, configure);
185a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
186a4af0ceeSJacob Faibussowitsch }
187a4af0ceeSJacob Faibussowitsch 
188a4af0ceeSJacob Faibussowitsch /*@C
189811af0c4SBarry Smith   PetscDeviceView - View a `PetscDevice`
190a4af0ceeSJacob Faibussowitsch 
1910e6b6b59SJacob Faibussowitsch   Collective on viewer
192a4af0ceeSJacob Faibussowitsch 
19391e63d38SStefano Zampini   Input Parameters:
194811af0c4SBarry Smith + device - The `PetscDevice` to view
1950e6b6b59SJacob Faibussowitsch - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
196a4af0ceeSJacob Faibussowitsch 
197a4af0ceeSJacob Faibussowitsch   Level: beginner
198a4af0ceeSJacob Faibussowitsch 
1990e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
2000e6b6b59SJacob Faibussowitsch `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
201a4af0ceeSJacob Faibussowitsch @*/
202d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer)
203d71ae5a4SJacob Faibussowitsch {
2040e6b6b59SJacob Faibussowitsch   auto      sub = viewer;
2050e6b6b59SJacob Faibussowitsch   PetscBool iascii;
2060e6b6b59SJacob Faibussowitsch 
207a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
208a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
2090e6b6b59SJacob Faibussowitsch   if (viewer) {
210a4af0ceeSJacob Faibussowitsch     PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2110e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
2120e6b6b59SJacob Faibussowitsch   } else {
2130e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
2140e6b6b59SJacob Faibussowitsch     iascii = PETSC_TRUE;
2150e6b6b59SJacob Faibussowitsch   }
2160e6b6b59SJacob Faibussowitsch 
2170e6b6b59SJacob Faibussowitsch   if (iascii) {
2180e6b6b59SJacob Faibussowitsch     auto        dtype = PETSC_DEVICE_HOST;
2190e6b6b59SJacob Faibussowitsch     MPI_Comm    comm;
2200e6b6b59SJacob Faibussowitsch     PetscMPIInt size;
2210e6b6b59SJacob Faibussowitsch     PetscInt    id = 0;
2220e6b6b59SJacob Faibussowitsch 
2230e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
2240e6b6b59SJacob Faibussowitsch     PetscCallMPI(MPI_Comm_size(comm, &size));
2250e6b6b59SJacob Faibussowitsch 
2260e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetDeviceId(device, &id));
2270e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetType(device, &dtype));
2280e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
2290e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
2300e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPushTab(sub));
2310e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
2320e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
2330e6b6b59SJacob Faibussowitsch   }
2340e6b6b59SJacob Faibussowitsch 
2350e6b6b59SJacob Faibussowitsch   // see if impls has extra viewer stuff
2360e6b6b59SJacob Faibussowitsch   PetscTryTypeMethod(device, view, sub);
2370e6b6b59SJacob Faibussowitsch 
2380e6b6b59SJacob Faibussowitsch   if (iascii) {
2390e6b6b59SJacob Faibussowitsch     // undo the ASCII specific stuff
2400e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPopTab(sub));
2410e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
2420e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerFlush(viewer));
2430e6b6b59SJacob Faibussowitsch   }
244a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
245a4af0ceeSJacob Faibussowitsch }
246a4af0ceeSJacob Faibussowitsch 
24791e63d38SStefano Zampini /*@C
2480e6b6b59SJacob Faibussowitsch   PetscDeviceGetType - Get the type of device
24991e63d38SStefano Zampini 
2500e6b6b59SJacob Faibussowitsch   Not Collective
25191e63d38SStefano Zampini 
25291e63d38SStefano Zampini   Input Parameter:
253811af0c4SBarry Smith . device - The `PetscDevice`
25491e63d38SStefano Zampini 
25591e63d38SStefano Zampini   Output Parameter:
2560e6b6b59SJacob Faibussowitsch . type - The `PetscDeviceType`
25791e63d38SStefano Zampini 
25891e63d38SStefano Zampini   Level: beginner
25991e63d38SStefano Zampini 
2600e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
2610e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
2620e6b6b59SJacob Faibussowitsch `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
2630e6b6b59SJacob Faibussowitsch @*/
264d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type)
265d71ae5a4SJacob Faibussowitsch {
2660e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
2670e6b6b59SJacob Faibussowitsch   PetscValidDevice(device, 1);
2680e6b6b59SJacob Faibussowitsch   PetscValidPointer(type, 2);
2690e6b6b59SJacob Faibussowitsch   *type = device->type;
2700e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
2710e6b6b59SJacob Faibussowitsch }
2720e6b6b59SJacob Faibussowitsch 
2730e6b6b59SJacob Faibussowitsch /*@C
2740e6b6b59SJacob Faibussowitsch   PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
2750e6b6b59SJacob Faibussowitsch 
2760e6b6b59SJacob Faibussowitsch   Not Collective
2770e6b6b59SJacob Faibussowitsch 
2780e6b6b59SJacob Faibussowitsch   Input Parameter:
2790e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
2800e6b6b59SJacob Faibussowitsch 
2810e6b6b59SJacob Faibussowitsch   Output Parameter:
2820e6b6b59SJacob Faibussowitsch . id - The id
2830e6b6b59SJacob Faibussowitsch 
2840e6b6b59SJacob Faibussowitsch   Notes:
2850e6b6b59SJacob Faibussowitsch   The returned ID may have been assigned by the underlying device backend. For example if the
2860e6b6b59SJacob Faibussowitsch   backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
2870e6b6b59SJacob Faibussowitsch   this device was configured.
2880e6b6b59SJacob Faibussowitsch 
2890e6b6b59SJacob Faibussowitsch   Level: beginner
2900e6b6b59SJacob Faibussowitsch 
2910e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
29291e63d38SStefano Zampini @*/
293d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
294d71ae5a4SJacob Faibussowitsch {
29591e63d38SStefano Zampini   PetscFunctionBegin;
29691e63d38SStefano Zampini   PetscValidDevice(device, 1);
29791e63d38SStefano Zampini   PetscValidIntPointer(id, 2);
29891e63d38SStefano Zampini   *id = device->deviceId;
29991e63d38SStefano Zampini   PetscFunctionReturn(0);
30091e63d38SStefano Zampini }
30191e63d38SStefano Zampini 
3020e6b6b59SJacob Faibussowitsch struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
3030e6b6b59SJacob Faibussowitsch   PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
3040e6b6b59SJacob Faibussowitsch 
305d71ae5a4SJacob Faibussowitsch   PETSC_NODISCARD PetscErrorCode finalize_() noexcept
306d71ae5a4SJacob Faibussowitsch   {
3070e6b6b59SJacob Faibussowitsch     PetscFunctionBegin;
3080e6b6b59SJacob Faibussowitsch     type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
3090e6b6b59SJacob Faibussowitsch     PetscFunctionReturn(0);
3100e6b6b59SJacob Faibussowitsch   }
3110e6b6b59SJacob Faibussowitsch };
3120e6b6b59SJacob Faibussowitsch 
3130e6b6b59SJacob Faibussowitsch static auto default_device_type = DefaultDeviceType();
3140e6b6b59SJacob Faibussowitsch 
3150e6b6b59SJacob Faibussowitsch /*@C
3160e6b6b59SJacob Faibussowitsch   PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
3170e6b6b59SJacob Faibussowitsch 
3180e6b6b59SJacob Faibussowitsch   Not Collective
3190e6b6b59SJacob Faibussowitsch 
3200e6b6b59SJacob Faibussowitsch   Notes:
3210e6b6b59SJacob Faibussowitsch   Unless selected by the user, the default device is selected in the following order\:
3220e6b6b59SJacob Faibussowitsch   `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
3230e6b6b59SJacob Faibussowitsch 
3240e6b6b59SJacob Faibussowitsch   Level: beginner
3250e6b6b59SJacob Faibussowitsch 
3260e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
3270e6b6b59SJacob Faibussowitsch @*/
328d71ae5a4SJacob Faibussowitsch PetscDeviceType PETSC_DEVICE_DEFAULT(void)
329d71ae5a4SJacob Faibussowitsch {
3300e6b6b59SJacob Faibussowitsch   return default_device_type.type;
3310e6b6b59SJacob Faibussowitsch }
3320e6b6b59SJacob Faibussowitsch 
3330e6b6b59SJacob Faibussowitsch /*@C
3340e6b6b59SJacob Faibussowitsch   PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
3350e6b6b59SJacob Faibussowitsch 
3360e6b6b59SJacob Faibussowitsch   Not Collective
3370e6b6b59SJacob Faibussowitsch 
3380e6b6b59SJacob Faibussowitsch   Input Parameter:
3390e6b6b59SJacob Faibussowitsch . type - the new default device type
3400e6b6b59SJacob Faibussowitsch 
3410e6b6b59SJacob Faibussowitsch   Notes:
3420e6b6b59SJacob Faibussowitsch   This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
3430e6b6b59SJacob Faibussowitsch 
3440e6b6b59SJacob Faibussowitsch   Level: beginner
3450e6b6b59SJacob Faibussowitsch 
3460e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
3470e6b6b59SJacob Faibussowitsch @*/
348d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
349d71ae5a4SJacob Faibussowitsch {
3500e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
3510e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3520e6b6b59SJacob Faibussowitsch   if (default_device_type.type != type) {
3530e6b6b59SJacob Faibussowitsch     // no need to waster a PetscRegisterFinalize() slot if we don't change it
3540e6b6b59SJacob Faibussowitsch     default_device_type.type = type;
3550e6b6b59SJacob Faibussowitsch     PetscCall(default_device_type.register_finalize());
3560e6b6b59SJacob Faibussowitsch   }
3570e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
3580e6b6b59SJacob Faibussowitsch }
3590e6b6b59SJacob Faibussowitsch 
3600e6b6b59SJacob Faibussowitsch static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
3610e6b6b59SJacob Faibussowitsch 
3620e6b6b59SJacob Faibussowitsch /*
3630e6b6b59SJacob Faibussowitsch   Actual intialization function; any functions claiming to initialize PetscDevice or
3640e6b6b59SJacob Faibussowitsch   PetscDeviceContext will have to run through this one
3650e6b6b59SJacob Faibussowitsch */
366d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId)
367d71ae5a4SJacob Faibussowitsch {
3680e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
3690e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3700e6b6b59SJacob Faibussowitsch   if (PetscUnlikely(!PetscDeviceInitialized(type))) {
3710e6b6b59SJacob Faibussowitsch     auto &dev  = defaultDevices[type].first;
3720e6b6b59SJacob Faibussowitsch     auto &init = defaultDevices[type].second;
3730e6b6b59SJacob Faibussowitsch 
3740e6b6b59SJacob Faibussowitsch     PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
3750e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
3760e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceConfigure(dev));
3770e6b6b59SJacob Faibussowitsch     init = true;
3780e6b6b59SJacob Faibussowitsch   }
3790e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
3800e6b6b59SJacob Faibussowitsch }
381a4af0ceeSJacob Faibussowitsch 
382a4af0ceeSJacob Faibussowitsch /*@C
383811af0c4SBarry Smith   PetscDeviceInitialize - Initialize `PetscDevice`
384a4af0ceeSJacob Faibussowitsch 
3850e6b6b59SJacob Faibussowitsch   Not Collective
386a4af0ceeSJacob Faibussowitsch 
387a4af0ceeSJacob Faibussowitsch   Input Parameter:
388811af0c4SBarry Smith . type - The `PetscDeviceType` to initialize
389a4af0ceeSJacob Faibussowitsch 
3900e6b6b59SJacob Faibussowitsch   Notes:
3910e6b6b59SJacob Faibussowitsch   Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
3920e6b6b59SJacob Faibussowitsch   result in device synchronization.
393a4af0ceeSJacob Faibussowitsch 
394a4af0ceeSJacob Faibussowitsch   Level: beginner
395a4af0ceeSJacob Faibussowitsch 
3960e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
3970e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
398a4af0ceeSJacob Faibussowitsch @*/
399d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
400d71ae5a4SJacob Faibussowitsch {
401a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
402a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
4039566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
404a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
405a4af0ceeSJacob Faibussowitsch }
406a4af0ceeSJacob Faibussowitsch 
407a4af0ceeSJacob Faibussowitsch /*@C
408811af0c4SBarry Smith   PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
409811af0c4SBarry Smith   `PetscDeviceType`
410a4af0ceeSJacob Faibussowitsch 
4110e6b6b59SJacob Faibussowitsch   Not Collective
412a4af0ceeSJacob Faibussowitsch 
413a4af0ceeSJacob Faibussowitsch   Input Parameter:
414811af0c4SBarry Smith . type - The `PetscDeviceType` to check
415a4af0ceeSJacob Faibussowitsch 
4160e6b6b59SJacob Faibussowitsch   Notes:
4170e6b6b59SJacob Faibussowitsch   Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
418a4af0ceeSJacob Faibussowitsch 
419811af0c4SBarry Smith   If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
420811af0c4SBarry Smith   return `PETSC_FALSE` for that `PetscDeviceType`.
421a4af0ceeSJacob Faibussowitsch 
422a4af0ceeSJacob Faibussowitsch   Level: beginner
423a4af0ceeSJacob Faibussowitsch 
4240e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
4250e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
426a4af0ceeSJacob Faibussowitsch @*/
427d71ae5a4SJacob Faibussowitsch PetscBool PetscDeviceInitialized(PetscDeviceType type)
428d71ae5a4SJacob Faibussowitsch {
4290e6b6b59SJacob Faibussowitsch   return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
4300e6b6b59SJacob Faibussowitsch }
4310e6b6b59SJacob Faibussowitsch 
4320e6b6b59SJacob Faibussowitsch /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
433d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device)
434d71ae5a4SJacob Faibussowitsch {
4350e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
4360e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 2);
4370e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitialize(type));
4380e6b6b59SJacob Faibussowitsch   *device = defaultDevices[type].first;
4390e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
440a4af0ceeSJacob Faibussowitsch }
441a4af0ceeSJacob Faibussowitsch 
442a16fd2c9SJacob Faibussowitsch /*@C
443a16fd2c9SJacob Faibussowitsch   PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
444a16fd2c9SJacob Faibussowitsch 
4450e6b6b59SJacob Faibussowitsch   Not Collective
446a16fd2c9SJacob Faibussowitsch 
447a16fd2c9SJacob Faibussowitsch   Input Parameters:
448a16fd2c9SJacob Faibussowitsch + device - The `PetscDevice`
449a16fd2c9SJacob Faibussowitsch - attr   - The attribute
450a16fd2c9SJacob Faibussowitsch 
451a16fd2c9SJacob Faibussowitsch   Output Parameter:
452a16fd2c9SJacob Faibussowitsch . value - The value of the attribute
453a16fd2c9SJacob Faibussowitsch 
454a16fd2c9SJacob Faibussowitsch   Notes:
455a16fd2c9SJacob Faibussowitsch   Since different attributes are often different types `value` is a `void *` to accommodate
456a16fd2c9SJacob Faibussowitsch   them all. The underlying type of the attribute is therefore included in the name of the
457a16fd2c9SJacob Faibussowitsch   `PetscDeviceAttribute` reponsible for querying it. For example,
458a16fd2c9SJacob Faibussowitsch   `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
459a16fd2c9SJacob Faibussowitsch 
4600e6b6b59SJacob Faibussowitsch   Level: intermediate
4610e6b6b59SJacob Faibussowitsch 
462a16fd2c9SJacob Faibussowitsch .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
463a16fd2c9SJacob Faibussowitsch @*/
464d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
465d71ae5a4SJacob Faibussowitsch {
466a16fd2c9SJacob Faibussowitsch   PetscFunctionBegin;
467a16fd2c9SJacob Faibussowitsch   PetscValidDevice(device, 1);
468a16fd2c9SJacob Faibussowitsch   PetscValidDeviceAttribute(attr, 2);
469a16fd2c9SJacob Faibussowitsch   PetscValidPointer(value, 3);
470a16fd2c9SJacob Faibussowitsch   PetscUseTypeMethod(device, getattribute, attr, value);
471a16fd2c9SJacob Faibussowitsch   PetscFunctionReturn(0);
472a16fd2c9SJacob Faibussowitsch }
473a16fd2c9SJacob Faibussowitsch 
474d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType)
475d71ae5a4SJacob Faibussowitsch {
476a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
477a4af0ceeSJacob Faibussowitsch   if (!PetscDeviceConfiguredFor_Internal(type)) {
4780e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
4790e6b6b59SJacob Faibussowitsch     defaultDevices[type].first = nullptr;
480a4af0ceeSJacob Faibussowitsch     PetscFunctionReturn(0);
481a4af0ceeSJacob Faibussowitsch   }
4820e6b6b59SJacob Faibussowitsch   PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
483a4af0ceeSJacob Faibussowitsch   /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
484a4af0ceeSJacob Faibussowitsch   switch (type) {
4850e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4860e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4870e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4880e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
489d71ae5a4SJacob Faibussowitsch   default:
490d71ae5a4SJacob Faibussowitsch     SETERRQ(comm, PETSC_ERR_PLIB, "PETSc was seemingly configured for PetscDeviceType %s but we've fallen through all cases in a switch", PetscDeviceTypes[type]);
491a4af0ceeSJacob Faibussowitsch   }
4920e6b6b59SJacob 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)]));
493cf3a2253SJacob Faibussowitsch   /*
4940e6b6b59SJacob Faibussowitsch     defaultInitType, defaultView  and defaultDeviceId now represent what the individual TYPES
4950e6b6b59SJacob Faibussowitsch     have decided to initialize as
496cf3a2253SJacob Faibussowitsch   */
4970e6b6b59SJacob Faibussowitsch   if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
4980e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
4999566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
5000e6b6b59SJacob Faibussowitsch     if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
5010e6b6b59SJacob Faibussowitsch   }
5020e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
5030e6b6b59SJacob Faibussowitsch }
504a4af0ceeSJacob Faibussowitsch 
505d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView)
506d71ae5a4SJacob Faibussowitsch {
5070e6b6b59SJacob Faibussowitsch   PetscInt initIdx       = PETSC_DEVICE_INIT_LAZY;
5080e6b6b59SJacob Faibussowitsch   auto     initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
5090e6b6b59SJacob Faibussowitsch   auto     flg           = PETSC_FALSE;
5100e6b6b59SJacob Faibussowitsch 
5110e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
5120e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
5130e6b6b59SJacob Faibussowitsch   if (flg) PetscCall(PetscLogGpuTime());
5140e6b6b59SJacob Faibussowitsch 
5150e6b6b59SJacob Faibussowitsch   PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
5160e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
5170e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
5180e6b6b59SJacob 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));
5190e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
5200e6b6b59SJacob Faibussowitsch   PetscOptionsEnd();
5210e6b6b59SJacob Faibussowitsch 
5220e6b6b59SJacob Faibussowitsch   if (initIdx == PETSC_DEVICE_INIT_NONE) {
5230e6b6b59SJacob Faibussowitsch     /* disabled all device initialization if devices are globally disabled */
5240e6b6b59SJacob 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");
5250e6b6b59SJacob Faibussowitsch     *defaultView  = PETSC_FALSE;
5260e6b6b59SJacob Faibussowitsch     initDeviceIdx = PETSC_DEVICE_HOST;
5270e6b6b59SJacob Faibussowitsch   } else {
5280e6b6b59SJacob Faibussowitsch     *defaultView = static_cast<PetscBool>(*defaultView && flg);
5290e6b6b59SJacob Faibussowitsch     if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
530a4af0ceeSJacob Faibussowitsch   }
5310e6b6b59SJacob Faibussowitsch   *defaultInitType         = PetscDeviceInitTypeCast(initIdx);
5320e6b6b59SJacob Faibussowitsch   *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
533030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
534030f984aSJacob Faibussowitsch }
535030f984aSJacob Faibussowitsch 
536030f984aSJacob Faibussowitsch /* called from PetscFinalize() do not call yourself! */
537d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceFinalize_Private()
538d71ae5a4SJacob Faibussowitsch {
539030f984aSJacob Faibussowitsch   PetscFunctionBegin;
540a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
541bf025ffbSJacob Faibussowitsch     const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] {
542a4af0ceeSJacob Faibussowitsch       PetscFunctionBegin;
5430e6b6b59SJacob Faibussowitsch       for (auto &&device : defaultDevices) {
5440e6b6b59SJacob Faibussowitsch         const auto dev = device.first;
5450e6b6b59SJacob Faibussowitsch 
5460e6b6b59SJacob 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);
5470e6b6b59SJacob Faibussowitsch       }
548a4af0ceeSJacob Faibussowitsch       PetscFunctionReturn(0);
549a4af0ceeSJacob Faibussowitsch     };
550bf025ffbSJacob Faibussowitsch     /*
551bf025ffbSJacob Faibussowitsch       you might be thinking, why on earth are you registered yet another finalizer in a
552bf025ffbSJacob Faibussowitsch       function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
553bf025ffbSJacob Faibussowitsch       because it is.
554bf025ffbSJacob Faibussowitsch 
555bf025ffbSJacob Faibussowitsch       The crux of the problem is that the initializer (and therefore the ~finalizer~) of
556bf025ffbSJacob Faibussowitsch       PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
557bf025ffbSJacob Faibussowitsch       a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
558bf025ffbSJacob Faibussowitsch       won't be destroyed yet. So we need to repeat the check that all devices have been
559bf025ffbSJacob Faibussowitsch       destroyed again ~after~ the global context is destroyed. In summary:
560bf025ffbSJacob Faibussowitsch 
561bf025ffbSJacob Faibussowitsch       1. This finalizer runs and destroys all devices, except it may not because the global
562bf025ffbSJacob Faibussowitsch          context may still hold a reference!
563bf025ffbSJacob Faibussowitsch       2. The global context finalizer runs and does the final reference count decrement
564bf025ffbSJacob Faibussowitsch          required, which actually destroys the held device.
565bf025ffbSJacob Faibussowitsch       3. Our newly added finalizer runs and checks that all is well.
566a4af0ceeSJacob Faibussowitsch     */
5670e6b6b59SJacob Faibussowitsch     PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize)));
568a4af0ceeSJacob Faibussowitsch   }
5690e6b6b59SJacob Faibussowitsch   for (auto &&device : defaultDevices) {
5700e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceDestroy(&device.first));
5710e6b6b59SJacob Faibussowitsch     device.second = false;
5720e6b6b59SJacob Faibussowitsch   }
573030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
574030f984aSJacob Faibussowitsch }
575030f984aSJacob Faibussowitsch 
576cf3a2253SJacob Faibussowitsch /*
577cf3a2253SJacob Faibussowitsch   Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
578cf3a2253SJacob Faibussowitsch   initialization types:
579cf3a2253SJacob Faibussowitsch 
580a4af0ceeSJacob Faibussowitsch   1. defaultInitType - how does PetscDevice as a whole expect to initialize?
581a4af0ceeSJacob Faibussowitsch   2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
582a4af0ceeSJacob Faibussowitsch      e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
583a4af0ceeSJacob Faibussowitsch      have all CUDA devices still initialize.
584a4af0ceeSJacob Faibussowitsch 
585a4af0ceeSJacob Faibussowitsch   All told the following happens:
586cf3a2253SJacob Faibussowitsch 
587a4af0ceeSJacob Faibussowitsch   0. defaultInitType -> LAZY
588a4af0ceeSJacob Faibussowitsch   1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
589a4af0ceeSJacob Faibussowitsch   2. PetscDevice initializes each sub type with deviceDefaultInitType.
590a4af0ceeSJacob Faibussowitsch   2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
591a4af0ceeSJacob Faibussowitsch       to checking for specific device init. if view or specific device init
592a4af0ceeSJacob Faibussowitsch       subTypeDefaultInitType -> EAGER. disabled once again overrides all.
593a4af0ceeSJacob Faibussowitsch */
5940e6b6b59SJacob Faibussowitsch 
595d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm)
596d71ae5a4SJacob Faibussowitsch {
5977a101e5eSJacob Faibussowitsch   auto defaultView                    = PETSC_FALSE;
5987a101e5eSJacob Faibussowitsch   auto initializeDeviceContextEagerly = PETSC_FALSE;
5990e6b6b59SJacob Faibussowitsch   auto defaultDeviceSet               = PETSC_FALSE;
6007a101e5eSJacob Faibussowitsch   auto defaultDevice                  = PetscInt{PETSC_DECIDE};
6010e6b6b59SJacob Faibussowitsch   auto deviceContextInitDevice        = PETSC_DEVICE_DEFAULT();
6020e6b6b59SJacob Faibussowitsch   auto defaultInitType                = PETSC_DEVICE_INIT_LAZY;
603a4af0ceeSJacob Faibussowitsch 
604a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
605a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
606a4af0ceeSJacob Faibussowitsch     int result;
607a4af0ceeSJacob Faibussowitsch 
6089566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
609a4af0ceeSJacob Faibussowitsch     /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
610a4af0ceeSJacob Faibussowitsch      * global space */
611a4af0ceeSJacob Faibussowitsch     if (PetscUnlikely(result != MPI_IDENT)) {
612a4af0ceeSJacob Faibussowitsch       char name[MPI_MAX_OBJECT_NAME] = {};
613a4af0ceeSJacob Faibussowitsch       int  len; /* unused */
614a4af0ceeSJacob Faibussowitsch 
6159566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
61698921bdaSJacob Faibussowitsch       SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
617a4af0ceeSJacob Faibussowitsch     }
618a4af0ceeSJacob Faibussowitsch   }
619a4af0ceeSJacob Faibussowitsch   comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
6209566063dSJacob Faibussowitsch   PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
621a4af0ceeSJacob Faibussowitsch 
6220e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
6237a101e5eSJacob Faibussowitsch 
6240e6b6b59SJacob Faibussowitsch   // the precise values don't matter here, so long as they are sequential
6250e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, "");
6260e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, "");
6270e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, "");
6280e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, "");
6290e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_MAX) == 4, "");
6300e6b6b59SJacob Faibussowitsch   for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
6310e6b6b59SJacob Faibussowitsch     const auto deviceType = PetscDeviceTypeCast(i);
632a4af0ceeSJacob Faibussowitsch     auto       initType   = defaultInitType;
633a4af0ceeSJacob Faibussowitsch 
6349566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
6350e6b6b59SJacob Faibussowitsch     if (PetscDeviceConfiguredFor_Internal(deviceType)) {
6360e6b6b59SJacob Faibussowitsch       if (initType == PETSC_DEVICE_INIT_EAGER) {
637a4af0ceeSJacob Faibussowitsch         initializeDeviceContextEagerly = PETSC_TRUE;
6380e6b6b59SJacob Faibussowitsch         // only update the default device if the user hasn't set it previously
6390e6b6b59SJacob Faibussowitsch         if (!defaultDeviceSet) {
640a4af0ceeSJacob Faibussowitsch           deviceContextInitDevice = deviceType;
6410e6b6b59SJacob Faibussowitsch           PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
6420e6b6b59SJacob Faibussowitsch         }
6430e6b6b59SJacob Faibussowitsch       } else if (initType == PETSC_DEVICE_INIT_NONE) {
644*1015a2a4SJacob Faibussowitsch         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]);
645a4af0ceeSJacob Faibussowitsch       }
646a4af0ceeSJacob Faibussowitsch     }
6470e6b6b59SJacob Faibussowitsch   }
6480e6b6b59SJacob Faibussowitsch 
6490e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
6500e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
6510e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
6520e6b6b59SJacob Faibussowitsch   /*                       PetscDevice is now fully initialized                          */
6530e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
6540e6b6b59SJacob Faibussowitsch   {
6550e6b6b59SJacob Faibussowitsch     /*
6560e6b6b59SJacob Faibussowitsch       query the options db to get the root settings from the user (if any).
6570e6b6b59SJacob Faibussowitsch 
6580e6b6b59SJacob Faibussowitsch       This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
6590e6b6b59SJacob Faibussowitsch       PetscDeviceContextSetFromOptions() before we even have one, then set a few static
6600e6b6b59SJacob Faibussowitsch       variables in that file with the results.
6610e6b6b59SJacob Faibussowitsch     */
6620e6b6b59SJacob Faibussowitsch     auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
6630e6b6b59SJacob Faibussowitsch     auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
6640e6b6b59SJacob Faibussowitsch 
6650e6b6b59SJacob Faibussowitsch     PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
6660e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
6670e6b6b59SJacob Faibussowitsch     PetscOptionsEnd();
6680e6b6b59SJacob Faibussowitsch 
6690e6b6b59SJacob Faibussowitsch     if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
6700e6b6b59SJacob Faibussowitsch     if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
6710e6b6b59SJacob Faibussowitsch   }
6720e6b6b59SJacob Faibussowitsch 
673a4af0ceeSJacob Faibussowitsch   if (initializeDeviceContextEagerly) {
674a4af0ceeSJacob Faibussowitsch     PetscDeviceContext dctx;
675a4af0ceeSJacob Faibussowitsch 
6760e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
6770e6b6b59SJacob Faibussowitsch     /* instantiates the device context */
6789566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
6799566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextSetUp(dctx));
680a4af0ceeSJacob Faibussowitsch   }
681a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
682a4af0ceeSJacob Faibussowitsch }
683