numba_mpi.api.requests

  1"""Utilities for handling MPI_Requests with respective Wait and Test function
  2wrappers.
  3"""
  4
  5import ctypes
  6
  7import numba
  8import numpy as np
  9from numba.core import types
 10from numba.core.extending import overload
 11
 12from numba_mpi.common import (
 13    RequestType,
 14    _MpiRequestPtr,
 15    _MpiStatusPtr,
 16    create_status_buffer,
 17    libmpi,
 18)
 19
 20
 21@numba.njit
 22def _allocate_numpy_array_of_request_handles(count=1):
 23    """Helper function for creating numpy array storing pointers to MPI_Request handles."""
 24    return np.empty(count, dtype=RequestType)
 25
 26
 27_MPI_Wait = libmpi.MPI_Wait
 28_MPI_Wait.restype = ctypes.c_int
 29_MPI_Wait.argtypes = [_MpiRequestPtr, _MpiStatusPtr]
 30
 31
 32@numba.njit
 33def wait(request):
 34    """Wrapper for MPI_Wait. Returns integer status code (0 == MPI_SUCCESS).
 35    Status is currently not handled. Requires 'request' parameter to be a
 36    c-style pointer to MPI_Request (such as returned by 'isend'/'irecv').
 37
 38    Uninitialized contents of 'request' (e.g., from numpy.empty()) may
 39    cause invalid pointer dereference and segmentation faults.
 40    """
 41
 42    status_buffer = create_status_buffer()
 43    status = _MPI_Wait(request.ctypes.data, status_buffer.ctypes.data)
 44
 45    return status
 46
 47
 48_MPI_Waitall = libmpi.MPI_Waitall
 49_MPI_Waitall.restype = ctypes.c_int
 50_MPI_Waitall.argtypes = [ctypes.c_int, _MpiRequestPtr, _MpiStatusPtr]
 51
 52
 53@numba.njit
 54def _waitall_array_impl(requests):
 55    """MPI_Waitall implementation for contiguous numpy arrays of MPI_Requests"""
 56
 57    status_buffer = create_status_buffer(requests.size)
 58
 59    status = _MPI_Waitall(
 60        requests.size, requests.ctypes.data, status_buffer.ctypes.data
 61    )
 62
 63    return status
 64
 65
 66def waitall(requests):
 67    """Wrapper for MPI_Waitall. Returns integer status code (0 == MPI_SUCCESS).
 68    Status is currently not handled. Requires 'requests' parameter to be an
 69    array or tuple of MPI_Request objects.
 70
 71    Uninitialized contents of 'requests' (e.g., from numpy.empty()) may
 72    cause invalid pointer dereference and segmentation faults.
 73    """
 74    if isinstance(requests, np.ndarray):
 75        return _waitall_array_impl(requests)
 76
 77    if isinstance(requests, (list, tuple)):
 78        request_buffer = np.hstack(requests)
 79        return _waitall_array_impl(request_buffer)
 80
 81    raise TypeError("Invalid type for array of MPI_Request objects")
 82
 83
 84@overload(waitall)
 85def _waitall_impl(requests):
 86    """List of overloads for MPI_Waitall implementation"""
 87
 88    if isinstance(requests, types.Array):
 89
 90        def impl(requests):
 91            return _waitall_array_impl(requests)
 92
 93    elif isinstance(requests, types.UniTuple):
 94
 95        def impl(requests):
 96            req_buffer = np.hstack(requests)
 97            return _waitall_array_impl(req_buffer)
 98
 99    else:
100        raise TypeError("Invalid type for array of MPI_Request objects")
101
102    return impl
103
104
105_MPI_Waitany = libmpi.MPI_Waitany
106_MPI_Waitany.restype = ctypes.c_int
107_MPI_Waitany.argtypes = [ctypes.c_int, _MpiRequestPtr, ctypes.c_void_p, _MpiStatusPtr]
108
109
110@numba.njit
111def _waitany_array_impl(requests):
112    """MPI_Waitany implementation for contiguous numpy arrays of MPI_Requests"""
113
114    status_buffer = create_status_buffer()
115    index = np.empty(1, dtype=np.intc)
116
117    status = _MPI_Waitany(
118        requests.size,
119        requests.ctypes.data,
120        index.ctypes.data,
121        status_buffer.ctypes.data,
122    )
123
124    return status, index[0]
125
126
127def waitany(requests):
128    """Wrapper for MPI_Waitany. Returns tuple of integers, first representing
129    status; second - the index of request that was completed. Status is
130    currently not handled. Requires 'requests' parameter to be an array
131    or tuple of MPI_Request objects.
132
133    Uninitialized contents of 'requests' (e.g., from numpy.empty()) may
134    cause invalid pointer dereference and segmentation faults.
135    """
136
137    if isinstance(requests, np.ndarray):
138        return _waitany_array_impl(requests)
139
140    if isinstance(requests, (list, tuple)):
141        request_buffer = np.hstack(requests)
142        return _waitany_array_impl(request_buffer)
143
144    raise TypeError("Invalid type for array of MPI_Request objects")
145
146
147@overload(waitany)
148def _waitany_impl(requests):
149    """List of overloads for MPI_Waitany implementation"""
150
151    if isinstance(requests, types.Array):
152
153        def impl(requests):
154            return _waitany_array_impl(requests)
155
156    elif isinstance(requests, types.UniTuple):
157
158        def impl(requests):
159            req_buffer = np.hstack(requests)
160            return _waitany_array_impl(req_buffer)
161
162    else:
163        raise TypeError("Invalid type for array of MPI_Request objects")
164
165    return impl
166
167
168_MPI_Test = libmpi.MPI_Test
169_MPI_Test.restype = ctypes.c_int
170_MPI_Test.argtypes = [_MpiRequestPtr, ctypes.c_void_p, _MpiStatusPtr]
171
172
173@numba.njit
174def test(request):
175    """Wrapper for MPI_Test. Returns tuple containing both status and boolean
176    flag that indicates whether given request is completed. Status is currently
177    not handled. Requires 'request' parameter to be a c-style pointer to
178    MPI_Request (such as returned by 'isend'/'irecv').
179
180    Uninitialized contents of 'request' (e.g., from numpy.empty()) may
181    cause invalid pointer dereference and segmentation faults.
182    """
183
184    status_buffer = create_status_buffer()
185    flag = np.empty(1, dtype=np.intc)
186
187    status = _MPI_Test(request.ctypes.data, flag.ctypes.data, status_buffer.ctypes.data)
188
189    return status, flag[0] != 0
190
191
192_MPI_Testall = libmpi.MPI_Testall
193_MPI_Testall.restype = ctypes.c_int
194_MPI_Testall.argtypes = [ctypes.c_int, _MpiRequestPtr, ctypes.c_void_p, _MpiStatusPtr]
195
196
197@numba.njit
198def _testall_array_impl(requests):
199    """MPI_Testall implementation for contiguous numpy arrays of MPI_Requests"""
200
201    status_buffer = create_status_buffer(requests.size)
202    flag = np.empty(1, dtype=np.intc)
203    status = _MPI_Testall(
204        requests.size,
205        requests.ctypes.data,
206        flag.ctypes.data,
207        status_buffer.ctypes.data,
208    )
209
210    return status, flag[0] != 0
211
212
213def testall(requests):
214    """Wrapper for MPI_Testall. Returns tuple containing both status and boolean
215    flag that indicates whether given request is completed. Status is currently
216    not handled. Requires 'requests' parameter to be an array or tuple of
217    MPI_Request objects.
218
219    Uninitialized contents of 'requests' (e.g., from numpy.empty()) may
220    cause invalid pointer dereference and segmentation faults.
221    """
222    if isinstance(requests, np.ndarray):
223        return _testall_array_impl(requests)
224
225    if isinstance(requests, (list, tuple)):
226        request_buffer = np.hstack(requests)
227        return _testall_array_impl(request_buffer)
228
229    raise TypeError("Invalid type for array of MPI_Request objects")
230
231
232@overload(testall)
233def _testall_impl(requests):
234    """List of overloads for MPI_Testall implementation"""
235
236    if isinstance(requests, types.Array):
237
238        def impl(requests):
239            return _testall_array_impl(requests)
240
241    elif isinstance(requests, types.UniTuple):
242
243        def impl(requests):
244            req_buffer = np.hstack(requests)
245            return _testall_array_impl(req_buffer)
246
247    else:
248        raise TypeError("Invalid type for array of MPI_Request objects")
249
250    return impl
251
252
253_MPI_Testany = libmpi.MPI_Testany
254_MPI_Testany.restype = ctypes.c_int
255_MPI_Testany.argtypes = [
256    ctypes.c_int,
257    _MpiRequestPtr,
258    ctypes.c_void_p,
259    ctypes.c_void_p,
260    _MpiStatusPtr,
261]
262
263
264@numba.njit
265def _testany_array_impl(requests):
266    """MPI_Testany implementation for contiguous numpy arrays of MPI_Requests"""
267
268    status_buffer = create_status_buffer()
269    flag = np.empty(1, dtype=np.intc)
270    index = np.empty(1, dtype=np.intc)
271    status = _MPI_Testany(
272        requests.size,
273        requests.ctypes.data,
274        index.ctypes.data,
275        flag.ctypes.data,
276        status_buffer.ctypes.data,
277    )
278
279    return status, flag[0] != 0, index[0]
280
281
282def testany(requests):
283    """Wrapper for MPI_Testany. Returns tuple containing status, boolean flag
284    that indicates whether any of requests is completed, and index of request
285    that is guaranteed to be completed. Requires 'requests' parameter to be an
286    array or tuple of MPI_Request objects.
287
288    Uninitialized contents of 'requests' (e.g., from numpy.empty()) may
289    cause invalid pointer dereference and segmentation faults.
290    """
291
292    if isinstance(requests, np.ndarray):
293        return _testany_array_impl(requests)
294
295    if isinstance(requests, (list, tuple)):
296        request_buffer = np.hstack(requests)
297        return _testany_array_impl(request_buffer)
298
299    raise TypeError("Invalid type for array of MPI_Request objects")
300
301
302@overload(testany)
303def _testany_impl(requests):
304    """List of overloads for MPI_Testany implementation"""
305
306    if isinstance(requests, types.Array):
307
308        def impl(requests):
309            return _testany_array_impl(requests)
310
311    elif isinstance(requests, types.UniTuple):
312
313        def impl(requests):
314            req_buffer = np.hstack(requests)
315            return _testany_array_impl(req_buffer)
316
317    else:
318        raise TypeError("Invalid type for array of MPI_Request objects")
319
320    return impl