compiler: imporve handling of dtype in cse

Author: mloubout

devito/finite_differences/differentiable.py

@@ -654,27 +654,29 @@ def __new__(cls, *args, **kwargs):
             raise ValueError(f"{cls.__name__} is constructed with exactly one arg;"
                              f" {len(args)} were supplied.")
 
-        # Diffify any Add, Mul, etc which might be in the expression
-        new_args = (diffify(args[0]),)
-
-        if not np.issubdtype(new_args[0].dtype, np.complexfloating):
-            raise ValueError(f"{cls.__name__} requires a complex dtype,"
-                             f" not {new_args[0].dtype.__name__}.")
-
-        return super().__new__(cls, *new_args, **kwargs)
+        return super().__new__(cls, *args, **kwargs)
 
     def __str__(self):
         return f"{self.__class__.__name__}({self.args[0]})"
 
     __repr__ = __str__
 
 
-class Real(ComplexPart):
+class RealComplexPart(ComplexPart):
+
+    @cached_property
+    def dtype(self):
+        dtypes = {getattr(e, 'dtype', None) for e in self.free_symbols}
+        dtype = infer_dtype(dtypes - {None})
+        return dtype(0).real.__class__
+
+
+class Real(RealComplexPart):
     """Get the real part of an expression"""
     _name = 'real'
 
 
-class Imag(ComplexPart):
+class Imag(RealComplexPart):
     """Get the imaginary part of an expression"""
     _name = 'imag'
 

devito/ir/clusters/cluster.py

@@ -304,12 +304,13 @@ def dtype(self):
         """
         dtypes = set()
         for i in self.exprs:
-            try:
-                if np.issubdtype(i.dtype, np.generic):
-                    dtypes.add(i.dtype)
-            except TypeError:
-                # E.g. `i.dtype` is a ctypes pointer, which has no dtype equivalent
-                pass
+            # try:
+            if np.issubdtype(i.dtype, np.generic):
+                dtypes.add(i.dtype)
+            # except TypeError:
+            #     print(i, type(i), i.dtype, np.issubdtype(i.dtype, np.generic))
+            #     # E.g. `i.dtype` is a ctypes pointer, which has no dtype equivalent
+            #     pass
 
         return infer_dtype(dtypes)
 

devito/passes/clusters/aliases.py

@@ -888,7 +888,6 @@ def lower_schedule(schedule, meta, sregistry, ftemps, min_dtype):
         else:
             # Degenerate case: scalar expression
             assert writeto.size == 0
-
             dtype = sympy_dtype(pivot, base=meta.dtype, smin=min_dtype)
             obj = Temp(name=name, dtype=dtype)
             expression = Eq(obj, uxreplace(pivot, subs))
@@ -1354,7 +1353,7 @@ def cost(self):
         # Not just the sum for the individual items' cost! There might be
         # redundancies, which we factor out here...
         counter = generator()
-        make = lambda: Symbol(name='dummy%d' % counter(), dtype=np.float32)
+        make = lambda _: Symbol(name='dummy%d' % counter(), dtype=np.float32)
 
         tot = 0
         for v in as_mapper(self, lambda i: i.ispace).values():

devito/passes/clusters/cse.py

@@ -14,7 +14,7 @@
 from devito.ir import Cluster, Scope, cluster_pass
 from devito.symbolics import estimate_cost, q_leaf, q_terminal
 from devito.symbolics.manipulation import _uxreplace
-from devito.tools import DAG, as_list, as_tuple, frozendict
+from devito.tools import DAG, as_list, as_tuple, frozendict, infer_dtype
 from devito.types import Eq, Symbol, Temp
 
 __all__ = ['cse']
@@ -78,7 +78,8 @@ def cse(cluster, sregistry=None, options=None, **kwargs):
     if cluster.is_fence:
         return cluster
 
-    make = lambda: CTemp(name=sregistry.make_name(), dtype=dtype)
+    make_dtype = lambda e: np.promote_types(e.dtype, dtype).type
+    make = lambda e: CTemp(name=sregistry.make_name(), dtype=make_dtype(e))
 
     exprs = _cse(cluster, make, min_cost=min_cost, mode=mode)
 
@@ -118,7 +119,7 @@ def _cse(maybe_exprs, make, min_cost=1, mode='basic'):
             exprs = maybe_exprs
             scope = Scope(maybe_exprs)
         else:
-            exprs = [Eq(make(), e) for e in maybe_exprs]
+            exprs = [Eq(make(e), e) for e in maybe_exprs]
             scope = Scope([])
 
     # Some sub-expressions aren't really "common" -- that's the case of Dimension-
@@ -155,7 +156,7 @@ def _cse(maybe_exprs, make, min_cost=1, mode='basic'):
         candidates = [c for c in candidates if c.cost == cost]
 
         # Apply replacements
-        chosen = [(c, scheduled.get(c.key) or make()) for c in candidates]
+        chosen = [(c, scheduled.get(c.key) or make(c)) for c in candidates]
         exprs = _inject(exprs, chosen, scheduled)
 
     # Drop useless temporaries (e.g., r0=r1)
@@ -275,6 +276,12 @@ def __new__(cls, expr, conditionals=None, sources=()):
     def expr(self):
         return self[0]
 
+    @property
+    def dtype(self):
+        dtypes = {getattr(e, 'dtype', None)
+                  for e in self.expr.free_symbols}
+        return infer_dtype(dtypes - {None})
+
     @property
     def conditionals(self):
         return self[1]

tests/test_cse.py

@@ -109,7 +109,7 @@ def test_default_algo(exprs, expected, min_cost):
         exprs[i] = DummyEq(indexify(diffify(eval(e).evaluate)))
 
     counter = generator()
-    make = lambda: CTemp(name='r%d' % counter()).indexify()
+    make = lambda _: CTemp(name='r%d' % counter()).indexify()
     processed = _cse(exprs, make, min_cost)
 
     assert len(processed) == len(expected)
@@ -241,7 +241,7 @@ def test_advanced_algo(exprs, expected):
         exprs[i] = DummyEq(indexify(diffify(eval(e).evaluate)))
 
     counter = generator()
-    make = lambda: CTemp(name='r%d' % counter(), dtype=np.float32).indexify()
+    make = lambda _: CTemp(name='r%d' % counter(), dtype=np.float32).indexify()
     processed = _cse(exprs, make, mode='advanced')
 
     assert len(processed) == len(expected)