Merge pull request #2742 from devitocodes/improve-guard-expr

compiler: Implement auto-simplification of GuardExpr

Author: mloubout

devito/ir/support/guards.py

@@ -4,6 +4,9 @@
 (e.g., Eq, Cluster, ...) should be evaluated at runtime.
 """
 
+from operator import ge, gt, le, lt
+
+from functools import singledispatch
 from sympy import And, Ge, Gt, Le, Lt, Mul, true
 from sympy.logic.boolalg import BooleanFunction
 import numpy as np
@@ -284,33 +287,54 @@ class GuardExpr(LocalObject, BooleanFunction):
     Being a LocalObject, a GuardExpr may carry an `initvalue`, which is
     the value that the guard assumes at the beginning of the scope where
     it is defined.
-
-    Through the `supersets` argument, a GuardExpr may also carry a set of
-    GuardExprs that are known to be more restrictive than itself. This is
-    usesful, e.g., to avoid redundant checks when chaining multiple guards
-    together (see `simplify_and`).
     """
 
     dtype = np.bool
 
-    def __init__(self, name, liveness='eager', supersets=None, **kwargs):
+    def __init__(self, name, liveness='eager', **kwargs):
         super().__init__(name, liveness=liveness, **kwargs)
 
-        self.supersets = frozenset(as_tuple(supersets))
+    @singledispatch
+    def _handle_boolean(obj, mapper):
+        raise NotImplementedError(f"Cannot handle boolean of type {type(obj)}")
 
-    def _hashable_content(self):
-        return super()._hashable_content() + (self.supersets,)
+    @_handle_boolean.register(And)
+    def _(obj, mapper):
+        for a in obj.args:
+            GuardExpr._handle_boolean(a, mapper)
 
-    __hash__ = LocalObject.__hash__
+    @_handle_boolean.register(Le)
+    @_handle_boolean.register(Ge)
+    @_handle_boolean.register(Lt)
+    @_handle_boolean.register(Gt)
+    def _(obj, mapper):
+        d, v = obj.args
+        k = obj.__class__
+        mapper.setdefault(k, {})[d] = v
+
+    @property
+    def as_mapper(self):
+        mapper = {}
+        GuardExpr._handle_boolean(self.initvalue, mapper)
+        return frozendict(mapper)
 
-    def __eq__(self, other):
-        return (isinstance(other, GuardExpr) and
-                super().__eq__(other) and
-                self.supersets == other.supersets)
+    def sort_key(self, order=None):
+        # Use the overarching LocalObject name for arguments ordering
+        class_key, args, exp, coeff = super().sort_key(order=order)
+        args = (len(args[1]) + 1, (self.name,) + args[1])
+        return class_key, args, exp, coeff
 
 
 # *** Utils
 
