mpi: fix distribution dtype

Author: mloubout

devito/data/distributed/transport.py

@@ -17,7 +17,7 @@
 import numpy as np
 
 from devito.mpi import MPI
-from devito.tools import dtype_to_mpidtype
+from devito.tools import dtype_to_mpidtype, mpi4py_mapper
 
 __all__ = ['sparse_exchange']
 
@@ -53,7 +53,13 @@ def sparse_exchange(comm, sendbufs, dtype, tag=0):
     """
     rank = comm.Get_rank()
     nprocs = comm.Get_size()
-    mpitype = dtype_to_mpidtype(dtype)
+
+    # Some MPI builds lack a native datatype for `dtype` (e.g. `float16`); send
+    # over a same-size byte-equivalent wire type and view back on receipt, just
+    # as the halo exchange does via `comm_dtype`. A no-op for mapped-to-self
+    # types.
+    wire = np.dtype(mpi4py_mapper.get(np.dtype(dtype).type, dtype))
+    mpitype = dtype_to_mpidtype(wire)
 
     recvd = {}
 
@@ -79,7 +85,7 @@ def sparse_exchange(comm, sendbufs, dtype, tag=0):
     for peer, buf in sendbufs.items():
         if peer == rank or buf.size == 0:
             continue
-        buf = np.ascontiguousarray(buf)
+        buf = np.ascontiguousarray(buf).view(wire)
         live_bufs.append(buf)
         sends.append(comm.Isend([buf, mpitype], dest=peer, tag=tag))
 
@@ -89,9 +95,9 @@ def sparse_exchange(comm, sendbufs, dtype, tag=0):
         comm.Probe(source=MPI.ANY_SOURCE, tag=tag, status=status)
         src = status.Get_source()
         count = status.Get_count(mpitype)
-        buf = np.empty(count, dtype=dtype)
+        buf = np.empty(count, dtype=wire)
         comm.Recv([buf, mpitype], source=src, tag=tag)
-        recvd[src] = buf
+        recvd[src] = buf.view(dtype)
 
     MPI.Request.Waitall(sends)
     return recvd