diff --git a/devito/passes/clusters/derivatives.py b/devito/passes/clusters/derivatives.py index b8c4108b83..ded4a8661e 100644 --- a/devito/passes/clusters/derivatives.py +++ b/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) diff --git a/devito/passes/iet/engine.py b/devito/passes/iet/engine.py index c879595e83..37334cea9b 100644 --- a/devito/passes/iet/engine.py +++ b/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 diff --git a/devito/tools/dtypes_lowering.py b/devito/tools/dtypes_lowering.py index 5e6ccf82e0..2a8f5d7800 100644 --- a/devito/tools/dtypes_lowering.py +++ b/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 diff --git a/tests/test_mpi.py b/tests/test_mpi.py index a9b464aef0..26c7bbf90c 100644 --- a/tests/test_mpi.py +++ b/tests/test_mpi.py @@ -2067,8 +2067,9 @@ 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 @@ -2076,10 +2077,14 @@ def test_allreduce_time(self, mode): 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):