api: fixup complex dtype handling for mpi

Author: mloubout

devito/passes/clusters/derivatives.py

@@ -7,6 +7,7 @@
 from devito.ir import Backward, Forward, Interval, IterationSpace, Queue
 from devito.passes.clusters.fusion import fuse
 from devito.symbolics import BasicWrapperMixin, reuse_if_untouched, search, uxreplace
+from devito.symbolics.inspection import sympy_dtype
 from devito.tools import infer_dtype, timed_pass
 from devito.types import Eq, Inc, Indexed, Symbol
 
@@ -170,12 +171,13 @@ def _(expr, c, ispace, weights, reusables, mapper, **kwargs):
 
     # The Symbol that will hold the result of the IndexDerivative computation
     # NOTE: created before recursing so that we ultimately get a sound ordering
+    dtype = sympy_dtype(ideriv)
     try:
         s = reusables.pop()
-        assert np.can_cast(s.dtype, w.dtype)
+        assert np.can_cast(s.dtype, dtype)
     except KeyError:
         name = sregistry.make_name(prefix='r')
-        s = Symbol(name=name, dtype=w.dtype)
+        s = Symbol(name=name, dtype=dtype)
 
     # Go inside `expr` and recursively lower any nested IndexDerivatives
     expr, processed = _core(expr, c, ispace1, weights, reusables, mapper, **kwargs)

devito/passes/iet/engine.py

@@ -259,11 +259,13 @@ def wrapper(*args, **kwargs):
         try:
             # Pure function case
             graph, = args
-            return maybe_timed(call(graph), func.__name__)(func, **kwargs)
+            ietfunc = func
         except ValueError:
             # Instance method case
             self, graph = args
-            return maybe_timed(call(graph), func.__name__)(partial(func, self), **kwargs)
+            ietfunc = partial(func, self)
+
+        return maybe_timed(call(graph), func.__name__)(ietfunc, **kwargs)
 
     return wrapper
 

devito/tools/dtypes_lowering.py

@@ -26,7 +26,7 @@
 dtype_mapper = {
     "int": np.int32,
     "float": np.float32,
-    "double": np.float64
+    "double": np.float64,
 }
 
 
@@ -98,6 +98,9 @@ def get_base_dtype(self, v, default=None):
 # Fallbacks
 dtypes_vector_mapper.update({(v, 1): v for v in dtype_mapper.values()})
 
+# Complex
+dtypes_vector_mapper.update({(np.complex64, 1): np.complex64})
+dtypes_vector_mapper.update({(np.complex128, 1): np.complex128})
 
 # *** Custom types escaping both the numpy and ctypes namespaces
 
@@ -238,7 +241,11 @@ class c_restrict_void_p(ctypes.c_void_p):
 
 ctypes_vector_mapper = {}
 for base_name, base_dtype in dtype_mapper.items():
-    base_ctype = dtype_to_ctype(base_dtype)
+    try:
+        base_ctype = dtype_to_ctype(base_dtype)
+    except NotImplementedError:
+        # E.g., np.complex64
+        continue
 
     for count in counts:
         dtype = dtypes_vector_mapper[(base_dtype, count)]
@@ -250,7 +257,6 @@ class c_restrict_void_p(ctypes.c_void_p):
 
         ctypes_vector_mapper[dtype] = ctype
 
-
 # *** ctypes lowering
 
 

tests/test_mpi.py

@@ -2067,19 +2067,24 @@ def test_merge_smart_if_within_conditional(self, mode):
         for n in FindNodes(Conditional).visit(op1):
             assert len(FindNodes(HaloUpdateCall).visit(n)) == 0
 
+    @pytest.mark.parametrize('dtype', [np.float32, np.complex64])
     @pytest.mark.parallel(mode=2)
-    def test_allreduce_time(self, mode):
+    def test_allreduce_time(self, dtype, mode):
         space_order = 8
         nx, ny = 11, 11
 
         grid = Grid(shape=(nx, ny))
         tt = grid.time_dim
         nt = 10
 
-        ux = TimeFunction(name="ux", grid=grid, time_order=1, space_order=space_order)
-        g = TimeFunction(name="g", grid=grid, dimensions=(tt, ), shape=(nt,))
+        c = tt if dtype == np.float32 else tt + tt * 1j
+
+        ux = TimeFunction(name="ux", grid=grid, time_order=1, space_order=space_order,
+                          dtype=dtype)
+        g = TimeFunction(name="g", grid=grid, dimensions=(tt, ), shape=(nt,),
+                         dtype=dtype)
 
-        op = Operator([Eq(ux.forward, ux + tt), Inc(g, ux)], name="Op")
+        op = Operator([Eq(ux.forward, ux + c), Inc(g, ux)], name="Op")
         assert_structure(op, ['t,x,y', 't'], 'txy')
 
         # Reduce should be in time loop but not in space loop
@@ -2091,7 +2096,10 @@ def test_allreduce_time(self, mode):
                 assert len(FindNodes(Call).visit(i)) == 0
 
         op.apply(time_m=0, time_M=nt-1)
-        assert np.isclose(np.max(g.data), 4356.0)
+        if dtype == np.float32:
+            assert np.isclose(np.max(g.data), 4356.0)
+        else:
+            assert np.isclose(np.max(g.data), 4356.0 + 4356.0j)
 
     @pytest.mark.parallel(mode=2)
     def test_multi_allreduce_time(self, mode):