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