compiler: add rudimentary support for multi-cond buffering

Author: mloubout

devito/ir/clusters/algorithms.py

@@ -254,6 +254,7 @@ def guard(clusters):
 
             # Chain together all `cds` conditions from all expressions in `c`
             guards = {}
+            mode = sympy.Or
             for cd in cds:
                 # `BOTTOM` parent implies a guard that lives outside of
                 # any iteration space, which corresponds to the placeholder None
@@ -270,6 +271,7 @@ def guard(clusters):
 
                 # Pull `cd` from any expr
                 condition = guards.setdefault(k, [])
+                mode = mode and cd.relation
                 for e in exprs:
                     try:
                         condition.append(e.conditionals[cd])
@@ -284,7 +286,8 @@ def guard(clusters):
                     conditionals.pop(cd, None)
                     exprs[i] = e.func(*e.args, conditionals=conditionals)
 
-            guards = {d: sympy.And(*v, evaluate=False) for d, v in guards.items()}
+            # Combination mode is And by default and Or if all conditions are
+            guards = {d: mode(*v, evaluate=False) for d, v in guards.items()}
 
             # Construct a guarded Cluster
             processed.append(c.rebuild(exprs=exprs, guards=guards))

devito/ir/equations/equation.py

@@ -223,7 +223,7 @@ def __new__(cls, *args, **kwargs):
             else:
                 cond = diff2sympy(lower_exprs(d.condition))
                 if d._factor is not None:
-                    cond = sympy.And(cond, GuardFactor(d))
+                    cond = d.relation(cond, GuardFactor(d))
                 conditionals[d] = cond
             # Replace dimension with index
             index = d.index

devito/types/dimension.py

@@ -860,6 +860,8 @@ class ConditionalDimension(DerivedDimension):
         If True, use `self`, rather than the parent Dimension, to
         index into arrays. A typical use case is when arrays are accessed
         indirectly via the ``condition`` expression.
+    relation: Or/And, default=And
+        How this ConditionalDimension will be combined with other ones.
 
     Examples
     --------
@@ -913,10 +915,10 @@ class ConditionalDimension(DerivedDimension):
     is_Conditional = True
 
     __rkwargs__ = DerivedDimension.__rkwargs__ + \
-        ('factor', 'condition', 'indirect')
+        ('factor', 'condition', 'indirect', 'relation')
 
     def __init_finalize__(self, name, parent=None, factor=None, condition=None,
-                          indirect=False, **kwargs):
+                          indirect=False, relation=sympy.And, **kwargs):
         # `parent=None` degenerates to a ConditionalDimension outside of
         # any iteration space
         if parent is None:
@@ -937,6 +939,7 @@ def __init_finalize__(self, name, parent=None, factor=None, condition=None,
 
         self._condition = condition
         self._indirect = indirect
+        self._relation = relation
 
     @property
     def uses_symbolic_factor(self):
@@ -978,6 +981,10 @@ def condition(self):
     def indirect(self):
         return self._indirect
 
+    @property
+    def relation(self):
+        return self._relation
+
     @cached_property
     def free_symbols(self):
         retval = set(super().free_symbols)

tests/test_buffering.py

@@ -1,9 +1,10 @@
 import numpy as np
 import pytest
+from sympy import Or
 
 from conftest import skipif
 from devito import (
-    ConditionalDimension, Constant, Eq, Grid, Operator, SubDimension, SubDomain,
+    CondEq, ConditionalDimension, Constant, Eq, Grid, Operator, SubDimension, SubDomain,
     TimeFunction, configuration, switchconfig
 )
 from devito.arch.archinfo import AppleArm
@@ -751,3 +752,36 @@ def test_buffer_reuse():
 
     assert all(np.all(usave.data[i-1] == i) for i in range(1, nt + 1))
     assert all(np.all(vsave.data[i-1] == i + 1) for i in range(1, nt + 1))
+
+
+def test_multi_cond():
+    grid = Grid((3, 3))
+    nt = 5
+
+    x, y = grid.dimensions
+
+    factor = 2
+    ntmod = (nt - 1) * factor + 1
+
+    ct1 = ConditionalDimension(name="ct1", parent=grid.time_dim,
+                               factor=factor, relation=Or)
+    ctend = ConditionalDimension(name="ctend", parent=grid.time_dim,
+                                 condition=CondEq(grid.time_dim, ntmod - 2),
+                                 relation=Or)
+
+    f = TimeFunction(grid=grid, name='f', time_order=0,
+                     space_order=0, save=nt, time_dim=ct1)
+    T = TimeFunction(grid=grid, name='T', time_order=0, space_order=0)
+
+    eqs = [Eq(T, grid.time_dim)]
+    # this to save times from 0 to nt - 2
+    eqs.append(Eq(f, T))
+    # this to save the last time sample nt - 1
+    eqs.append(Eq(f.forward, T+1, implicit_dims=ctend))
+
+    # run operator with buffering
+    op = Operator(eqs, opt=('streaming', 'buffering'))
+    op.apply(time_m=0, time_M=ntmod-2)
+
+    for i in range(nt):
+        assert np.allclose(f.data[i], i*2)