compiler: Introduce extract_dtype

Author: EdCaunt

devito/finite_differences/differentiable.py

@@ -17,7 +17,7 @@
 from devito.finite_differences.tools import make_shift_x0, coeff_priority
 from devito.logger import warning
 from devito.tools import (as_tuple, filter_ordered, flatten, frozendict,
-                          infer_dtype, is_integer, split, is_number)
+                          infer_dtype, extract_dtype, is_integer, split, is_number)
 from devito.types import Array, DimensionTuple, Evaluable, StencilDimension
 from devito.types.basic import AbstractFunction
 
@@ -665,8 +665,7 @@ class RealComplexPart(ComplexPart):
 
     @cached_property
     def dtype(self):
-        dtypes = {getattr(e, 'dtype', None) for e in self.free_symbols}
-        dtype = infer_dtype(dtypes - {None})
+        dtype = extract_dtype(self)
         return dtype(0).real.__class__
 
 

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, infer_dtype
+from devito.tools import DAG, as_list, as_tuple, frozendict, extract_dtype
 from devito.types import Eq, Symbol, Temp
 
 __all__ = ['cse']
@@ -278,9 +278,7 @@ def expr(self):
 
     @property
     def dtype(self):
-        dtypes = {getattr(e, 'dtype', None)
-                  for e in self.expr.free_symbols}
-        return infer_dtype(dtypes - {None})
+        return extract_dtype(self.expr)
 
     @property
     def conditionals(self):

devito/tools/dtypes_lowering.py

@@ -15,7 +15,8 @@
            'double3', 'double4', 'dtypes_vector_mapper', 'dtype_to_mpidtype',
            'dtype_to_cstr', 'dtype_to_ctype', 'infer_datasize', 'dtype_to_mpitype',
            'dtype_len', 'ctypes_to_cstr', 'c_restrict_void_p', 'ctypes_vector_mapper',
-           'is_external_ctype', 'infer_dtype', 'CustomDtype', 'mpi4py_mapper']
+           'is_external_ctype', 'infer_dtype', 'extract_dtype', 'CustomDtype',
+           'mpi4py_mapper']
 
 
 # *** Custom np.dtypes
@@ -365,3 +366,10 @@ def infer_dtype(dtypes):
     else:
         # E.g., mixed integer arithmetic
         return max(dtypes, key=lambda i: np.dtype(i).itemsize, default=None)
+
+
+def extract_dtype(expr):
+    """Extract the "winning" dtype from an expression"""
+    dtypes = {getattr(e, 'dtype', None)
+              for e in expr.free_symbols}
+    return infer_dtype(dtypes - {None})