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 @*/ 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()); 930e6b6b59SJacob Faibussowitsch PetscCall(PetscNew(device)); 940e6b6b59SJacob Faibussowitsch (*device)->id = PetscDeviceCounter++; 950e6b6b59SJacob Faibussowitsch (*device)->type = type; 960e6b6b59SJacob 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) { 1020e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid); 1030e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid); 1040e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid); 1050e6b6b59SJacob 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 1170e6b6b59SJacob Faibussowitsch Not Collective 118030f984aSJacob Faibussowitsch 119030f984aSJacob Faibussowitsch Input Parameter: 1200e6b6b59SJacob Faibussowitsch . device - The `PetscDevice` 121030f984aSJacob Faibussowitsch 122030f984aSJacob Faibussowitsch Level: beginner 123030f984aSJacob Faibussowitsch 1240e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`, 1250e6b6b59SJacob Faibussowitsch `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()` 126030f984aSJacob Faibussowitsch @*/ 1279371c9d4SSatish Balay PetscErrorCode PetscDeviceDestroy(PetscDevice *device) { 128a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 1290e6b6b59SJacob 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) { 1340e6b6b59SJacob 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 1450e6b6b59SJacob Faibussowitsch Not Collective 146a4af0ceeSJacob Faibussowitsch 147a4af0ceeSJacob Faibussowitsch Input Parameter: 148811af0c4SBarry Smith . device - The `PetscDevice` to configure 149a4af0ceeSJacob Faibussowitsch 1500e6b6b59SJacob Faibussowitsch Notes: 1510e6b6b59SJacob Faibussowitsch The user should not assume that this is a cheap operation. 152a4af0ceeSJacob Faibussowitsch 153a4af0ceeSJacob Faibussowitsch Level: beginner 154a4af0ceeSJacob Faibussowitsch 1550e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`, 1560e6b6b59SJacob 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 */ 1650e6b6b59SJacob Faibussowitsch switch (const auto dtype = device->type) { 1660e6b6b59SJacob Faibussowitsch case PETSC_DEVICE_HOST: 1670e6b6b59SJacob Faibussowitsch if (PetscDefined(HAVE_HOST)) break; // always true 1689371c9d4SSatish Balay case PETSC_DEVICE_CUDA: 1699371c9d4SSatish Balay if (PetscDefined(HAVE_CUDA)) break; 1700e6b6b59SJacob Faibussowitsch goto error; 1719371c9d4SSatish Balay case PETSC_DEVICE_HIP: 1729371c9d4SSatish Balay if (PetscDefined(HAVE_HIP)) break; 1730e6b6b59SJacob Faibussowitsch goto error; 1749371c9d4SSatish Balay case PETSC_DEVICE_SYCL: 1759371c9d4SSatish Balay if (PetscDefined(HAVE_SYCL)) break; 176*f4d061e9SPierre Jolivet goto error; 1770e6b6b59SJacob Faibussowitsch default: 1780e6b6b59SJacob Faibussowitsch error: 1790e6b6b59SJacob Faibussowitsch SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]); 180a4af0ceeSJacob Faibussowitsch } 181dbbe0bcdSBarry Smith PetscUseTypeMethod(device, configure); 182a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 183a4af0ceeSJacob Faibussowitsch } 184a4af0ceeSJacob Faibussowitsch 185a4af0ceeSJacob Faibussowitsch /*@C 186811af0c4SBarry Smith PetscDeviceView - View a `PetscDevice` 187a4af0ceeSJacob Faibussowitsch 1880e6b6b59SJacob Faibussowitsch Collective on viewer 189a4af0ceeSJacob Faibussowitsch 19091e63d38SStefano Zampini Input Parameters: 191811af0c4SBarry Smith + device - The `PetscDevice` to view 1920e6b6b59SJacob Faibussowitsch - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`) 193a4af0ceeSJacob Faibussowitsch 194a4af0ceeSJacob Faibussowitsch Level: beginner 195a4af0ceeSJacob Faibussowitsch 1960e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, 1970e6b6b59SJacob Faibussowitsch `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()` 198a4af0ceeSJacob Faibussowitsch @*/ 1999371c9d4SSatish Balay PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer) { 2000e6b6b59SJacob Faibussowitsch auto sub = viewer; 2010e6b6b59SJacob Faibussowitsch PetscBool iascii; 2020e6b6b59SJacob Faibussowitsch 203a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 204a4af0ceeSJacob Faibussowitsch PetscValidDevice(device, 1); 2050e6b6b59SJacob Faibussowitsch if (viewer) { 206a4af0ceeSJacob Faibussowitsch PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 2070e6b6b59SJacob Faibussowitsch PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii)); 2080e6b6b59SJacob Faibussowitsch } else { 2090e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer)); 2100e6b6b59SJacob Faibussowitsch iascii = PETSC_TRUE; 2110e6b6b59SJacob Faibussowitsch } 2120e6b6b59SJacob Faibussowitsch 2130e6b6b59SJacob Faibussowitsch if (iascii) { 2140e6b6b59SJacob Faibussowitsch auto dtype = PETSC_DEVICE_HOST; 2150e6b6b59SJacob Faibussowitsch MPI_Comm comm; 2160e6b6b59SJacob Faibussowitsch PetscMPIInt size; 2170e6b6b59SJacob Faibussowitsch PetscInt id = 0; 2180e6b6b59SJacob Faibussowitsch 2190e6b6b59SJacob Faibussowitsch PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm)); 2200e6b6b59SJacob Faibussowitsch PetscCallMPI(MPI_Comm_size(comm, &size)); 2210e6b6b59SJacob Faibussowitsch 2220e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceGetDeviceId(device, &id)); 2230e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceGetType(device, &dtype)); 2240e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub)); 2250e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes")); 2260e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerASCIIPushTab(sub)); 2270e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype])); 2280e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id)); 2290e6b6b59SJacob Faibussowitsch } 2300e6b6b59SJacob Faibussowitsch 2310e6b6b59SJacob Faibussowitsch // see if impls has extra viewer stuff 2320e6b6b59SJacob Faibussowitsch PetscTryTypeMethod(device, view, sub); 2330e6b6b59SJacob Faibussowitsch 2340e6b6b59SJacob Faibussowitsch if (iascii) { 2350e6b6b59SJacob Faibussowitsch // undo the ASCII specific stuff 2360e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerASCIIPopTab(sub)); 2370e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub)); 2380e6b6b59SJacob Faibussowitsch PetscCall(PetscViewerFlush(viewer)); 2390e6b6b59SJacob Faibussowitsch } 240a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 241a4af0ceeSJacob Faibussowitsch } 242a4af0ceeSJacob Faibussowitsch 24391e63d38SStefano Zampini /*@C 2440e6b6b59SJacob Faibussowitsch PetscDeviceGetType - Get the type of device 24591e63d38SStefano Zampini 2460e6b6b59SJacob Faibussowitsch Not Collective 24791e63d38SStefano Zampini 24891e63d38SStefano Zampini Input Parameter: 249811af0c4SBarry Smith . device - The `PetscDevice` 25091e63d38SStefano Zampini 25191e63d38SStefano Zampini Output Parameter: 2520e6b6b59SJacob Faibussowitsch . type - The `PetscDeviceType` 25391e63d38SStefano Zampini 25491e63d38SStefano Zampini Level: beginner 25591e63d38SStefano Zampini 2560e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, 2570e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`, 2580e6b6b59SJacob Faibussowitsch `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()` 2590e6b6b59SJacob Faibussowitsch @*/ 2600e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type) { 2610e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 2620e6b6b59SJacob Faibussowitsch PetscValidDevice(device, 1); 2630e6b6b59SJacob Faibussowitsch PetscValidPointer(type, 2); 2640e6b6b59SJacob Faibussowitsch *type = device->type; 2650e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 2660e6b6b59SJacob Faibussowitsch } 2670e6b6b59SJacob Faibussowitsch 2680e6b6b59SJacob Faibussowitsch /*@C 2690e6b6b59SJacob Faibussowitsch PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice` 2700e6b6b59SJacob Faibussowitsch 2710e6b6b59SJacob Faibussowitsch Not Collective 2720e6b6b59SJacob Faibussowitsch 2730e6b6b59SJacob Faibussowitsch Input Parameter: 2740e6b6b59SJacob Faibussowitsch . device - The `PetscDevice` 2750e6b6b59SJacob Faibussowitsch 2760e6b6b59SJacob Faibussowitsch Output Parameter: 2770e6b6b59SJacob Faibussowitsch . id - The id 2780e6b6b59SJacob Faibussowitsch 2790e6b6b59SJacob Faibussowitsch Notes: 2800e6b6b59SJacob Faibussowitsch The returned ID may have been assigned by the underlying device backend. For example if the 2810e6b6b59SJacob Faibussowitsch backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when 2820e6b6b59SJacob Faibussowitsch this device was configured. 2830e6b6b59SJacob Faibussowitsch 2840e6b6b59SJacob Faibussowitsch Level: beginner 2850e6b6b59SJacob Faibussowitsch 2860e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()` 28791e63d38SStefano Zampini @*/ 2889371c9d4SSatish Balay PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id) { 28991e63d38SStefano Zampini PetscFunctionBegin; 29091e63d38SStefano Zampini PetscValidDevice(device, 1); 29191e63d38SStefano Zampini PetscValidIntPointer(id, 2); 29291e63d38SStefano Zampini *id = device->deviceId; 29391e63d38SStefano Zampini PetscFunctionReturn(0); 29491e63d38SStefano Zampini } 29591e63d38SStefano Zampini 2960e6b6b59SJacob Faibussowitsch struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> { 2970e6b6b59SJacob Faibussowitsch PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE; 2980e6b6b59SJacob Faibussowitsch 2990e6b6b59SJacob Faibussowitsch PETSC_NODISCARD PetscErrorCode finalize_() noexcept { 3000e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 3010e6b6b59SJacob Faibussowitsch type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE; 3020e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 3030e6b6b59SJacob Faibussowitsch } 3040e6b6b59SJacob Faibussowitsch }; 3050e6b6b59SJacob Faibussowitsch 3060e6b6b59SJacob Faibussowitsch static auto default_device_type = DefaultDeviceType(); 3070e6b6b59SJacob Faibussowitsch 3080e6b6b59SJacob Faibussowitsch /*@C 3090e6b6b59SJacob Faibussowitsch PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType` 3100e6b6b59SJacob Faibussowitsch 3110e6b6b59SJacob Faibussowitsch Not Collective 3120e6b6b59SJacob Faibussowitsch 3130e6b6b59SJacob Faibussowitsch Notes: 3140e6b6b59SJacob Faibussowitsch Unless selected by the user, the default device is selected in the following order\: 3150e6b6b59SJacob Faibussowitsch `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`. 3160e6b6b59SJacob Faibussowitsch 3170e6b6b59SJacob Faibussowitsch Level: beginner 3180e6b6b59SJacob Faibussowitsch 3190e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()` 3200e6b6b59SJacob Faibussowitsch @*/ 3210e6b6b59SJacob Faibussowitsch PetscDeviceType PETSC_DEVICE_DEFAULT(void) { 3220e6b6b59SJacob Faibussowitsch return default_device_type.type; 3230e6b6b59SJacob Faibussowitsch } 3240e6b6b59SJacob Faibussowitsch 3250e6b6b59SJacob Faibussowitsch /*@C 3260e6b6b59SJacob Faibussowitsch PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice` 3270e6b6b59SJacob Faibussowitsch 3280e6b6b59SJacob Faibussowitsch Not Collective 3290e6b6b59SJacob Faibussowitsch 3300e6b6b59SJacob Faibussowitsch Input Parameter: 3310e6b6b59SJacob Faibussowitsch . type - the new default device type 3320e6b6b59SJacob Faibussowitsch 3330e6b6b59SJacob Faibussowitsch Notes: 3340e6b6b59SJacob Faibussowitsch This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`. 3350e6b6b59SJacob Faibussowitsch 3360e6b6b59SJacob Faibussowitsch Level: beginner 3370e6b6b59SJacob Faibussowitsch 3380e6b6b59SJacob Faibussowitsch .seealso: `PetscDeviceType`, `PetscDeviceGetType`, 3390e6b6b59SJacob Faibussowitsch @*/ 3400e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type) { 3410e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 3420e6b6b59SJacob Faibussowitsch PetscValidDeviceType(type, 1); 3430e6b6b59SJacob Faibussowitsch if (default_device_type.type != type) { 3440e6b6b59SJacob Faibussowitsch // no need to waster a PetscRegisterFinalize() slot if we don't change it 3450e6b6b59SJacob Faibussowitsch default_device_type.type = type; 3460e6b6b59SJacob Faibussowitsch PetscCall(default_device_type.register_finalize()); 3470e6b6b59SJacob Faibussowitsch } 3480e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 3490e6b6b59SJacob Faibussowitsch } 3500e6b6b59SJacob Faibussowitsch 3510e6b6b59SJacob Faibussowitsch static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {}; 3520e6b6b59SJacob Faibussowitsch 3530e6b6b59SJacob Faibussowitsch /* 3540e6b6b59SJacob Faibussowitsch Actual intialization function; any functions claiming to initialize PetscDevice or 3550e6b6b59SJacob Faibussowitsch PetscDeviceContext will have to run through this one 3560e6b6b59SJacob Faibussowitsch */ 3570e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId) { 3580e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 3590e6b6b59SJacob Faibussowitsch PetscValidDeviceType(type, 1); 3600e6b6b59SJacob Faibussowitsch if (PetscUnlikely(!PetscDeviceInitialized(type))) { 3610e6b6b59SJacob Faibussowitsch auto &dev = defaultDevices[type].first; 3620e6b6b59SJacob Faibussowitsch auto &init = defaultDevices[type].second; 3630e6b6b59SJacob Faibussowitsch 3640e6b6b59SJacob Faibussowitsch PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]); 3650e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev)); 3660e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceConfigure(dev)); 3670e6b6b59SJacob Faibussowitsch init = true; 3680e6b6b59SJacob Faibussowitsch } 3690e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 3700e6b6b59SJacob Faibussowitsch } 371a4af0ceeSJacob Faibussowitsch 372a4af0ceeSJacob Faibussowitsch /*@C 373811af0c4SBarry Smith PetscDeviceInitialize - Initialize `PetscDevice` 374a4af0ceeSJacob Faibussowitsch 3750e6b6b59SJacob Faibussowitsch Not Collective 376a4af0ceeSJacob Faibussowitsch 377a4af0ceeSJacob Faibussowitsch Input Parameter: 378811af0c4SBarry Smith . type - The `PetscDeviceType` to initialize 379a4af0ceeSJacob Faibussowitsch 3800e6b6b59SJacob Faibussowitsch Notes: 3810e6b6b59SJacob Faibussowitsch Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may 3820e6b6b59SJacob Faibussowitsch result in device synchronization. 383a4af0ceeSJacob Faibussowitsch 384a4af0ceeSJacob Faibussowitsch Level: beginner 385a4af0ceeSJacob Faibussowitsch 3860e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`, 3870e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()` 388a4af0ceeSJacob Faibussowitsch @*/ 3899371c9d4SSatish Balay PetscErrorCode PetscDeviceInitialize(PetscDeviceType type) { 390a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 391a4af0ceeSJacob Faibussowitsch PetscValidDeviceType(type, 1); 3929566063dSJacob Faibussowitsch PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE)); 393a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 394a4af0ceeSJacob Faibussowitsch } 395a4af0ceeSJacob Faibussowitsch 396a4af0ceeSJacob Faibussowitsch /*@C 397811af0c4SBarry Smith PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular 398811af0c4SBarry Smith `PetscDeviceType` 399a4af0ceeSJacob Faibussowitsch 4000e6b6b59SJacob Faibussowitsch Not Collective 401a4af0ceeSJacob Faibussowitsch 402a4af0ceeSJacob Faibussowitsch Input Parameter: 403811af0c4SBarry Smith . type - The `PetscDeviceType` to check 404a4af0ceeSJacob Faibussowitsch 4050e6b6b59SJacob Faibussowitsch Notes: 4060e6b6b59SJacob Faibussowitsch Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise. 407a4af0ceeSJacob Faibussowitsch 408811af0c4SBarry Smith If one has not configured PETSc for a particular `PetscDeviceType` then this routine will 409811af0c4SBarry Smith return `PETSC_FALSE` for that `PetscDeviceType`. 410a4af0ceeSJacob Faibussowitsch 411a4af0ceeSJacob Faibussowitsch Level: beginner 412a4af0ceeSJacob Faibussowitsch 4130e6b6b59SJacob Faibussowitsch .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`, 4140e6b6b59SJacob Faibussowitsch `PetscDeviceCreate()`, `PetscDeviceDestroy()` 415a4af0ceeSJacob Faibussowitsch @*/ 4169371c9d4SSatish Balay PetscBool PetscDeviceInitialized(PetscDeviceType type) { 4170e6b6b59SJacob Faibussowitsch return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second); 4180e6b6b59SJacob Faibussowitsch } 4190e6b6b59SJacob Faibussowitsch 4200e6b6b59SJacob Faibussowitsch /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */ 4210e6b6b59SJacob Faibussowitsch PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device) { 4220e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4230e6b6b59SJacob Faibussowitsch PetscValidPointer(device, 2); 4240e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceInitialize(type)); 4250e6b6b59SJacob Faibussowitsch *device = defaultDevices[type].first; 4260e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 427a4af0ceeSJacob Faibussowitsch } 428a4af0ceeSJacob Faibussowitsch 429a16fd2c9SJacob Faibussowitsch /*@C 430a16fd2c9SJacob Faibussowitsch PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice` 431a16fd2c9SJacob Faibussowitsch 4320e6b6b59SJacob Faibussowitsch Not Collective 433a16fd2c9SJacob Faibussowitsch 434a16fd2c9SJacob Faibussowitsch Input Parameters: 435a16fd2c9SJacob Faibussowitsch + device - The `PetscDevice` 436a16fd2c9SJacob Faibussowitsch - attr - The attribute 437a16fd2c9SJacob Faibussowitsch 438a16fd2c9SJacob Faibussowitsch Output Parameter: 439a16fd2c9SJacob Faibussowitsch . value - The value of the attribute 440a16fd2c9SJacob Faibussowitsch 441a16fd2c9SJacob Faibussowitsch Notes: 442a16fd2c9SJacob Faibussowitsch Since different attributes are often different types `value` is a `void *` to accommodate 443a16fd2c9SJacob Faibussowitsch them all. The underlying type of the attribute is therefore included in the name of the 444a16fd2c9SJacob Faibussowitsch `PetscDeviceAttribute` reponsible for querying it. For example, 445a16fd2c9SJacob Faibussowitsch `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`. 446a16fd2c9SJacob Faibussowitsch 4470e6b6b59SJacob Faibussowitsch Level: intermediate 4480e6b6b59SJacob Faibussowitsch 449a16fd2c9SJacob Faibussowitsch .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice` 450a16fd2c9SJacob Faibussowitsch @*/ 451a16fd2c9SJacob Faibussowitsch PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value) { 452a16fd2c9SJacob Faibussowitsch PetscFunctionBegin; 453a16fd2c9SJacob Faibussowitsch PetscValidDevice(device, 1); 454a16fd2c9SJacob Faibussowitsch PetscValidDeviceAttribute(attr, 2); 455a16fd2c9SJacob Faibussowitsch PetscValidPointer(value, 3); 456a16fd2c9SJacob Faibussowitsch PetscUseTypeMethod(device, getattribute, attr, value); 457a16fd2c9SJacob Faibussowitsch PetscFunctionReturn(0); 458a16fd2c9SJacob Faibussowitsch } 459a16fd2c9SJacob Faibussowitsch 4609371c9d4SSatish Balay static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType) { 461a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 462a4af0ceeSJacob Faibussowitsch if (!PetscDeviceConfiguredFor_Internal(type)) { 4630e6b6b59SJacob Faibussowitsch PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type])); 4640e6b6b59SJacob Faibussowitsch defaultDevices[type].first = nullptr; 465a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 466a4af0ceeSJacob Faibussowitsch } 4670e6b6b59SJacob Faibussowitsch PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type])); 468a4af0ceeSJacob Faibussowitsch /* ugly switch needed to pick the right global variable... could maybe do this as a union? */ 469a4af0ceeSJacob Faibussowitsch switch (type) { 4700e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 4710e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 4720e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 4730e6b6b59SJacob Faibussowitsch PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 4749371c9d4SSatish 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]); 475a4af0ceeSJacob Faibussowitsch } 4760e6b6b59SJacob 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)])); 477cf3a2253SJacob Faibussowitsch /* 4780e6b6b59SJacob Faibussowitsch defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES 4790e6b6b59SJacob Faibussowitsch have decided to initialize as 480cf3a2253SJacob Faibussowitsch */ 4810e6b6b59SJacob Faibussowitsch if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) { 4820e6b6b59SJacob Faibussowitsch PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type])); 4839566063dSJacob Faibussowitsch PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId)); 4840e6b6b59SJacob Faibussowitsch if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr)); 4850e6b6b59SJacob Faibussowitsch } 4860e6b6b59SJacob Faibussowitsch PetscFunctionReturn(0); 4870e6b6b59SJacob Faibussowitsch } 488a4af0ceeSJacob Faibussowitsch 4890e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView) { 4900e6b6b59SJacob Faibussowitsch PetscInt initIdx = PETSC_DEVICE_INIT_LAZY; 4910e6b6b59SJacob Faibussowitsch auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice); 4920e6b6b59SJacob Faibussowitsch auto flg = PETSC_FALSE; 4930e6b6b59SJacob Faibussowitsch 4940e6b6b59SJacob Faibussowitsch PetscFunctionBegin; 4950e6b6b59SJacob Faibussowitsch PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg)); 4960e6b6b59SJacob Faibussowitsch if (flg) PetscCall(PetscLogGpuTime()); 4970e6b6b59SJacob Faibussowitsch 4980e6b6b59SJacob Faibussowitsch PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys"); 4990e6b6b59SJacob Faibussowitsch PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr)); 5000e6b6b59SJacob Faibussowitsch PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet)); 5010e6b6b59SJacob 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)); 5020e6b6b59SJacob Faibussowitsch PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg)); 5030e6b6b59SJacob Faibussowitsch PetscOptionsEnd(); 5040e6b6b59SJacob Faibussowitsch 5050e6b6b59SJacob Faibussowitsch if (initIdx == PETSC_DEVICE_INIT_NONE) { 5060e6b6b59SJacob Faibussowitsch /* disabled all device initialization if devices are globally disabled */ 5070e6b6b59SJacob 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"); 5080e6b6b59SJacob Faibussowitsch *defaultView = PETSC_FALSE; 5090e6b6b59SJacob Faibussowitsch initDeviceIdx = PETSC_DEVICE_HOST; 5100e6b6b59SJacob Faibussowitsch } else { 5110e6b6b59SJacob Faibussowitsch *defaultView = static_cast<PetscBool>(*defaultView && flg); 5120e6b6b59SJacob Faibussowitsch if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER; 513a4af0ceeSJacob Faibussowitsch } 5140e6b6b59SJacob Faibussowitsch *defaultInitType = PetscDeviceInitTypeCast(initIdx); 5150e6b6b59SJacob Faibussowitsch *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx); 516030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 517030f984aSJacob Faibussowitsch } 518030f984aSJacob Faibussowitsch 519030f984aSJacob Faibussowitsch /* called from PetscFinalize() do not call yourself! */ 5200e6b6b59SJacob Faibussowitsch static PetscErrorCode PetscDeviceFinalize_Private() { 521030f984aSJacob Faibussowitsch PetscFunctionBegin; 522a4af0ceeSJacob Faibussowitsch if (PetscDefined(USE_DEBUG)) { 523bf025ffbSJacob Faibussowitsch const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] { 524a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 5250e6b6b59SJacob Faibussowitsch for (auto &&device : defaultDevices) { 5260e6b6b59SJacob Faibussowitsch const auto dev = device.first; 5270e6b6b59SJacob Faibussowitsch 5280e6b6b59SJacob 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); 5290e6b6b59SJacob Faibussowitsch } 530a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 531a4af0ceeSJacob Faibussowitsch }; 532bf025ffbSJacob Faibussowitsch /* 533bf025ffbSJacob Faibussowitsch you might be thinking, why on earth are you registered yet another finalizer in a 534bf025ffbSJacob Faibussowitsch function already called during PetscRegisterFinalizeAll()? If this seems stupid it's 535bf025ffbSJacob Faibussowitsch because it is. 536bf025ffbSJacob Faibussowitsch 537bf025ffbSJacob Faibussowitsch The crux of the problem is that the initializer (and therefore the ~finalizer~) of 538bf025ffbSJacob Faibussowitsch PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had 539bf025ffbSJacob Faibussowitsch a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence 540bf025ffbSJacob Faibussowitsch won't be destroyed yet. So we need to repeat the check that all devices have been 541bf025ffbSJacob Faibussowitsch destroyed again ~after~ the global context is destroyed. In summary: 542bf025ffbSJacob Faibussowitsch 543bf025ffbSJacob Faibussowitsch 1. This finalizer runs and destroys all devices, except it may not because the global 544bf025ffbSJacob Faibussowitsch context may still hold a reference! 545bf025ffbSJacob Faibussowitsch 2. The global context finalizer runs and does the final reference count decrement 546bf025ffbSJacob Faibussowitsch required, which actually destroys the held device. 547bf025ffbSJacob Faibussowitsch 3. Our newly added finalizer runs and checks that all is well. 548a4af0ceeSJacob Faibussowitsch */ 5490e6b6b59SJacob Faibussowitsch PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize))); 550a4af0ceeSJacob Faibussowitsch } 5510e6b6b59SJacob Faibussowitsch for (auto &&device : defaultDevices) { 5520e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceDestroy(&device.first)); 5530e6b6b59SJacob Faibussowitsch device.second = false; 5540e6b6b59SJacob Faibussowitsch } 555030f984aSJacob Faibussowitsch PetscFunctionReturn(0); 556030f984aSJacob Faibussowitsch } 557030f984aSJacob Faibussowitsch 558cf3a2253SJacob Faibussowitsch /* 559cf3a2253SJacob Faibussowitsch Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of 560cf3a2253SJacob Faibussowitsch initialization types: 561cf3a2253SJacob Faibussowitsch 562a4af0ceeSJacob Faibussowitsch 1. defaultInitType - how does PetscDevice as a whole expect to initialize? 563a4af0ceeSJacob Faibussowitsch 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize? 564a4af0ceeSJacob Faibussowitsch e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but 565a4af0ceeSJacob Faibussowitsch have all CUDA devices still initialize. 566a4af0ceeSJacob Faibussowitsch 567a4af0ceeSJacob Faibussowitsch All told the following happens: 568cf3a2253SJacob Faibussowitsch 569a4af0ceeSJacob Faibussowitsch 0. defaultInitType -> LAZY 570a4af0ceeSJacob Faibussowitsch 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER 571a4af0ceeSJacob Faibussowitsch 2. PetscDevice initializes each sub type with deviceDefaultInitType. 572a4af0ceeSJacob Faibussowitsch 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition 573a4af0ceeSJacob Faibussowitsch to checking for specific device init. if view or specific device init 574a4af0ceeSJacob Faibussowitsch subTypeDefaultInitType -> EAGER. disabled once again overrides all. 575a4af0ceeSJacob Faibussowitsch */ 5760e6b6b59SJacob Faibussowitsch 5779371c9d4SSatish Balay PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm) { 5787a101e5eSJacob Faibussowitsch auto defaultView = PETSC_FALSE; 5797a101e5eSJacob Faibussowitsch auto initializeDeviceContextEagerly = PETSC_FALSE; 5800e6b6b59SJacob Faibussowitsch auto defaultDeviceSet = PETSC_FALSE; 5817a101e5eSJacob Faibussowitsch auto defaultDevice = PetscInt{PETSC_DECIDE}; 5820e6b6b59SJacob Faibussowitsch auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT(); 5830e6b6b59SJacob Faibussowitsch auto defaultInitType = PETSC_DEVICE_INIT_LAZY; 584a4af0ceeSJacob Faibussowitsch 585a4af0ceeSJacob Faibussowitsch PetscFunctionBegin; 586a4af0ceeSJacob Faibussowitsch if (PetscDefined(USE_DEBUG)) { 587a4af0ceeSJacob Faibussowitsch int result; 588a4af0ceeSJacob Faibussowitsch 5899566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result)); 590a4af0ceeSJacob Faibussowitsch /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the 591a4af0ceeSJacob Faibussowitsch * global space */ 592a4af0ceeSJacob Faibussowitsch if (PetscUnlikely(result != MPI_IDENT)) { 593a4af0ceeSJacob Faibussowitsch char name[MPI_MAX_OBJECT_NAME] = {}; 594a4af0ceeSJacob Faibussowitsch int len; /* unused */ 595a4af0ceeSJacob Faibussowitsch 5969566063dSJacob Faibussowitsch PetscCallMPI(MPI_Comm_get_name(comm, name, &len)); 59798921bdaSJacob Faibussowitsch SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name); 598a4af0ceeSJacob Faibussowitsch } 599a4af0ceeSJacob Faibussowitsch } 600a4af0ceeSJacob Faibussowitsch comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */ 6019566063dSJacob Faibussowitsch PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private)); 602a4af0ceeSJacob Faibussowitsch 6030e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView)); 6047a101e5eSJacob Faibussowitsch 6050e6b6b59SJacob Faibussowitsch // the precise values don't matter here, so long as they are sequential 6060e6b6b59SJacob Faibussowitsch static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, ""); 6070e6b6b59SJacob Faibussowitsch static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, ""); 6080e6b6b59SJacob Faibussowitsch static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, ""); 6090e6b6b59SJacob Faibussowitsch static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, ""); 6100e6b6b59SJacob Faibussowitsch static_assert(Petsc::util::integral_value(PETSC_DEVICE_MAX) == 4, ""); 6110e6b6b59SJacob Faibussowitsch for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) { 6120e6b6b59SJacob Faibussowitsch const auto deviceType = PetscDeviceTypeCast(i); 613a4af0ceeSJacob Faibussowitsch auto initType = defaultInitType; 614a4af0ceeSJacob Faibussowitsch 6159566063dSJacob Faibussowitsch PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType)); 6160e6b6b59SJacob Faibussowitsch if (PetscDeviceConfiguredFor_Internal(deviceType)) { 6170e6b6b59SJacob Faibussowitsch if (initType == PETSC_DEVICE_INIT_EAGER) { 618a4af0ceeSJacob Faibussowitsch initializeDeviceContextEagerly = PETSC_TRUE; 6190e6b6b59SJacob Faibussowitsch // only update the default device if the user hasn't set it previously 6200e6b6b59SJacob Faibussowitsch if (!defaultDeviceSet) { 621a4af0ceeSJacob Faibussowitsch deviceContextInitDevice = deviceType; 6220e6b6b59SJacob Faibussowitsch PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType])); 6230e6b6b59SJacob Faibussowitsch } 6240e6b6b59SJacob Faibussowitsch } else if (initType == PETSC_DEVICE_INIT_NONE) { 6250e6b6b59SJacob 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]); 626a4af0ceeSJacob Faibussowitsch } 627a4af0ceeSJacob Faibussowitsch } 6280e6b6b59SJacob Faibussowitsch } 6290e6b6b59SJacob Faibussowitsch 6300e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice)); 6310e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT())); 6320e6b6b59SJacob Faibussowitsch /* ----------------------------------------------------------------------------------- */ 6330e6b6b59SJacob Faibussowitsch /* PetscDevice is now fully initialized */ 6340e6b6b59SJacob Faibussowitsch /* ----------------------------------------------------------------------------------- */ 6350e6b6b59SJacob Faibussowitsch { 6360e6b6b59SJacob Faibussowitsch /* 6370e6b6b59SJacob Faibussowitsch query the options db to get the root settings from the user (if any). 6380e6b6b59SJacob Faibussowitsch 6390e6b6b59SJacob Faibussowitsch This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call 6400e6b6b59SJacob Faibussowitsch PetscDeviceContextSetFromOptions() before we even have one, then set a few static 6410e6b6b59SJacob Faibussowitsch variables in that file with the results. 6420e6b6b59SJacob Faibussowitsch */ 6430e6b6b59SJacob Faibussowitsch auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE); 6440e6b6b59SJacob Faibussowitsch auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE); 6450e6b6b59SJacob Faibussowitsch 6460e6b6b59SJacob Faibussowitsch PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys"); 6470e6b6b59SJacob Faibussowitsch PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype)); 6480e6b6b59SJacob Faibussowitsch PetscOptionsEnd(); 6490e6b6b59SJacob Faibussowitsch 6500e6b6b59SJacob Faibussowitsch if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first)); 6510e6b6b59SJacob Faibussowitsch if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first)); 6520e6b6b59SJacob Faibussowitsch } 6530e6b6b59SJacob Faibussowitsch 654a4af0ceeSJacob Faibussowitsch if (initializeDeviceContextEagerly) { 655a4af0ceeSJacob Faibussowitsch PetscDeviceContext dctx; 656a4af0ceeSJacob Faibussowitsch 6570e6b6b59SJacob Faibussowitsch PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice])); 6580e6b6b59SJacob Faibussowitsch /* instantiates the device context */ 6599566063dSJacob Faibussowitsch PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 6609566063dSJacob Faibussowitsch PetscCall(PetscDeviceContextSetUp(dctx)); 661a4af0ceeSJacob Faibussowitsch } 662a4af0ceeSJacob Faibussowitsch PetscFunctionReturn(0); 663a4af0ceeSJacob Faibussowitsch } 664