numba_mpi.api.irecv

MPI_Irecv() implementation

View Source

 1# pylint: disable=duplicate-code
 2
 3"""MPI_Irecv() implementation"""
 4
 5import ctypes
 6
 7import numba
 8from mpi4py.MPI import ANY_SOURCE, ANY_TAG
 9
10from numba_mpi.api.requests import _allocate_numpy_array_of_request_handles
11from numba_mpi.common import _MPI_Comm_World_ptr, libmpi, send_recv_async_args
12from numba_mpi.utils import _mpi_addr, _mpi_dtype
13
14_MPI_Irecv = libmpi.MPI_Irecv
15_MPI_Irecv.restype = ctypes.c_int
16_MPI_Irecv.argtypes = send_recv_async_args
17
18
19@numba.njit()
20def irecv(data, source=ANY_SOURCE, tag=ANY_TAG):
21    """Wrapper for MPI_Irecv (only handles contiguous arrays, at least for now).
22    If successful (i.e. result is MPI_SUCCESS), returns c-style pointer to valid
23    MPI_Request handle that may be used with appropriate wait and test functions."""
24
25    assert data.flags.c_contiguous  # TODO #60
26
27    request_buffer = _allocate_numpy_array_of_request_handles()
28
29    status = _MPI_Irecv(
30        data.ctypes.data,
31        data.size,
32        _mpi_dtype(data),
33        source,
34        tag,
35        _mpi_addr(_MPI_Comm_World_ptr),
36        request_buffer.ctypes.data,
37    )
38
39    return status, request_buffer

@numba.njit()

def irecv(data, source=-2, tag=-1): View Source

20@numba.njit()
21def irecv(data, source=ANY_SOURCE, tag=ANY_TAG):
22    """Wrapper for MPI_Irecv (only handles contiguous arrays, at least for now).
23    If successful (i.e. result is MPI_SUCCESS), returns c-style pointer to valid
24    MPI_Request handle that may be used with appropriate wait and test functions."""
25
26    assert data.flags.c_contiguous  # TODO #60
27
28    request_buffer = _allocate_numpy_array_of_request_handles()
29
30    status = _MPI_Irecv(
31        data.ctypes.data,
32        data.size,
33        _mpi_dtype(data),
34        source,
35        tag,
36        _mpi_addr(_MPI_Comm_World_ptr),
37        request_buffer.ctypes.data,
38    )
39
40    return status, request_buffer

Wrapper for MPI_Irecv (only handles contiguous arrays, at least for now). If successful (i.e. result is MPI_SUCCESS), returns c-style pointer to valid MPI_Request handle that may be used with appropriate wait and test functions.