xref: /petsc/src/binding/petsc4py/src/petsc4py/PETSc/petscdevice.pxi (revision 5dafb6f581977b074a43082461f4a734d7d3976b) 
1# --------------------------------------------------------------------
2
3cdef extern from * nogil:
4
5    ctypedef enum PetscOffloadMask:
6        PETSC_OFFLOAD_UNALLOCATED
7        PETSC_OFFLOAD_CPU
8        PETSC_OFFLOAD_GPU
9        PETSC_OFFLOAD_BOTH
10        PETSC_OFFLOAD_KOKKOS
11
12    ctypedef enum PetscMemType:
13        PETSC_MEMTYPE_HOST
14        PETSC_MEMTYPE_CUDA
15        PETSC_MEMTYPE_HIP
16        PETSC_MEMTYPE_SYCL
17
18    ctypedef enum PetscDeviceType:
19        PETSC_DEVICE_HOST
20        PETSC_DEVICE_CUDA
21        PETSC_DEVICE_HIP
22        PETSC_DEVICE_SYCL
23
24    ctypedef enum PetscStreamType:
25        PETSC_STREAM_DEFAULT
26        PETSC_STREAM_NONBLOCKING
27        PETSC_STREAM_DEFAULT_WITH_BARRIER
28        PETSC_STREAM_NONBLOCKING_WITH_BARRIER
29
30    ctypedef enum PetscDeviceContextJoinMode:
31        PETSC_DEVICE_CONTEXT_JOIN_DESTROY
32        PETSC_DEVICE_CONTEXT_JOIN_SYNC
33        PETSC_DEVICE_CONTEXT_JOIN_NO_SYNC
34
35    PetscErrorCode PetscDeviceCreate(PetscDeviceType, PetscInt, PetscDevice *)
36    PetscErrorCode PetscDeviceDestroy(PetscDevice *)
37    PetscErrorCode PetscDeviceConfigure(PetscDevice)
38    PetscErrorCode PetscDeviceView(PetscDevice, PetscViewer)
39    PetscErrorCode PetscDeviceGetType(PetscDevice, PetscDeviceType *)
40    PetscErrorCode PetscDeviceGetDeviceId(PetscDevice, PetscInt *)
41    PetscDeviceType PETSC_DEVICE_DEFAULT()
42    PetscErrorCode PetscDeviceSetDefaultDeviceType(PetscDeviceType)
43    PetscErrorCode PetscDeviceInitialize(PetscDeviceType)
44    PetscBool PetscDeviceInitialized(PetscDeviceType)
45
46    PetscErrorCode PetscDeviceContextCreate(PetscDeviceContext *)
47    PetscErrorCode PetscDeviceContextDestroy(PetscDeviceContext *)
48    PetscErrorCode PetscDeviceContextSetStreamType(PetscDeviceContext, PetscStreamType)
49    PetscErrorCode PetscDeviceContextGetStreamType(PetscDeviceContext, PetscStreamType *)
50    PetscErrorCode PetscDeviceContextSetDevice(PetscDeviceContext, PetscDevice)
51    PetscErrorCode PetscDeviceContextGetDevice(PetscDeviceContext, PetscDevice *)
52    PetscErrorCode PetscDeviceContextGetDeviceType(PetscDeviceContext, PetscDeviceType *)
53    PetscErrorCode PetscDeviceContextSetUp(PetscDeviceContext)
54    PetscErrorCode PetscDeviceContextDuplicate(PetscDeviceContext, PetscDeviceContext *)
55    PetscErrorCode PetscDeviceContextQueryIdle(PetscDeviceContext, PetscBool *)
56    PetscErrorCode PetscDeviceContextWaitForContext(PetscDeviceContext, PetscDeviceContext)
57    PetscErrorCode PetscDeviceContextForkWithStreamType(PetscDeviceContext, PetscStreamType, PetscInt, PetscDeviceContext **)
58    PetscErrorCode PetscDeviceContextFork(PetscDeviceContext, PetscInt, PetscDeviceContext **)
59    PetscErrorCode PetscDeviceContextJoin(PetscDeviceContext, PetscInt, PetscDeviceContextJoinMode, PetscDeviceContext **)
60    PetscErrorCode PetscDeviceContextSynchronize(PetscDeviceContext)
61    PetscErrorCode PetscDeviceContextSetFromOptions(MPI_Comm, PetscDeviceContext)
62    PetscErrorCode PetscDeviceContextView(PetscDeviceContext, PetscViewer)
63    PetscErrorCode PetscDeviceContextViewFromOptions(PetscDeviceContext, PetscObject, const char name[])
64    PetscErrorCode PetscDeviceContextGetCurrentContext(PetscDeviceContext *)
65    PetscErrorCode PetscDeviceContextSetCurrentContext(PetscDeviceContext)
66
67cdef extern from * nogil: # custom.h
68    PetscErrorCode PetscDeviceReference(PetscDevice)
69
70cdef inline PetscDeviceType asDeviceType(object dtype) except <PetscDeviceType>(-1):
71    if isinstance(dtype, str):
72        dtype = dtype.upper()
73        try:
74            return getattr(Device.Type, dtype)
75        except AttributeError:
76            raise ValueError("unknown device type: %s" % dtype)
77    return dtype
78
79cdef inline str toDeviceType(PetscDeviceType dtype):
80    try:
81        return Device.Type.__enum2str[dtype]
82    except KeyError:
83        raise NotImplementedError("unhandled PetscDeviceType %d" % <int>dtype)
84
85cdef inline PetscStreamType asStreamType(object stype) except <PetscStreamType>(-1):
86    if isinstance(stype, str):
87        stype = stype.upper()
88        try:
89            return getattr(DeviceContext.StreamType, stype)
90        except AttributeError:
91            raise ValueError("unknown stream type: %s" % stype)
92    return stype
93
94cdef inline str toStreamType(PetscStreamType stype):
95    try:
96        return DeviceContext.StreamType.__enum2str[stype]
97    except KeyError:
98        raise NotImplementedError("unhandled PetscStreamType %d" % <int>stype)
99
100cdef inline PetscDeviceContextJoinMode asJoinMode(object jmode) except <PetscDeviceContextJoinMode>(-1):
101    if isinstance(jmode, str):
102        jmode = jmode.upper()
103        try:
104            return getattr(DeviceContext.JoinMode, jmode)
105        except AttributeError:
106            raise ValueError("unknown join mode: %s" % jmode)
107    return jmode
108
109cdef inline str toJoinMode(PetscDeviceContextJoinMode jmode):
110    try:
111        return DeviceContext.JoinMode.__enum2str[jmode]
112    except KeyError:
113        raise NotImplementedError("unhandled PetscDeviceContextJoinMode %d" % <int>jmode)
114