api: revamp subsampling factor

Author: mloubout

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

@@ -9,11 +9,9 @@
 from devito.data import LEFT, RIGHT
 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.caching import Cached
-from devito.types.constant import Constant
 
 
 __all__ = ['Dimension', 'SpaceDimension', 'TimeDimension', 'DefaultDimension',
@@ -823,10 +821,8 @@ def bound_symbols(self):
         return self.parent.bound_symbols
 
 
-class SubsamplingFactor(Constant, Cached):
-
-    __hash__ = sympy.Symbol.__hash__
-    _cache_key = Symbol._cache_key
+class SubsamplingFactor(Scalar):
+    pass
 
 
 class ConditionalDimension(DerivedDimension):
@@ -905,40 +901,52 @@ class ConditionalDimension(DerivedDimension):
     is_NonlinearDerived = True
     is_Conditional = True
 
-    __rkwargs__ = DerivedDimension.__rkwargs__ + ('factor', 'condition', 'indirect')
+    __rkwargs__ = DerivedDimension.__rkwargs__ + \
+        ('symbolic_factor', 'factor', 'condition', 'indirect')
 
     def __init_finalize__(self, name, parent=None, factor=None, condition=None,
-                          indirect=False, **kwargs):
+                          indirect=False, symbolic_factor=None, **kwargs):
         # `parent=None` degenerates to a ConditionalDimension outside of
         # any iteration space
         if parent is None:
             parent = BOTTOM
 
         super().__init_finalize__(name, parent)
 
-        # Always make the factor symbolic to allow overrides with different factor.
-        if factor is None or factor == 1:
+        # Process subsampling factor
+        fname = f"{name}f"
+        if factor is None:
             self._factor = None
-        elif is_integer(factor):
-            self._factor = SubsamplingFactor(name=f"{name}f", value=factor,
-                                             dtype=np.int32)
-        elif factor.is_Constant and is_integer(factor.data):
-            self._factor = factor
+        elif is_number(factor):
+            self._factor = int(factor)
+        elif factor.is_Constant:
+            self._factor = factor.data
+            fname = factor.name
         else:
             raise ValueError("factor must be an integer or integer Constant")
 
+        if self._factor is not None:
+            # Always make the factor symbolic to allow overrides with different factor.
+            self._symbolic_factor = symbolic_factor or \
+                SubsamplingFactor(name=fname, dtype=np.int32, is_const=True)
+        else:
+            self._symbolic_factor = None
+
         self._condition = condition
         self._indirect = indirect
 
     @property
     def spacing(self):
-        s = self._factor.data if self._factor is not None else 1
-        return s * self.parent.spacing
+        return self.factor * self.parent.spacing
 
     @property
     def factor(self):
         return self._factor if self._factor is not None else 1
 
+    @property
+    def symbolic_factor(self):
+        return self._symbolic_factor
+
     @property
     def condition(self):
         return self._condition
@@ -960,7 +968,7 @@ def free_symbols(self):
 
     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
+        fact = self.factor
         toint = lambda x: math.ceil(x / fact)
         vals = {}
         try:
@@ -984,12 +992,8 @@ def _arg_defaults(self, _min=None, size=None, alias=None):
         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
 
+        factor = defaults[dim.symbolic_factor.name] = self.factor
         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