Merge pull request #2575 from devitocodes/fix-subsampling-rebuild

api: revamp subsampling factors to avoid duplicates

Author: FabioLuporini

devito/deprecations.py

@@ -27,5 +27,12 @@ def subdomain_warn(self):
              DeprecationWarning, stacklevel=2)
         return
 
+    @cached_property
+    def constant_factor_warn(self):
+        warn("Using a `Constant` as a factor when creating a ConditionalDimension"
+             " is deprecated. Use an integer instead.",
+             DeprecationWarning, stacklevel=2)
+        return
+
 
 deprecations = DevitoDeprecation()

devito/ir/equations/equation.py

@@ -220,7 +220,7 @@ def __new__(cls, *args, **kwargs):
             index = d.index
             if d.condition is not None and d in expr.free_symbols:
                 index = index - relational_min(d.condition, d.parent)
-            expr = uxreplace(expr, {d: IntDiv(index, d.factor)})
+            expr = uxreplace(expr, {d: IntDiv(index, d.symbolic_factor)})
 
         conditionals = frozendict(conditionals)
 

devito/ir/support/guards.py

@@ -47,7 +47,7 @@ class GuardFactor(Guard, CondEq, Pickable):
     def __new__(cls, d, **kwargs):
         assert d.is_Conditional
 
-        obj = super().__new__(cls, d.parent % d.factor, 0)
+        obj = super().__new__(cls, d.parent % d.symbolic_factor, 0)
         obj.d = d
 
         return obj
@@ -129,7 +129,7 @@ def __new__(cls, d, direction, **kwargs):
             p1 = d.root.symbolic_max
 
             if d.is_Conditional:
-                v = d.factor
+                v = d.symbolic_factor
                 # Round `p0 + 1` up to the nearest multiple of `v`
                 p0 = Mul((((p0 + 1) + v - 1) / v), v, evaluate=False)
             else:
@@ -140,7 +140,7 @@ def __new__(cls, d, direction, **kwargs):
             p1 = d.root
 
             if d.is_Conditional:
-                v = d.factor
+                v = d.symbolic_factor
                 # Round `p1 - 1` down to the nearest sub-multiple of `v`
                 # NOTE: we use ABS to make sure we handle negative values properly.
                 # Once `p1 - 1` is negative (e.g. `iteration=time - 1` and `time=0`),

devito/symbolics/extended_sympy.py

@@ -84,7 +84,7 @@ class IntDiv(sympy.Expr):
     def __new__(cls, lhs, rhs, params=None):
         if rhs == 0:
             raise ValueError("Cannot divide by 0")
-        elif rhs == 1:
+        elif rhs == 1 or rhs is None:
             return lhs
 
         if not is_integer(rhs):

devito/types/dimension.py

@@ -7,9 +7,10 @@
 import numpy as np
 
 from devito.data import LEFT, RIGHT
+from devito.deprecations import deprecations
 from devito.exceptions import InvalidArgument
 from devito.logger import debug
-from devito.tools import Pickable, is_integer, memoized_meth
+from devito.tools import Pickable, is_integer, is_number, memoized_meth
 from devito.types.args import ArgProvider
 from devito.types.basic import Symbol, DataSymbol, Scalar
 from devito.types.constant import Constant
@@ -822,6 +823,10 @@ def bound_symbols(self):
         return self.parent.bound_symbols
 
 