+op_mapper = {
+    Le: le,
+    Lt: lt,
+    Ge: ge,
+    Gt: gt
+}
+
+
 def simplify_and(relation, v):
     """
     Given `x = And(*relation.args, v)`, return `relation` if `x ≡ relation`,
@@ -327,19 +351,44 @@ def simplify_and(relation, v):
     else:
         candidates, other = [], [relation, v]
 
-    # Quick check based on GuardExpr.supersets to avoid adding `v` to `relation`
-    # if `relation` already includes a more restrictive guard than `v`
-    if isinstance(v, GuardExpr):
-        if any(a in v.supersets for a in candidates):
-            return relation
-
     covered = False
     new_args = []
     for a in candidates:
-        if isinstance(a, GuardExpr) or a.lhs is not v.lhs:
+        if isinstance(v, GuardExpr) and isinstance(a, GuardExpr):
+            # Attempt optimizing guards in GuardExpr form
+            covered = True
+
+            m0 = v.as_mapper
+            m1 = a.as_mapper
+
+            for cls, op in op_mapper.items():
+                if cls in m0 and cls in m1:
+                    try:
+                        if set(m0[cls]) != set(m1[cls]):
+                            new_args.extend([a, v])
+                        elif all(op(m0[cls][d], m1[cls][d]) for d in m0[cls]):
+                            new_args.append(v)
+                        elif all(op(m1[cls][d], m0[cls][d]) for d in m1[cls]):
+                            new_args.append(a)
+                        else:
+                            new_args.extend([a, v])
+                    except TypeError:
+                        # E.g., `cls = Le`, then `z <= 2` and `z <= z_M + 1`
+                        new_args.extend([a, v])
+
+                elif cls in m0:
+                    new_args.append(v)
+
+                elif cls in m1:
+                    new_args.append(a)
+
+        elif a.lhs is not v.lhs:
             new_args.append(a)
+
         else:
+            # Attempt optimizing guards in relational form
             covered = True
+
             try:
                 if type(a) in (Gt, Ge) and v.rhs > a.rhs:
                     new_args.append(v)

tests/test_symbolics.py

@@ -4,12 +4,13 @@
 import pytest
 import numpy as np
 
-from sympy import Expr, Number, Symbol
+from sympy import And, Expr, Number, Symbol
 from devito import (Constant, Dimension, Grid, Function, solve, TimeFunction, Eq,  # noqa
                     Operator, SubDimension, norm, Le, Ge, Gt, Lt, Abs, sin, cos,
                     Min, Max, Real, Imag, Conj, SubDomain, configuration)
 from devito.finite_differences.differentiable import SafeInv, Weights, Mul
 from devito.ir import Expression, FindNodes, ccode
+from devito.ir.support.guards import GuardExpr, simplify_and
 from devito.mpi.halo_scheme import HaloTouch
 from devito.symbolics import (
     retrieve_functions, retrieve_indexed, evalrel, CallFromPointer, Cast, # noqa
@@ -543,6 +544,133 @@ def test_halo_touch():
     assert hash(ht0) == hash(ht0._rebuild())
 
 
+def test_guard_expr_Le():
+    grid = Grid(shape=(3, 3, 3))
+    _, y, z = grid.dimensions
+
+    y_M = y.symbolic_max
+    z_M = z.symbolic_max
+
+    g0 = GuardExpr('g0', initvalue=And(y <= y_M + 3, z <= z_M + 3))
+    g1 = GuardExpr('g1', initvalue=And(y <= y_M + 3, z <= z_M + 4))
+
+    v0 = simplify_and(g0, g1)
+    v1 = simplify_and(g1, g0)
+    assert v0 is g0
+    assert v0 is v1
+
+    v2 = simplify_and(And(g0, g1), g0)
+    assert v2 is g0
+
+    g3 = GuardExpr('g3', initvalue=And(y <= y_M + 2, z <= z_M + 3))
+    v3 = simplify_and(g0, g3)
+    v4 = simplify_and(g3, g0)
+    assert v3 is g3
+    assert v4 is v3
+
+    g4 = GuardExpr('g4', initvalue=And(y <= y_M + 4, z <= z_M + 2))
+    v5 = simplify_and(g0, g4)
+    assert v5 == And(g0, g4)
+
+    g5 = GuardExpr('g5', initvalue=And(y <= y_M))
+    g6 = GuardExpr('g6', initvalue=And(z <= z_M))
+    v6 = simplify_and(g0, g5)
+    v7 = simplify_and(g0, g6)
+    v8 = simplify_and(g5, g6)
+    assert v6 == And(g0, g5)
+    assert v7 == And(g0, g6)
+    assert v8 == And(g5, g6)
+
+
+def test_guard_expr_Ge():
+    grid = Grid(shape=(3, 3, 3))
+    _, y, z = grid.dimensions
+
+    y_m = y.symbolic_min
+    z_m = z.symbolic_min
+
+    g0 = GuardExpr('g0', initvalue=And(y >= y_m - 3, z >= z_m - 3))
+    g1 = GuardExpr('g1', initvalue=And(y >= y_m - 3, z >= z_m - 4))
+    v0 = simplify_and(g0, g1)
+    v1 = simplify_and(g1, g0)
+    assert v0 is g0
+    assert v0 is v1
+
+    v2 = simplify_and(And(g0, g1), g0)
+    assert v2 is g0
+
+    g3 = GuardExpr('g3', initvalue=And(y >= y_m - 2, z >= z_m - 3))
+    v3 = simplify_and(g0, g3)
+    v4 = simplify_and(g3, g0)
+    assert v3 is g3
+    assert v4 is v3
+
+    g4 = GuardExpr('g4', initvalue=And(y >= y_m - 4, z >= z_m - 2))
+    v5 = simplify_and(g0, g4)
+    assert v5 == And(g0, g4)
+    g5 = GuardExpr('g5', initvalue=And(y >= y_m))
+    g6 = GuardExpr('g6', initvalue=And(z >= z_m))
+    v6 = simplify_and(g0, g5)
+    v7 = simplify_and(g0, g6)
+    v8 = simplify_and(g5, g6)
+    assert v6 == And(g0, g5)
+    assert v7 == And(g0, g6)
+    assert v8 == And(g5, g6)
+
+    g7 = GuardExpr('g7', initvalue=And(y >= -2, z >= -2))
+    v9 = simplify_and(g0, g7)
+    assert v9 == And(g0, g7)
+
+
+def test_guard_expr_Le_Ge_mixed():
+    grid = Grid(shape=(3, 3, 3))
+    _, y, z = grid.dimensions
+
+    y_m = y.symbolic_min
+    y_M = y.symbolic_max
+    z_m = z.symbolic_min
+    z_M = z.symbolic_max
+
+    g0 = GuardExpr('g0', initvalue=And(y <= y_M + 3, z <= z_M + 3))
+    g1 = GuardExpr('g1', initvalue=And(y >= y_m - 3, z >= z_m - 3))
+    v0 = simplify_and(g0, g1)
+    v1 = simplify_and(g1, g0)
+    assert v0 == And(g0, g1)
+    assert v1 == And(g0, g1)
+
+    g2 = GuardExpr('g2', initvalue=And(y <= y_M + 2, z >= z_m - 2))
+    v2 = simplify_and(g0, g2)
+    v3 = simplify_and(g2, g0)
+    assert v2 == And(g0, g2)
+    assert v3 == And(g0, g2)
+
+    g3 = GuardExpr('g3', initvalue=And(y >= y_m - 2, z <= z_M + 2))
+    v4 = simplify_and(g0, g3)
+    v5 = simplify_and(g3, g0)
+    assert v4 == And(g0, g3)
+    assert v5 == And(g0, g3)
+
+    g4 = GuardExpr('g4', initvalue=And(y <= y_M + 2, z >= z_m, z <= z_M + 2))
+    v6 = simplify_and(g0, g4)
+    v7 = simplify_and(g4, g0)
+    assert v6 is g4
+    assert v7 is g4
+
+    g5 = GuardExpr('g5', initvalue=And(y >= y_m - 2, y <= y_M + 2,
+                                       z >= z_m - 2, z <= z_M + 2))
+    v8 = simplify_and(g0, g5)
+    v9 = simplify_and(g5, g0)
+    assert v8 is g5
+    assert v9 is g5
+
+    g6 = GuardExpr('g6', initvalue=And(y >= y_m, y <= y_M,
+                                       z >= z_m - 3, z <= z_M))
+    v10 = simplify_and(g5, g6)
+    v11 = simplify_and(g6, g5)
+    assert v10 is And(g5, g6)
+    assert v11 is And(g5, g6)
+
+
 def test_canonical_ordering_of_weights():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions