dsl: Continue refactoring SubDimensions

devito/ir/equations/algorithms.py

@@ -267,11 +267,9 @@ def _(d, mapper, rebuilt, sregistry):
         # Already have a substitution for this dimension
         return
 
-    abstract_tkns = d._symbolic_thickness
-    concrete_tkns = tuple(tkn._rebuild(name=sregistry.make_name(prefix=tkn.name))
-                          for tkn in abstract_tkns)
-
-    kwargs = {'thickness': concrete_tkns}
+    tkns = tuple(tkn._rebuild(name=sregistry.make_name(prefix=tkn.name))
+                 for tkn in d.thickness)
+    kwargs = {'thickness': tkns}
     fkwargs = {}
 
     idim0 = d.implicit_dimension

devito/types/dimension.py

@@ -9,7 +9,7 @@
 from devito.data import LEFT, RIGHT
 from devito.exceptions import InvalidArgument
 from devito.logger import debug
-from devito.tools import Pickable, is_integer, flatten, memoized_meth
+from devito.tools import Pickable, is_integer, memoized_meth
 from devito.types.args import ArgProvider
 from devito.types.basic import Symbol, DataSymbol, Scalar
 from devito.types.constant import Constant
@@ -542,18 +542,9 @@ class SubDimensionThickness(DataSymbol):
     __rkwargs__ = DataSymbol.__rkwargs__ + ('root', 'side', 'local', 'value')
 
     def __new__(cls, *args, **kwargs):
-        try:
-            root = kwargs.pop('root')
-        except KeyError:
-            raise ValueError("No root dimension provided")
-        try:
-            side = kwargs.pop('side')
-        except KeyError:
-            raise ValueError("No side specified")
-        try:
-            local = kwargs.pop('local')
-        except KeyError:
-            raise ValueError("Locality of thickness not specified")
+        root = kwargs.pop('root')
+        side = kwargs.pop('side')
+        local = kwargs.pop('local')
 
         newobj = super().__new__(cls, *args, **kwargs)
         newobj._root = root
@@ -595,7 +586,6 @@ def _arg_values(self, grid=None, **kwargs):
         # So overriding the thickness to a nonzero value should not cause
         # boundaries to exist between ranks where they did not before
         r_tkn = kwargs.get(self.name, self.value)
-
         if grid is not None and grid.is_distributed(self.root):
             # Get local thickness
             if self.local:
@@ -638,21 +628,18 @@ class AbstractSubDimension(DerivedDimension):
     __rargs__ = DerivedDimension.__rargs__ + ('thickness',)
     __rkwargs__ = ()
 
-    _thickness_type = SubDimensionThickness
+    _thickness_type = Symbol
 
     def __init_finalize__(self, name, parent, thickness, **kwargs):
         super().__init_finalize__(name, parent)
-        self._process_thicknesses(thickness)
-
-    def _process_thicknesses(self, thickness):
-        if any(isinstance(tkn, SubDimensionThickness) for tkn in thickness):
-            ltkn, rtkn = thickness
+        thickness = thickness or (None, None)
+        if any(isinstance(tkn, self._thickness_type) for tkn in thickness):
+            self._thickness = Thickness(*thickness)
         else:
-            ltkn, rtkn = self._symbolic_thickness(thickness=thickness)
+            self._thickness = self._symbolic_thickness(thickness=thickness)
 
-        self._thickness = Thickness(ltkn, rtkn)
-
-        # Set up the interval
+    @cached_property
+    def _interval(self):
         if self.thickness.right.value is None:  # Left SubDimension
             left = self.parent.symbolic_min
             right = self.parent.symbolic_min + self.ltkn - 1
@@ -663,21 +650,14 @@ def _process_thicknesses(self, thickness):
             left = self.parent.symbolic_min + self.ltkn
             right = self.parent.symbolic_max - self.rtkn
 
-        self._interval = sympy.Interval(left, right)
+        return sympy.Interval(left, right)
 
     @memoized_meth
-    def _symbolic_thickness(self, thickness=None):
-        thickness = thickness or (None, None)
-        kwargs = {'dtype': np.int64, 'is_const': True, 'nonnegative': True,
-                  'root': self.root, 'local': self.local}
-
-        ltkn = self._thickness_type(name="%s_ltkn" % self.parent.name,
-                                    side=LEFT, value=thickness[0], **kwargs)
+    def _symbolic_thickness(self, **kwargs):
+        kwargs = {'dtype': np.int32, 'is_const': True, 'nonnegative': True}
 
-        rtkn = self._thickness_type(name="%s_rtkn" % self.parent.name,
-                                    side=RIGHT, value=thickness[1], **kwargs)
-
-        return (ltkn, rtkn)
+        names = ["%s_%stkn" % (self.parent.name, s) for s in ('l', 'r')]
+        return Thickness(*[Symbol(name=n, **kwargs) for n in names])
 
     @cached_property
     def symbolic_min(self):
@@ -696,9 +676,7 @@ def symbolic_size(self):
     def thickness(self):
         return self._thickness
 
-    @property
-    def is_abstract(self):
-        return all(i is None for i in flatten(self.thickness))
+    tkns = thickness  # Shortcut for thickness
 
     @property
     def ltkn(self):
@@ -710,12 +688,6 @@ def rtkn(self):
         # Shortcut for the right thickness symbol
         return self.thickness.right
 
