1 #include "petscdevice_interface_internal.hpp" /*I <petscdevice.h> I*/ 2 #include <petsc/private/petscadvancedmacros.h> 3 4 #include "../impls/host/hostdevice.hpp" 5 #include "../impls/cupm/cupmdevice.hpp" 6 #include "../impls/sycl/sycldevice.hpp" 7 8 #include <limits> // std::numeric_limits 9 #include <utility> // std::make_pair 10 11 using namespace Petsc::device; 12 13 /* 14 note to anyone adding more classes, the name must be ALL_CAPS_SHORT_NAME + Device exactly to 15 be picked up by the switch-case macros below 16 */ 17 static host::Device HOSTDevice{PetscDeviceContextCreate_HOST}; 18 #if PetscDefined(HAVE_CUDA) 19 static cupm::Device<cupm::DeviceType::CUDA> CUDADevice{PetscDeviceContextCreate_CUDA}; 20 #endif 21 #if PetscDefined(HAVE_HIP) 22 static cupm::Device<cupm::DeviceType::HIP> HIPDevice{PetscDeviceContextCreate_HIP}; 23 #endif 24 #if PetscDefined(HAVE_SYCL) 25 static sycl::Device SYCLDevice{PetscDeviceContextCreate_SYCL}; 26 #endif 27 28 #define PETSC_DEVICE_CASE(IMPLS, func, ...) \ 29 case PetscConcat_(PETSC_DEVICE_, IMPLS): { \ 30 PetscCall(PetscConcat_(IMPLS, Device).func(__VA_ARGS__)); \ 31 } break 32 33 /* 34 Suppose you have: 35 36 CUDADevice.myFunction(arg1,arg2) 37 38 that you would like to conditionally define and call in a switch-case: 39 40 switch(PetscDeviceType) { 41 #if PetscDefined(HAVE_CUDA) 42 case PETSC_DEVICE_CUDA: { 43 PetscCall(CUDADevice.myFunction(arg1,arg2)); 44 } break; 45 #endif 46 } 47 48 then calling this macro: 49 50 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA,myFunction,arg1,arg2) 51 52 will expand to the following case statement: 53 54 case PETSC_DEVICE_CUDA: { 55 PetscCall(CUDADevice.myFunction(arg1,arg2)); 56 } break 57 58 if PetscDefined(HAVE_CUDA) evaluates to 1, and expand to nothing otherwise 59 */ 60 #define PETSC_DEVICE_CASE_IF_PETSC_DEFINED(IMPLS, func, ...) PetscIfPetscDefined(PetscConcat_(HAVE_, IMPLS), PETSC_DEVICE_CASE, PetscExpandToNothing)(IMPLS, func, __VA_ARGS__) 61 62 /*@C 63 PetscDeviceCreate - Get a new handle for a particular device (often a GPU) type 64 65 Not Collective 66 67 Input Parameters: 68 + type - The type of `PetscDevice` 69 - devid - The numeric ID# of the device (pass `PETSC_DECIDE` to assign automatically) 70 71 Output Parameter: 72 . device - The `PetscDevice` 73 74 Notes: 75 This routine may initialize `PetscDevice`. If this is the case, it may cause some sort of 76 device synchronization. 77 78 `devid` is what you might pass to `cudaSetDevice()` for example. 79 80 Level: beginner 81 82 .seealso: `PetscDevice`, `PetscDeviceInitType`, 83 `PetscDeviceInitialize()`,`PetscDeviceInitialized()`, `PetscDeviceConfigure()`, 84 `PetscDeviceView()`, `PetscDeviceDestroy()` 85 @*/ 86 PetscErrorCode PetscDeviceCreate(PetscDeviceType type, PetscInt devid, PetscDevice *device) { 87 static PetscInt PetscDeviceCounter = 0; 88 89 PetscFunctionBegin; 90 PetscValidDeviceType(type, 1); 91 PetscValidPointer(device, 3); 92 PetscCall(PetscDeviceInitializePackage()); 93 PetscCall(PetscNew(device)); 94 (*device)->id = PetscDeviceCounter++; 95 (*device)->type = type; 96 (*device)->refcnt = 1; 97 /* 98 if you are adding a device, you also need to add it's initialization in 99 PetscDeviceInitializeTypeFromOptions_Private() below 100 */ 101 switch (type) { 102 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, getDevice, *device, devid); 103 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, getDevice, *device, devid); 104 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, getDevice, *device, devid); 105 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, getDevice, *device, devid); 106 default: 107 /* in case the above macros expand to nothing this silences any unused variable warnings */ 108 (void)(devid); 109 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]); 110 } 111 PetscFunctionReturn(0); 112 } 113 114 /*@C 115 PetscDeviceDestroy - Free a `PetscDevice` 116 117 Not Collective 118 119 Input Parameter: 120 . device - The `PetscDevice` 121 122 Level: beginner 123 124 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceView()`, 125 `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()` 126 @*/ 127 PetscErrorCode PetscDeviceDestroy(PetscDevice *device) { 128 PetscFunctionBegin; 129 PetscValidPointer(device, 1); 130 if (!*device) PetscFunctionReturn(0); 131 PetscValidDevice(*device, 1); 132 PetscCall(PetscDeviceDereference_Internal(*device)); 133 if ((*device)->refcnt) { 134 *device = nullptr; 135 PetscFunctionReturn(0); 136 } 137 PetscCall(PetscFree((*device)->data)); 138 PetscCall(PetscFree(*device)); 139 PetscFunctionReturn(0); 140 } 141 142 /*@C 143 PetscDeviceConfigure - Configure a particular `PetscDevice` 144 145 Not Collective 146 147 Input Parameter: 148 . device - The `PetscDevice` to configure 149 150 Notes: 151 The user should not assume that this is a cheap operation. 152 153 Level: beginner 154 155 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceView()`, `PetscDeviceDestroy()`, 156 `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()` 157 @*/ 158 PetscErrorCode PetscDeviceConfigure(PetscDevice device) { 159 PetscFunctionBegin; 160 PetscValidDevice(device, 1); 161 /* 162 if no available configuration is available, this cascades all the way down to default 163 and error 164 */ 165 switch (const auto dtype = device->type) { 166 case PETSC_DEVICE_HOST: 167 if (PetscDefined(HAVE_HOST)) break; // always true 168 case PETSC_DEVICE_CUDA: 169 if (PetscDefined(HAVE_CUDA)) break; 170 goto error; 171 case PETSC_DEVICE_HIP: 172 if (PetscDefined(HAVE_HIP)) break; 173 goto error; 174 case PETSC_DEVICE_SYCL: 175 if (PetscDefined(HAVE_SYCL)) break; 176 goto error; 177 default: 178 error: 179 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]); 180 } 181 PetscUseTypeMethod(device, configure); 182 PetscFunctionReturn(0); 183 } 184 185 /*@C 186 PetscDeviceView - View a `PetscDevice` 187 188 Collective on viewer 189 190 Input Parameters: 191 + device - The `PetscDevice` to view 192 - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`) 193 194 Level: beginner 195 196 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, 197 `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()` 198 @*/ 199 PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer) { 200 auto sub = viewer; 201 PetscBool iascii; 202 203 PetscFunctionBegin; 204 PetscValidDevice(device, 1); 205 if (viewer) { 206 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 207 PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii)); 208 } else { 209 PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer)); 210 iascii = PETSC_TRUE; 211 } 212 213 if (iascii) { 214 auto dtype = PETSC_DEVICE_HOST; 215 MPI_Comm comm; 216 PetscMPIInt size; 217 PetscInt id = 0; 218 219 PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm)); 220 PetscCallMPI(MPI_Comm_size(comm, &size)); 221 222 PetscCall(PetscDeviceGetDeviceId(device, &id)); 223 PetscCall(PetscDeviceGetType(device, &dtype)); 224 PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub)); 225 PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes")); 226 PetscCall(PetscViewerASCIIPushTab(sub)); 227 PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype])); 228 PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id)); 229 } 230 231 // see if impls has extra viewer stuff 232 PetscTryTypeMethod(device, view, sub); 233 234 if (iascii) { 235 // undo the ASCII specific stuff 236 PetscCall(PetscViewerASCIIPopTab(sub)); 237 PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub)); 238 PetscCall(PetscViewerFlush(viewer)); 239 } 240 PetscFunctionReturn(0); 241 } 242 243 /*@C 244 PetscDeviceGetType - Get the type of device 245 246 Not Collective 247 248 Input Parameter: 249 . device - The `PetscDevice` 250 251 Output Parameter: 252 . type - The `PetscDeviceType` 253 254 Level: beginner 255 256 .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, 257 `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`, 258 `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()` 259 @*/ 260 PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type) { 261 PetscFunctionBegin; 262 PetscValidDevice(device, 1); 263 PetscValidPointer(type, 2); 264 *type = device->type; 265 PetscFunctionReturn(0); 266 } 267 268 /*@C 269 PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice` 270 271 Not Collective 272 273 Input Parameter: 274 . device - The `PetscDevice` 275 276 Output Parameter: 277 . id - The id 278 279 Notes: 280 The returned ID may have been assigned by the underlying device backend. For example if the 281 backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when 282 this device was configured. 283 284 Level: beginner 285 286 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()` 287 @*/ 288 PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id) { 289 PetscFunctionBegin; 290 PetscValidDevice(device, 1); 291 PetscValidIntPointer(id, 2); 292 *id = device->deviceId; 293 PetscFunctionReturn(0); 294 } 295 296 struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> { 297 PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE; 298 299 PETSC_NODISCARD PetscErrorCode finalize_() noexcept { 300 PetscFunctionBegin; 301 type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE; 302 PetscFunctionReturn(0); 303 } 304 }; 305 306 static auto default_device_type = DefaultDeviceType(); 307 308 /*@C 309 PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType` 310 311 Not Collective 312 313 Notes: 314 Unless selected by the user, the default device is selected in the following order\: 315 `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`. 316 317 Level: beginner 318 319 .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()` 320 @*/ 321 PetscDeviceType PETSC_DEVICE_DEFAULT(void) { 322 return default_device_type.type; 323 } 324 325 /*@C 326 PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice` 327 328 Not Collective 329 330 Input Parameter: 331 . type - the new default device type 332 333 Notes: 334 This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`. 335 336 Level: beginner 337 338 .seealso: `PetscDeviceType`, `PetscDeviceGetType`, 339 @*/ 340 PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type) { 341 PetscFunctionBegin; 342 PetscValidDeviceType(type, 1); 343 if (default_device_type.type != type) { 344 // no need to waster a PetscRegisterFinalize() slot if we don't change it 345 default_device_type.type = type; 346 PetscCall(default_device_type.register_finalize()); 347 } 348 PetscFunctionReturn(0); 349 } 350 351 static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {}; 352 353 /* 354 Actual intialization function; any functions claiming to initialize PetscDevice or 355 PetscDeviceContext will have to run through this one 356 */ 357 static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId) { 358 PetscFunctionBegin; 359 PetscValidDeviceType(type, 1); 360 if (PetscUnlikely(!PetscDeviceInitialized(type))) { 361 auto &dev = defaultDevices[type].first; 362 auto &init = defaultDevices[type].second; 363 364 PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]); 365 PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev)); 366 PetscCall(PetscDeviceConfigure(dev)); 367 init = true; 368 } 369 PetscFunctionReturn(0); 370 } 371 372 /*@C 373 PetscDeviceInitialize - Initialize `PetscDevice` 374 375 Not Collective 376 377 Input Parameter: 378 . type - The `PetscDeviceType` to initialize 379 380 Notes: 381 Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may 382 result in device synchronization. 383 384 Level: beginner 385 386 .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`, 387 `PetscDeviceCreate()`, `PetscDeviceDestroy()` 388 @*/ 389 PetscErrorCode PetscDeviceInitialize(PetscDeviceType type) { 390 PetscFunctionBegin; 391 PetscValidDeviceType(type, 1); 392 PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE)); 393 PetscFunctionReturn(0); 394 } 395 396 /*@C 397 PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular 398 `PetscDeviceType` 399 400 Not Collective 401 402 Input Parameter: 403 . type - The `PetscDeviceType` to check 404 405 Notes: 406 Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise. 407 408 If one has not configured PETSc for a particular `PetscDeviceType` then this routine will 409 return `PETSC_FALSE` for that `PetscDeviceType`. 410 411 Level: beginner 412 413 .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`, 414 `PetscDeviceCreate()`, `PetscDeviceDestroy()` 415 @*/ 416 PetscBool PetscDeviceInitialized(PetscDeviceType type) { 417 return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second); 418 } 419 420 /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */ 421 PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device) { 422 PetscFunctionBegin; 423 PetscValidPointer(device, 2); 424 PetscCall(PetscDeviceInitialize(type)); 425 *device = defaultDevices[type].first; 426 PetscFunctionReturn(0); 427 } 428 429 /*@C 430 PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice` 431 432 Not Collective 433 434 Input Parameters: 435 + device - The `PetscDevice` 436 - attr - The attribute 437 438 Output Parameter: 439 . value - The value of the attribute 440 441 Notes: 442 Since different attributes are often different types `value` is a `void *` to accommodate 443 them all. The underlying type of the attribute is therefore included in the name of the 444 `PetscDeviceAttribute` reponsible for querying it. For example, 445 `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`. 446 447 Level: intermediate 448 449 .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice` 450 @*/ 451 PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value) { 452 PetscFunctionBegin; 453 PetscValidDevice(device, 1); 454 PetscValidDeviceAttribute(attr, 2); 455 PetscValidPointer(value, 3); 456 PetscUseTypeMethod(device, getattribute, attr, value); 457 PetscFunctionReturn(0); 458 } 459 460 static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType) { 461 PetscFunctionBegin; 462 if (!PetscDeviceConfiguredFor_Internal(type)) { 463 PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type])); 464 defaultDevices[type].first = nullptr; 465 PetscFunctionReturn(0); 466 } 467 PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type])); 468 /* ugly switch needed to pick the right global variable... could maybe do this as a union? */ 469 switch (type) { 470 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 471 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 472 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 473 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 474 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]); 475 } 476 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)])); 477 /* 478 defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES 479 have decided to initialize as 480 */ 481 if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) { 482 PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type])); 483 PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId)); 484 if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr)); 485 } 486 PetscFunctionReturn(0); 487 } 488 489 static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView) { 490 PetscInt initIdx = PETSC_DEVICE_INIT_LAZY; 491 auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice); 492 auto flg = PETSC_FALSE; 493 494 PetscFunctionBegin; 495 PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg)); 496 if (flg) PetscCall(PetscLogGpuTime()); 497 498 PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys"); 499 PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr)); 500 PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet)); 501 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)); 502 PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg)); 503 PetscOptionsEnd(); 504 505 if (initIdx == PETSC_DEVICE_INIT_NONE) { 506 /* disabled all device initialization if devices are globally disabled */ 507 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"); 508 *defaultView = PETSC_FALSE; 509 initDeviceIdx = PETSC_DEVICE_HOST; 510 } else { 511 *defaultView = static_cast<PetscBool>(*defaultView && flg); 512 if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER; 513 } 514 *defaultInitType = PetscDeviceInitTypeCast(initIdx); 515 *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx); 516 PetscFunctionReturn(0); 517 } 518 519 /* called from PetscFinalize() do not call yourself! */ 520 static PetscErrorCode PetscDeviceFinalize_Private() { 521 PetscFunctionBegin; 522 if (PetscDefined(USE_DEBUG)) { 523 const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] { 524 PetscFunctionBegin; 525 for (auto &&device : defaultDevices) { 526 const auto dev = device.first; 527 528 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); 529 } 530 PetscFunctionReturn(0); 531 }; 532 /* 533 you might be thinking, why on earth are you registered yet another finalizer in a 534 function already called during PetscRegisterFinalizeAll()? If this seems stupid it's 535 because it is. 536 537 The crux of the problem is that the initializer (and therefore the ~finalizer~) of 538 PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had 539 a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence 540 won't be destroyed yet. So we need to repeat the check that all devices have been 541 destroyed again ~after~ the global context is destroyed. In summary: 542 543 1. This finalizer runs and destroys all devices, except it may not because the global 544 context may still hold a reference! 545 2. The global context finalizer runs and does the final reference count decrement 546 required, which actually destroys the held device. 547 3. Our newly added finalizer runs and checks that all is well. 548 */ 549 PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize))); 550 } 551 for (auto &&device : defaultDevices) { 552 PetscCall(PetscDeviceDestroy(&device.first)); 553 device.second = false; 554 } 555 PetscFunctionReturn(0); 556 } 557 558 /* 559 Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of 560 initialization types: 561 562 1. defaultInitType - how does PetscDevice as a whole expect to initialize? 563 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize? 564 e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but 565 have all CUDA devices still initialize. 566 567 All told the following happens: 568 569 0. defaultInitType -> LAZY 570 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER 571 2. PetscDevice initializes each sub type with deviceDefaultInitType. 572 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition 573 to checking for specific device init. if view or specific device init 574 subTypeDefaultInitType -> EAGER. disabled once again overrides all. 575 */ 576 577 PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm) { 578 auto defaultView = PETSC_FALSE; 579 auto initializeDeviceContextEagerly = PETSC_FALSE; 580 auto defaultDeviceSet = PETSC_FALSE; 581 auto defaultDevice = PetscInt{PETSC_DECIDE}; 582 auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT(); 583 auto defaultInitType = PETSC_DEVICE_INIT_LAZY; 584 585 PetscFunctionBegin; 586 if (PetscDefined(USE_DEBUG)) { 587 int result; 588 589 PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result)); 590 /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the 591 * global space */ 592 if (PetscUnlikely(result != MPI_IDENT)) { 593 char name[MPI_MAX_OBJECT_NAME] = {}; 594 int len; /* unused */ 595 596 PetscCallMPI(MPI_Comm_get_name(comm, name, &len)); 597 SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name); 598 } 599 } 600 comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */ 601 PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private)); 602 603 PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView)); 604 605 // the precise values don't matter here, so long as they are sequential 606 static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, ""); 607 static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, ""); 608 static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, ""); 609 static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, ""); 610 static_assert(Petsc::util::integral_value(PETSC_DEVICE_MAX) == 4, ""); 611 for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) { 612 const auto deviceType = PetscDeviceTypeCast(i); 613 auto initType = defaultInitType; 614 615 PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType)); 616 if (PetscDeviceConfiguredFor_Internal(deviceType)) { 617 if (initType == PETSC_DEVICE_INIT_EAGER) { 618 initializeDeviceContextEagerly = PETSC_TRUE; 619 // only update the default device if the user hasn't set it previously 620 if (!defaultDeviceSet) { 621 deviceContextInitDevice = deviceType; 622 PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType])); 623 } 624 } else if (initType == PETSC_DEVICE_INIT_NONE) { 625 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]); 626 } 627 } 628 } 629 630 PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice)); 631 PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT())); 632 /* ----------------------------------------------------------------------------------- */ 633 /* PetscDevice is now fully initialized */ 634 /* ----------------------------------------------------------------------------------- */ 635 { 636 /* 637 query the options db to get the root settings from the user (if any). 638 639 This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call 640 PetscDeviceContextSetFromOptions() before we even have one, then set a few static 641 variables in that file with the results. 642 */ 643 auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE); 644 auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE); 645 646 PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys"); 647 PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype)); 648 PetscOptionsEnd(); 649 650 if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first)); 651 if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first)); 652 } 653 654 if (initializeDeviceContextEagerly) { 655 PetscDeviceContext dctx; 656 657 PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice])); 658 /* instantiates the device context */ 659 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 660 PetscCall(PetscDeviceContextSetUp(dctx)); 661 } 662 PetscFunctionReturn(0); 663 } 664