+class SubsamplingFactor(Scalar):
+    pass
+
+
 class ConditionalDimension(DerivedDimension):
 
     """
@@ -898,7 +903,8 @@ class ConditionalDimension(DerivedDimension):
     is_NonlinearDerived = True
     is_Conditional = True
 
-    __rkwargs__ = DerivedDimension.__rkwargs__ + ('factor', 'condition', 'indirect')
+    __rkwargs__ = DerivedDimension.__rkwargs__ + \
+        ('factor', 'condition', 'indirect')
 
     def __init_finalize__(self, name, parent=None, factor=None, condition=None,
                           indirect=False, **kwargs):
@@ -909,27 +915,49 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
 
         super().__init_finalize__(name, parent)
 
-        # Always make the factor symbolic to allow overrides with different factor.
+        # Process subsampling factor
         if factor is None or factor == 1:
             self._factor = None
-        elif is_integer(factor):
-            self._factor = Constant(name="%sf" % name, value=factor, dtype=np.int32)
-        elif factor.is_Constant and is_integer(factor.data):
+        elif is_number(factor):
+            self._factor = int(factor)
+        elif factor.is_Constant:
+            deprecations.constant_factor_warn
             self._factor = factor
         else:
-            raise ValueError("factor must be an integer or integer Constant")
+            raise ValueError("factor must be an integer")
 
         self._condition = condition
         self._indirect = indirect
 
+    @property
+    def uses_symbolic_factor(self):
+        return self._factor is not None
+
+    @property
+    def factor_data(self):
+        if isinstance(self.factor, Constant):
+            return self.factor.data
+        else:
+            return self.factor
+
     @property
     def spacing(self):
-        s = self._factor.data if self._factor is not None else 1
-        return s * self.parent.spacing
+        return self.factor_data * self.parent.spacing
 
     @property
     def factor(self):
-        return self._factor if self._factor is not None else 1
+        return self._factor if self.uses_symbolic_factor else 1
+
+    @cached_property
+    def symbolic_factor(self):
+        if not self.uses_symbolic_factor:
+            return None
+        elif isinstance(self.factor, Constant):
+            return self.factor
+        else:
+            return SubsamplingFactor(
+                name=f'{self.name}f', dtype=np.int32, is_const=True
+            )
 
     @property
     def condition(self):
@@ -950,9 +978,14 @@ def free_symbols(self):
             pass
         return retval
 
-    def _arg_values(self, interval, grid=None, **kwargs):
-        # Parent dimension define the interval
-        fact = self._factor.data if self._factor is not None else 1
+    def _arg_values(self, interval, grid=None, args=None, **kwargs):
+        if not self.uses_symbolic_factor:
+            return {}
+
+        args = args or {}
+        fname = self.symbolic_factor.name
+        fact = kwargs.get(fname, args.get(fname, self.factor_data))
+
         toint = lambda x: math.ceil(x / fact)
         vals = {}
         try:
@@ -965,6 +998,8 @@ def _arg_values(self, interval, grid=None, **kwargs):
         except (KeyError, TypeError):
             pass
 
+        vals[self.symbolic_factor.name] = fact
+
         return vals
 
     def _arg_defaults(self, _min=None, size=None, alias=None):
@@ -974,15 +1009,9 @@ def _arg_defaults(self, _min=None, size=None, alias=None):
         # `factor` endpoints are legal, so we return them all. It's then
         # up to the caller to decide which one to pick upon reduction
         dim = alias or self
-        if dim.condition is not None or size is None or dim._factor is None:
-            return defaults
-        try:
-            # Is it a symbolic factor?
-            factor = defaults[dim._factor.name] = self._factor.data
-        except AttributeError:
-            factor = dim._factor
-
-        defaults[dim.parent.max_name] = range(0, factor*size - 1)
+        if dim.uses_symbolic_factor:
+            factor = defaults[dim.symbolic_factor.name] = self.factor_data
+            defaults[dim.parent.max_name] = range(0, factor*size - 1)
 
         return defaults
 

devito/types/grid.py

@@ -305,7 +305,7 @@ def spacing_map(self):
                 # Special case subsampling: `Grid.dimensions` -> (xb, yb, zb)`
                 # where `xb, yb, zb` are ConditionalDimensions whose parents
                 # are SpaceDimensions
-                mapper[d.root.spacing] = s/self.dtype(d.factor.data)
+                mapper[d.root.spacing] = s/self.dtype(d.factor)
             elif d.is_Space:
                 # Typical case: `Grid.dimensions` -> (x, y, z)` where `x, y, z` are
                 # the SpaceDimensions

tests/test_dimension.py

@@ -1,4 +1,5 @@
 from itertools import product
+from copy import deepcopy
 
 import numpy as np
 from sympy import And, Or
@@ -16,6 +17,7 @@
 from devito.ir import SymbolRegistry
 from devito.symbolics import indexify, retrieve_functions, IntDiv, INT
 from devito.types import Array, StencilDimension, Symbol
+from devito.types.basic import Scalar
 from devito.types.dimension import AffineIndexAccessFunction, Thickness
 
 
@@ -1011,9 +1013,9 @@ def test_issue_1592(self):
         op = Operator(Eq(v.forward, v.dx))
         op.apply(time=6)
         exprs = FindNodes(Expression).visit(op)
-        assert exprs[-1].expr.lhs.indices[0] == IntDiv(time, time_sub.factor) + 1
-        assert time_sub.factor.data == 2
-        assert time_sub.factor.is_Constant
+        assert exprs[-1].expr.lhs.indices[0] == IntDiv(time, time_sub.symbolic_factor) + 1
+        assert time_sub.factor == 2
+        assert isinstance(time_sub.symbolic_factor, Scalar)
 
     def test_issue_1753(self):
         grid = Grid(shape=(3, 3, 3))
@@ -1896,6 +1898,41 @@ def test_cond_notime(self):
         op(time_m=1, time_M=nt-1, dt=1)
         assert norm(g, order=1) == norm(sum(usaved, dims=time_under), order=1)
 
+    def test_cond_copy(self):
+        grid = Grid((11, 11, 11))
+        time = grid.time_dim
+
+        cd = ConditionalDimension(name='tsub', parent=time, factor=5)
+        u = TimeFunction(name='u', grid=grid, space_order=4, time_order=2, save=Buffer(2))
+        u1 = TimeFunction(name='u1', grid=grid, space_order=0,
+                          time_order=0, save=5, time_dim=cd)
+        u2 = TimeFunction(name='u2', grid=grid, space_order=0,
+                          time_order=0, save=5, time_dim=cd)
+
+        # Mimic what happens when an operator is copied
+        u12 = deepcopy(u1)
+        u22 = deepcopy(u2)
+
+        op = Operator([Eq(u.forward, u.laplace), Eq(u12, u), Eq(u22, u)])
+        assert len([p for p in op.parameters if p.name == 'tsubf']) == 1
+
+    def test_const_factor(self):
+        grid = Grid(shape=(4, 4))
+        time = grid.time_dim
+
+        f1 = 4
+        f2 = Constant(name='f2', dtype=np.int32, value=4)
+        t1 = ConditionalDimension('t_sub', parent=time, factor=f1)
+        t2 = ConditionalDimension('t_sub2', parent=time, factor=f2)
+
+        assert isinstance(t1.symbolic_factor, Scalar)
+        assert t1.factor == f1
+
+        assert t2.symbolic_factor.is_Constant
+        assert t2.factor == f2
+        assert t2.factor.data == f1
+        assert t2.spacing == t1.spacing
+
 
 class TestCustomDimension:
 

tests/test_pickle.py

@@ -320,6 +320,7 @@ def test_conditional_dimension(self, pickle):
         assert cd.name == new_cd.name
         assert cd.parent.name == new_cd.parent.name
         assert cd.factor == new_cd.factor
+        assert cd.symbolic_factor == new_cd.symbolic_factor
         assert cd.condition == new_cd.condition
 
     def test_incr_dimension(self, pickle):
@@ -603,7 +604,7 @@ def test_equation(self, pickle):
         assert new_eq.lhs.name == f.name
         assert str(new_eq.rhs) == 'f(x) + 1'
         assert new_eq.implicit_dims[0].name == 'xs'
-        assert new_eq.implicit_dims[0].factor.data == 4
+        assert new_eq.implicit_dims[0].factor == 4
 
     @pytest.mark.parametrize('typ', [ctypes.c_float, 'struct truct'])
     def test_Cast(self, pickle, typ):