numba_mpi.api.recv

MPI_Recv() implementation

 1"""MPI_Recv() implementation"""
 2
 3import ctypes
 4
 5import numba
 6import numpy as np
 7from mpi4py.MPI import ANY_SOURCE, ANY_TAG
 8
 9from numba_mpi.common import (
10    _MPI_Comm_World_ptr,
11    _MpiStatusPtr,
12    create_status_buffer,
13    libmpi,
14    send_recv_args,
15)
16from numba_mpi.utils import _mpi_addr, _mpi_dtype
17
18_MPI_Recv = libmpi.MPI_Recv
19_MPI_Recv.restype = ctypes.c_int
20_MPI_Recv.argtypes = send_recv_args + [
21    _MpiStatusPtr,
22]
23
24
25@numba.njit()
26def recv(data, source=ANY_SOURCE, tag=ANY_TAG):
27    """file containing wrapper for MPI_Recv (writes data directly if `data` is contiguous, otherwise
28    allocates a buffer and later copies the data into non-contiguous `data` array).
29    Returns integer status code (0 == MPI_SUCCESS)"""
30    status_buffer = create_status_buffer()
31
32    buffer = (
33        data
34        if data.flags.c_contiguous
35        else np.empty(
36            data.shape, data.dtype
37        )  # np.empty_like(data, order='C') fails with Numba
38    )
39
40    status = _MPI_Recv(
41        buffer.ctypes.data,
42        buffer.size,
43        _mpi_dtype(data),
44        source,
45        tag,
46        _mpi_addr(_MPI_Comm_World_ptr),
47        status_buffer.ctypes.data,
48    )
49
50    if not data.flags.c_contiguous:
51        data[...] = buffer
52
53    return status
@numba.njit()
def recv(data, source=-2, tag=-1): 
26@numba.njit()
27def recv(data, source=ANY_SOURCE, tag=ANY_TAG):
28    """file containing wrapper for MPI_Recv (writes data directly if `data` is contiguous, otherwise
29    allocates a buffer and later copies the data into non-contiguous `data` array).
30    Returns integer status code (0 == MPI_SUCCESS)"""
31    status_buffer = create_status_buffer()
32
33    buffer = (
34        data
35        if data.flags.c_contiguous
36        else np.empty(
37            data.shape, data.dtype
38        )  # np.empty_like(data, order='C') fails with Numba
39    )
40
41    status = _MPI_Recv(
42        buffer.ctypes.data,
43        buffer.size,
44        _mpi_dtype(data),
45        source,
46        tag,
47        _mpi_addr(_MPI_Comm_World_ptr),
48        status_buffer.ctypes.data,
49    )
50
51    if not data.flags.c_contiguous:
52        data[...] = buffer
53
54    return status

file containing wrapper for MPI_Recv (writes data directly if data is contiguous, otherwise allocates a buffer and later copies the data into non-contiguous data array). Returns integer status code (0 == MPI_SUCCESS)