xref: /petsc/src/sys/objects/device/interface/device.cxx (revision 0e6b6b5985dd9b1172860d21fb88bd3966bf7c54)
1*0e6b6b59SJacob Faibussowitsch #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/
217f48955SJacob Faibussowitsch #include <petsc/private/petscadvancedmacros.h>
3030f984aSJacob Faibussowitsch 
4*0e6b6b59SJacob Faibussowitsch #include "../impls/host/hostdevice.hpp"
5*0e6b6b59SJacob Faibussowitsch #include "../impls/cupm/cupmdevice.hpp"
6*0e6b6b59SJacob Faibussowitsch #include "../impls/sycl/sycldevice.hpp"
7*0e6b6b59SJacob Faibussowitsch 
8*0e6b6b59SJacob Faibussowitsch #include <limits>  // std::numeric_limits
9*0e6b6b59SJacob Faibussowitsch #include <utility> // std::make_pair
10*0e6b6b59SJacob Faibussowitsch 
11*0e6b6b59SJacob 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 */
17*0e6b6b59SJacob Faibussowitsch static host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
18030f984aSJacob Faibussowitsch #if PetscDefined(HAVE_CUDA)
19*0e6b6b59SJacob Faibussowitsch static cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
20030f984aSJacob Faibussowitsch #endif
21030f984aSJacob Faibussowitsch #if PetscDefined(HAVE_HIP)
22*0e6b6b59SJacob Faibussowitsch static cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
23030f984aSJacob Faibussowitsch #endif
24a2158755SJunchao Zhang #if PetscDefined(HAVE_SYCL)
25*0e6b6b59SJacob 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 
65*0e6b6b59SJacob 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:
75*0e6b6b59SJacob Faibussowitsch   This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
76*0e6b6b59SJacob 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 
82*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`,
83*0e6b6b59SJacob Faibussowitsch `PetscDeviceInitialize()`,`PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
84*0e6b6b59SJacob 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());
93*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscNew(device));
94*0e6b6b59SJacob Faibussowitsch   (*device)->id     = PetscDeviceCounter++;
95*0e6b6b59SJacob Faibussowitsch   (*device)->type   = type;
96*0e6b6b59SJacob 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) {
102*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
103*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
104*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
105*0e6b6b59SJacob 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 
117*0e6b6b59SJacob Faibussowitsch   Not Collective
118030f984aSJacob Faibussowitsch 
119030f984aSJacob Faibussowitsch   Input Parameter:
120*0e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
121030f984aSJacob Faibussowitsch 
122030f984aSJacob Faibussowitsch   Level: beginner
123030f984aSJacob Faibussowitsch 
124*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
125*0e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
126030f984aSJacob Faibussowitsch @*/
1279371c9d4SSatish Balay PetscErrorCode PetscDeviceDestroy(PetscDevice *device) {
128a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
129*0e6b6b59SJacob 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) {
134*0e6b6b59SJacob 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 
145*0e6b6b59SJacob Faibussowitsch   Not Collective
146a4af0ceeSJacob Faibussowitsch 
147a4af0ceeSJacob Faibussowitsch   Input Parameter:
148811af0c4SBarry Smith . device - The `PetscDevice` to configure
149a4af0ceeSJacob Faibussowitsch 
150*0e6b6b59SJacob Faibussowitsch   Notes:
151*0e6b6b59SJacob Faibussowitsch   The user should not assume that this is a cheap operation.
152a4af0ceeSJacob Faibussowitsch 
153a4af0ceeSJacob Faibussowitsch   Level: beginner
154a4af0ceeSJacob Faibussowitsch 
155*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
156*0e6b6b59SJacob 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   */
165*0e6b6b59SJacob Faibussowitsch   switch (const auto dtype = device->type) {
166*0e6b6b59SJacob Faibussowitsch   case PETSC_DEVICE_HOST:
167*0e6b6b59SJacob Faibussowitsch     if (PetscDefined(HAVE_HOST)) break; // always true
1689371c9d4SSatish Balay   case PETSC_DEVICE_CUDA:
1699371c9d4SSatish Balay     if (PetscDefined(HAVE_CUDA)) break;
170*0e6b6b59SJacob Faibussowitsch     goto error;
1719371c9d4SSatish Balay   case PETSC_DEVICE_HIP:
1729371c9d4SSatish Balay     if (PetscDefined(HAVE_HIP)) break;
173*0e6b6b59SJacob Faibussowitsch     goto error;
1749371c9d4SSatish Balay   case PETSC_DEVICE_SYCL:
1759371c9d4SSatish Balay     if (PetscDefined(HAVE_SYCL)) break;
176*0e6b6b59SJacob Faibussowitsch   default:
177*0e6b6b59SJacob Faibussowitsch   error:
178*0e6b6b59SJacob Faibussowitsch     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
179a4af0ceeSJacob Faibussowitsch   }
180dbbe0bcdSBarry Smith   PetscUseTypeMethod(device, configure);
181a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
182a4af0ceeSJacob Faibussowitsch }
183a4af0ceeSJacob Faibussowitsch 
184a4af0ceeSJacob Faibussowitsch /*@C
185811af0c4SBarry Smith   PetscDeviceView - View a `PetscDevice`
186a4af0ceeSJacob Faibussowitsch 
187*0e6b6b59SJacob Faibussowitsch   Collective on viewer
188a4af0ceeSJacob Faibussowitsch 
18991e63d38SStefano Zampini   Input Parameters:
190811af0c4SBarry Smith + device - The `PetscDevice` to view
191*0e6b6b59SJacob Faibussowitsch - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
192a4af0ceeSJacob Faibussowitsch 
193a4af0ceeSJacob Faibussowitsch   Level: beginner
194a4af0ceeSJacob Faibussowitsch 
195*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
196*0e6b6b59SJacob Faibussowitsch `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
197a4af0ceeSJacob Faibussowitsch @*/
1989371c9d4SSatish Balay PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer) {
199*0e6b6b59SJacob Faibussowitsch   auto      sub = viewer;
200*0e6b6b59SJacob Faibussowitsch   PetscBool iascii;
201*0e6b6b59SJacob Faibussowitsch 
202a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
203a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
204*0e6b6b59SJacob Faibussowitsch   if (viewer) {
205a4af0ceeSJacob Faibussowitsch     PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
206*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
207*0e6b6b59SJacob Faibussowitsch   } else {
208*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
209*0e6b6b59SJacob Faibussowitsch     iascii = PETSC_TRUE;
210*0e6b6b59SJacob Faibussowitsch   }
211*0e6b6b59SJacob Faibussowitsch 
212*0e6b6b59SJacob Faibussowitsch   if (iascii) {
213*0e6b6b59SJacob Faibussowitsch     auto        dtype = PETSC_DEVICE_HOST;
214*0e6b6b59SJacob Faibussowitsch     MPI_Comm    comm;
215*0e6b6b59SJacob Faibussowitsch     PetscMPIInt size;
216*0e6b6b59SJacob Faibussowitsch     PetscInt    id = 0;
217*0e6b6b59SJacob Faibussowitsch 
218*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
219*0e6b6b59SJacob Faibussowitsch     PetscCallMPI(MPI_Comm_size(comm, &size));
220*0e6b6b59SJacob Faibussowitsch 
221*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetDeviceId(device, &id));
222*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetType(device, &dtype));
223*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
224*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
225*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPushTab(sub));
226*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
227*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
228*0e6b6b59SJacob Faibussowitsch   }
229*0e6b6b59SJacob Faibussowitsch 
230*0e6b6b59SJacob Faibussowitsch   // see if impls has extra viewer stuff
231*0e6b6b59SJacob Faibussowitsch   PetscTryTypeMethod(device, view, sub);
232*0e6b6b59SJacob Faibussowitsch 
233*0e6b6b59SJacob Faibussowitsch   if (iascii) {
234*0e6b6b59SJacob Faibussowitsch     // undo the ASCII specific stuff
235*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPopTab(sub));
236*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
237*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerFlush(viewer));
238*0e6b6b59SJacob Faibussowitsch   }
239a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
240a4af0ceeSJacob Faibussowitsch }
241a4af0ceeSJacob Faibussowitsch 
24291e63d38SStefano Zampini /*@C
243*0e6b6b59SJacob Faibussowitsch   PetscDeviceGetType - Get the type of device
24491e63d38SStefano Zampini 
245*0e6b6b59SJacob Faibussowitsch   Not Collective
24691e63d38SStefano Zampini 
24791e63d38SStefano Zampini   Input Parameter:
248811af0c4SBarry Smith . device - The `PetscDevice`
24991e63d38SStefano Zampini 
25091e63d38SStefano Zampini   Output Parameter:
251*0e6b6b59SJacob Faibussowitsch . type - The `PetscDeviceType`
25291e63d38SStefano Zampini 
25391e63d38SStefano Zampini   Level: beginner
25491e63d38SStefano Zampini 
255*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
256*0e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
257*0e6b6b59SJacob Faibussowitsch `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
258*0e6b6b59SJacob Faibussowitsch @*/
259*0e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type) {
260*0e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
261*0e6b6b59SJacob Faibussowitsch   PetscValidDevice(device, 1);
262*0e6b6b59SJacob Faibussowitsch   PetscValidPointer(type, 2);
263*0e6b6b59SJacob Faibussowitsch   *type = device->type;
264*0e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
265*0e6b6b59SJacob Faibussowitsch }
266*0e6b6b59SJacob Faibussowitsch 
267*0e6b6b59SJacob Faibussowitsch /*@C
268*0e6b6b59SJacob Faibussowitsch   PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
269*0e6b6b59SJacob Faibussowitsch 
270*0e6b6b59SJacob Faibussowitsch   Not Collective
271*0e6b6b59SJacob Faibussowitsch 
272*0e6b6b59SJacob Faibussowitsch   Input Parameter:
273*0e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
274*0e6b6b59SJacob Faibussowitsch 
275*0e6b6b59SJacob Faibussowitsch   Output Parameter:
276*0e6b6b59SJacob Faibussowitsch . id - The id
277*0e6b6b59SJacob Faibussowitsch 
278*0e6b6b59SJacob Faibussowitsch   Notes:
279*0e6b6b59SJacob Faibussowitsch   The returned ID may have been assigned by the underlying device backend. For example if the
280*0e6b6b59SJacob Faibussowitsch   backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
281*0e6b6b59SJacob Faibussowitsch   this device was configured.
282*0e6b6b59SJacob Faibussowitsch 
283*0e6b6b59SJacob Faibussowitsch   Level: beginner
284*0e6b6b59SJacob Faibussowitsch 
285*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
28691e63d38SStefano Zampini @*/
2879371c9d4SSatish Balay PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id) {
28891e63d38SStefano Zampini   PetscFunctionBegin;
28991e63d38SStefano Zampini   PetscValidDevice(device, 1);
29091e63d38SStefano Zampini   PetscValidIntPointer(id, 2);
29191e63d38SStefano Zampini   *id = device->deviceId;
29291e63d38SStefano Zampini   PetscFunctionReturn(0);
29391e63d38SStefano Zampini }
29491e63d38SStefano Zampini 
295*0e6b6b59SJacob Faibussowitsch struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
296*0e6b6b59SJacob Faibussowitsch   PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
297*0e6b6b59SJacob Faibussowitsch 
298*0e6b6b59SJacob Faibussowitsch   PETSC_NODISCARD PetscErrorCode finalize_() noexcept {
299*0e6b6b59SJacob Faibussowitsch     PetscFunctionBegin;
300*0e6b6b59SJacob Faibussowitsch     type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
301*0e6b6b59SJacob Faibussowitsch     PetscFunctionReturn(0);
302*0e6b6b59SJacob Faibussowitsch   }
303*0e6b6b59SJacob Faibussowitsch };
304*0e6b6b59SJacob Faibussowitsch 
305*0e6b6b59SJacob Faibussowitsch static auto default_device_type = DefaultDeviceType();
306*0e6b6b59SJacob Faibussowitsch 
307*0e6b6b59SJacob Faibussowitsch /*@C
308*0e6b6b59SJacob Faibussowitsch   PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
309*0e6b6b59SJacob Faibussowitsch 
310*0e6b6b59SJacob Faibussowitsch   Not Collective
311*0e6b6b59SJacob Faibussowitsch 
312*0e6b6b59SJacob Faibussowitsch   Notes:
313*0e6b6b59SJacob Faibussowitsch   Unless selected by the user, the default device is selected in the following order\:
314*0e6b6b59SJacob Faibussowitsch   `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
315*0e6b6b59SJacob Faibussowitsch 
316*0e6b6b59SJacob Faibussowitsch   Level: beginner
317*0e6b6b59SJacob Faibussowitsch 
318*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
319*0e6b6b59SJacob Faibussowitsch @*/
320*0e6b6b59SJacob Faibussowitsch PetscDeviceType PETSC_DEVICE_DEFAULT(void) {
321*0e6b6b59SJacob Faibussowitsch   return default_device_type.type;
322*0e6b6b59SJacob Faibussowitsch }
323*0e6b6b59SJacob Faibussowitsch 
324*0e6b6b59SJacob Faibussowitsch /*@C
325*0e6b6b59SJacob Faibussowitsch   PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
326*0e6b6b59SJacob Faibussowitsch 
327*0e6b6b59SJacob Faibussowitsch   Not Collective
328*0e6b6b59SJacob Faibussowitsch 
329*0e6b6b59SJacob Faibussowitsch   Input Parameter:
330*0e6b6b59SJacob Faibussowitsch . type - the new default device type
331*0e6b6b59SJacob Faibussowitsch 
332*0e6b6b59SJacob Faibussowitsch   Notes:
333*0e6b6b59SJacob Faibussowitsch   This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
334*0e6b6b59SJacob Faibussowitsch 
335*0e6b6b59SJacob Faibussowitsch   Level: beginner
336*0e6b6b59SJacob Faibussowitsch 
337*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
338*0e6b6b59SJacob Faibussowitsch @*/
339*0e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type) {
340*0e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
341*0e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
342*0e6b6b59SJacob Faibussowitsch   if (default_device_type.type != type) {
343*0e6b6b59SJacob Faibussowitsch     // no need to waster a PetscRegisterFinalize() slot if we don't change it
344*0e6b6b59SJacob Faibussowitsch     default_device_type.type = type;
345*0e6b6b59SJacob Faibussowitsch     PetscCall(default_device_type.register_finalize());
346*0e6b6b59SJacob Faibussowitsch   }
347*0e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
348*0e6b6b59SJacob Faibussowitsch }
349*0e6b6b59SJacob Faibussowitsch 
350*0e6b6b59SJacob Faibussowitsch static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
351*0e6b6b59SJacob Faibussowitsch 
352*0e6b6b59SJacob Faibussowitsch /*
353*0e6b6b59SJacob Faibussowitsch   Actual intialization function; any functions claiming to initialize PetscDevice or
354*0e6b6b59SJacob Faibussowitsch   PetscDeviceContext will have to run through this one
355*0e6b6b59SJacob Faibussowitsch */
356*0e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId) {
357*0e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
358*0e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
359*0e6b6b59SJacob Faibussowitsch   if (PetscUnlikely(!PetscDeviceInitialized(type))) {
360*0e6b6b59SJacob Faibussowitsch     auto &dev  = defaultDevices[type].first;
361*0e6b6b59SJacob Faibussowitsch     auto &init = defaultDevices[type].second;
362*0e6b6b59SJacob Faibussowitsch 
363*0e6b6b59SJacob Faibussowitsch     PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
364*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
365*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceConfigure(dev));
366*0e6b6b59SJacob Faibussowitsch     init = true;
367*0e6b6b59SJacob Faibussowitsch   }
368*0e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
369*0e6b6b59SJacob Faibussowitsch }
370a4af0ceeSJacob Faibussowitsch 
371a4af0ceeSJacob Faibussowitsch /*@C
372811af0c4SBarry Smith   PetscDeviceInitialize - Initialize `PetscDevice`
373a4af0ceeSJacob Faibussowitsch 
374*0e6b6b59SJacob Faibussowitsch   Not Collective
375a4af0ceeSJacob Faibussowitsch 
376a4af0ceeSJacob Faibussowitsch   Input Parameter:
377811af0c4SBarry Smith . type - The `PetscDeviceType` to initialize
378a4af0ceeSJacob Faibussowitsch 
379*0e6b6b59SJacob Faibussowitsch   Notes:
380*0e6b6b59SJacob Faibussowitsch   Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
381*0e6b6b59SJacob Faibussowitsch   result in device synchronization.
382a4af0ceeSJacob Faibussowitsch 
383a4af0ceeSJacob Faibussowitsch   Level: beginner
384a4af0ceeSJacob Faibussowitsch 
385*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
386*0e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
387a4af0ceeSJacob Faibussowitsch @*/
3889371c9d4SSatish Balay PetscErrorCode PetscDeviceInitialize(PetscDeviceType type) {
389a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
390a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3919566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
392a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
393a4af0ceeSJacob Faibussowitsch }
394a4af0ceeSJacob Faibussowitsch 
395a4af0ceeSJacob Faibussowitsch /*@C
396811af0c4SBarry Smith   PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
397811af0c4SBarry Smith   `PetscDeviceType`
398a4af0ceeSJacob Faibussowitsch 
399*0e6b6b59SJacob Faibussowitsch   Not Collective
400a4af0ceeSJacob Faibussowitsch 
401a4af0ceeSJacob Faibussowitsch   Input Parameter:
402811af0c4SBarry Smith . type - The `PetscDeviceType` to check
403a4af0ceeSJacob Faibussowitsch 
404*0e6b6b59SJacob Faibussowitsch   Notes:
405*0e6b6b59SJacob Faibussowitsch   Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
406a4af0ceeSJacob Faibussowitsch 
407811af0c4SBarry Smith   If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
408811af0c4SBarry Smith   return `PETSC_FALSE` for that `PetscDeviceType`.
409a4af0ceeSJacob Faibussowitsch 
410a4af0ceeSJacob Faibussowitsch   Level: beginner
411a4af0ceeSJacob Faibussowitsch 
412*0e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
413*0e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
414a4af0ceeSJacob Faibussowitsch @*/
4159371c9d4SSatish Balay PetscBool PetscDeviceInitialized(PetscDeviceType type) {
416*0e6b6b59SJacob Faibussowitsch   return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
417*0e6b6b59SJacob Faibussowitsch }
418*0e6b6b59SJacob Faibussowitsch 
419*0e6b6b59SJacob Faibussowitsch /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
420*0e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device) {
421*0e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
422*0e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 2);
423*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitialize(type));
424*0e6b6b59SJacob Faibussowitsch   *device = defaultDevices[type].first;
425*0e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
426a4af0ceeSJacob Faibussowitsch }
427a4af0ceeSJacob Faibussowitsch 
428a16fd2c9SJacob Faibussowitsch /*@C
429a16fd2c9SJacob Faibussowitsch   PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
430a16fd2c9SJacob Faibussowitsch 
431*0e6b6b59SJacob Faibussowitsch   Not Collective
432a16fd2c9SJacob Faibussowitsch 
433a16fd2c9SJacob Faibussowitsch   Input Parameters:
434a16fd2c9SJacob Faibussowitsch + device - The `PetscDevice`
435a16fd2c9SJacob Faibussowitsch - attr   - The attribute
436a16fd2c9SJacob Faibussowitsch 
437a16fd2c9SJacob Faibussowitsch   Output Parameter:
438a16fd2c9SJacob Faibussowitsch . value - The value of the attribute
439a16fd2c9SJacob Faibussowitsch 
440a16fd2c9SJacob Faibussowitsch   Notes:
441a16fd2c9SJacob Faibussowitsch   Since different attributes are often different types `value` is a `void *` to accommodate
442a16fd2c9SJacob Faibussowitsch   them all. The underlying type of the attribute is therefore included in the name of the
443a16fd2c9SJacob Faibussowitsch   `PetscDeviceAttribute` reponsible for querying it. For example,
444a16fd2c9SJacob Faibussowitsch   `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
445a16fd2c9SJacob Faibussowitsch 
446*0e6b6b59SJacob Faibussowitsch   Level: intermediate
447*0e6b6b59SJacob Faibussowitsch 
448a16fd2c9SJacob Faibussowitsch .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
449a16fd2c9SJacob Faibussowitsch @*/
450a16fd2c9SJacob Faibussowitsch PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value) {
451a16fd2c9SJacob Faibussowitsch   PetscFunctionBegin;
452a16fd2c9SJacob Faibussowitsch   PetscValidDevice(device, 1);
453a16fd2c9SJacob Faibussowitsch   PetscValidDeviceAttribute(attr, 2);
454a16fd2c9SJacob Faibussowitsch   PetscValidPointer(value, 3);
455a16fd2c9SJacob Faibussowitsch   PetscUseTypeMethod(device, getattribute, attr, value);
456a16fd2c9SJacob Faibussowitsch   PetscFunctionReturn(0);
457a16fd2c9SJacob Faibussowitsch }
458a16fd2c9SJacob Faibussowitsch 
4599371c9d4SSatish Balay static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType) {
460a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
461a4af0ceeSJacob Faibussowitsch   if (!PetscDeviceConfiguredFor_Internal(type)) {
462*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
463*0e6b6b59SJacob Faibussowitsch     defaultDevices[type].first = nullptr;
464a4af0ceeSJacob Faibussowitsch     PetscFunctionReturn(0);
465a4af0ceeSJacob Faibussowitsch   }
466*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
467a4af0ceeSJacob Faibussowitsch   /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
468a4af0ceeSJacob Faibussowitsch   switch (type) {
469*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
470*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
471*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
472*0e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4739371c9d4SSatish 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]);
474a4af0ceeSJacob Faibussowitsch   }
475*0e6b6b59SJacob 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)]));
476cf3a2253SJacob Faibussowitsch   /*
477*0e6b6b59SJacob Faibussowitsch     defaultInitType, defaultView  and defaultDeviceId now represent what the individual TYPES
478*0e6b6b59SJacob Faibussowitsch     have decided to initialize as
479cf3a2253SJacob Faibussowitsch   */
480*0e6b6b59SJacob Faibussowitsch   if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
481*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
4829566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
483*0e6b6b59SJacob Faibussowitsch     if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
484*0e6b6b59SJacob Faibussowitsch   }
485*0e6b6b59SJacob Faibussowitsch   PetscFunctionReturn(0);
486*0e6b6b59SJacob Faibussowitsch }
487a4af0ceeSJacob Faibussowitsch 
488*0e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView) {
489*0e6b6b59SJacob Faibussowitsch   PetscInt initIdx       = PETSC_DEVICE_INIT_LAZY;
490*0e6b6b59SJacob Faibussowitsch   auto     initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
491*0e6b6b59SJacob Faibussowitsch   auto     flg           = PETSC_FALSE;
492*0e6b6b59SJacob Faibussowitsch 
493*0e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
494*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
495*0e6b6b59SJacob Faibussowitsch   if (flg) PetscCall(PetscLogGpuTime());
496*0e6b6b59SJacob Faibussowitsch 
497*0e6b6b59SJacob Faibussowitsch   PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
498*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
499*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
500*0e6b6b59SJacob 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));
501*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
502*0e6b6b59SJacob Faibussowitsch   PetscOptionsEnd();
503*0e6b6b59SJacob Faibussowitsch 
504*0e6b6b59SJacob Faibussowitsch   if (initIdx == PETSC_DEVICE_INIT_NONE) {
505*0e6b6b59SJacob Faibussowitsch     /* disabled all device initialization if devices are globally disabled */
506*0e6b6b59SJacob 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");
507*0e6b6b59SJacob Faibussowitsch     *defaultView  = PETSC_FALSE;
508*0e6b6b59SJacob Faibussowitsch     initDeviceIdx = PETSC_DEVICE_HOST;
509*0e6b6b59SJacob Faibussowitsch   } else {
510*0e6b6b59SJacob Faibussowitsch     *defaultView = static_cast<PetscBool>(*defaultView && flg);
511*0e6b6b59SJacob Faibussowitsch     if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
512a4af0ceeSJacob Faibussowitsch   }
513*0e6b6b59SJacob Faibussowitsch   *defaultInitType         = PetscDeviceInitTypeCast(initIdx);
514*0e6b6b59SJacob Faibussowitsch   *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
515030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
516030f984aSJacob Faibussowitsch }
517030f984aSJacob Faibussowitsch 
518030f984aSJacob Faibussowitsch /* called from PetscFinalize() do not call yourself! */
519*0e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceFinalize_Private() {
520030f984aSJacob Faibussowitsch   PetscFunctionBegin;
521a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
522bf025ffbSJacob Faibussowitsch     const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] {
523a4af0ceeSJacob Faibussowitsch       PetscFunctionBegin;
524*0e6b6b59SJacob Faibussowitsch       for (auto &&device : defaultDevices) {
525*0e6b6b59SJacob Faibussowitsch         const auto dev = device.first;
526*0e6b6b59SJacob Faibussowitsch 
527*0e6b6b59SJacob 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);
528*0e6b6b59SJacob Faibussowitsch       }
529a4af0ceeSJacob Faibussowitsch       PetscFunctionReturn(0);
530a4af0ceeSJacob Faibussowitsch     };
531bf025ffbSJacob Faibussowitsch     /*
532bf025ffbSJacob Faibussowitsch       you might be thinking, why on earth are you registered yet another finalizer in a
533bf025ffbSJacob Faibussowitsch       function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
534bf025ffbSJacob Faibussowitsch       because it is.
535bf025ffbSJacob Faibussowitsch 
536bf025ffbSJacob Faibussowitsch       The crux of the problem is that the initializer (and therefore the ~finalizer~) of
537bf025ffbSJacob Faibussowitsch       PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
538bf025ffbSJacob Faibussowitsch       a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
539bf025ffbSJacob Faibussowitsch       won't be destroyed yet. So we need to repeat the check that all devices have been
540bf025ffbSJacob Faibussowitsch       destroyed again ~after~ the global context is destroyed. In summary:
541bf025ffbSJacob Faibussowitsch 
542bf025ffbSJacob Faibussowitsch       1. This finalizer runs and destroys all devices, except it may not because the global
543bf025ffbSJacob Faibussowitsch          context may still hold a reference!
544bf025ffbSJacob Faibussowitsch       2. The global context finalizer runs and does the final reference count decrement
545bf025ffbSJacob Faibussowitsch          required, which actually destroys the held device.
546bf025ffbSJacob Faibussowitsch       3. Our newly added finalizer runs and checks that all is well.
547a4af0ceeSJacob Faibussowitsch     */
548*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize)));
549a4af0ceeSJacob Faibussowitsch   }
550*0e6b6b59SJacob Faibussowitsch   for (auto &&device : defaultDevices) {
551*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceDestroy(&device.first));
552*0e6b6b59SJacob Faibussowitsch     device.second = false;
553*0e6b6b59SJacob Faibussowitsch   }
554030f984aSJacob Faibussowitsch   PetscFunctionReturn(0);
555030f984aSJacob Faibussowitsch }
556030f984aSJacob Faibussowitsch 
557cf3a2253SJacob Faibussowitsch /*
558cf3a2253SJacob Faibussowitsch   Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
559cf3a2253SJacob Faibussowitsch   initialization types:
560cf3a2253SJacob Faibussowitsch 
561a4af0ceeSJacob Faibussowitsch   1. defaultInitType - how does PetscDevice as a whole expect to initialize?
562a4af0ceeSJacob Faibussowitsch   2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
563a4af0ceeSJacob Faibussowitsch      e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
564a4af0ceeSJacob Faibussowitsch      have all CUDA devices still initialize.
565a4af0ceeSJacob Faibussowitsch 
566a4af0ceeSJacob Faibussowitsch   All told the following happens:
567cf3a2253SJacob Faibussowitsch 
568a4af0ceeSJacob Faibussowitsch   0. defaultInitType -> LAZY
569a4af0ceeSJacob Faibussowitsch   1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
570a4af0ceeSJacob Faibussowitsch   2. PetscDevice initializes each sub type with deviceDefaultInitType.
571a4af0ceeSJacob Faibussowitsch   2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
572a4af0ceeSJacob Faibussowitsch       to checking for specific device init. if view or specific device init
573a4af0ceeSJacob Faibussowitsch       subTypeDefaultInitType -> EAGER. disabled once again overrides all.
574a4af0ceeSJacob Faibussowitsch */
575*0e6b6b59SJacob Faibussowitsch 
5769371c9d4SSatish Balay PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm) {
5777a101e5eSJacob Faibussowitsch   auto defaultView                    = PETSC_FALSE;
5787a101e5eSJacob Faibussowitsch   auto initializeDeviceContextEagerly = PETSC_FALSE;
579*0e6b6b59SJacob Faibussowitsch   auto defaultDeviceSet               = PETSC_FALSE;
5807a101e5eSJacob Faibussowitsch   auto defaultDevice                  = PetscInt{PETSC_DECIDE};
581*0e6b6b59SJacob Faibussowitsch   auto deviceContextInitDevice        = PETSC_DEVICE_DEFAULT();
582*0e6b6b59SJacob Faibussowitsch   auto defaultInitType                = PETSC_DEVICE_INIT_LAZY;
583a4af0ceeSJacob Faibussowitsch 
584a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
585a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
586a4af0ceeSJacob Faibussowitsch     int result;
587a4af0ceeSJacob Faibussowitsch 
5889566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
589a4af0ceeSJacob Faibussowitsch     /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
590a4af0ceeSJacob Faibussowitsch      * global space */
591a4af0ceeSJacob Faibussowitsch     if (PetscUnlikely(result != MPI_IDENT)) {
592a4af0ceeSJacob Faibussowitsch       char name[MPI_MAX_OBJECT_NAME] = {};
593a4af0ceeSJacob Faibussowitsch       int  len; /* unused */
594a4af0ceeSJacob Faibussowitsch 
5959566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
59698921bdaSJacob Faibussowitsch       SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
597a4af0ceeSJacob Faibussowitsch     }
598a4af0ceeSJacob Faibussowitsch   }
599a4af0ceeSJacob Faibussowitsch   comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
6009566063dSJacob Faibussowitsch   PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
601a4af0ceeSJacob Faibussowitsch 
602*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
6037a101e5eSJacob Faibussowitsch 
604*0e6b6b59SJacob Faibussowitsch   // the precise values don't matter here, so long as they are sequential
605*0e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, "");
606*0e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, "");
607*0e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, "");
608*0e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, "");
609*0e6b6b59SJacob Faibussowitsch   static_assert(Petsc::util::integral_value(PETSC_DEVICE_MAX) == 4, "");
610*0e6b6b59SJacob Faibussowitsch   for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
611*0e6b6b59SJacob Faibussowitsch     const auto deviceType = PetscDeviceTypeCast(i);
612a4af0ceeSJacob Faibussowitsch     auto       initType   = defaultInitType;
613a4af0ceeSJacob Faibussowitsch 
6149566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
615*0e6b6b59SJacob Faibussowitsch     if (PetscDeviceConfiguredFor_Internal(deviceType)) {
616*0e6b6b59SJacob Faibussowitsch       if (initType == PETSC_DEVICE_INIT_EAGER) {
617a4af0ceeSJacob Faibussowitsch         initializeDeviceContextEagerly = PETSC_TRUE;
618*0e6b6b59SJacob Faibussowitsch         // only update the default device if the user hasn't set it previously
619*0e6b6b59SJacob Faibussowitsch         if (!defaultDeviceSet) {
620a4af0ceeSJacob Faibussowitsch           deviceContextInitDevice = deviceType;
621*0e6b6b59SJacob Faibussowitsch           PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
622*0e6b6b59SJacob Faibussowitsch         }
623*0e6b6b59SJacob Faibussowitsch       } else if (initType == PETSC_DEVICE_INIT_NONE) {
624*0e6b6b59SJacob 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]);
625a4af0ceeSJacob Faibussowitsch       }
626a4af0ceeSJacob Faibussowitsch     }
627*0e6b6b59SJacob Faibussowitsch   }
628*0e6b6b59SJacob Faibussowitsch 
629*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
630*0e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
631*0e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
632*0e6b6b59SJacob Faibussowitsch   /*                       PetscDevice is now fully initialized                          */
633*0e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
634*0e6b6b59SJacob Faibussowitsch   {
635*0e6b6b59SJacob Faibussowitsch     /*
636*0e6b6b59SJacob Faibussowitsch       query the options db to get the root settings from the user (if any).
637*0e6b6b59SJacob Faibussowitsch 
638*0e6b6b59SJacob Faibussowitsch       This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
639*0e6b6b59SJacob Faibussowitsch       PetscDeviceContextSetFromOptions() before we even have one, then set a few static
640*0e6b6b59SJacob Faibussowitsch       variables in that file with the results.
641*0e6b6b59SJacob Faibussowitsch     */
642*0e6b6b59SJacob Faibussowitsch     auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
643*0e6b6b59SJacob Faibussowitsch     auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
644*0e6b6b59SJacob Faibussowitsch 
645*0e6b6b59SJacob Faibussowitsch     PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
646*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
647*0e6b6b59SJacob Faibussowitsch     PetscOptionsEnd();
648*0e6b6b59SJacob Faibussowitsch 
649*0e6b6b59SJacob Faibussowitsch     if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
650*0e6b6b59SJacob Faibussowitsch     if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
651*0e6b6b59SJacob Faibussowitsch   }
652*0e6b6b59SJacob Faibussowitsch 
653a4af0ceeSJacob Faibussowitsch   if (initializeDeviceContextEagerly) {
654a4af0ceeSJacob Faibussowitsch     PetscDeviceContext dctx;
655a4af0ceeSJacob Faibussowitsch 
656*0e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
657*0e6b6b59SJacob Faibussowitsch     /* instantiates the device context */
6589566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
6599566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextSetUp(dctx));
660a4af0ceeSJacob Faibussowitsch   }
661a4af0ceeSJacob Faibussowitsch   PetscFunctionReturn(0);
662a4af0ceeSJacob Faibussowitsch }
663