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 default: 177 error: 178 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "PETSc was not configured for PetscDeviceType %s", PetscDeviceTypes[dtype]); 179 } 180 PetscUseTypeMethod(device, configure); 181 PetscFunctionReturn(0); 182 } 183 184 /*@C 185 PetscDeviceView - View a `PetscDevice` 186 187 Collective on viewer 188 189 Input Parameters: 190 + device - The `PetscDevice` to view 191 - viewer - The `PetscViewer` to view the device with (`NULL` for `PETSC_VIEWER_STDOUT_WORLD`) 192 193 Level: beginner 194 195 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceConfigure()`, 196 `PetscDeviceDestroy()`, `PetscDeviceGetType()`, `PetscDeviceGetDeviceId()` 197 @*/ 198 PetscErrorCode PetscDeviceView(PetscDevice device, PetscViewer viewer) { 199 auto sub = viewer; 200 PetscBool iascii; 201 202 PetscFunctionBegin; 203 PetscValidDevice(device, 1); 204 if (viewer) { 205 PetscValidHeaderSpecific(viewer, PETSC_VIEWER_CLASSID, 2); 206 PetscCall(PetscObjectTypeCompare(PetscObjectCast(viewer), PETSCVIEWERASCII, &iascii)); 207 } else { 208 PetscCall(PetscViewerASCIIGetStdout(PETSC_COMM_WORLD, &viewer)); 209 iascii = PETSC_TRUE; 210 } 211 212 if (iascii) { 213 auto dtype = PETSC_DEVICE_HOST; 214 MPI_Comm comm; 215 PetscMPIInt size; 216 PetscInt id = 0; 217 218 PetscCall(PetscObjectGetComm(PetscObjectCast(viewer), &comm)); 219 PetscCallMPI(MPI_Comm_size(comm, &size)); 220 221 PetscCall(PetscDeviceGetDeviceId(device, &id)); 222 PetscCall(PetscDeviceGetType(device, &dtype)); 223 PetscCall(PetscViewerGetSubViewer(viewer, PETSC_COMM_SELF, &sub)); 224 PetscCall(PetscViewerASCIIPrintf(sub, "PetscDevice Object: %d MPI %s\n", size, size == 1 ? "process" : "processes")); 225 PetscCall(PetscViewerASCIIPushTab(sub)); 226 PetscCall(PetscViewerASCIIPrintf(sub, "type: %s\n", PetscDeviceTypes[dtype])); 227 PetscCall(PetscViewerASCIIPrintf(sub, "id: %" PetscInt_FMT "\n", id)); 228 } 229 230 // see if impls has extra viewer stuff 231 PetscTryTypeMethod(device, view, sub); 232 233 if (iascii) { 234 // undo the ASCII specific stuff 235 PetscCall(PetscViewerASCIIPopTab(sub)); 236 PetscCall(PetscViewerRestoreSubViewer(viewer, PETSC_COMM_SELF, &sub)); 237 PetscCall(PetscViewerFlush(viewer)); 238 } 239 PetscFunctionReturn(0); 240 } 241 242 /*@C 243 PetscDeviceGetType - Get the type of device 244 245 Not Collective 246 247 Input Parameter: 248 . device - The `PetscDevice` 249 250 Output Parameter: 251 . type - The `PetscDeviceType` 252 253 Level: beginner 254 255 .seealso: `PetscDevice`, `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, 256 `PetscDeviceCreate()`, `PetscDeviceConfigure()`, `PetscDeviceDestroy()`, 257 `PetscDeviceGetDeviceId()`, `PETSC_DEVICE_DEFAULT()` 258 @*/ 259 PetscErrorCode PetscDeviceGetType(PetscDevice device, PetscDeviceType *type) { 260 PetscFunctionBegin; 261 PetscValidDevice(device, 1); 262 PetscValidPointer(type, 2); 263 *type = device->type; 264 PetscFunctionReturn(0); 265 } 266 267 /*@C 268 PetscDeviceGetDeviceId - Get the device ID for a `PetscDevice` 269 270 Not Collective 271 272 Input Parameter: 273 . device - The `PetscDevice` 274 275 Output Parameter: 276 . id - The id 277 278 Notes: 279 The returned ID may have been assigned by the underlying device backend. For example if the 280 backend is CUDA then `id` is exactly the value returned by `cudaGetDevice()` at the time when 281 this device was configured. 282 283 Level: beginner 284 285 .seealso: `PetscDevice`, `PetscDeviceCreate()`, `PetscDeviceGetType()` 286 @*/ 287 PetscErrorCode PetscDeviceGetDeviceId(PetscDevice device, PetscInt *id) { 288 PetscFunctionBegin; 289 PetscValidDevice(device, 1); 290 PetscValidIntPointer(id, 2); 291 *id = device->deviceId; 292 PetscFunctionReturn(0); 293 } 294 295 struct DefaultDeviceType : public Petsc::RegisterFinalizeable<DefaultDeviceType> { 296 PetscDeviceType type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE; 297 298 PETSC_NODISCARD PetscErrorCode finalize_() noexcept { 299 PetscFunctionBegin; 300 type = PETSC_DEVICE_HARDWARE_DEFAULT_TYPE; 301 PetscFunctionReturn(0); 302 } 303 }; 304 305 static auto default_device_type = DefaultDeviceType(); 306 307 /*@C 308 PETSC_DEVICE_DEFAULT - Retrieve the current default `PetscDeviceType` 309 310 Not Collective 311 312 Notes: 313 Unless selected by the user, the default device is selected in the following order\: 314 `PETSC_DEVICE_HIP`, `PETSC_DEVICE_CUDA`, `PETSC_DEVICE_SYCL`, `PETSC_DEVICE_HOST`. 315 316 Level: beginner 317 318 .seealso: `PetscDeviceType`, `PetscDeviceSetDefaultDeviceType()`, `PetscDeviceGetType()` 319 @*/ 320 PetscDeviceType PETSC_DEVICE_DEFAULT(void) { 321 return default_device_type.type; 322 } 323 324 /*@C 325 PetscDeviceSetDefaultDeviceType - Set the default device type for `PetscDevice` 326 327 Not Collective 328 329 Input Parameter: 330 . type - the new default device type 331 332 Notes: 333 This sets the `PetscDeviceType` returned by `PETSC_DEVICE_DEFAULT()`. 334 335 Level: beginner 336 337 .seealso: `PetscDeviceType`, `PetscDeviceGetType`, 338 @*/ 339 PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType type) { 340 PetscFunctionBegin; 341 PetscValidDeviceType(type, 1); 342 if (default_device_type.type != type) { 343 // no need to waster a PetscRegisterFinalize() slot if we don't change it 344 default_device_type.type = type; 345 PetscCall(default_device_type.register_finalize()); 346 } 347 PetscFunctionReturn(0); 348 } 349 350 static std::array<std::pair<PetscDevice, bool>, PETSC_DEVICE_MAX> defaultDevices = {}; 351 352 /* 353 Actual intialization function; any functions claiming to initialize PetscDevice or 354 PetscDeviceContext will have to run through this one 355 */ 356 static PetscErrorCode PetscDeviceInitializeDefaultDevice_Internal(PetscDeviceType type, PetscInt defaultDeviceId) { 357 PetscFunctionBegin; 358 PetscValidDeviceType(type, 1); 359 if (PetscUnlikely(!PetscDeviceInitialized(type))) { 360 auto &dev = defaultDevices[type].first; 361 auto &init = defaultDevices[type].second; 362 363 PetscAssert(!dev, PETSC_COMM_SELF, PETSC_ERR_MEM, "Trying to overwrite existing default device of type %s", PetscDeviceTypes[type]); 364 PetscCall(PetscDeviceCreate(type, defaultDeviceId, &dev)); 365 PetscCall(PetscDeviceConfigure(dev)); 366 init = true; 367 } 368 PetscFunctionReturn(0); 369 } 370 371 /*@C 372 PetscDeviceInitialize - Initialize `PetscDevice` 373 374 Not Collective 375 376 Input Parameter: 377 . type - The `PetscDeviceType` to initialize 378 379 Notes: 380 Eagerly initializes the corresponding `PetscDeviceType` if needed. If this is the case it may 381 result in device synchronization. 382 383 Level: beginner 384 385 .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialized()`, 386 `PetscDeviceCreate()`, `PetscDeviceDestroy()` 387 @*/ 388 PetscErrorCode PetscDeviceInitialize(PetscDeviceType type) { 389 PetscFunctionBegin; 390 PetscValidDeviceType(type, 1); 391 PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, PETSC_DECIDE)); 392 PetscFunctionReturn(0); 393 } 394 395 /*@C 396 PetscDeviceInitialized - Determines whether `PetscDevice` is initialized for a particular 397 `PetscDeviceType` 398 399 Not Collective 400 401 Input Parameter: 402 . type - The `PetscDeviceType` to check 403 404 Notes: 405 Returns `PETSC_TRUE` if `type` is initialized, `PETSC_FALSE` otherwise. 406 407 If one has not configured PETSc for a particular `PetscDeviceType` then this routine will 408 return `PETSC_FALSE` for that `PetscDeviceType`. 409 410 Level: beginner 411 412 .seealso: `PetscDevice`, `PetscDeviceInitType`, `PetscDeviceInitialize()`, 413 `PetscDeviceCreate()`, `PetscDeviceDestroy()` 414 @*/ 415 PetscBool PetscDeviceInitialized(PetscDeviceType type) { 416 return static_cast<PetscBool>(PetscDeviceConfiguredFor_Internal(type) && defaultDevices[type].second); 417 } 418 419 /* Get the default PetscDevice for a particular type and constructs them if lazily initialized. */ 420 PetscErrorCode PetscDeviceGetDefaultForType_Internal(PetscDeviceType type, PetscDevice *device) { 421 PetscFunctionBegin; 422 PetscValidPointer(device, 2); 423 PetscCall(PetscDeviceInitialize(type)); 424 *device = defaultDevices[type].first; 425 PetscFunctionReturn(0); 426 } 427 428 /*@C 429 PetscDeviceGetAttribute - Query a particular attribute of a `PetscDevice` 430 431 Not Collective 432 433 Input Parameters: 434 + device - The `PetscDevice` 435 - attr - The attribute 436 437 Output Parameter: 438 . value - The value of the attribute 439 440 Notes: 441 Since different attributes are often different types `value` is a `void *` to accommodate 442 them all. The underlying type of the attribute is therefore included in the name of the 443 `PetscDeviceAttribute` reponsible for querying it. For example, 444 `PETSC_DEVICE_ATTR_SIZE_T_SHARED_MEM_PER_BLOCK` is of type `size_t`. 445 446 Level: intermediate 447 448 .seealso: `PetscDeviceAtrtibute`, `PetscDeviceConfigure()`, `PetscDevice` 449 @*/ 450 PetscErrorCode PetscDeviceGetAttribute(PetscDevice device, PetscDeviceAttribute attr, void *value) { 451 PetscFunctionBegin; 452 PetscValidDevice(device, 1); 453 PetscValidDeviceAttribute(attr, 2); 454 PetscValidPointer(value, 3); 455 PetscUseTypeMethod(device, getattribute, attr, value); 456 PetscFunctionReturn(0); 457 } 458 459 static PetscErrorCode PetscDeviceInitializeTypeFromOptions_Private(MPI_Comm comm, PetscDeviceType type, PetscInt defaultDeviceId, PetscBool defaultView, PetscDeviceInitType *defaultInitType) { 460 PetscFunctionBegin; 461 if (!PetscDeviceConfiguredFor_Internal(type)) { 462 PetscCall(PetscInfo(nullptr, "PetscDeviceType %s not available\n", PetscDeviceTypes[type])); 463 defaultDevices[type].first = nullptr; 464 PetscFunctionReturn(0); 465 } 466 PetscCall(PetscInfo(nullptr, "PetscDeviceType %s available, initializing\n", PetscDeviceTypes[type])); 467 /* ugly switch needed to pick the right global variable... could maybe do this as a union? */ 468 switch (type) { 469 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HOST, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 470 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(CUDA, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 471 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(HIP, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 472 PETSC_DEVICE_CASE_IF_PETSC_DEFINED(SYCL, initialize, comm, &defaultDeviceId, &defaultView, defaultInitType); 473 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]); 474 } 475 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)])); 476 /* 477 defaultInitType, defaultView and defaultDeviceId now represent what the individual TYPES 478 have decided to initialize as 479 */ 480 if ((*defaultInitType == PETSC_DEVICE_INIT_EAGER) || defaultView) { 481 PetscCall(PetscInfo(nullptr, "Eagerly initializing %s PetscDevice\n", PetscDeviceTypes[type])); 482 PetscCall(PetscDeviceInitializeDefaultDevice_Internal(type, defaultDeviceId)); 483 if (defaultView) PetscCall(PetscDeviceView(defaultDevices[type].first, nullptr)); 484 } 485 PetscFunctionReturn(0); 486 } 487 488 static PetscErrorCode PetscDeviceInitializeQueryOptions_Private(MPI_Comm comm, PetscDeviceType *deviceContextInitDevice, PetscDeviceInitType *defaultInitType, PetscInt *defaultDevice, PetscBool *defaultDeviceSet, PetscBool *defaultView) { 489 PetscInt initIdx = PETSC_DEVICE_INIT_LAZY; 490 auto initDeviceIdx = static_cast<PetscInt>(*deviceContextInitDevice); 491 auto flg = PETSC_FALSE; 492 493 PetscFunctionBegin; 494 PetscCall(PetscOptionsHasName(nullptr, nullptr, "-log_view_gpu_time", &flg)); 495 if (flg) PetscCall(PetscLogGpuTime()); 496 497 PetscOptionsBegin(comm, nullptr, "PetscDevice Options", "Sys"); 498 PetscCall(PetscOptionsEList("-device_enable", "How (or whether) to initialize PetscDevices", "PetscDeviceInitialize()", PetscDeviceInitTypes, 3, PetscDeviceInitTypes[initIdx], &initIdx, nullptr)); 499 PetscCall(PetscOptionsEList("-default_device_type", "Set the PetscDeviceType returned by PETSC_DEVICE_DEFAULT()", "PetscDeviceSetDefaultDeviceType()", PetscDeviceTypes, PETSC_DEVICE_MAX, PetscDeviceTypes[initDeviceIdx], &initDeviceIdx, defaultDeviceSet)); 500 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 PetscCall(PetscOptionsBool("-device_view", "Display device information and assignments (forces eager initialization)", "PetscDeviceView()", *defaultView, defaultView, &flg)); 502 PetscOptionsEnd(); 503 504 if (initIdx == PETSC_DEVICE_INIT_NONE) { 505 /* disabled all device initialization if devices are globally disabled */ 506 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 *defaultView = PETSC_FALSE; 508 initDeviceIdx = PETSC_DEVICE_HOST; 509 } else { 510 *defaultView = static_cast<PetscBool>(*defaultView && flg); 511 if (*defaultView) initIdx = PETSC_DEVICE_INIT_EAGER; 512 } 513 *defaultInitType = PetscDeviceInitTypeCast(initIdx); 514 *deviceContextInitDevice = PetscDeviceTypeCast(initDeviceIdx); 515 PetscFunctionReturn(0); 516 } 517 518 /* called from PetscFinalize() do not call yourself! */ 519 static PetscErrorCode PetscDeviceFinalize_Private() { 520 PetscFunctionBegin; 521 if (PetscDefined(USE_DEBUG)) { 522 const auto PetscDeviceCheckAllDestroyedAfterFinalize = [] { 523 PetscFunctionBegin; 524 for (auto &&device : defaultDevices) { 525 const auto dev = device.first; 526 527 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 } 529 PetscFunctionReturn(0); 530 }; 531 /* 532 you might be thinking, why on earth are you registered yet another finalizer in a 533 function already called during PetscRegisterFinalizeAll()? If this seems stupid it's 534 because it is. 535 536 The crux of the problem is that the initializer (and therefore the ~finalizer~) of 537 PetscDeviceContext is guaranteed to run after PetscDevice's. So if the global context had 538 a default PetscDevice attached, that PetscDevice will have a reference count >0 and hence 539 won't be destroyed yet. So we need to repeat the check that all devices have been 540 destroyed again ~after~ the global context is destroyed. In summary: 541 542 1. This finalizer runs and destroys all devices, except it may not because the global 543 context may still hold a reference! 544 2. The global context finalizer runs and does the final reference count decrement 545 required, which actually destroys the held device. 546 3. Our newly added finalizer runs and checks that all is well. 547 */ 548 PetscCall(PetscRegisterFinalize(std::move(PetscDeviceCheckAllDestroyedAfterFinalize))); 549 } 550 for (auto &&device : defaultDevices) { 551 PetscCall(PetscDeviceDestroy(&device.first)); 552 device.second = false; 553 } 554 PetscFunctionReturn(0); 555 } 556 557 /* 558 Begins the init proceeedings for the entire PetscDevice stack. there are 3 stages of 559 initialization types: 560 561 1. defaultInitType - how does PetscDevice as a whole expect to initialize? 562 2. subTypeDefaultInitType - how does each PetscDevice implementation expect to initialize? 563 e.g. you may want to blanket disable PetscDevice init (and disable say Kokkos init), but 564 have all CUDA devices still initialize. 565 566 All told the following happens: 567 568 0. defaultInitType -> LAZY 569 1. Check for log_view/log_summary, if yes defaultInitType -> EAGER 570 2. PetscDevice initializes each sub type with deviceDefaultInitType. 571 2.1 Each enabled PetscDevice sub-type then does the above disable or view check in addition 572 to checking for specific device init. if view or specific device init 573 subTypeDefaultInitType -> EAGER. disabled once again overrides all. 574 */ 575 576 PetscErrorCode PetscDeviceInitializeFromOptions_Internal(MPI_Comm comm) { 577 auto defaultView = PETSC_FALSE; 578 auto initializeDeviceContextEagerly = PETSC_FALSE; 579 auto defaultDeviceSet = PETSC_FALSE; 580 auto defaultDevice = PetscInt{PETSC_DECIDE}; 581 auto deviceContextInitDevice = PETSC_DEVICE_DEFAULT(); 582 auto defaultInitType = PETSC_DEVICE_INIT_LAZY; 583 584 PetscFunctionBegin; 585 if (PetscDefined(USE_DEBUG)) { 586 int result; 587 588 PetscCallMPI(MPI_Comm_compare(comm, PETSC_COMM_WORLD, &result)); 589 /* in order to accurately assign ranks to gpus we need to get the MPI_Comm_rank of the 590 * global space */ 591 if (PetscUnlikely(result != MPI_IDENT)) { 592 char name[MPI_MAX_OBJECT_NAME] = {}; 593 int len; /* unused */ 594 595 PetscCallMPI(MPI_Comm_get_name(comm, name, &len)); 596 SETERRQ(comm, PETSC_ERR_MPI, "Default devices being initialized on MPI_Comm '%s' not PETSC_COMM_WORLD", name); 597 } 598 } 599 comm = PETSC_COMM_WORLD; /* from this point on we assume we're on PETSC_COMM_WORLD */ 600 PetscCall(PetscRegisterFinalize(PetscDeviceFinalize_Private)); 601 602 PetscCall(PetscDeviceInitializeQueryOptions_Private(comm, &deviceContextInitDevice, &defaultInitType, &defaultDevice, &defaultDeviceSet, &defaultView)); 603 604 // the precise values don't matter here, so long as they are sequential 605 static_assert(Petsc::util::integral_value(PETSC_DEVICE_HOST) == 0, ""); 606 static_assert(Petsc::util::integral_value(PETSC_DEVICE_CUDA) == 1, ""); 607 static_assert(Petsc::util::integral_value(PETSC_DEVICE_HIP) == 2, ""); 608 static_assert(Petsc::util::integral_value(PETSC_DEVICE_SYCL) == 3, ""); 609 static_assert(Petsc::util::integral_value(PETSC_DEVICE_MAX) == 4, ""); 610 for (int i = PETSC_DEVICE_HOST; i < PETSC_DEVICE_MAX; ++i) { 611 const auto deviceType = PetscDeviceTypeCast(i); 612 auto initType = defaultInitType; 613 614 PetscCall(PetscDeviceInitializeTypeFromOptions_Private(comm, deviceType, defaultDevice, defaultView, &initType)); 615 if (PetscDeviceConfiguredFor_Internal(deviceType)) { 616 if (initType == PETSC_DEVICE_INIT_EAGER) { 617 initializeDeviceContextEagerly = PETSC_TRUE; 618 // only update the default device if the user hasn't set it previously 619 if (!defaultDeviceSet) { 620 deviceContextInitDevice = deviceType; 621 PetscCall(PetscInfo(nullptr, "PetscDevice %s set as default device type due to eager initialization\n", PetscDeviceTypes[deviceType])); 622 } 623 } else if (initType == PETSC_DEVICE_INIT_NONE) { 624 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]); 625 } 626 } 627 } 628 629 PetscCall(PetscDeviceSetDefaultDeviceType(deviceContextInitDevice)); 630 PetscCall(PetscDeviceContextSetRootDeviceType_Internal(PETSC_DEVICE_DEFAULT())); 631 /* ----------------------------------------------------------------------------------- */ 632 /* PetscDevice is now fully initialized */ 633 /* ----------------------------------------------------------------------------------- */ 634 { 635 /* 636 query the options db to get the root settings from the user (if any). 637 638 This section is a bit of a hack. We have to reach across to dcontext.cxx to all but call 639 PetscDeviceContextSetFromOptions() before we even have one, then set a few static 640 variables in that file with the results. 641 */ 642 auto dtype = std::make_pair(PETSC_DEVICE_DEFAULT(), PETSC_FALSE); 643 auto stype = std::make_pair(PETSC_DEVICE_CONTEXT_DEFAULT_STREAM_TYPE, PETSC_FALSE); 644 645 PetscOptionsBegin(comm, "root_", "Root PetscDeviceContext Options", "Sys"); 646 PetscCall(PetscDeviceContextQueryOptions_Internal(PetscOptionsObject, dtype, stype)); 647 PetscOptionsEnd(); 648 649 if (dtype.second) PetscCall(PetscDeviceContextSetRootDeviceType_Internal(dtype.first)); 650 if (stype.second) PetscCall(PetscDeviceContextSetRootStreamType_Internal(stype.first)); 651 } 652 653 if (initializeDeviceContextEagerly) { 654 PetscDeviceContext dctx; 655 656 PetscCall(PetscInfo(nullptr, "Eagerly initializing PetscDeviceContext with %s device\n", PetscDeviceTypes[deviceContextInitDevice])); 657 /* instantiates the device context */ 658 PetscCall(PetscDeviceContextGetCurrentContext(&dctx)); 659 PetscCall(PetscDeviceContextSetUp(dctx)); 660 } 661 PetscFunctionReturn(0); 662 } 663