tests: Add tests for specialising ConditionalDimension factors

Author: EdCaunt

devito/ir/equations/equation.py

@@ -5,8 +5,9 @@
 
 from devito.ir.equations.algorithms import dimension_sort, lower_exprs
 from devito.finite_differences.differentiable import diff2sympy
-from devito.ir.support import (GuardFactor, Interval, IntervalGroup, IterationSpace,
+from devito.ir.support import (Interval, IntervalGroup, IterationSpace,
                                Stencil, detect_io, detect_accesses)
+from devito.ir.support.guards import GuardFactorEq
 from devito.symbolics import IntDiv, limits_mapper, uxreplace
 from devito.tools import Pickable, Tag, frozendict
 from devito.types import Eq, Inc, ReduceMax, ReduceMin, relational_min
@@ -210,11 +211,11 @@ def __new__(cls, *args, **kwargs):
             if not d.is_Conditional:
                 continue
             if d.condition is None:
-                conditionals[d] = GuardFactor(d)
+                conditionals[d] = GuardFactorEq.new_from_dim(d)
             else:
                 cond = diff2sympy(lower_exprs(d.condition))
                 if d._factor is not None:
-                    cond = sympy.And(cond, GuardFactor(d))
+                    cond = sympy.And(cond, GuardFactorEq.new_from_dim(d))
                 conditionals[d] = cond
             # Replace dimension with index
             index = d.index

devito/ir/iet/visitors.py

@@ -15,6 +15,7 @@
 from sympy.core.function import Application
 
 from devito.exceptions import CompilationError
+from devito.symbolics import IndexedPointer
 from devito.ir.iet.nodes import (Node, Iteration, Expression, ExpressionBundle,
                                  Call, Lambda, BlankLine, Section, ListMajor)
 from devito.ir.support.space import Backward
@@ -1516,9 +1517,7 @@ def __init__(self, mapper, nested=False):
 
         # Sanity check
         for k, v in self.mapper.items():
-            # FIXME: Erronously blocks f_vec->size[1]
-            # Apparently this is an IndexedPointer
-            if not isinstance(k, AbstractSymbol):
+            if not isinstance(k, (AbstractSymbol, IndexedPointer)):
                 raise ValueError(f"Attempted to specialize non-scalar symbol: {k}")
 
             if not isinstance(v, Number):
@@ -1531,7 +1530,10 @@ def visit_Operator(self, o, **kwargs):
         # is the intended use case
         body = self._visit(o.body)
 
-        not_params = tuple(i for i in self.mapper if i not in o.parameters)
+        # NOTE: IndexedPointers that want replacing with a hardcoded value won't appear in
+        # the Operator parameters. Perhaps this check wants relaxing.
+        not_params = tuple(i for i in self.mapper
+                           if i not in o.parameters and isinstance(i, AbstractSymbol))
         if not_params:
             raise ValueError(f"Attempted to specialize symbols {not_params} which are not"
                              " found in the Operator parameters")

devito/ir/support/guards.py

@@ -38,35 +38,34 @@ def canonical(self):
 
     @property
     def negated(self):
-        return negations[self.__class__](*self._args_rebuild, evaluate=False)
+        try:
+            return negations[self.__class__](*self._args_rebuild, evaluate=False)
+        except KeyError:
+            raise ValueError(f"Class {self.__class__.__name__} does not have a negation")
 
 
 # *** GuardFactor
 
 
-class GuardFactor(Guard, CondEq, Pickable):
+class GuardFactor(Guard, Pickable):
 
     """
     A guard for factor-based ConditionalDimensions.
 
-    Given the ConditionalDimension `d` with factor `k`, create the
-    symbolic relational `d.parent % k == 0`.
+    Introduces a constructor where, given the ConditionalDimension `d` with factor `k`,
+    the symbolic relational `d.parent % k == 0` is created.
     """
 
-    __rargs__ = ('d',)
+    __rargs__ = ('lhs', 'rhs')
 
-    def __new__(cls, d, **kwargs):
+    @classmethod
+    def new_from_dim(cls, d, **kwargs):
         assert d.is_Conditional
 
         obj = super().__new__(cls, d.parent % d.symbolic_factor, 0)
-        obj.d = d
 
         return obj
 
-    @property
-    def _args_rebuild(self):
-        return (self.d,)
-
 
 class GuardFactorEq(GuardFactor, CondEq):
     pass
@@ -76,9 +75,6 @@ class GuardFactorNe(GuardFactor, CondNe):
     pass
 
 
-GuardFactor = GuardFactorEq
-
-
 # *** GuardBound
 
 

devito/symbolics/extended_sympy.py

@@ -35,7 +35,8 @@ class CondEq(sympy.Eq):
     """
 
     def __new__(cls, *args, **kwargs):
-        return sympy.Eq.__new__(cls, *args, evaluate=False)
+        kwargs['evaluate'] = False
+        return sympy.Eq.__new__(cls, *args, **kwargs)
 
     @property
     def canonical(self):

devito/symbolics/manipulation.py

@@ -136,9 +136,9 @@ def _(mapper, rule):
 @singledispatch
 def _uxreplace_handle(expr, args, kwargs):
     try:
-        return expr.func(*args, evaluate=False)
+        return expr.func(*args, evaluate=False, **kwargs)
     except TypeError:
-        return expr.func(*args)
+        return expr.func(*args, **kwargs)
 
 
 @_uxreplace_handle.register(Min)

tests/test_pickle.py

@@ -10,7 +10,8 @@
                     Dimension, SubDimension, ConditionalDimension, IncrDimension,
                     TimeDimension, SteppingDimension, Operator, MPI, Min, solve,
                     PrecomputedSparseTimeFunction, SubDomain)
-from devito.ir import Backward, Forward, GuardFactor, GuardBound, GuardBoundNext
+from devito.ir import Backward, Forward, GuardBound, GuardBoundNext
+from devito.ir.support.guards import GuardFactorEq
 from devito.data import LEFT, OWNED
 from devito.finite_differences.tools import direct, transpose, left, right, centered
 from devito.mpi.halo_scheme import Halo
@@ -500,7 +501,7 @@ def test_guard_factor(self, pickle):
         d = Dimension(name='d')
         cd = ConditionalDimension(name='cd', parent=d, factor=4)
 
-        gf = GuardFactor(cd)
+        gf = GuardFactorEq.new_from_dim(cd)
 
         pkl_gf = pickle.dumps(gf)
         new_gf = pickle.loads(pkl_gf)

tests/test_specialization.py

@@ -120,27 +120,25 @@ def define(self, dimensions):
         for m, o in zip(mappers, ops):
             check_op(m, o)
 
-    # FIXME: Currently throws an error - probably a missing handler for GuardFactor
-    # in Uxreplace
-    # def test_factor(self):
-    #     """Test that ConditionalDimensions can have their symbolic factors specialized"""
-    #     size = 16
-    #     factor = 4
-    #     i = Dimension(name='i')
-    #     ci = ConditionalDimension(name='ci', parent=i, factor=factor)
+    def test_factor(self):
+        """Test that ConditionalDimensions can have their symbolic factors specialized"""
+        size = 16
+        factor = 4
+        i = Dimension(name='i')
+        ci = ConditionalDimension(name='ci', parent=i, factor=factor)
 
-    #     g = Function(name='g', shape=(size,), dimensions=(i,))
-    #     f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))
+        g = Function(name='g', shape=(size,), dimensions=(i,))
+        f = Function(name='f', shape=(int(size/factor),), dimensions=(ci,))
 
-    #     op0 = Operator([Eq(f, g)])
+        op0 = Operator([Eq(f, g)])
 
-    #     mapper = {ci.symbolic_factor: sympy.Integer(factor)}
+        mapper = {ci.symbolic_factor: sympy.Integer(factor)}
 
-    #     op1 = Specializer(mapper).visit(op0)
+        op1 = Specializer(mapper).visit(op0)
 
-    #     assert ci.symbolic_factor not in op1.parameters
-    #     assert ci.symbolic_factor.name not in str(op1.ccode)
-    #     assert "if ((i)%(4) == 0)" in str(op1.ccode)
+        assert ci.symbolic_factor not in op1.parameters
+        assert ci.symbolic_factor.name not in str(op1.ccode)
+        assert "if ((i)%(4) == 0)" in str(op1.ccode)
 
     def test_spacing(self):
         """Test that grid spacings can be specialized"""
@@ -156,14 +154,21 @@ def test_spacing(self):
         assert grid.dimensions[0].spacing.name not in str(op1.ccode)
         assert "/1.0e-1F;" in str(op1.ccode)
 
-    # Strides/sizes
-    def test_strides(self):
-        """Test that strides and sizes generated for linearization can be specialized"""
+    def test_sizes(self):
+        """Test that strides generated for linearization can be specialized"""
         grid = Grid(shape=(11, 11))
 
         f = TimeFunction(name='f', grid=grid, space_order=2)
 
         op0 = Operator(Eq(f.forward, f.dx2),
                        opt=('advanced', {'expand': True, 'linearize': True}))
 
-        from IPython import embed; embed()
+        mapper = {f.symbolic_shape[1]: sympy.Integer(11),
+                  f.symbolic_shape[2]: sympy.Integer(11)}
+
+        op1 = Specializer(mapper).visit(op0)
+
+        assert "const int x_fsz0 = 11;" in str(op1.ccode)
+        assert "const int y_fsz0 = 11;" in str(op1.ccode)
+
+    # TODO: Should strides get linearized? If so, how?