compiler: Add Guards.pairwise_or

Author: FabioLuporini

devito/ir/support/guards.py

@@ -4,6 +4,7 @@
 (e.g., Eq, Cluster, ...) should be evaluated at runtime.
 """
 
+from collections import Counter, defaultdict
 from operator import ge, gt, le, lt
 
 from functools import singledispatch
@@ -254,6 +255,20 @@ def xandg(self, d, guard):
 
         return Guards(m)
 
+    def pairwise_or(self, d, *guards):
+        m = dict(self)
+
+        if d in m:
+            guards.append(m[d])
+
+        g = pairwise_or(*guards)
+        if g is true:
+            m.pop(d, None)
+        else:
+            m[d] = g
+
+        return Guards(m)
+
     def impose(self, d, guard):
         m = dict(self)
 
@@ -405,3 +420,72 @@ def simplify_and(relation, v):
         new_args.append(v)
 
     return And(*(new_args + other))
+
+
+def pairwise_or(*guards):
+    """
+    Given a series of guards, create a new guard that combines them by taking
+    the OR of each subset of homogeneous components. Here, "homogeneous" means
+    that they apply to the same variable with the same operator (e.g., given
+    `y >= 2`, `y >= 3` is homogeneous, while `z >= 3` and `y <= 4` are not).
+
+    Examples
+    --------
+    Given:
+
+        g0 = {flag0 and z >= 2 and z <= 10 and y >= 3}
+        g1 = {z >= 4 and z <= 8}
+        g2 = {y >= 2 and y <= 5}
+
+    Where `flag0` is of type GuardExpr, then:
+
+    Return:
+
+        {z >= 2 and z <= 10 and y >= 2}
+    """
+    errmsg = lambda g: f"Cannot handle guard component of type {type(g)}"
+
+    flags = Counter()
+    mapper = defaultdict(list)
+
+    # Analysis
+    for guard in guards:
+        if guard is true or guard is None:
+            continue
+        elif isinstance(guard, And):
+            components = guard.args
+        elif isinstance(guard, GuardExpr) or guard.is_Relational:
+            components = [guard]
+        else:
+            raise NotImplementedError(errmsg(guard))
+
+        for g in components:
+            if isinstance(g, GuardExpr):
+                flags[g] += 1
+            elif g.is_Relational and g.lhs.is_Symbol and g.rhs.is_Number:
+                mapper[(g.lhs, type(g))].append(g.rhs)
+            else:
+                raise NotImplementedError(errmsg(g))
+
+    # Synthesis
+    guard = true
+    for (s, op), v in mapper.items():
+        if len(v) < len(guards):
+            # Not all guards contributed to this component; cannot simplify
+            pass
+        elif op in (Ge, Gt):
+            guard = And(guard, op(s, min(v)))
+        else:
+            guard = And(guard, op(s, max(v)))
+
+    for flag, v in flags.items():
+        if v == len(guards):
+            guard = And(guard, flag)
+        elif flag.initvalue.free_symbols & guard.free_symbols:
+            # We still lack the logic to properly handle this case
+            raise NotImplementedError(errmsg(flag))
+        else:
+            # Safe to ignore
+            pass
+
+    return guard

tests/test_symbolics.py

@@ -4,13 +4,13 @@
 import pytest
 import numpy as np
 
-from sympy import And, Expr, Number, Symbol
+from sympy import And, Expr, Number, Symbol, true
 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.ir.support.guards import GuardExpr, simplify_and, pairwise_or
 from devito.mpi.halo_scheme import HaloTouch
 from devito.symbolics import (
     retrieve_functions, retrieve_indexed, evalrel, CallFromPointer, Cast, # noqa
@@ -671,6 +671,33 @@ def test_guard_expr_Le_Ge_mixed():
     assert v11 is And(g5, g6)
 
 
+def test_guard_pairwise_or():
+    grid = Grid(shape=(3, 3, 3))
+    x, y, z = grid.dimensions
+
+    flag = GuardExpr('flag', initvalue=And(x >= 4, x <= 14))
+
+    g0 = And(flag, z >= 8, z <= 39)
+    g1 = And(z >= 8, z <= 40)
+    g2 = y >= 9
+    v0 = pairwise_or(g0, g1, g2)
+    assert v0 is true
+
+    g3 = And(z >= 7, z <= 40)
+    g4 = And(z >= 8, z <= 42, y >= 9)
+    v1 = pairwise_or(g0, g3, g4)
+    assert v1 == And(z >= 7, z <= 42)
+
+    # Some unsupported cases
+    g5 = And(flag, z >= 9, x >= 5)
+    g6 = x >= 3
+    with pytest.raises(NotImplementedError):
+        pairwise_or(g5, g6)
+    g7 = And(z <= y)
+    with pytest.raises(NotImplementedError):
+        pairwise_or(g0, g7)
+
+
 def test_canonical_ordering_of_weights():
     grid = Grid(shape=(3, 3, 3))
     x, y, z = grid.dimensions