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