compiler: prevent hosted per-thread arrays are dereferenced within partree at read

Author: mloubout

devito/finite_differences/differentiable.py

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

devito/passes/iet/parpragma.py

@@ -317,6 +317,9 @@ def _make_parregion(self, partree, parrays):
             i = n.write
             if not (i.is_Array or i.is_TempFunction):
                 continue
+            elif partree.dim in i.dimensions:
+                # Non-local Array (full iteration space): no need to vector-expand
+                continue
             elif i in parrays:
                 pi = parrays[i]
             else:

devito/symbolics/inspection.py

@@ -4,6 +4,7 @@
 from sympy import (Function, Indexed, Integer, Mul, Number,
                    Pow, S, Symbol, Tuple)
 from sympy.core.numbers import ImaginaryUnit
+from sympy.core.function import Application
 
 from devito.finite_differences import Derivative
 from devito.finite_differences.differentiable import IndexDerivative
@@ -116,7 +117,7 @@ def estimate_cost(exprs, estimate=False):
 estimate_values = {
     'elementary': 100,
     'pow': 50,
-    'SafeInv': 10,
+    'SafeInv': 50,
     'div': 5,
     'Abs': 5,
     'floor': 1,
@@ -211,6 +212,7 @@ def _(expr, estimate, seen):
 
 
 @_estimate_cost.register(Function)
+@_estimate_cost.register(Application)
 def _(expr, estimate, seen):
     if q_routine(expr):
         flops, _ = zip(*[_estimate_cost(a, estimate, seen) for a in expr.args])
@@ -227,6 +229,7 @@ def _(expr, estimate, seen):
             flops += 1
     else:
         flops = 0
+
     return flops, False
 
 

devito/types/basic.py

@@ -849,7 +849,7 @@ def __init_finalize__(self, *args, **kwargs):
 
         # Averaging mode for off the grid evaluation
         self._avg_mode = kwargs.get('avg_mode', 'arithmetic')
-        if self._avg_mode not in ['arithmetic', 'harmonic']:
+        if self._avg_mode not in ['arithmetic', 'harmonic', 'safe_harmonic']:
             raise ValueError("Invalid averaging mode_mode %s, accepted values are"
                              " arithmetic or harmonic" % self._avg_mode)
 
@@ -989,7 +989,7 @@ def c0(self):
     def _eval_deriv(self):
         return self
 
-    @cached_property
+    @property
     def _grid_map(self):
         """
         Mapper of off-grid interpolation points indices for each dimension.
@@ -1049,14 +1049,13 @@ def _evaluate(self, **kwargs):
             return self
 
         io = self.interp_order
-        if self._avg_mode == 'harmonic':
-            retval = 1 / self
-        else:
-            retval = self
+        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)
 
         # Apply interpolation from inner most dim
         for d, i in mapper.items():
-            retval = retval._subs(i.subs(subs), self.indices_ref[d])
             retval = retval.diff(d, deriv_order=0, fd_order=io, x0={d: i})
 
         # Evaluate. Since we used `self.function` it will be on the grid when
@@ -1065,9 +1064,9 @@ def _evaluate(self, **kwargs):
         retval = retval.subs(subs)
 
         # If harmonic averaging, invert at the end
-        if self._avg_mode == 'harmonic':
-            from devito.finite_differences.differentiable import SafeInv
-            retval = SafeInv(retval, self.function.subs(subs))
+        if 'harmonic' in self._avg_mode:
+            retval = retval.safe_inv(self.function.subs(subs),
+                                     safe='safe' in self._avg_mode)
 
         return retval
 

devito/types/dense.py

@@ -1535,6 +1535,15 @@ def _time_buffering(self):
     def _time_buffering_default(self):
         return self._time_buffering and not isinstance(self.save, Buffer)
 
+    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.
+            return retval.evaluate
+        else:
+            return retval
+
     def _arg_check(self, args, intervals, **kwargs):
         super()._arg_check(args, intervals, **kwargs)
 

tests/test_differentiable.py

@@ -119,7 +119,7 @@ def test_avg_mode(ndim, io):
 
     a0 = Function(name="a0", grid=grid, **kw)
     a = Function(name="a", grid=grid, **kw)
-    b = Function(name="b", grid=grid, avg_mode='harmonic', **kw)
+    b = Function(name="b", grid=grid, avg_mode='safe_harmonic', **kw)
 
     a0_avg = a0._eval_at(v)
     a_avg = a._eval_at(v).evaluate.simplify()
@@ -141,7 +141,8 @@ def test_avg_mode(ndim, io):
     assert sympy.simplify(a_avg - expected) == 0
 
     # Harmonic average, h(a[.5]) = 1/(.5/a[0] + .5/a[1])
-    expected = (sum(c / b.subs(arg) for c, arg in zip(ndcoeffs.flatten(), args)))
+    expected = (sum(c * SafeInv(b.subs(arg), b.subs(arg))
+                    for c, arg in zip(ndcoeffs.flatten(), args)))
     assert sympy.simplify(b_avg.args[0] - expected) == 0
     assert isinstance(b_avg, SafeInv)
     assert b_avg.base == b