def one_norm(in_buf, inout_buf, datatype): # This mapping MPI datatype -> NumPy dtype # only works for built-in MPI datatypes typecode = MPI._typecode(datatype) assert typecode is not None # check MPI datatype is built-in dtype = np.dtype(typecode) # `in_buf`, `inout_buf` are buffer-like objects exposing raw bytes # don't use them directly but convert them to numpy arrays in_array = np.frombuffer(in_buf, dtype) # `in_array` and `in_buf` share memory inout_array = np.frombuffer(inout_buf, dtype) # `inout_array` and `inout_buf` share memory
def fn_sum(buffer_a, buffer_b, t): tc = MPI._typecode(t) # map MPI datatype -> Python typecode array_a = np.frombuffer(buffer_a, dtype=tc) array_b = np.frombuffer(buffer_b, dtype=tc) array_b += array_a