numba_mpi.api.size

MPI_Comm_size() implementation

 1"""MPI_Comm_size() implementation"""
 2
 3import ctypes
 4
 5import numba
 6import numpy as np
 7
 8from numba_mpi.common import _MPI_Comm_World_ptr, _MpiComm, libmpi
 9from numba_mpi.utils import _mpi_addr
10
11_MPI_Comm_size = libmpi.MPI_Comm_size
12_MPI_Comm_size.restype = ctypes.c_int
13_MPI_Comm_size.argtypes = [_MpiComm, ctypes.c_void_p]
14
15
16@numba.njit()
17def size():
18    """wrapper for MPI_Comm_size(), in case of failure returns 0"""
19    value = np.empty(1, dtype=np.intc)
20    status = _MPI_Comm_size(_mpi_addr(_MPI_Comm_World_ptr), value.ctypes.data)
21    if status != 0:
22        value[0] = 0
23    return value[0]
@numba.njit()
def size(): 
17@numba.njit()
18def size():
19    """wrapper for MPI_Comm_size(), in case of failure returns 0"""
20    value = np.empty(1, dtype=np.intc)
21    status = _MPI_Comm_size(_mpi_addr(_MPI_Comm_World_ptr), value.ctypes.data)
22    if status != 0:
23        value[0] = 0
24    return value[0]

wrapper for MPI_Comm_size(), in case of failure returns 0