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