compiler: patch multi-subdim isapce handling

Author: mloubout

devito/finite_differences/differentiable.py

@@ -259,7 +259,7 @@ def __rfloordiv__(self, other):
         from .elementary import floor
         return floor(other / self)
 
-    def safe_inv(self, ref, safe=False):
+    def _inv(self, ref, safe=False):
         if safe:
             return SafeInv(self, ref or self)
         else:

devito/passes/clusters/implicit.py

@@ -175,13 +175,17 @@ def callback(self, clusters, prefix):
 
             # Make sure the "implicit expressions" are scheduled in
             # the innermost loop such that the thicknesses can be computed
-            edims = set(retrieve_dimensions(mapper.values(), deep=True))
-            if dim not in edims or not edims.issubset(prefix.dimensions):
+            def key(tkn):
+                edims = set(retrieve_dimensions(tkn, deep=True))
+                return dim._defines & edims and edims.issubset(prefix.dimensions)
+
+            mapper = {k: v for k, v in mapper.items() if key(v)}
+            if not mapper:
                 continue
 
             found[d.functions].clusters.append(c)
             found[d.functions].mapper = reduce(found[d.functions].mapper,
-                                               mapper, edims, prefix)
+                                               mapper, {dim}, prefix)
 
         # Turn the reduced mapper into a list of equations
         processed = []
@@ -262,7 +266,7 @@ def reduce(m0, m1, edims, prefix):
     def key(i):
         try:
             return i.indices[d]
-        except AttributeError:
+        except (KeyError, AttributeError):
             return i
 
     mapper = {}

devito/passes/iet/misc.py

@@ -6,7 +6,6 @@
 
 from devito.finite_differences import Max, Min
 from devito.finite_differences.differentiable import SafeInv
-from devito.logger import warning
 from devito.ir import (Any, Forward, DummyExpr, Iteration, EmptyList, Prodder,
                        FindApplications, FindNodes, FindSymbols, Transformer,
                        Uxreplace, filter_iterations, retrieve_iteration_tree,
@@ -155,7 +154,7 @@ def _generate_macros(iet, tracker=None, langbb=None, printer=CPrinter, **kwargs)
                      for define, expr in headers)
 
     # Generate Macros from higher-level SymPy objects
-    mheaders, includes = _generate_macros_math(iet, langbb=langbb)
+    mheaders, includes = _generate_macros_math(iet, langbb=langbb, printer=printer)
     includes = sorted(includes, key=str)
     headers.extend(sorted(mheaders, key=str))
 
@@ -199,25 +198,25 @@ def _generate_macros_findexeds(iet, sregistry=None, tracker=None, **kwargs):
     return iet
 
 
-def _generate_macros_math(iet, langbb=None):
+def _generate_macros_math(iet, langbb=None, printer=CPrinter):
     headers = []
     includes = []
     for i in FindApplications().visit(iet):
-        header, include = _lower_macro_math(i, langbb)
+        header, include = _lower_macro_math(i, langbb, printer)
         headers.extend(header)
         includes.extend(include)
 
     return headers, set(includes) - {None}
 
 
 @singledispatch
-def _lower_macro_math(expr, langbb):
+def _lower_macro_math(expr, langbb, printer):
     return (), {}
 
 
 @_lower_macro_math.register(Min)
 @_lower_macro_math.register(sympy.Min)
-def _(expr, langbb):
+def _(expr, langbb, printer):
     if has_integer_args(*expr.args):
         return (('MIN(a,b)', ('(((a) < (b)) ? (a) : (b))')),), {}
     else:
@@ -226,23 +225,28 @@ def _(expr, langbb):
 
 @_lower_macro_math.register(Max)
 @_lower_macro_math.register(sympy.Max)
-def _(expr, langbb):
+def _(expr, langbb, printer):
     if has_integer_args(*expr.args):
         return (('MAX(a,b)', ('(((a) > (b)) ? (a) : (b))')),), {}
     else:
         return (), as_tuple(langbb.get('header-math'))
 
 
 @_lower_macro_math.register(SafeInv)
-def _(expr, langbb):
+def _(expr, langbb, printer):
     try:
-        eps = np.finfo(expr.base.dtype).resolution**2
-    except ValueError:
-        warning(f"dtype not recognized in SafeInv for {expr.base}, assuming float32")
-        eps = np.finfo(np.float32).resolution**2
-    b = Cast('b', dtype=np.float32)
+        dtype = expr.base.dtype
+        eps = np.finfo(dtype).resolution**2
+    except (AttributeError, ValueError):
+        dtype = np.float32
+        eps = np.finfo(dtype).resolution**2
+
+    b = printer()._print(Cast('b', dtype=dtype))
+    ext = 'F' if dtype is np.float32 else ''
+
     return (('SAFEINV(a, b)',
-             f'(((a) < {eps}F || ({b}) < {eps}F) ? (0.0F) : ((1.0F) / (a)))'),), {}
+             f'(((a) < {eps}{ext} || ({b}) < {eps}{ext}) ? '
+             f'(0.0{ext}) : ((1.0{ext}) / (a)))'),), {}
 
 
 @iet_pass

devito/symbolics/inspection.py

@@ -13,7 +13,7 @@
 from devito.symbolics.extended_sympy import (CallFromPointer, Cast,
                                              DefFunction, ReservedWord)
 from devito.symbolics.queries import q_routine
-from devito.tools import as_tuple, prod
+from devito.tools import as_tuple, prod, is_integer
 from devito.tools.dtypes_lowering import infer_dtype
 
 __all__ = ['compare_ops', 'estimate_cost', 'has_integer_args', 'sympy_dtype']
@@ -287,12 +287,14 @@ def has_integer_args(*args):
         try:
             return np.issubdtype(args[0].dtype, np.integer)
         except AttributeError:
-            return args[0].is_integer
+            return is_integer(args[0])
 
     res = True
     for a in args:
         try:
-            if isinstance(a, INT):
+            if isinstance(a, INT) or \
+               is_integer(a) or \
+               has_integer_args(a):
                 res = res and True
             elif len(a.args) > 0:
                 res = res and has_integer_args(*a.args)

devito/types/basic.py

@@ -1028,6 +1028,14 @@ def _grid_map(self):
 
         return mapper
 
+    @property
+    def is_harmonic(self):
+        return self.avg_mode == 'harmonic' or self.avg_mode == 'safe_harmonic'
+
+    @property
+    def is_harmonic_safe(self):
+        return self.avg_mode == 'safe_harmonic'
+
     def _evaluate(self, **kwargs):
         """
         Evaluate off the grid with 2nd order interpolation.
@@ -1046,8 +1054,8 @@ def _evaluate(self, **kwargs):
         io = self.interp_order
         retval = self.subs({i.subs(subs): self.indices_ref[d]
                             for d, i in mapper.items()})
-        if 'harmonic' in self._avg_mode:
-            retval = retval.safe_inv(retval, safe='safe' in self._avg_mode)
+        if self.is_harmonic:
+            retval = retval._inv(retval, safe=self.is_harmonic_safe)
 
         # Apply interpolation from inner most dim
         for d, i in mapper.items():
@@ -1059,9 +1067,8 @@ def _evaluate(self, **kwargs):
         retval = retval.subs(subs)
 
         # If harmonic averaging, invert at the end
-        if 'harmonic' in self._avg_mode:
-            retval = retval.safe_inv(self.function.subs(subs),
-                                     safe='safe' in self._avg_mode)
+        if self.is_harmonic:
+            retval = retval._inv(self.function.subs(subs), safe=self.is_harmonic_safe)
 
         return retval
 

devito/types/dense.py

@@ -1541,7 +1541,7 @@ def _evaluate(self, **kwargs):
         retval = super()._evaluate(**kwargs)
         if not self._time_buffering and not retval.is_Function:
             # Saved TimeFunction might need streaming, expand interpolations
-            # for easier processing.
+            # for easier processing
             return retval.evaluate
         else:
             return retval