1# -------------------------------------------------------------------- 2 3cdef extern from * nogil: 4 5 const MPI_Comm MPI_COMM_NULL 6 const MPI_Comm MPI_COMM_SELF 7 const MPI_Comm MPI_COMM_WORLD 8 9 const MPI_Datatype MPI_DATATYPE_NULL 10 const MPI_Op MPI_OP_NULL 11 12 enum: MPI_IDENT 13 enum: MPI_CONGRUENT 14 int MPI_Comm_compare(MPI_Comm, MPI_Comm, int*) 15 16 int MPI_Comm_size(MPI_Comm, int*) 17 int MPI_Comm_rank(MPI_Comm, int*) 18 int MPI_Barrier(MPI_Comm) 19 20 int MPI_Initialized(int*) 21 int MPI_Finalized(int*) 22 23 ctypedef int MPI_Fint 24 MPI_Fint MPI_Comm_c2f(MPI_Comm) 25 26cdef extern from * nogil: 27 28 MPI_Comm PETSC_COMM_SELF 29 MPI_Comm PETSC_COMM_WORLD 30 31 PetscErrorCode PetscCommDuplicate(MPI_Comm, MPI_Comm*, int*) 32 PetscErrorCode PetscCommDestroy(MPI_Comm*) 33 34# -------------------------------------------------------------------- 35 36cdef extern from "cython.h": 37 void *Cython_ImportFunction(object, char[], char[]) except? NULL 38 39ctypedef MPI_Comm* PyMPICommGet(object) except NULL 40ctypedef object PyMPICommNew(MPI_Comm) 41ctypedef MPI_Datatype* PyMPIDatatypeGet(object) except NULL 42ctypedef MPI_Op* PyMPIOpGet(object) except NULL 43 44cdef inline MPI_Comm mpi4py_Comm_Get( 45 object comm, 46) except? MPI_COMM_NULL: 47 from mpi4py import MPI 48 cdef PyMPICommGet *commget = \ 49 <PyMPICommGet*> Cython_ImportFunction(MPI, b"PyMPIComm_Get", b"MPI_Comm *(PyObject *)") 50 if commget == NULL: return MPI_COMM_NULL 51 cdef MPI_Comm *ptr = commget(comm) 52 if ptr == NULL: return MPI_COMM_NULL 53 return ptr[0] 54 55cdef inline object mpi4py_Comm_New(MPI_Comm comm): 56 from mpi4py import MPI 57 cdef PyMPICommNew *commnew = \ 58 <PyMPICommNew*> Cython_ImportFunction(MPI, b"PyMPIComm_New", b"PyObject *(MPI_Comm)") 59 if commnew == NULL: return None 60 return commnew(comm) 61 62cdef inline MPI_Datatype mpi4py_Datatype_Get( 63 object datatype, 64) except? MPI_DATATYPE_NULL: 65 from mpi4py import MPI 66 cdef PyMPIDatatypeGet *datatypeget = \ 67 <PyMPIDatatypeGet*> Cython_ImportFunction(MPI, b"PyMPIDatatype_Get", b"MPI_Datatype *(PyObject *)") 68 if datatypeget == NULL: return MPI_DATATYPE_NULL 69 cdef MPI_Datatype *ptr = datatypeget(datatype) 70 if ptr == NULL: return MPI_DATATYPE_NULL 71 return ptr[0] 72 73cdef inline MPI_Op mpi4py_Op_Get( 74 object op, 75) except? MPI_OP_NULL: 76 from mpi4py import MPI 77 cdef PyMPIOpGet *opget = \ 78 <PyMPIOpGet*> Cython_ImportFunction(MPI, b"PyMPIOp_Get", b"MPI_Op *(PyObject *)") 79 if opget == NULL: return MPI_OP_NULL 80 cdef MPI_Op *ptr = opget(op) 81 if ptr == NULL: return MPI_OP_NULL 82 return ptr[0] 83 84# -------------------------------------------------------------------- 85 86cdef inline PetscErrorCode PetscCommDEALLOC(MPI_Comm* comm): 87 if comm == NULL: return PETSC_SUCCESS 88 cdef MPI_Comm tmp = comm[0] 89 if tmp == MPI_COMM_NULL: return PETSC_SUCCESS 90 comm[0] = MPI_COMM_NULL 91 if not (<int>PetscInitializeCalled): return PETSC_SUCCESS 92 if (<int>PetscFinalizeCalled): return PETSC_SUCCESS 93 return PetscCommDestroy(&tmp) 94 95cdef inline MPI_Comm def_Comm( 96 object comm, MPI_Comm defv, 97) except? MPI_COMM_NULL: 98 cdef MPI_Comm retv = MPI_COMM_NULL 99 if comm is None: 100 retv = defv 101 elif isinstance(comm, Comm): 102 retv = (<Comm>comm).comm 103 elif type(comm).__module__ == 'mpi4py.MPI': 104 retv = mpi4py_Comm_Get(comm) 105 else: 106 retv = (<Comm?>comm).comm 107 return retv 108 109cdef inline Comm new_Comm(MPI_Comm comm): 110 cdef Comm ob = <Comm> Comm() 111 ob.comm = comm 112 return ob 113 114# -------------------------------------------------------------------- 115 116cdef inline int comm_size(MPI_Comm comm) except ? -1: 117 if comm == MPI_COMM_NULL: raise ValueError("null communicator") 118 cdef int size = 0 119 CHKERR(<PetscErrorCode>MPI_Comm_size(comm, &size)) 120 return size 121 122cdef inline int comm_rank(MPI_Comm comm) except ? -1: 123 if comm == MPI_COMM_NULL: raise ValueError("null communicator") 124 cdef int rank = 0 125 CHKERR(<PetscErrorCode>MPI_Comm_rank(comm, &rank)) 126 return rank 127 128# -------------------------------------------------------------------- 129