-    @property
-    def tkns(self):
-        # Shortcut for both thickness symbols
-        # FIXME: Maybe redundant now?
-        return self.ltkn, self.rtkn
-
     def __hash__(self):
         return id(self)
 
@@ -772,31 +744,33 @@ class SubDimension(AbstractSubDimension):
 
     __rargs__ = AbstractSubDimension.__rargs__ + ('local',)
 
+    _thickness_type = SubDimensionThickness
+
     def __init_finalize__(self, name, parent, thickness, local,
                           **kwargs):
         self._local = local
         super().__init_finalize__(name, parent, thickness)
 
     @classmethod
     def left(cls, name, parent, thickness, local=True):
-        return cls(name, parent,
-                   thickness=(thickness, None),
-                   # thickness=((lst, thickness), (rst, None)),
-                   local=local)
+        return cls(name, parent, thickness=(thickness, None), local=local)
 
     @classmethod
     def right(cls, name, parent, thickness, local=True):
-        return cls(name, parent,
-                   thickness=(None, thickness),
-                   # thickness=((lst, None), (rst, thickness)),
-                   local=local)
+        return cls(name, parent, thickness=(None, thickness), local=local)
 
     @classmethod
     def middle(cls, name, parent, thickness_left, thickness_right, local=False):
-        return cls(name, parent,
-                   thickness=(thickness_left, thickness_right),
-                   # thickness=((lst, thickness_left), (rst, thickness_right)),
-                   local=local)
+        return cls(name, parent, thickness=(thickness_left, thickness_right), local=local)
+
+    @memoized_meth
+    def _symbolic_thickness(self, thickness=None):
+        kwargs = {'dtype': np.int64, 'is_const': True, 'nonnegative': True,
+                  'root': self.root, 'local': self.local}
+
+        names = ["%s_%stkn" % (self.parent.name, s) for s in ('l', 'r')]
+        return Thickness(*[SubDimensionThickness(name=n, side=s, value=t, **kwargs)
+                           for n, s, t in zip(names, (LEFT, RIGHT), thickness)])
 
     @property
     def local(self):
@@ -846,39 +820,6 @@ class MultiSubDimension(AbstractSubDimension):
 
     __rkwargs__ = ('functions', 'bounds_indices', 'implicit_dimension')
 
-    # FIXME: Is this still necessary?
-    _thickness_type = Symbol
-
-    def _process_thicknesses(self, thickness):
-        # FIXME: Requires homogensisation with SubDimension._process_thickness
-        # Canonicalize thickness
-        if thickness is None:
-            # Using dummy left/right is the only thing we can do for such
-            # an abstract MultiSubDimension
-            thickness = ((None, None), (None, None))
-
-            left = sympy.S.NegativeInfinity
-            right = sympy.S.Infinity
-        elif isinstance(thickness, tuple):
-            if all(isinstance(i, Symbol) for i in thickness):
-                ltkn, rtkn = thickness
-                thickness = ((ltkn, None), (rtkn, None))
-            elif all(isinstance(i, tuple) and len(i) == 2 for i in thickness):
-                (ltkn, _), (rtkn, _) = thickness
-
-            try:
-                left = self.parent.symbolic_min + ltkn
-                right = self.parent.symbolic_max - rtkn
-            except TypeError:
-                # May end up here after a reconstruction
-                left = sympy.S.NegativeInfinity
-                right = sympy.S.Infinity
-        else:
-            raise ValueError("MultiSubDimension expects a tuple of thicknesses")
-
-        self._thickness = Thickness(*thickness)
-        self._interval = sympy.Interval(left, right)
-
     def __init_finalize__(self, name, parent, thickness, functions=None,
                           bounds_indices=None, implicit_dimension=None):
 
@@ -887,6 +828,12 @@ def __init_finalize__(self, name, parent, thickness, functions=None,
         self.bounds_indices = bounds_indices
         self.implicit_dimension = implicit_dimension
 
+    @cached_property
+    def _interval(self):
+        left = self.parent.symbolic_min + self.ltkn
+        right = self.parent.symbolic_max - self.rtkn
+        return sympy.Interval(left, right)
+
     @cached_property
     def bound_symbols(self):
         return self.parent.bound_symbols

tests/test_caching.py

@@ -842,9 +842,9 @@ def test_solve(self, operate_on_empty_cache):
         # created by the finite difference (u.dt, u.dx2). We would have had
         # three extra references to u(t + dt), u(x - h_x) and u(x + h_x).
         # But this is not the case anymore!
-        assert len(_SymbolCache) == 10
+        assert len(_SymbolCache) == 8
         clear_cache()
-        assert len(_SymbolCache) == 6
+        assert len(_SymbolCache) == 4
         clear_cache()
         assert len(_SymbolCache) == 2
         clear_cache()

tests/test_dimension.py

@@ -787,6 +787,18 @@ def test_expandingbox_like(self, opt):
         assert np.all(u.data[:, :, :2] == 0.)
         assert np.all(u.data[:, :, -2:] == 0.)
 
+    def test_standalone_thickness(self):
+        x = Dimension('x')
+        ix = SubDimension.left('ix', x, 2)
+        f = Function(name="f", dimensions=(ix,), shape=(5,), dtype=np.int32)
+
+        eqns = Eq(f[ix.symbolic_max], 1)
+
+        op = Operator(eqns)
+        assert 'x_ltkn0' in str(op.ccode)
+        op(x_m=0)
+        assert np.all(f.data == np.array([0, 1, 0, 0, 0]))
+
 
 class TestConditionalDimension: