compiler: Add TempArray.shift

Author: FabioLuporini

devito/passes/clusters/aliases.py

@@ -913,7 +913,8 @@ def lower_schedule(schedule, meta, sregistry, opt_ftemps, opt_min_dtype,
                     indices.append(i.dim - i.lower + s)
 
             dtype = sympy_dtype(pivot, base=meta.dtype)
-            obj = make(name=name, dimensions=dimensions, halo=halo, dtype=dtype)
+            obj = make(name=name, dimensions=dimensions, halo=halo, dtype=dtype,
+                       shift=shift)
             expression = Eq(obj[indices], uxreplace(pivot, subs))
 
             callback = lambda idx: obj[[i + s for i, s in zip(idx, shift)]]

devito/types/misc.py

@@ -10,7 +10,7 @@
 
 from devito.types import Array, CompositeObject, Indexed, Symbol, LocalObject
 from devito.types.basic import IndexedData
-from devito.tools import CustomDtype, Pickable, frozendict
+from devito.tools import CustomDtype, Pickable, as_tuple, frozendict
 
 __all__ = ['Timer', 'Pointer', 'VolatileInt', 'FIndexed', 'Wildcard', 'Fence',
            'Global', 'Hyperplane', 'Indirection', 'Temp', 'TempArray', 'Jump',
@@ -235,12 +235,25 @@ class TempArray(Array):
 
     is_autopaddable = True
 
+    __rkwargs__ = (Array.__rkwargs__ + ('shift',))
+
+    def __init_finalize__(self, *args, shift=None, **kwargs):
+        super().__init_finalize__(*args, **kwargs)
+
+        # An integer for each Dimension representing the shift applied to the halo
+        # for homogeneity reasons
+        self._shift = as_tuple(shift)
+
     def __padding_setup__(self, **kwargs):
         padding = kwargs.pop('padding', None)
         if padding is None:
             padding = self.__padding_setup_smart__(**kwargs)
         return super().__padding_setup__(padding=padding, **kwargs)
 
+    @property
+    def shift(self):
+        return self._shift
+
 
 class Fence: