1# -------------------------------------------------------------------- 2 3class staticproperty(property): 4 def __get__(self, *args, **kwargs): 5 return self.fget.__get__(*args, **kwargs)() 6 7 8cdef object make_enum_class(str class_name, str class_docstring, tuple args): 9 cdef dict enum2str = {} 10 cdef dict attrs = {} 11 12 for name, c_enum in args: 13 enum2str[c_enum] = name 14 attrs[name] = c_enum 15 16 attrs['__enum2str'] = enum2str 17 attrs['__doc__'] = class_docstring 18 return type(class_name, (object,), attrs) 19 20DeviceType = make_enum_class( 21 "DeviceType", 22 """The type of device. 23 24 See Also 25 -------- 26 Device, Device.create, Device.getDeviceType, Device.type, petsc.PetscDeviceType 27 28 """, 29 ( 30 ("HOST" , PETSC_DEVICE_HOST), 31 ("CUDA" , PETSC_DEVICE_CUDA), 32 ("HIP" , PETSC_DEVICE_HIP), 33 ("SYCL" , PETSC_DEVICE_SYCL), 34 ("DEFAULT" , staticproperty(lambda *_, **__: PETSC_DEVICE_DEFAULT())) 35 ) 36) 37 38StreamType = make_enum_class( 39 "StreamType", 40 """The type of stream. 41 42 See Also 43 -------- 44 DeviceContext, DeviceContext.getStreamType 45 DeviceContext.setStreamType, petsc.PetscStreamType 46 47 """, 48 ( 49 ("DEFAULT" , PETSC_STREAM_DEFAULT), 50 ("NONBLOCKING" , PETSC_STREAM_NONBLOCKING), 51 ("DEFAULT_WITH_BARRIER" , PETSC_STREAM_DEFAULT_WITH_BARRIER), 52 ("NONBLOCKING_WITH_BARRIER" , PETSC_STREAM_NONBLOCKING_WITH_BARRIER), 53 ) 54) 55 56DeviceJoinMode = make_enum_class( 57 "JoinMode", 58 """The type of join to perform. 59 60 See Also 61 -------- 62 DeviceContext, DeviceContext.join, DeviceContext.fork 63 petsc.PetscDeviceContextJoinMode 64 65 """, 66 ( 67 ("DESTROY" , PETSC_DEVICE_CONTEXT_JOIN_DESTROY), 68 ("SYNC" , PETSC_DEVICE_CONTEXT_JOIN_SYNC), 69 ("NO_SYNC" , PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC), 70 ) 71) 72 73# -------------------------------------------------------------------- 74 75cdef class Device: 76 """The device object. 77 78 Represents a handle to an accelerator (which may be the host). 79 80 See Also 81 -------- 82 DeviceContext, petsc.PetscDevice 83 84 """ 85 86 Type = DeviceType 87 88 def __cinit__(self): 89 self.device = NULL 90 91 def __dealloc__(self): 92 self.destroy() 93 94 @classmethod 95 def create(cls, dtype: Type | None = None, device_id: int = DECIDE) -> Device: 96 """Create a device object. 97 98 Not collective. 99 100 Parameters 101 ---------- 102 dtype 103 The type of device to create (or `None` for the default). 104 105 device_id 106 The numeric id of the device to create. 107 108 See Also 109 -------- 110 destroy, petsc.PetscDeviceCreate 111 112 """ 113 cdef PetscInt cdevice_id = asInt(device_id) 114 cdef PetscDeviceType cdevice_type = asDeviceType(dtype if dtype is not None else cls.Type.DEFAULT) 115 cdef Device device = cls() 116 117 CHKERR(PetscDeviceCreate(cdevice_type, cdevice_id, &device.device)) 118 return device 119 120 def destroy(self) -> None: 121 """Destroy a device object. 122 123 Not collective. 124 125 See Also 126 -------- 127 create, petsc.PetscDeviceDestroy 128 129 """ 130 CHKERR(PetscDeviceDestroy(&self.device)) 131 132 def configure(self) -> None: 133 """Configure and setup a device object. 134 135 Not collective. 136 137 See Also 138 -------- 139 create, petsc.PetscDeviceConfigure 140 141 """ 142 CHKERR(PetscDeviceConfigure(self.device)) 143 144 def view(self, Viewer viewer=None) -> None: 145 """View a device object. 146 147 Collective. 148 149 Parameters 150 ---------- 151 viewer 152 A `Viewer` instance or `None` for the default viewer. 153 154 See Also 155 -------- 156 petsc.PetscDeviceView 157 158 """ 159 cdef PetscViewer vwr = NULL 160 161 if viewer is not None: 162 vwr = viewer.vwr 163 CHKERR(PetscDeviceView(self.device, vwr)) 164 165 def getDeviceType(self) -> str: 166 """Return the type of the device. 167 168 Not collective. 169 170 See Also 171 -------- 172 type, petsc.PetscDeviceGetType 173 174 """ 175 cdef PetscDeviceType cdtype = PETSC_DEVICE_HOST 176 177 CHKERR(PetscDeviceGetType(self.device, &cdtype)) 178 return toDeviceType(cdtype) 179 180 def getDeviceId(self) -> int: 181 """Return the device id. 182 183 Not collective. 184 185 See Also 186 -------- 187 create, petsc.PetscDeviceGetDeviceId 188 189 """ 190 cdef PetscInt cdevice_id = 0 191 192 CHKERR(PetscDeviceGetDeviceId(self.device, &cdevice_id)) 193 return toInt(cdevice_id) 194 195 @staticmethod 196 def setDefaultType(device_type: Type | str) -> None: 197 """Set the device type to be used as the default in subsequent calls to `create`. 198 199 Not collective. 200 201 See Also 202 -------- 203 create, petsc.PetscDeviceSetDefaultDeviceType 204 205 """ 206 cdef PetscDeviceType cdevice_type = asDeviceType(device_type) 207 208 CHKERR(PetscDeviceSetDefaultDeviceType(cdevice_type)) 209 210 property type: 211 """The device type.""" 212 def __get__(self) -> str: 213 return self.getDeviceType() 214 215 property device_id: 216 """The device id.""" 217 def __get__(self) -> int: 218 return self.getDeviceId() 219 220 221# -------------------------------------------------------------------- 222 223cdef class DeviceContext(Object): 224 """DeviceContext object. 225 226 Represents an abstract handle to a device context. 227 228 See Also 229 -------- 230 Device, petsc.PetscDeviceContext 231 232 """ 233 JoinMode = DeviceJoinMode 234 StreamType = StreamType 235 236 def __cinit__(self): 237 self.obj = <PetscObject*> &self.dctx 238 self.dctx = NULL 239 240 def create(self) -> Self: 241 """Create an empty DeviceContext. 242 243 Not collective. 244 245 See Also 246 -------- 247 destroy, Device, petsc.PetscDeviceContextCreate 248 249 """ 250 cdef PetscDeviceContext dctx = NULL 251 CHKERR(PetscDeviceContextCreate(&dctx)) 252 CHKERR(PetscCLEAR(self.obj)); self.dctx = dctx 253 return self 254 255 def destroy(self) -> Self: 256 """Destroy a device context. 257 258 Not collective. 259 260 See Also 261 -------- 262 create, petsc.PetscDeviceContextDestroy 263 264 """ 265 CHKERR(PetscDeviceContextDestroy(&self.dctx)) 266 return self 267 268 def getStreamType(self) -> str: 269 """Return the `StreamType`. 270 271 Not collective. 272 273 See Also 274 -------- 275 stream_type, setStreamType, petsc.PetscDeviceContextGetStreamType 276 277 """ 278 cdef PetscStreamType cstream_type = PETSC_STREAM_DEFAULT 279 280 CHKERR(PetscDeviceContextGetStreamType(self.dctx, &cstream_type)) 281 return toStreamType(cstream_type) 282 283 def setStreamType(self, stream_type: StreamType | str) -> None: 284 """Set the `StreamType`. 285 286 Not collective. 287 288 Parameters 289 ---------- 290 stream_type 291 The type of stream to set 292 293 See Also 294 -------- 295 stream_type, getStreamType, petsc.PetscDeviceContextSetStreamType 296 297 """ 298 cdef PetscStreamType cstream_type = asStreamType(stream_type) 299 300 CHKERR(PetscDeviceContextSetStreamType(self.dctx, cstream_type)) 301 302 def getDevice(self) -> Device: 303 """Get the `Device` which this instance is attached to. 304 305 Not collective. 306 307 See Also 308 -------- 309 setDevice, device, Device, petsc.PetscDeviceContextGetDevice 310 311 """ 312 cdef PetscDevice device = NULL 313 314 CHKERR(PetscDeviceContextGetDevice(self.dctx, &device)) 315 return PyPetscDevice_New(device) 316 317 def setDevice(self, Device device not None) -> None: 318 """Set the `Device` which this `DeviceContext` is attached to. 319 320 Collective. 321 322 Parameters 323 ---------- 324 device 325 The `Device` to which this instance is attached to. 326 327 See Also 328 -------- 329 getDevice, device, Device, petsc.PetscDeviceContextSetDevice 330 331 """ 332 cdef PetscDevice cdevice = PyPetscDevice_Get(device) 333 334 CHKERR(PetscDeviceContextSetDevice(self.dctx, cdevice)) 335 336 def setUp(self) -> None: 337 """Set up the internal data structures for using the device context. 338 339 Not collective. 340 341 See Also 342 -------- 343 create, petsc.PetscDeviceContextSetUp 344 345 """ 346 CHKERR(PetscDeviceContextSetUp(self.dctx)) 347 348 def duplicate(self) -> DeviceContext: 349 """Duplicate a the device context. 350 351 Not collective. 352 353 See Also 354 -------- 355 create, petsc.PetscDeviceContextDuplicate 356 357 """ 358 cdef DeviceContext octx = type(self)() 359 360 CHKERR(PetscDeviceContextDuplicate(self.dctx, &octx.dctx)) 361 return octx 362 363 def idle(self) -> bool: 364 """Return whether the underlying stream for the device context is idle. 365 366 Not collective. 367 368 See Also 369 -------- 370 synchronize, petsc.PetscDeviceContextQueryIdle 371 372 """ 373 cdef PetscBool is_idle = PETSC_FALSE 374 375 CHKERR(PetscDeviceContextQueryIdle(self.dctx, &is_idle)) 376 return toBool(is_idle) 377 378 def waitFor(self, other: DeviceContext | None) -> None: 379 """Make this instance wait for ``other``. 380 381 Not collective. 382 383 Parameters 384 ---------- 385 other 386 The other `DeviceContext` to wait for 387 388 See Also 389 -------- 390 fork, join, petsc.PetscDeviceContextWaitForContext 391 392 """ 393 cdef PetscDeviceContext cother = NULL 394 395 if other is not None: 396 cother = PyPetscDeviceContext_Get(other) 397 CHKERR(PetscDeviceContextWaitForContext(self.dctx, cother)) 398 399 def fork(self, n: int, stream_type: DeviceContext.StreamType | str | None = None) -> list[DeviceContext]: 400 """Create multiple device contexts which are all logically dependent on this one. 401 402 Not collective. 403 404 Parameters 405 ---------- 406 n 407 The number of device contexts to create. 408 stream_type 409 The type of stream of the forked device context. 410 411 Examples 412 -------- 413 The device contexts created must be destroyed using `join`. 414 415 >>> dctx = PETSc.DeviceContext().getCurrent() 416 >>> dctxs = dctx.fork(4) 417 >>> ... # perform computations 418 >>> # we can mix various join modes 419 >>> dctx.join(PETSc.DeviceContext.JoinMode.SYNC, dctxs[0:2]) 420 >>> dctx.join(PETSc.DeviceContext.JoinMode.SYNC, dctxs[2:]) 421 >>> ... # some more computations and joins 422 >>> # dctxs must be all destroyed with joinMode.DESTROY 423 >>> dctx.join(PETSc.DeviceContext.JoinMode.DESTROY, dctxs) 424 425 See Also 426 -------- 427 join, waitFor, petsc.PetscDeviceContextFork 428 429 """ 430 cdef PetscDeviceContext *csubctxs = NULL 431 cdef PetscStreamType cstream_type = PETSC_STREAM_DEFAULT 432 cdef PetscInt cn = asInt(n) 433 cdef list subctxs = [] 434 if stream_type is None: 435 CHKERR(PetscDeviceContextFork(self.dctx, cn, &csubctxs)) 436 else: 437 cstream_type = asStreamType(stream_type) 438 CHKERR(PetscDeviceContextForkWithStreamType(self.dctx, cstream_type, cn, &csubctxs)) 439 # FIXME: without CXX compiler, csubctxs is NULL 440 if csubctxs: 441 subctxs = [None] * cn 442 for i from 0 <= i < cn: 443 subctxs[i] = DeviceContext() 444 (<DeviceContext?>subctxs[i]).dctx = csubctxs[i] 445 CHKERR(PetscFree(csubctxs)) 446 return subctxs 447 448 def join(self, join_mode: JoinMode | str, py_sub_ctxs: list[DeviceContext]) -> None: 449 """Join a set of device contexts on this one. 450 451 Not collective. 452 453 Parameters 454 ---------- 455 join_mode 456 The type of join to perform. 457 py_sub_ctxs 458 The list of device contexts to join. 459 460 See Also 461 -------- 462 fork, waitFor, petsc.PetscDeviceContextJoin 463 464 """ 465 cdef PetscDeviceContext *np_subctx = NULL 466 cdef PetscDeviceContextJoinMode cjoin_mode = asJoinMode(join_mode) 467 cdef Py_ssize_t nctxs = len(py_sub_ctxs) 468 469 CHKERR(PetscMalloc(<size_t>(nctxs) * sizeof(PetscDeviceContext *), &np_subctx)) 470 for i from 0 <= i < nctxs: 471 dctx = py_sub_ctxs[i] 472 np_subctx[i] = (<DeviceContext?>dctx).dctx if dctx is not None else NULL 473 CHKERR(PetscDeviceContextJoin(self.dctx, <PetscInt>nctxs, cjoin_mode, &np_subctx)) 474 475 if cjoin_mode == PETSC_DEVICE_CONTEXT_JOIN_DESTROY: 476 # in this case, PETSc destroys the contexts and frees the array 477 for i in range(nctxs): 478 (<DeviceContext?>py_sub_ctxs[i]).dctx = NULL 479 else: 480 # we need to free the temporary array 481 CHKERR(PetscFree(np_subctx)) 482 483 def synchronize(self) -> None: 484 """Synchronize a device context. 485 486 Not collective. 487 488 Notes 489 ----- 490 The underlying stream is considered idle after this routine returns, 491 i.e. `idle` will return ``True``. 492 493 See Also 494 -------- 495 idle, petsc.PetscDeviceContextSynchronize 496 497 """ 498 CHKERR(PetscDeviceContextSynchronize(self.dctx)) 499 500 def setFromOptions(self, comm: Comm | None = None) -> None: 501 """Configure the `DeviceContext` from the options database. 502 503 Collective. 504 505 Parameters 506 ---------- 507 comm 508 MPI communicator, defaults to `Sys.getDefaultComm`. 509 510 See Also 511 -------- 512 Sys.getDefaultComm, petsc.PetscDeviceContextSetFromOptions 513 514 """ 515 cdef MPI_Comm ccomm = def_Comm(comm, PETSC_COMM_DEFAULT) 516 517 CHKERR(PetscDeviceContextSetFromOptions(ccomm, self.dctx)) 518 519 @staticmethod 520 def getCurrent() -> DeviceContext: 521 """Return the current device context. 522 523 Not collective. 524 525 See Also 526 -------- 527 current, setCurrent, petsc.PetscDeviceContextGetCurrentContext 528 529 """ 530 cdef PetscDeviceContext dctx = NULL 531 532 CHKERR(PetscDeviceContextGetCurrentContext(&dctx)) 533 return PyPetscDeviceContext_New(dctx) 534 535 @staticmethod 536 def setCurrent(dctx: DeviceContext | None) -> None: 537 """Set the current device context. 538 539 Not collective. 540 541 Parameters 542 ---------- 543 dctx 544 The `DeviceContext` to set as current (or `None` to use 545 the default context). 546 547 See Also 548 -------- 549 current, getCurrent, petsc.PetscDeviceContextSetCurrentContext 550 551 """ 552 cdef PetscDeviceContext cdctx = NULL 553 554 if dctx is not None: 555 cdctx = PyPetscDeviceContext_Get(dctx) 556 CHKERR(PetscDeviceContextSetCurrentContext(cdctx)) 557 558 property stream_type: 559 """The stream type.""" 560 def __get__(self) -> str: 561 return self.getStreamType() 562 563 def __set__(self, stype: StreamType | str) -> None: 564 self.setStreamType(stype) 565 566 property device: 567 """The device associated to the device context.""" 568 def __get__(self) -> Device: 569 return self.getDevice() 570 571 def __set__(self, Device device) -> None: 572 self.setDevice(device) 573 574 property current: 575 """The current global device context.""" 576 def __get__(self) -> DeviceContext: 577 return self.getCurrent() 578 579 def __set__(self, dctx: DeviceContext | None) -> None: 580 self.setCurrent(dctx) 581 582# -------------------------------------------------------------------- 583 584del DeviceType 585del DeviceJoinMode 586del StreamType 587del staticproperty 588