xref: /petsc/src/sys/objects/device/interface/device.cxx (revision 394bf645184363d868b70bf802ab31bc3e6c5573)
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 <utility> // std::make_pair
90e6b6b59SJacob Faibussowitsch 
100e6b6b59SJacob Faibussowitsch using namespace Petsc::device;
11030f984aSJacob Faibussowitsch 
12cf3a2253SJacob Faibussowitsch /*
13cf3a2253SJacob Faibussowitsch   note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to
14cf3a2253SJacob Faibussowitsch   be picked up by the switch-case macros below
15cf3a2253SJacob Faibussowitsch */
160e6b6b59SJacob Faibussowitsch static host::Device HOSTDevice{PetscDeviceContextCreate_HOST};
17030f984aSJacob Faibussowitsch #if PetscDefined(HAVE_CUDA)
180e6b6b59SJacob Faibussowitsch static cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA};
19030f984aSJacob Faibussowitsch #endif
20030f984aSJacob Faibussowitsch #if PetscDefined(HAVE_HIP)
210e6b6b59SJacob Faibussowitsch static cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP};
22030f984aSJacob Faibussowitsch #endif
23a2158755SJunchao Zhang #if PetscDefined(HAVE_SYCL)
240e6b6b59SJacob Faibussowitsch static sycl::Device SYCLDevice{PetscDeviceContextCreate_SYCL};
25a2158755SJunchao Zhang #endif
26030f984aSJacob Faibussowitsch 
2717f48955SJacob Faibussowitsch #define PETSC_DEVICE_CASE(IMPLS, func, ...) \
2817f48955SJacob Faibussowitsch   case PetscConcat_(PETSC_DEVICE_, IMPLS): { \
299566063dSJacob Faibussowitsch     PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \
3017f48955SJacob Faibussowitsch   } break
31a4af0ceeSJacob Faibussowitsch 
32*394bf645SJacob Faibussowitsch #define PETSC_VOID_0(...) ((void)0)
33*394bf645SJacob Faibussowitsch 
34cf3a2253SJacob Faibussowitsch /*
35cf3a2253SJacob Faibussowitsch   Suppose you have:
36cf3a2253SJacob Faibussowitsch 
37cf3a2253SJacob Faibussowitsch   CUDADevice.myFunction(arg1,arg2)
38cf3a2253SJacob Faibussowitsch 
39cf3a2253SJacob Faibussowitsch   that you would like to conditionally define and call in a switch-case:
40cf3a2253SJacob Faibussowitsch 
41cf3a2253SJacob Faibussowitsch   switch(PetscDeviceType) {
42cf3a2253SJacob Faibussowitsch   #if PetscDefined(HAVE_CUDA)
43cf3a2253SJacob Faibussowitsch   case PETSC_DEVICE_CUDA: {
449566063dSJacob Faibussowitsch     PetscCall(CUDADevice.myFunction(arg1,arg2));
45cf3a2253SJacob Faibussowitsch   } break;
46cf3a2253SJacob Faibussowitsch   #endif
47cf3a2253SJacob Faibussowitsch   }
48cf3a2253SJacob Faibussowitsch 
49cf3a2253SJacob Faibussowitsch   then calling this macro:
50cf3a2253SJacob Faibussowitsch 
51cf3a2253SJacob Faibussowitsch   PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2)
52cf3a2253SJacob Faibussowitsch 
53cf3a2253SJacob Faibussowitsch   will expand to the following case statement:
54cf3a2253SJacob Faibussowitsch 
55cf3a2253SJacob Faibussowitsch   case PETSC_DEVICE_CUDA: {
569566063dSJacob Faibussowitsch     PetscCall(CUDADevice.myFunction(arg1,arg2));
57cf3a2253SJacob Faibussowitsch   } break
58cf3a2253SJacob Faibussowitsch 
59cf3a2253SJacob Faibussowitsch   if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise
6017f48955SJacob Faibussowitsch */
61*394bf645SJacob Faibussowitsch #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(PetscConcat_(HAVE_, IMPLS), PETSC_DEVICE_CASE, PETSC_VOID_0)(IMPLS, func, __VA_ARGS__)
62030f984aSJacob Faibussowitsch 
63030f984aSJacob Faibussowitsch /*@C
64811af0c4SBarry Smith   PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type
65030f984aSJacob Faibussowitsch 
660e6b6b59SJacob Faibussowitsch   Not Collective
67030f984aSJacob Faibussowitsch 
68f1a722f8SMatthew G. Knepley   Input Parameters:
69811af0c4SBarry Smith + type  - The type of `PetscDevice`
70811af0c4SBarry Smith - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically)
71030f984aSJacob Faibussowitsch 
72030f984aSJacob Faibussowitsch   Output Parameter:
73811af0c4SBarry Smith . device - The `PetscDevice`
74030f984aSJacob Faibussowitsch 
752fe279fdSBarry Smith   Level: beginner
762fe279fdSBarry Smith 
77030f984aSJacob Faibussowitsch   Notes:
780e6b6b59SJacob Faibussowitsch   This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of
790e6b6b59SJacob Faibussowitsch   device synchronization.
80a4af0ceeSJacob Faibussowitsch 
81811af0c4SBarry Smith   `devid` is what you might pass to `cudaSetDevice()` for example.
82030f984aSJacob Faibussowitsch 
830e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`,
840e6b6b59SJacob Faibussowitsch `PetscDeviceInitialize()`, `PetscDeviceInitialized()`, `PetscDeviceConfigure()`,
850e6b6b59SJacob Faibussowitsch `PetscDeviceView()`, `PetscDeviceDestroy()`
86030f984aSJacob Faibussowitsch @*/
87d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device)
88d71ae5a4SJacob Faibussowitsch {
89030f984aSJacob Faibussowitsch   static PetscInt PetscDeviceCounter = 0;
90030f984aSJacob Faibussowitsch 
91030f984aSJacob Faibussowitsch   PetscFunctionBegin;
92a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
93a4af0ceeSJacob Faibussowitsch   PetscValidPointer(device, 3);
949566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializePackage());
950e6b6b59SJacob Faibussowitsch   PetscCall(PetscNew(device));
960e6b6b59SJacob Faibussowitsch   (*device)->id     = PetscDeviceCounter++;
970e6b6b59SJacob Faibussowitsch   (*device)->type   = type;
980e6b6b59SJacob Faibussowitsch   (*device)->refcnt = 1;
99cf3a2253SJacob Faibussowitsch   /*
10091c35059SPierre Jolivet     if you are adding a device, you also need to add its initialization in
101cf3a2253SJacob Faibussowitsch     PetscDeviceInitializeTypeFromOptions_Private() below
102cf3a2253SJacob Faibussowitsch   */
103a4af0ceeSJacob Faibussowitsch   switch (type) {
1040e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid);
1050e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid);
1060e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid);
1070e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid);
108030f984aSJacob Faibussowitsch   default:
10917f48955SJacob Faibussowitsch     /* in case the above macros expand to nothing this silences any unused variable warnings */
11017f48955SJacob Faibussowitsch     (void)(devid);
11198921bdaSJacob 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]);
112030f984aSJacob Faibussowitsch   }
1133ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
114030f984aSJacob Faibussowitsch }
115030f984aSJacob Faibussowitsch 
116030f984aSJacob Faibussowitsch /*@C
117811af0c4SBarry Smith   PetscDeviceDestroy - Free a `PetscDevice`
118030f984aSJacob Faibussowitsch 
1190e6b6b59SJacob Faibussowitsch   Not Collective
120030f984aSJacob Faibussowitsch 
121030f984aSJacob Faibussowitsch   Input Parameter:
1220e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
123030f984aSJacob Faibussowitsch 
124030f984aSJacob Faibussowitsch   Level: beginner
125030f984aSJacob Faibussowitsch 
1260e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`,
1270e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
128030f984aSJacob Faibussowitsch @*/
129d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceDestroy(PetscDevice *device)
130d71ae5a4SJacob Faibussowitsch {
131a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
1320e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 1);
1333ba16761SJacob Faibussowitsch   if (!*device) PetscFunctionReturn(PETSC_SUCCESS);
134a4af0ceeSJacob Faibussowitsch   PetscValidDevice(*device, 1);
1359566063dSJacob Faibussowitsch   PetscCall(PetscDeviceDereference_Internal(*device));
136a4af0ceeSJacob Faibussowitsch   if ((*device)->refcnt) {
1370e6b6b59SJacob Faibussowitsch     *device = nullptr;
1383ba16761SJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
139030f984aSJacob Faibussowitsch   }
1409566063dSJacob Faibussowitsch   PetscCall(PetscFree((*device)->data));
1419566063dSJacob Faibussowitsch   PetscCall(PetscFree(*device));
1423ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
143030f984aSJacob Faibussowitsch }
144030f984aSJacob Faibussowitsch 
145a4af0ceeSJacob Faibussowitsch /*@C
146811af0c4SBarry Smith   PetscDeviceConfigure - Configure a particular `PetscDevice`
147030f984aSJacob Faibussowitsch 
1480e6b6b59SJacob Faibussowitsch   Not Collective
149a4af0ceeSJacob Faibussowitsch 
150a4af0ceeSJacob Faibussowitsch   Input Parameter:
151811af0c4SBarry Smith . device - The `PetscDevice` to configure
152a4af0ceeSJacob Faibussowitsch 
1532fe279fdSBarry Smith   Level: beginner
1542fe279fdSBarry Smith 
1550e6b6b59SJacob Faibussowitsch   Notes:
1560e6b6b59SJacob Faibussowitsch   The user should not assume that this is a cheap operation.
157a4af0ceeSJacob Faibussowitsch 
1580e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`,
1590e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
160a4af0ceeSJacob Faibussowitsch @*/
161d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceConfigure(PetscDevice device)
162d71ae5a4SJacob Faibussowitsch {
163030f984aSJacob Faibussowitsch   PetscFunctionBegin;
164a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
165cf3a2253SJacob Faibussowitsch   /*
166cf3a2253SJacob Faibussowitsch     if no available configuration is available, this cascades all the way down to default
167cf3a2253SJacob Faibussowitsch     and error
168cf3a2253SJacob Faibussowitsch   */
1690e6b6b59SJacob Faibussowitsch   switch (const auto dtype = device->type) {
1700e6b6b59SJacob Faibussowitsch   case PETSC_DEVICE_HOST:
1710e6b6b59SJacob Faibussowitsch     if (PetscDefined(HAVE_HOST)) break; // always true
1729371c9d4SSatish Balay   case PETSC_DEVICE_CUDA:
1739371c9d4SSatish Balay     if (PetscDefined(HAVE_CUDA)) break;
1740e6b6b59SJacob Faibussowitsch     goto error;
1759371c9d4SSatish Balay   case PETSC_DEVICE_HIP:
1769371c9d4SSatish Balay     if (PetscDefined(HAVE_HIP)) break;
1770e6b6b59SJacob Faibussowitsch     goto error;
1789371c9d4SSatish Balay   case PETSC_DEVICE_SYCL:
1799371c9d4SSatish Balay     if (PetscDefined(HAVE_SYCL)) break;
180f4d061e9SPierre Jolivet     goto error;
1810e6b6b59SJacob Faibussowitsch   default:
1820e6b6b59SJacob Faibussowitsch   error:
1830e6b6b59SJacob Faibussowitsch     SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]);
184a4af0ceeSJacob Faibussowitsch   }
185dbbe0bcdSBarry Smith   PetscUseTypeMethod(device, configure);
1863ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
187a4af0ceeSJacob Faibussowitsch }
188a4af0ceeSJacob Faibussowitsch 
189a4af0ceeSJacob Faibussowitsch /*@C
190811af0c4SBarry Smith   PetscDeviceView - View a `PetscDevice`
191a4af0ceeSJacob Faibussowitsch 
1920e6b6b59SJacob Faibussowitsch   Collective on viewer
193a4af0ceeSJacob Faibussowitsch 
19491e63d38SStefano Zampini   Input Parameters:
195811af0c4SBarry Smith + device - The `PetscDevice` to view
1960e6b6b59SJacob Faibussowitsch - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`)
197a4af0ceeSJacob Faibussowitsch 
198a4af0ceeSJacob Faibussowitsch   Level: beginner
199a4af0ceeSJacob Faibussowitsch 
2000e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`,
2010e6b6b59SJacob Faibussowitsch `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()`
202a4af0ceeSJacob Faibussowitsch @*/
203d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer)
204d71ae5a4SJacob Faibussowitsch {
2050e6b6b59SJacob Faibussowitsch   auto      sub = viewer;
2060e6b6b59SJacob Faibussowitsch   PetscBool iascii;
2070e6b6b59SJacob Faibussowitsch 
208a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
209a4af0ceeSJacob Faibussowitsch   PetscValidDevice(device, 1);
2100e6b6b59SJacob Faibussowitsch   if (viewer) {
211a4af0ceeSJacob Faibussowitsch     PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2);
2120e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii));
2130e6b6b59SJacob Faibussowitsch   } else {
2140e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer));
2150e6b6b59SJacob Faibussowitsch     iascii = PETSC_TRUE;
2160e6b6b59SJacob Faibussowitsch   }
2170e6b6b59SJacob Faibussowitsch 
2180e6b6b59SJacob Faibussowitsch   if (iascii) {
2190e6b6b59SJacob Faibussowitsch     auto        dtype = PETSC_DEVICE_HOST;
2200e6b6b59SJacob Faibussowitsch     MPI_Comm    comm;
2210e6b6b59SJacob Faibussowitsch     PetscMPIInt size;
2220e6b6b59SJacob Faibussowitsch     PetscInt    id = 0;
2230e6b6b59SJacob Faibussowitsch 
2240e6b6b59SJacob Faibussowitsch     PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm));
2250e6b6b59SJacob Faibussowitsch     PetscCallMPI(MPI_Comm_size(comm, &size));
2260e6b6b59SJacob Faibussowitsch 
2270e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetDeviceId(device, &id));
2280e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceGetType(device, &dtype));
2290e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub));
2300e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes"));
2310e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPushTab(sub));
2320e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype]));
2330e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id));
2340e6b6b59SJacob Faibussowitsch   }
2350e6b6b59SJacob Faibussowitsch 
2360e6b6b59SJacob Faibussowitsch   // see if impls has extra viewer stuff
2370e6b6b59SJacob Faibussowitsch   PetscTryTypeMethod(device, view, sub);
2380e6b6b59SJacob Faibussowitsch 
2390e6b6b59SJacob Faibussowitsch   if (iascii) {
2400e6b6b59SJacob Faibussowitsch     // undo the ASCII specific stuff
2410e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerASCIIPopTab(sub));
2420e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub));
2430e6b6b59SJacob Faibussowitsch     PetscCall(PetscViewerFlush(viewer));
2440e6b6b59SJacob Faibussowitsch   }
2453ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
246a4af0ceeSJacob Faibussowitsch }
247a4af0ceeSJacob Faibussowitsch 
24891e63d38SStefano Zampini /*@C
2490e6b6b59SJacob Faibussowitsch   PetscDeviceGetType - Get the type of device
25091e63d38SStefano Zampini 
2510e6b6b59SJacob Faibussowitsch   Not Collective
25291e63d38SStefano Zampini 
25391e63d38SStefano Zampini   Input Parameter:
254811af0c4SBarry Smith . device - The `PetscDevice`
25591e63d38SStefano Zampini 
25691e63d38SStefano Zampini   Output Parameter:
2570e6b6b59SJacob Faibussowitsch . type - The `PetscDeviceType`
25891e63d38SStefano Zampini 
25991e63d38SStefano Zampini   Level: beginner
26091e63d38SStefano Zampini 
2610e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`,
2620e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`,
2630e6b6b59SJacob Faibussowitsch `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()`
2640e6b6b59SJacob Faibussowitsch @*/
265d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type)
266d71ae5a4SJacob Faibussowitsch {
2670e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
2680e6b6b59SJacob Faibussowitsch   PetscValidDevice(device, 1);
2690e6b6b59SJacob Faibussowitsch   PetscValidPointer(type, 2);
2700e6b6b59SJacob Faibussowitsch   *type = device->type;
2713ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
2720e6b6b59SJacob Faibussowitsch }
2730e6b6b59SJacob Faibussowitsch 
2740e6b6b59SJacob Faibussowitsch /*@C
2750e6b6b59SJacob Faibussowitsch   PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice`
2760e6b6b59SJacob Faibussowitsch 
2770e6b6b59SJacob Faibussowitsch   Not Collective
2780e6b6b59SJacob Faibussowitsch 
2790e6b6b59SJacob Faibussowitsch   Input Parameter:
2800e6b6b59SJacob Faibussowitsch . device - The `PetscDevice`
2810e6b6b59SJacob Faibussowitsch 
2820e6b6b59SJacob Faibussowitsch   Output Parameter:
2830e6b6b59SJacob Faibussowitsch . id - The id
2840e6b6b59SJacob Faibussowitsch 
2852fe279fdSBarry Smith   Level: beginner
2862fe279fdSBarry Smith 
2870e6b6b59SJacob Faibussowitsch   Notes:
2880e6b6b59SJacob Faibussowitsch   The returned ID may have been assigned by the underlying device backend. For example if the
2890e6b6b59SJacob Faibussowitsch   backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when
2900e6b6b59SJacob Faibussowitsch   this device was configured.
2910e6b6b59SJacob Faibussowitsch 
2920e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()`
29391e63d38SStefano Zampini @*/
294d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id)
295d71ae5a4SJacob Faibussowitsch {
29691e63d38SStefano Zampini   PetscFunctionBegin;
29791e63d38SStefano Zampini   PetscValidDevice(device, 1);
29891e63d38SStefano Zampini   PetscValidIntPointer(id, 2);
29991e63d38SStefano Zampini   *id = device->deviceId;
3003ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
30191e63d38SStefano Zampini }
30291e63d38SStefano Zampini 
3030e6b6b59SJacob Faibussowitsch struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> {
3040e6b6b59SJacob Faibussowitsch   PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
3050e6b6b59SJacob Faibussowitsch 
306089fb57cSJacob Faibussowitsch   PetscErrorCode finalize_() noexcept
307d71ae5a4SJacob Faibussowitsch   {
3080e6b6b59SJacob Faibussowitsch     PetscFunctionBegin;
3090e6b6b59SJacob Faibussowitsch     type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE;
3103ba16761SJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
3110e6b6b59SJacob Faibussowitsch   }
3120e6b6b59SJacob Faibussowitsch };
3130e6b6b59SJacob Faibussowitsch 
3140e6b6b59SJacob Faibussowitsch static auto default_device_type = DefaultDeviceType();
3150e6b6b59SJacob Faibussowitsch 
3160e6b6b59SJacob Faibussowitsch /*@C
3170e6b6b59SJacob Faibussowitsch   PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType`
3180e6b6b59SJacob Faibussowitsch 
3190e6b6b59SJacob Faibussowitsch   Not Collective
3200e6b6b59SJacob Faibussowitsch 
3212fe279fdSBarry Smith   Level: beginner
3222fe279fdSBarry Smith 
3230e6b6b59SJacob Faibussowitsch   Notes:
3240e6b6b59SJacob Faibussowitsch   Unless selected by the user, the default device is selected in the following order\:
3250e6b6b59SJacob Faibussowitsch   `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`.
3260e6b6b59SJacob Faibussowitsch 
3270e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()`
3280e6b6b59SJacob Faibussowitsch @*/
329d71ae5a4SJacob Faibussowitsch PetscDeviceType PETSC_DEVICE_DEFAULT(void)
330d71ae5a4SJacob Faibussowitsch {
3310e6b6b59SJacob Faibussowitsch   return default_device_type.type;
3320e6b6b59SJacob Faibussowitsch }
3330e6b6b59SJacob Faibussowitsch 
3340e6b6b59SJacob Faibussowitsch /*@C
3350e6b6b59SJacob Faibussowitsch   PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice`
3360e6b6b59SJacob Faibussowitsch 
3370e6b6b59SJacob Faibussowitsch   Not Collective
3380e6b6b59SJacob Faibussowitsch 
3390e6b6b59SJacob Faibussowitsch   Input Parameter:
3400e6b6b59SJacob Faibussowitsch . type - the new default device type
3410e6b6b59SJacob Faibussowitsch 
3422fe279fdSBarry Smith   Level: beginner
3432fe279fdSBarry Smith 
3440e6b6b59SJacob Faibussowitsch   Notes:
3450e6b6b59SJacob Faibussowitsch   This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`.
3460e6b6b59SJacob Faibussowitsch 
3470e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceGetType`,
3480e6b6b59SJacob Faibussowitsch @*/
349d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type)
350d71ae5a4SJacob Faibussowitsch {
3510e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
3520e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3530e6b6b59SJacob Faibussowitsch   if (default_device_type.type != type) {
3540e6b6b59SJacob Faibussowitsch     // no need to waster a PetscRegisterFinalize() slot if we don't change it
3550e6b6b59SJacob Faibussowitsch     default_device_type.type = type;
3560e6b6b59SJacob Faibussowitsch     PetscCall(default_device_type.register_finalize());
3570e6b6b59SJacob Faibussowitsch   }
3583ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3590e6b6b59SJacob Faibussowitsch }
3600e6b6b59SJacob Faibussowitsch 
3610e6b6b59SJacob Faibussowitsch static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {};
3620e6b6b59SJacob Faibussowitsch 
3630e6b6b59SJacob Faibussowitsch /*
364da81f932SPierre Jolivet   Actual initialization function; any functions claiming to initialize PetscDevice or
3650e6b6b59SJacob Faibussowitsch   PetscDeviceContext will have to run through this one
3660e6b6b59SJacob Faibussowitsch */
367d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId)
368d71ae5a4SJacob Faibussowitsch {
3690e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
3700e6b6b59SJacob Faibussowitsch   PetscValidDeviceType(type, 1);
3710e6b6b59SJacob Faibussowitsch   if (PetscUnlikely(!PetscDeviceInitialized(type))) {
3720e6b6b59SJacob Faibussowitsch     auto &dev  = defaultDevices[type].first;
3730e6b6b59SJacob Faibussowitsch     auto &init = defaultDevices[type].second;
3740e6b6b59SJacob Faibussowitsch 
3750e6b6b59SJacob Faibussowitsch     PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]);
3760e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev));
3770e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceConfigure(dev));
3780e6b6b59SJacob Faibussowitsch     init = true;
3790e6b6b59SJacob Faibussowitsch   }
3803ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
3810e6b6b59SJacob Faibussowitsch }
382a4af0ceeSJacob Faibussowitsch 
383a4af0ceeSJacob Faibussowitsch /*@C
384811af0c4SBarry Smith   PetscDeviceInitialize - Initialize `PetscDevice`
385a4af0ceeSJacob Faibussowitsch 
3860e6b6b59SJacob Faibussowitsch   Not Collective
387a4af0ceeSJacob Faibussowitsch 
388a4af0ceeSJacob Faibussowitsch   Input Parameter:
389811af0c4SBarry Smith . type - The `PetscDeviceType` to initialize
390a4af0ceeSJacob Faibussowitsch 
3912fe279fdSBarry Smith   Level: beginner
3922fe279fdSBarry Smith 
3930e6b6b59SJacob Faibussowitsch   Notes:
3940e6b6b59SJacob Faibussowitsch   Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may
3950e6b6b59SJacob Faibussowitsch   result in device synchronization.
396a4af0ceeSJacob Faibussowitsch 
3970e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`,
3980e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
399a4af0ceeSJacob Faibussowitsch @*/
400d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceInitialize(PetscDeviceType type)
401d71ae5a4SJacob Faibussowitsch {
402a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
403a4af0ceeSJacob Faibussowitsch   PetscValidDeviceType(type, 1);
4049566063dSJacob Faibussowitsch   PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE));
4053ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
406a4af0ceeSJacob Faibussowitsch }
407a4af0ceeSJacob Faibussowitsch 
408a4af0ceeSJacob Faibussowitsch /*@C
409811af0c4SBarry Smith   PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular
410811af0c4SBarry Smith   `PetscDeviceType`
411a4af0ceeSJacob Faibussowitsch 
4120e6b6b59SJacob Faibussowitsch   Not Collective
413a4af0ceeSJacob Faibussowitsch 
414a4af0ceeSJacob Faibussowitsch   Input Parameter:
415811af0c4SBarry Smith . type - The `PetscDeviceType` to check
416a4af0ceeSJacob Faibussowitsch 
4172fe279fdSBarry Smith   Level: beginner
4182fe279fdSBarry Smith 
4190e6b6b59SJacob Faibussowitsch   Notes:
4200e6b6b59SJacob Faibussowitsch   Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise.
421a4af0ceeSJacob Faibussowitsch 
422811af0c4SBarry Smith   If one has not configured PETSc for a particular `PetscDeviceType` then this routine will
423811af0c4SBarry Smith   return `PETSC_FALSE` for that `PetscDeviceType`.
424a4af0ceeSJacob Faibussowitsch 
4250e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`,
4260e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()`
427a4af0ceeSJacob Faibussowitsch @*/
428d71ae5a4SJacob Faibussowitsch PetscBool PetscDeviceInitialized(PetscDeviceType type)
429d71ae5a4SJacob Faibussowitsch {
4300e6b6b59SJacob Faibussowitsch   return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second);
4310e6b6b59SJacob Faibussowitsch }
4320e6b6b59SJacob Faibussowitsch 
4330e6b6b59SJacob Faibussowitsch /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */
434d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device)
435d71ae5a4SJacob Faibussowitsch {
4360e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
4370e6b6b59SJacob Faibussowitsch   PetscValidPointer(device, 2);
4380e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitialize(type));
4390e6b6b59SJacob Faibussowitsch   *device = defaultDevices[type].first;
4403ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
441a4af0ceeSJacob Faibussowitsch }
442a4af0ceeSJacob Faibussowitsch 
443a16fd2c9SJacob Faibussowitsch /*@C
444a16fd2c9SJacob Faibussowitsch   PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice`
445a16fd2c9SJacob Faibussowitsch 
4460e6b6b59SJacob Faibussowitsch   Not Collective
447a16fd2c9SJacob Faibussowitsch 
448a16fd2c9SJacob Faibussowitsch   Input Parameters:
449a16fd2c9SJacob Faibussowitsch + device - The `PetscDevice`
450a16fd2c9SJacob Faibussowitsch - attr   - The attribute
451a16fd2c9SJacob Faibussowitsch 
452a16fd2c9SJacob Faibussowitsch   Output Parameter:
453a16fd2c9SJacob Faibussowitsch . value - The value of the attribute
454a16fd2c9SJacob Faibussowitsch 
4552fe279fdSBarry Smith   Level: intermediate
4562fe279fdSBarry Smith 
457a16fd2c9SJacob Faibussowitsch   Notes:
458a16fd2c9SJacob Faibussowitsch   Since different attributes are often different types `value` is a `void *` to accommodate
459a16fd2c9SJacob Faibussowitsch   them all. The underlying type of the attribute is therefore included in the name of the
460da81f932SPierre Jolivet   `PetscDeviceAttribute` responsible for querying it. For example,
461a16fd2c9SJacob Faibussowitsch   `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`.
462a16fd2c9SJacob Faibussowitsch 
463a16fd2c9SJacob Faibussowitsch .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice`
464a16fd2c9SJacob Faibussowitsch @*/
465d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value)
466d71ae5a4SJacob Faibussowitsch {
467a16fd2c9SJacob Faibussowitsch   PetscFunctionBegin;
468a16fd2c9SJacob Faibussowitsch   PetscValidDevice(device, 1);
469a16fd2c9SJacob Faibussowitsch   PetscValidDeviceAttribute(attr, 2);
470a16fd2c9SJacob Faibussowitsch   PetscValidPointer(value, 3);
471a16fd2c9SJacob Faibussowitsch   PetscUseTypeMethod(device, getattribute, attr, value);
4723ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
473a16fd2c9SJacob Faibussowitsch }
474a16fd2c9SJacob Faibussowitsch 
475d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType)
476d71ae5a4SJacob Faibussowitsch {
477a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
478a4af0ceeSJacob Faibussowitsch   if (!PetscDeviceConfiguredFor_Internal(type)) {
4790e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type]));
4800e6b6b59SJacob Faibussowitsch     defaultDevices[type].first = nullptr;
4813ba16761SJacob Faibussowitsch     PetscFunctionReturn(PETSC_SUCCESS);
482a4af0ceeSJacob Faibussowitsch   }
4830e6b6b59SJacob Faibussowitsch   PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type]));
484a4af0ceeSJacob Faibussowitsch   /* ugly switch needed to pick the right global variable... could maybe do this as a union? */
485a4af0ceeSJacob Faibussowitsch   switch (type) {
4860e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4870e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4880e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
4890e6b6b59SJacob Faibussowitsch     PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType);
490d71ae5a4SJacob Faibussowitsch   default:
491d71ae5a4SJacob 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]);
492a4af0ceeSJacob Faibussowitsch   }
493bd2fcf0cSJacob 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::to_underlying(*defaultInitType)]));
494cf3a2253SJacob Faibussowitsch   /*
4950e6b6b59SJacob Faibussowitsch     defaultInitType, defaultView  and defaultDeviceId now represent what the individual TYPES
4960e6b6b59SJacob Faibussowitsch     have decided to initialize as
497cf3a2253SJacob Faibussowitsch   */
4980e6b6b59SJacob Faibussowitsch   if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) {
4990e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type]));
5009566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId));
5010e6b6b59SJacob Faibussowitsch     if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr));
5020e6b6b59SJacob Faibussowitsch   }
5033ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
5040e6b6b59SJacob Faibussowitsch }
505a4af0ceeSJacob Faibussowitsch 
506d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView)
507d71ae5a4SJacob Faibussowitsch {
5080e6b6b59SJacob Faibussowitsch   PetscInt initIdx       = PETSC_DEVICE_INIT_LAZY;
5090e6b6b59SJacob Faibussowitsch   auto     initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice);
5100e6b6b59SJacob Faibussowitsch   auto     flg           = PETSC_FALSE;
5110e6b6b59SJacob Faibussowitsch 
5120e6b6b59SJacob Faibussowitsch   PetscFunctionBegin;
5130e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg));
5140e6b6b59SJacob Faibussowitsch   if (flg) PetscCall(PetscLogGpuTime());
5150e6b6b59SJacob Faibussowitsch 
5160e6b6b59SJacob Faibussowitsch   PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys");
5170e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr));
5180e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet));
5190e6b6b59SJacob 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));
5200e6b6b59SJacob Faibussowitsch   PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg));
5210e6b6b59SJacob Faibussowitsch   PetscOptionsEnd();
5220e6b6b59SJacob Faibussowitsch 
5230e6b6b59SJacob Faibussowitsch   if (initIdx == PETSC_DEVICE_INIT_NONE) {
5240e6b6b59SJacob Faibussowitsch     /* disabled all device initialization if devices are globally disabled */
525da81f932SPierre Jolivet     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 exclusive");
5260e6b6b59SJacob Faibussowitsch     *defaultView  = PETSC_FALSE;
5270e6b6b59SJacob Faibussowitsch     initDeviceIdx = PETSC_DEVICE_HOST;
5280e6b6b59SJacob Faibussowitsch   } else {
5290e6b6b59SJacob Faibussowitsch     *defaultView = static_cast<PetscBool>(*defaultView && flg);
5300e6b6b59SJacob Faibussowitsch     if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER;
531a4af0ceeSJacob Faibussowitsch   }
5320e6b6b59SJacob Faibussowitsch   *defaultInitType         = PetscDeviceInitTypeCast(initIdx);
5330e6b6b59SJacob Faibussowitsch   *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx);
5343ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
535030f984aSJacob Faibussowitsch }
536030f984aSJacob Faibussowitsch 
537030f984aSJacob Faibussowitsch /* called from PetscFinalize() do not call yourself! */
538d71ae5a4SJacob Faibussowitsch static PetscErrorCode PetscDeviceFinalize_Private()
539d71ae5a4SJacob Faibussowitsch {
540030f984aSJacob Faibussowitsch   PetscFunctionBegin;
541a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
542bf025ffbSJacob Faibussowitsch     const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] {
543a4af0ceeSJacob Faibussowitsch       PetscFunctionBegin;
5440e6b6b59SJacob Faibussowitsch       for (auto &&device : defaultDevices) {
5450e6b6b59SJacob Faibussowitsch         const auto dev = device.first;
5460e6b6b59SJacob Faibussowitsch 
5470e6b6b59SJacob 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);
5480e6b6b59SJacob Faibussowitsch       }
5493ba16761SJacob Faibussowitsch       PetscFunctionReturn(PETSC_SUCCESS);
550a4af0ceeSJacob Faibussowitsch     };
551bf025ffbSJacob Faibussowitsch     /*
552bf025ffbSJacob Faibussowitsch       you might be thinking, why on earth are you registered yet another finalizer in a
553bf025ffbSJacob Faibussowitsch       function already called during PetscRegisterFinalizeAll()? If this seems stupid it's
554bf025ffbSJacob Faibussowitsch       because it is.
555bf025ffbSJacob Faibussowitsch 
556bf025ffbSJacob Faibussowitsch       The crux of the problem is that the initializer (and therefore the ~finalizer~) of
557bf025ffbSJacob Faibussowitsch       PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had
558bf025ffbSJacob Faibussowitsch       a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence
559bf025ffbSJacob Faibussowitsch       won't be destroyed yet. So we need to repeat the check that all devices have been
560bf025ffbSJacob Faibussowitsch       destroyed again ~after~ the global context is destroyed. In summary:
561bf025ffbSJacob Faibussowitsch 
562bf025ffbSJacob Faibussowitsch       1. This finalizer runs and destroys all devices, except it may not because the global
563bf025ffbSJacob Faibussowitsch          context may still hold a reference!
564bf025ffbSJacob Faibussowitsch       2. The global context finalizer runs and does the final reference count decrement
565bf025ffbSJacob Faibussowitsch          required, which actually destroys the held device.
566bf025ffbSJacob Faibussowitsch       3. Our newly added finalizer runs and checks that all is well.
567a4af0ceeSJacob Faibussowitsch     */
5680e6b6b59SJacob Faibussowitsch     PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize)));
569a4af0ceeSJacob Faibussowitsch   }
5700e6b6b59SJacob Faibussowitsch   for (auto &&device : defaultDevices) {
5710e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceDestroy(&device.first));
5720e6b6b59SJacob Faibussowitsch     device.second = false;
5730e6b6b59SJacob Faibussowitsch   }
5743ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
575030f984aSJacob Faibussowitsch }
576030f984aSJacob Faibussowitsch 
577cf3a2253SJacob Faibussowitsch /*
578cf3a2253SJacob Faibussowitsch   Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of
579cf3a2253SJacob Faibussowitsch   initialization types:
580cf3a2253SJacob Faibussowitsch 
581a4af0ceeSJacob Faibussowitsch   1. defaultInitType - how does PetscDevice as a whole expect to initialize?
582a4af0ceeSJacob Faibussowitsch   2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize?
583a4af0ceeSJacob Faibussowitsch      e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but
584a4af0ceeSJacob Faibussowitsch      have all CUDA devices still initialize.
585a4af0ceeSJacob Faibussowitsch 
586a4af0ceeSJacob Faibussowitsch   All told the following happens:
587cf3a2253SJacob Faibussowitsch 
588a4af0ceeSJacob Faibussowitsch   0. defaultInitType -> LAZY
589a4af0ceeSJacob Faibussowitsch   1. Check for log_view/log_summary, if yes defaultInitType -> EAGER
590a4af0ceeSJacob Faibussowitsch   2. PetscDevice initializes each sub type with deviceDefaultInitType.
591a4af0ceeSJacob Faibussowitsch   2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition
592a4af0ceeSJacob Faibussowitsch       to checking for specific device init. if view or specific device init
593a4af0ceeSJacob Faibussowitsch       subTypeDefaultInitType -> EAGER. disabled once again overrides all.
594a4af0ceeSJacob Faibussowitsch */
5950e6b6b59SJacob Faibussowitsch 
596d71ae5a4SJacob Faibussowitsch PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm)
597d71ae5a4SJacob Faibussowitsch {
5987a101e5eSJacob Faibussowitsch   auto defaultView                    = PETSC_FALSE;
5997a101e5eSJacob Faibussowitsch   auto initializeDeviceContextEagerly = PETSC_FALSE;
6000e6b6b59SJacob Faibussowitsch   auto defaultDeviceSet               = PETSC_FALSE;
6017a101e5eSJacob Faibussowitsch   auto defaultDevice                  = PetscInt{PETSC_DECIDE};
6020e6b6b59SJacob Faibussowitsch   auto deviceContextInitDevice        = PETSC_DEVICE_DEFAULT();
6030e6b6b59SJacob Faibussowitsch   auto defaultInitType                = PETSC_DEVICE_INIT_LAZY;
604a4af0ceeSJacob Faibussowitsch 
605a4af0ceeSJacob Faibussowitsch   PetscFunctionBegin;
606a4af0ceeSJacob Faibussowitsch   if (PetscDefined(USE_DEBUG)) {
607a4af0ceeSJacob Faibussowitsch     int result;
608a4af0ceeSJacob Faibussowitsch 
6099566063dSJacob Faibussowitsch     PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result));
610a4af0ceeSJacob Faibussowitsch     /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the
611a4af0ceeSJacob Faibussowitsch      * global space */
612a4af0ceeSJacob Faibussowitsch     if (PetscUnlikely(result != MPI_IDENT)) {
613a4af0ceeSJacob Faibussowitsch       char name[MPI_MAX_OBJECT_NAME] = {};
614a4af0ceeSJacob Faibussowitsch       int  len; /* unused */
615a4af0ceeSJacob Faibussowitsch 
6169566063dSJacob Faibussowitsch       PetscCallMPI(MPI_Comm_get_name(comm, name, &len));
61798921bdaSJacob Faibussowitsch       SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name);
618a4af0ceeSJacob Faibussowitsch     }
619a4af0ceeSJacob Faibussowitsch   }
620a4af0ceeSJacob Faibussowitsch   comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */
6219566063dSJacob Faibussowitsch   PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private));
622a4af0ceeSJacob Faibussowitsch 
6230e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView));
6247a101e5eSJacob Faibussowitsch 
6250e6b6b59SJacob Faibussowitsch   // the precise values don't matter here, so long as they are sequential
626bd2fcf0cSJacob Faibussowitsch   static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HOST) == 0, "");
627bd2fcf0cSJacob Faibussowitsch   static_assert(Petsc::util::to_underlying(PETSC_DEVICE_CUDA) == 1, "");
628bd2fcf0cSJacob Faibussowitsch   static_assert(Petsc::util::to_underlying(PETSC_DEVICE_HIP) == 2, "");
629bd2fcf0cSJacob Faibussowitsch   static_assert(Petsc::util::to_underlying(PETSC_DEVICE_SYCL) == 3, "");
630bd2fcf0cSJacob Faibussowitsch   static_assert(Petsc::util::to_underlying(PETSC_DEVICE_MAX) == 4, "");
6310e6b6b59SJacob Faibussowitsch   for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) {
6320e6b6b59SJacob Faibussowitsch     const auto deviceType = PetscDeviceTypeCast(i);
633a4af0ceeSJacob Faibussowitsch     auto       initType   = defaultInitType;
634a4af0ceeSJacob Faibussowitsch 
6359566063dSJacob Faibussowitsch     PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType));
6360e6b6b59SJacob Faibussowitsch     if (PetscDeviceConfiguredFor_Internal(deviceType)) {
6370e6b6b59SJacob Faibussowitsch       if (initType == PETSC_DEVICE_INIT_EAGER) {
638a4af0ceeSJacob Faibussowitsch         initializeDeviceContextEagerly = PETSC_TRUE;
6390e6b6b59SJacob Faibussowitsch         // only update the default device if the user hasn't set it previously
6400e6b6b59SJacob Faibussowitsch         if (!defaultDeviceSet) {
641a4af0ceeSJacob Faibussowitsch           deviceContextInitDevice = deviceType;
6420e6b6b59SJacob Faibussowitsch           PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType]));
6430e6b6b59SJacob Faibussowitsch         }
6440e6b6b59SJacob Faibussowitsch       } else if (initType == PETSC_DEVICE_INIT_NONE) {
6451015a2a4SJacob 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]);
646a4af0ceeSJacob Faibussowitsch       }
647a4af0ceeSJacob Faibussowitsch     }
6480e6b6b59SJacob Faibussowitsch   }
6490e6b6b59SJacob Faibussowitsch 
6500e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice));
6510e6b6b59SJacob Faibussowitsch   PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT()));
6520e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
6530e6b6b59SJacob Faibussowitsch   /*                       PetscDevice is now fully initialized                          */
6540e6b6b59SJacob Faibussowitsch   /* ----------------------------------------------------------------------------------- */
6550e6b6b59SJacob Faibussowitsch   {
6560e6b6b59SJacob Faibussowitsch     /*
6570e6b6b59SJacob Faibussowitsch       query the options db to get the root settings from the user (if any).
6580e6b6b59SJacob Faibussowitsch 
6590e6b6b59SJacob Faibussowitsch       This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call
6600e6b6b59SJacob Faibussowitsch       PetscDeviceContextSetFromOptions() before we even have one, then set a few static
6610e6b6b59SJacob Faibussowitsch       variables in that file with the results.
6620e6b6b59SJacob Faibussowitsch     */
6630e6b6b59SJacob Faibussowitsch     auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE);
6640e6b6b59SJacob Faibussowitsch     auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE);
6650e6b6b59SJacob Faibussowitsch 
6660e6b6b59SJacob Faibussowitsch     PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys");
6670e6b6b59SJacob Faibussowitsch     PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype));
6680e6b6b59SJacob Faibussowitsch     PetscOptionsEnd();
6690e6b6b59SJacob Faibussowitsch 
6700e6b6b59SJacob Faibussowitsch     if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first));
6710e6b6b59SJacob Faibussowitsch     if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first));
6720e6b6b59SJacob Faibussowitsch   }
6730e6b6b59SJacob Faibussowitsch 
674a4af0ceeSJacob Faibussowitsch   if (initializeDeviceContextEagerly) {
675a4af0ceeSJacob Faibussowitsch     PetscDeviceContext dctx;
676a4af0ceeSJacob Faibussowitsch 
6770e6b6b59SJacob Faibussowitsch     PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice]));
6780e6b6b59SJacob Faibussowitsch     /* instantiates the device context */
6799566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextGetCurrentContext(&dctx));
6809566063dSJacob Faibussowitsch     PetscCall(PetscDeviceContextSetUp(dctx));
681a4af0ceeSJacob Faibussowitsch   }
6823ba16761SJacob Faibussowitsch   PetscFunctionReturn(PETSC_SUCCESS);
683a4af0ceeSJacob Faibussowitsch }
684